Add support for Mage-Flow and Mage-Flow Edit

.gitignore

@@ -186,4 +186,5 @@ aitk_db.db-wal
 aitk_db.db-shm
 /notes.md
 /data
-.claude
+.claude
+original_repo

README.md

@@ -34,6 +34,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [circlestone-labs/Anima-Base-v1.0-Diffusers](https://huggingface.co/circlestone-labs/Anima-Base-v1.0-Diffusers) (Anima)
 - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2)
 - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo)
+- [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow)
 
 ### Instruction / Edit
 - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev)
@@ -44,6 +45,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [Boogu/Boogu-Image-0.1-Edit](https://huggingface.co/Boogu/Boogu-Image-0.1-Edit) (Boogu Image Edit)
 - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2 Edit Training)
 - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo Edit Training)
+- [microsoft/Mage-Flow-Edit-Base](https://huggingface.co/microsoft/Mage-Flow-Edit-Base) (Mage-Flow Edit)
 
 ### Video
 - [Wan-AI/Wan2.1-T2V-1.3B-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) (Wan 2.1 1.3B)

extensions_built_in/diffusion_models/__init__.py

@@ -18,6 +18,7 @@ from .ideogram4 import Ideogram4Model
 from .prx_pixel_t2i import PRXPixelT2IModel
 from .krea2 import Krea2Model
 from .boogu_image import BooguImageModel, BooguImageEditModel
+from .mageflow import MageFlowModel, MageFlowEditModel
 
 AI_TOOLKIT_MODELS = [
     # put a list of models here
@@ -51,4 +52,6 @@ AI_TOOLKIT_MODELS = [
     Krea2Model,
     BooguImageModel,
     BooguImageEditModel,
+    MageFlowModel,
+    MageFlowEditModel,
 ]

extensions_built_in/diffusion_models/mageflow/__init__.py

@@ -0,0 +1 @@
+from .mageflow import MageFlowModel, MageFlowEditModel

extensions_built_in/diffusion_models/mageflow/mageflow.py

@@ -0,0 +1,650 @@
+"""Mage-Flow (microsoft/Mage) for ai-toolkit.
+
+Two archs backed by the same NR-MMDiT stack:
+  - ``mageflow``       text-to-image  (microsoft/Mage-Flow-Base)
+  - ``mageflow_edit``  instruction edit (microsoft/Mage-Flow-Edit-Base)
+
+Components (diffusers-style repo layout):
+  - transformer: ``MageFlow`` dual-stream DiT (src/transformer.py) — packed
+    variable-length [text | image(+refs)] sequences, per-sample 2D RoPE,
+    joint varlen attention. 4B params, hidden 3072, 12 blocks.
+  - text encoder: Qwen3-VL (ships inside the repo under ``text_encoder/``);
+    conditioning is the final hidden states with the templated system prompt
+    dropped (34 tokens for t2i, 64 for edit).
+  - autoencoder: ``MageVAE`` (src/vae.py) — 128-channel, 16x downsample
+    one-step diffusion codec, no latent normalization.
+
+Flow matching convention matches ai-toolkit exactly (sigma=1 noise -> sigma=0
+clean, target = noise - clean), static sigma shift 6.0. Edit mode feeds the
+reference images in two places, as in the reference implementation: through
+the Qwen3-VL encoder alongside the instruction (long edge capped at 384) and
+as clean VAE latents sequence-appended after the noisy target tokens (at the
+target resolution).
+"""
+
+import json
+import math
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision.transforms.functional import to_tensor
+from safetensors.torch import load_file, save_file
+
+import huggingface_hub
+from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration
+from optimum.quanto import freeze
+
+from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from toolkit.models.base_model import BaseModel
+from toolkit.basic import flush
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+from toolkit.accelerator import unwrap_model
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.util.quantize import quantize, get_qtype, quantize_model
+from toolkit.memory_management import MemoryManager
+
+from .src.transformer import MageFlow, MageFlowParams
+from .src.vae import MageVAE
+from .src.text_encoder import (
+    encode_mageflow_prompt,
+    patch_qwen_vl_patch_embed,
+    resize_vl_images,
+)
+from .src.pipeline import MageFlowPipeline, predict_velocity
+
+if TYPE_CHECKING:
+    from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
+
+
+# Training/sampling timestep distribution: the released checkpoints use a
+# static sigma shift of 6.0 (transformer/config.json + scheduler config).
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "use_dynamic_shifting": False,
+    "shift": 6.0,
+}
+
+# Keys of the checkpoint transformer/config.json that are NOT MageFlowParams
+# constructor args (legacy/unused fields) — same filter as the reference
+# ``load_from_repo``. Everything else becomes the DiT structure.
+_CONFIG_META_KEYS = {
+    "_class_name",
+    "txt_max_length",
+    "max_sequence_length",
+    "param_dtype",
+    "packing",
+    "schedule_mode",
+    "static_shift",
+    "use_time_shift",
+    "rope_type",
+    "apply_text_rotary_emb",
+    "mlp_ratio",
+    "depth_single_blocks",
+    "theta",
+    "qkv_bias",
+    "guidance_embed",
+    "vec_in_dim",
+    "vec_type",
+    "time_type",
+    "double_block_type",
+}
+
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
+
+class MageFlowModel(BaseModel):
+    arch = "mageflow"
+    is_edit = False
+    use_old_lokr_format = False
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_flow_matching = True
+        self.is_transformer = True
+        self.target_lora_modules = ["MageFlow"]
+
+        # MageVAE is 16x downsampling; the DiT patch size is 1 (one token per
+        # latent pixel), so pixel sizes must be multiples of 16.
+        self.patch_size = 1
+        self.vae_scale_factor = 16
+        # Safety cap on prompt token length (truncation only); embeds are stored
+        # per-sample at natural length and packed varlen at the model call.
+        self.max_text_length = int(
+            self.model_config.model_kwargs.get("max_text_length", 2048)
+        )
+        # Qwen3-VL AutoProcessor for encoding reference images into the prompt
+        # (edit only).
+        self.vl_processor = None
+
+        if self.is_edit:
+            # Reference images feed the model in two places: through the
+            # Qwen3-VL encoder alongside the instruction, and as clean VAE
+            # latents sequence-appended after the noisy target tokens.
+            self.encode_control_in_text_embeddings = True
+            self.has_multiple_control_images = True
+            # References keep their own aspect/size in the dataloader; they are
+            # resized to the target resolution here (reference behavior).
+            self.use_raw_control_images = True
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+
+    def get_bucket_divisibility(self):
+        # 16 for the VAE downsample, patch size is 1.
+        return self.vae_scale_factor * self.patch_size
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def _get_model_file(self, relpath: str) -> str:
+        """Resolve a repo-relative file from a local directory or the HF hub."""
+        name_or_path = self.model_config.name_or_path
+        if os.path.isdir(name_or_path):
+            path = os.path.join(name_or_path, relpath)
+            if not os.path.exists(path):
+                raise FileNotFoundError(
+                    f"Could not find {relpath!r} in {name_or_path!r}"
+                )
+            return path
+        return huggingface_hub.hf_hub_download(
+            repo_id=name_or_path, filename=relpath, token=HF_TOKEN
+        )
+
+    def _load_transformer(self) -> MageFlow:
+        dtype = self.torch_dtype
+        self.print_and_status_update("Loading transformer (MageFlow NR-MMDiT)")
+
+        with open(self._get_model_file("transformer/config.json")) as f:
+            tcfg = json.load(f)
+        structure = {k: v for k, v in tcfg.items() if k not in _CONFIG_META_KEYS}
+        structure.update(self.model_config.model_kwargs.get("transformer_config", {}))
+        params = MageFlowParams(**structure)
+
+        # Build on meta, then materialize straight from the checkpoint.
+        with torch.device("meta"):
+            transformer = MageFlow(params)
+
+        self.print_and_status_update("  - fetching transformer weights")
+        state_dict = load_file(
+            self._get_model_file("transformer/diffusion_pytorch_model.safetensors")
+        )
+        state_dict = {
+            k: (v.to(dtype) if v.is_floating_point() else v)
+            for k, v in state_dict.items()
+        }
+        self.print_and_status_update("  - loading transformer state dict")
+        transformer.load_state_dict(state_dict, strict=True, assign=True)
+        # The RoPE tables are plain tensors (not buffers/params), so the meta
+        # init left them unmaterialized — rebuild them for real.
+        transformer.reset_rope()
+        del state_dict
+        flush()
+        return transformer
+
+    def _load_text_encoder(self):
+        dtype = self.torch_dtype
+        te_path = self.model_config.model_kwargs.get("text_encoder_path", None)
+        if te_path is not None:
+            te_kwargs = {}
+        else:
+            te_path = self.model_config.name_or_path
+            te_kwargs = {"subfolder": "text_encoder"}
+        self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
+
+        tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN, **te_kwargs)
+        text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+            te_path, torch_dtype=dtype, token=HF_TOKEN, **te_kwargs
+        )
+        vl_processor = None
+        if self.is_edit:
+            # Edit mode: reference images are encoded into the text embeddings,
+            # so the vision tower stays. Swap its Conv3d patch_embed for an
+            # equivalent GEMM (bf16 Conv3d has no fast cuDNN kernel).
+            vl_processor = AutoProcessor.from_pretrained(
+                te_path, token=HF_TOKEN, **te_kwargs
+            )
+            patch_qwen_vl_patch_embed(text_encoder)
+        else:
+            # We only ever encode text, so the vision tower is dead weight --
+            # drop it to free VRAM.
+            if getattr(text_encoder.model, "visual", None) is not None:
+                text_encoder.model.visual = None
+        text_encoder.eval()
+        text_encoder.requires_grad_(False)
+        flush()
+        return tokenizer, vl_processor, text_encoder
+
+    def _load_vae(self) -> MageVAE:
+        self.print_and_status_update("Loading MageVAE")
+        vae_path = self.model_config.model_kwargs.get("vae_path", None)
+        if vae_path is None:
+            vae_path = self._get_model_file("vae/diffusion_pytorch_model.safetensors")
+        vae = MageVAE(
+            ckpt_path=vae_path,
+            sample_posterior=bool(
+                self.model_config.model_kwargs.get("vae_sample_posterior", True)
+            ),
+        )
+        vae.eval()
+        vae.requires_grad_(False)
+        return vae
+
+    def get_quantization_exclude_modules(self):
+        # sensitive modules kept in full precision (fnmatch patterns on module
+        # names within MageFlow):
+        #   img_in / txt_in / txt_norm - input projections
+        #   time_text_embed*           - timestep embedder feeding every
+        #                                block's modulation
+        #   norm_out* / proj_out       - final adaptive norm / output projection
+        return [
+            "img_in",
+            "txt_in",
+            "txt_norm",
+            "time_text_embed*",
+            "norm_out*",
+            "proj_out",
+        ]
+
+    def load_model(self):
+        dtype = self.torch_dtype
+        self.print_and_status_update("Loading Mage-Flow model")
+
+        transformer = self._load_transformer()
+
+        if self.model_config.quantize:
+            self.print_and_status_update("Quantizing transformer")
+            quantize_model(self, transformer)
+            flush()
+
+        if (
+            self.model_config.layer_offloading
+            and self.model_config.layer_offloading_transformer_percent > 0
+        ):
+            MemoryManager.attach(
+                transformer,
+                self.device_torch,
+                offload_percent=self.model_config.layer_offloading_transformer_percent,
+            )
+
+        if self.model_config.low_vram:
+            self.print_and_status_update("Moving transformer to CPU")
+            transformer.to("cpu")
+        else:
+            transformer.to(self.device_torch, dtype=dtype)
+        flush()
+
+        tokenizer, vl_processor, text_encoder = self._load_text_encoder()
+        if self.model_config.quantize_te:
+            self.print_and_status_update("Quantizing text encoder")
+            text_encoder.to(self.device_torch)
+            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
+            freeze(text_encoder)
+            flush()
+        if (
+            self.model_config.layer_offloading
+            and self.model_config.layer_offloading_text_encoder_percent > 0
+        ):
+            MemoryManager.attach(
+                text_encoder,
+                self.device_torch,
+                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
+            )
+
+        if self.model_config.low_vram:
+            self.print_and_status_update("Moving text encoder to CPU")
+            text_encoder.to("cpu")
+        else:
+            text_encoder.to(self.device_torch)
+        flush()
+
+        vae = self._load_vae()
+        vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
+
+        self.noise_scheduler = MageFlowModel.get_train_scheduler()
+
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.vl_processor = vl_processor
+        self.model = transformer
+        self.pipeline = MageFlowPipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # Generation (training previews)
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return MageFlowPipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: MageFlowPipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: AdvancedPromptEmbeds,
+        unconditional_embeds: AdvancedPromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        sc = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // sc * sc)
+        gen_config.height = int(gen_config.height // sc * sc)
+
+        # Reference image(s) -> clean VAE latents appended to the sequence.
+        # The Qwen3-VL side already saw them (baked into the prompt embeds).
+        # ctrl_img_1 mirrors ctrl_img when unset, so use one or the other.
+        ctrl_paths = []
+        if self.is_edit:
+            if gen_config.ctrl_img is not None:
+                ctrl_paths.append(gen_config.ctrl_img)
+            elif gen_config.ctrl_img_1 is not None:
+                ctrl_paths.append(gen_config.ctrl_img_1)
+            if gen_config.ctrl_img_2 is not None:
+                ctrl_paths.append(gen_config.ctrl_img_2)
+            if gen_config.ctrl_img_3 is not None:
+                ctrl_paths.append(gen_config.ctrl_img_3)
+
+        ref_latents = None
+        if ctrl_paths:
+            ctrl_tensors = [
+                to_tensor(Image.open(path).convert("RGB")) for path in ctrl_paths
+            ]
+            # one batch item (preview batch size is 1) -> List[List[(128, h, w)]]
+            ref_latents = [
+                self._encode_ref_latents(
+                    ctrl_tensors, gen_config.height, gen_config.width
+                )
+            ]
+
+        img = pipeline(
+            conditional_embeds=conditional_embeds,
+            unconditional_embeds=unconditional_embeds,
+            height=gen_config.height,
+            width=gen_config.width,
+            num_inference_steps=gen_config.num_inference_steps,
+            guidance_scale=gen_config.guidance_scale,
+            latents=gen_config.latents,
+            generator=generator,
+            ref_latents=ref_latents,
+        )[0]
+        return img
+
+    # ------------------------------------------------------------------
+    # Reference-image helpers (edit)
+    # ------------------------------------------------------------------
+    def _encode_ref_latents(
+        self, control_tensors, target_height: int, target_width: int
+    ) -> List[torch.Tensor]:
+        """Encode ``[0, 1]`` reference image tensors to clean VAE latents.
+
+        Sizing (deliberate deviation from the reference ``generate_edits``,
+        which resizes every ref to exactly the target resolution and squishes
+        mismatched aspect ratios): each reference is resized to the TARGET's
+        total pixel count while keeping its OWN aspect ratio, snapped to the
+        16px divisibility. A ref that shares the target's aspect ratio still
+        lands on exactly the target size (identical to the reference
+        behavior); one with a different aspect ratio keeps its shape instead
+        of being distorted. Returns a list of ``(128, h, w)`` latents (one per
+        reference image). ``control_tensors`` is a list of ``(C, H, W)`` or
+        ``(1, C, H, W)`` tensors in ``[0, 1]``.
+        """
+        sc = self.get_bucket_divisibility()  # 16
+        target_area = target_height * target_width
+        latents = []
+        for img in control_tensors:
+            if img.dim() == 3:
+                img = img.unsqueeze(0)
+            img = img.to(self.device_torch, dtype=torch.float32)
+
+            h, w = img.shape[2], img.shape[3]
+            ratio = h / w
+            new_h = math.sqrt(target_area * ratio)
+            new_w = new_h / ratio
+            new_h = max(sc, int(round(new_h / sc)) * sc)
+            new_w = max(sc, int(round(new_w / sc)) * sc)
+            if (new_h, new_w) != (h, w):
+                img = F.interpolate(
+                    img,
+                    size=(new_h, new_w),
+                    mode="bicubic",
+                    antialias=True,
+                ).clamp(0, 1)
+            # encode_images expects [-1, 1]; control tensors arrive in [0, 1].
+            latent = self.encode_images(
+                img * 2 - 1, device=self.device_torch, dtype=self.torch_dtype
+            )
+            latents.append(latent[0])  # drop batch dim -> (128, h, w)
+        return latents
+
+    def _batch_ref_latents_from_batch(
+        self,
+        batch: "DataLoaderBatchDTO",
+        batch_size: int,
+        target_height: int,
+        target_width: int,
+    ) -> Optional[List[List[torch.Tensor]]]:
+        """Build predict_velocity's ``ref_latents`` from a train batch."""
+        control_list = batch.control_tensor_list
+        if control_list is None and batch.control_tensor is not None:
+            control_list = [batch.control_tensor[b : b + 1] for b in range(batch_size)]
+        if control_list is None:
+            return None
+        if len(control_list) != batch_size:
+            raise ValueError("Control tensor list length does not match batch size")
+        ref_latents = []
+        for controls in control_list:
+            if isinstance(controls, torch.Tensor):
+                controls = [controls]
+            ref_latents.append(
+                self._encode_ref_latents(controls, target_height, target_width)
+            )
+        return ref_latents
+
+    # ------------------------------------------------------------------
+    # Training hooks
+    # ------------------------------------------------------------------
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,  # (B, 128, h, w)
+        timestep: torch.Tensor,  # 0..1000 scale
+        text_embeddings: AdvancedPromptEmbeds,
+        batch: "DataLoaderBatchDTO" = None,
+        **kwargs,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        # Clean reference latents from the batch's control images (if any);
+        # they ride along in the sequence, never noised.
+        ref_latents = None
+        if batch is not None and self.is_edit:
+            with torch.no_grad():
+                _, _, lh, lw = latent_model_input.shape
+                ref_latents = self._batch_ref_latents_from_batch(
+                    batch,
+                    latent_model_input.shape[0],
+                    target_height=lh * self.vae_scale_factor,
+                    target_width=lw * self.vae_scale_factor,
+                )
+
+        # toolkit timestep (0..1000, 1000 = pure noise) -> flow sigma in
+        # [0, 1] with 1 = pure noise. Same convention -> straight divide.
+        t = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
+        if t.dim() == 0:
+            t = t.unsqueeze(0)
+        if t.shape[0] != latent_model_input.shape[0]:
+            t = t.expand(latent_model_input.shape[0])
+
+        pred = predict_velocity(
+            self.transformer,
+            latent_model_input.to(self.device_torch, self.torch_dtype),
+            t,
+            text_embeddings.text_embeds,
+            ref_latents=ref_latents,
+        )
+        return pred
+
+    def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+
+        # Normalize control images to a per-prompt list (List[List[Tensor]]).
+        # They arrive as a (B, C, H, W) batch tensor (control_tensor), a list of
+        # per-sample lists (control_tensor_list), or a flat list of (1, C, H, W)
+        # tensors for a single prompt (sampling / blank-embed caching).
+        if control_images is not None:
+            if isinstance(control_images, torch.Tensor):
+                control_images = [
+                    [control_images[i]] for i in range(control_images.shape[0])
+                ]
+            elif len(control_images) > 0 and not isinstance(control_images[0], list):
+                control_images = [control_images]
+            if len(control_images) == 1 and len(prompt) > 1:
+                control_images = control_images * len(prompt)
+            if len(control_images) != len(prompt):
+                raise ValueError(
+                    "Number of prompts must match number of control image sets"
+                )
+        else:
+            control_images = [None] * len(prompt)
+
+        template_name = "mage-flow-edit" if self.is_edit else "mage-flow"
+        vl_long_edge = int(self.model_config.model_kwargs.get("vl_cond_long_edge", 384))
+
+        # Encode each prompt at its natural length and store one (L, 2560)
+        # tensor per batch item. Padding is never needed: the packed varlen
+        # model call consumes the per-sample lengths directly.
+        features_list = []
+        for p, ctrl in zip(prompt, control_images):
+            images = None
+            if self.is_edit and ctrl is not None and len(ctrl) > 0:
+                images = resize_vl_images(ctrl, max_long_edge=vl_long_edge)
+            features = encode_mageflow_prompt(
+                self.text_encoder,
+                self.tokenizer,
+                p,
+                template_name=template_name,
+                max_length=self.max_text_length,
+                images=images,
+                processor=self.vl_processor,
+                dtype=self.torch_dtype,
+            )
+            features_list.append(features)
+
+        return AdvancedPromptEmbeds(text_embeds=features_list)
+
+    def get_loss_target(self, *args, **kwargs):
+        # Flow-matching velocity target: noise - clean.
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        return (noise - batch.latents).detach()
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # VAE (MageVAE -- raw latents, no normalization)
+    # ------------------------------------------------------------------
+    def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        self.vae.eval()
+
+        if isinstance(image_list, list):
+            images = torch.stack(
+                [img.squeeze(0) if img.dim() == 4 else img for img in image_list], dim=0
+            )
+        else:
+            images = image_list
+        images = images.to(self.vae.device, dtype=self.vae.dtype)
+
+        latents = self.vae.encode(images)  # (B, 128, H/16, W/16)
+        return latents.to(device, dtype=dtype)
+
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+
+        latents = latents.to(self.vae.device, dtype=self.vae.dtype)
+        images = self.vae.decode(latents)  # (B, 3, H, W) in [-1, 1]
+        return images.to(device, dtype=dtype)
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def save_model(self, output_path, meta, save_dtype):
+        from toolkit.util.quantize import dequantize_if_quantized
+
+        if not output_path.endswith(".safetensors"):
+            output_path = output_path + ".safetensors"
+        transformer: MageFlow = unwrap_model(self.model)
+        state_dict = transformer.state_dict()
+        save_dict = {}
+        for k, v in state_dict.items():
+            # dequantize any quantized (e.g. quanto/torchao) weights so we save plain full precision tensors
+            save_dict[k] = (
+                dequantize_if_quantized(v).clone().to("cpu", dtype=save_dtype)
+            )
+        meta = get_meta_for_safetensors(meta, name=self.arch)
+        save_file(save_dict, output_path, metadata=meta)
+
+    def get_base_model_version(self):
+        return "mageflow"
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["transformer_blocks"]
+
+    def convert_lora_weights_before_save(self, state_dict):
+        return {
+            k.replace("transformer.", "diffusion_model."): v
+            for k, v in state_dict.items()
+        }
+
+    def convert_lora_weights_before_load(self, state_dict):
+        return {
+            k.replace("diffusion_model.", "transformer."): v
+            for k, v in state_dict.items()
+        }
+
+
+class MageFlowEditModel(MageFlowModel):
+    arch = "mageflow_edit"
+    is_edit = True
+
+    def get_base_model_version(self):
+        return "mageflow_edit"

extensions_built_in/diffusion_models/mageflow/src/attn.py

@@ -0,0 +1,73 @@
+"""Varlen attention shim for Mage-Flow.
+
+The reference implementation (github.com/microsoft/Mage) runs every attention
+call through flash-attn's ``flash_attn_varlen_func`` so that several
+variable-length samples can be packed into one sequence and kept isolated via
+``cu_seqlens``. This shim exposes the same function with the FA2 calling
+convention, falling back to a per-sequence ``torch.scaled_dot_product_attention``
+loop when flash-attn is not installed (functionally equivalent for the dense /
+non-causal path the DiT uses, just slower).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import torch
+import torch.nn.functional as F
+
+_RESOLVED_FN: Callable[..., Any] | None = None
+
+
+def _sdpa_varlen(
+    q,
+    k,
+    v,
+    cu_seqlens_q=None,
+    cu_seqlens_k=None,
+    max_seqlen_q=None,
+    max_seqlen_k=None,
+    dropout_p: float = 0.0,
+    softmax_scale=None,
+    causal: bool = False,
+    **_unused: Any,
+):
+    """FA2-varlen-compatible SDPA fallback: one SDPA dispatch per sequence."""
+    cu_q = cu_seqlens_q.tolist()
+    cu_k = cu_seqlens_k.tolist()
+    outs = []
+    for qs, qe, ks, ke in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]):
+        # (s, h, d) -> (1, h, s, d)
+        q_i = q[qs:qe].transpose(0, 1).unsqueeze(0)
+        k_i = k[ks:ke].transpose(0, 1).unsqueeze(0)
+        v_i = v[ks:ke].transpose(0, 1).unsqueeze(0)
+        out_i = F.scaled_dot_product_attention(
+            q_i,
+            k_i,
+            v_i,
+            attn_mask=None,
+            dropout_p=dropout_p,
+            is_causal=causal,
+            scale=softmax_scale,
+        )
+        outs.append(out_i.squeeze(0).transpose(0, 1))
+    return torch.cat(outs, dim=0).contiguous()
+
+
+def _resolve() -> Callable[..., Any]:
+    global _RESOLVED_FN
+    if _RESOLVED_FN is None:
+        try:
+            from flash_attn import flash_attn_varlen_func as _fn
+
+            _RESOLVED_FN = _fn
+        except ImportError:
+            _RESOLVED_FN = _sdpa_varlen
+    return _RESOLVED_FN
+
+
+def flash_attn_varlen_func(*args, **kwargs):
+    return _resolve()(*args, **kwargs)
+
+
+__all__ = ["flash_attn_varlen_func"]

extensions_built_in/diffusion_models/mageflow/src/pipeline.py

@@ -0,0 +1,229 @@
+"""Packing / sampling helpers for Mage-Flow.
+
+Turns latents + Qwen3-VL text features into the packed variable-length
+sequences the ``MageFlow`` DiT consumes (batch dim 1, per-sample
+``cu_seqlens``, exactly mirroring the reference training/inference packing),
+and provides a minimal flow-matching sampler for ai-toolkit preview images.
+
+Time convention: rectified flow with ``x_t = (1 - sigma) * clean + sigma *
+noise``, the model predicts the velocity ``noise - clean``, and sampling
+integrates sigma from 1 (pure noise) down to 0 (clean). This is identical to
+ai-toolkit's convention, so the toolkit ``timestep / 1000`` flows straight
+through as sigma. The inference sigma schedule is the reference's static
+shift: ``shifted = shift * s / (1 + (shift - 1) * s)`` over ``s =
+linspace(1, 1/steps, steps)`` with a terminal 0 (shift 6.0).
+
+Edit conditioning: each sample's sequence is ``[target, ref_1, …, ref_N]`` —
+clean reference latents ride along after the noisy target tokens, sharing the
+sample's timestep modulation. RoPE gives every image segment in the pack its
+own "frame" index (segment order), so references land on later frame
+coordinates than their target, as in the reference ``generate_edits``.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import torch
+from einops import rearrange
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+from .transformer import MageFlow
+
+
+def lens_to_cu(lens: List[int], device) -> torch.Tensor:
+    """Sequence lengths -> cumulative cu_seqlens [0, l0, l0+l1, ...] (int32)."""
+    t = torch.tensor(lens, device=device, dtype=torch.int32)
+    return torch.cat(
+        [
+            torch.zeros(1, dtype=torch.int32, device=device),
+            torch.cumsum(t, dim=0, dtype=torch.int32),
+        ]
+    )
+
+
+def pack_text_features(
+    features_list: List[torch.Tensor], device, dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Concatenate per-sample ``(L_i, D)`` text features into the packed
+    ``[1, sum L, D]`` text stream + its ``cu_seqlens``. No padding — varlen
+    attention isolates the samples."""
+    feats = [f.to(device, dtype) for f in features_list]
+    txt = torch.cat(feats, dim=0).unsqueeze(0)
+    txt_cu = lens_to_cu([f.shape[0] for f in feats], device)
+    return txt, txt_cu
+
+
+def predict_velocity(
+    model: MageFlow,
+    latents: torch.Tensor,  # (B, C, h, w) noisy target latents
+    t: torch.Tensor,  # (B,) flow sigma in [0, 1] (1 = pure noise)
+    text_embeds: List[torch.Tensor],  # per-sample (L_i, D) Qwen3-VL features
+    ref_latents: Optional[
+        List[List[torch.Tensor]]
+    ] = None,  # per-sample clean (C, hr, wr) refs
+) -> torch.Tensor:
+    """Run the DiT on the packed [text | target(+refs)] sequences.
+
+    Packs the batch into one varlen forward (mirroring reference packing:
+    per-sample cu_seqlens isolate the samples inside the attention kernel).
+    Reference latents, when given, are appended clean after each sample's
+    target tokens; the returned velocity covers only the target tokens,
+    reshaped back to ``(B, C, h, w)``. No time flip / negation: Mage-Flow's
+    convention matches ai-toolkit's.
+    """
+    device = latents.device
+    dtype = latents.dtype
+    b, c, h, w = latents.shape
+
+    if ref_latents is not None and not any(len(r) > 0 for r in ref_latents):
+        ref_latents = None
+
+    img_parts, samp_lens, shape_seq, target_idx_parts = [], [], [], []
+    off = 0
+    for i in range(b):
+        tgt = rearrange(latents[i], "c h w -> (h w) c")
+        img_parts.append(tgt)
+        shape_seq.append((1, h, w))
+        target_idx_parts.append(torch.arange(off, off + h * w, device=device))
+        samp_len = h * w
+        off += h * w
+        if ref_latents is not None:
+            for ref in ref_latents[i]:
+                ref = ref.to(device, dtype)
+                _, rh, rw = ref.shape
+                img_parts.append(rearrange(ref, "c h w -> (h w) c"))
+                shape_seq.append((1, rh, rw))
+                samp_len += rh * rw
+                off += rh * rw
+        samp_lens.append(samp_len)
+
+    img = torch.cat(img_parts, dim=0).unsqueeze(0)  # [1, sum, C]
+    img_cu = lens_to_cu(samp_lens, device)
+    img_shapes = [shape_seq]
+    target_idx = torch.cat(target_idx_parts)
+
+    txt, txt_cu = pack_text_features(text_embeds, device, dtype)
+
+    out = model(
+        img=img,
+        txt=txt,
+        timesteps=t.to(device),
+        img_shapes=img_shapes,
+        img_cu_seqlens=img_cu,
+        txt_cu_seqlens=txt_cu,
+    )  # [1, sum, C_out]
+
+    pred = out[:, target_idx, :]  # [1, B*h*w, C]
+    return rearrange(pred, "1 (b h w) c -> b c h w", b=b, h=h, w=w)
+
+
+def build_shifted_sigmas(
+    num_steps: int, shift: float = 6.0, device=None
+) -> torch.Tensor:
+    """The reference inference sigma schedule: ``linspace(1, 1/steps, steps)``
+    run through the static shift ``shift*s/(1+(shift-1)*s)`` with a terminal 0."""
+    s = torch.linspace(1.0, 1.0 / num_steps, num_steps, dtype=torch.float64)
+    s = shift * s / (1 + (shift - 1) * s)
+    sigmas = torch.cat([s, torch.zeros(1, dtype=torch.float64)])
+    return sigmas.to(device=device, dtype=torch.float32)
+
+
+class MageFlowPipeline:
+    """Lightweight flow-matching sampler used by ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # ``model`` is the MageFlowModel (BaseModel) so we can reuse its
+        # encode/decode and config.
+        self.model = model
+
+    @property
+    def device(self):
+        return self.model.device_torch
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def set_progress_bar_config(self, **kwargs):
+        pass
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        conditional_embeds,
+        unconditional_embeds,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 5.0,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        ref_latents: Optional[List[List[torch.Tensor]]] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        model = self.model
+        device = model.device_torch
+        dtype = model.torch_dtype
+        transformer: MageFlow = model.transformer
+
+        shift = float(model.model_config.model_kwargs.get("static_shift", 6.0))
+        renorm = bool(model.model_config.model_kwargs.get("cfg_renormalization", False))
+
+        do_cfg = guidance_scale > 1.0 and unconditional_embeds is not None
+
+        gh = height // model.vae_scale_factor
+        gw = width // model.vae_scale_factor
+        latent_channels = transformer.in_channels
+
+        if latents is None:
+            latents = randn_tensor(
+                (1, latent_channels, gh, gw),
+                generator=generator,
+                device=device,
+                dtype=torch.float32,
+            )
+        latents = latents.to(device, dtype=torch.float32)
+
+        cond_feats = [f.to(device, dtype) for f in conditional_embeds.text_embeds]
+        if do_cfg:
+            uncond_feats = [
+                f.to(device, dtype) for f in unconditional_embeds.text_embeds
+            ]
+
+        sigmas = build_shifted_sigmas(num_inference_steps, shift=shift, device=device)
+
+        # Euler integration of the flow ODE, sigma 1 -> 0 (with optional CFG:
+        # v = uncond + cfg * (cond - uncond), reference convention).
+        for i in range(num_inference_steps):
+            s_cur = sigmas[i].item()
+            s_next = sigmas[i + 1].item()
+            t = torch.full((latents.shape[0],), s_cur, dtype=dtype, device=device)
+            v_cond = predict_velocity(
+                transformer, latents.to(dtype), t, cond_feats, ref_latents=ref_latents
+            )
+            if do_cfg:
+                v_uncond = predict_velocity(
+                    transformer,
+                    latents.to(dtype),
+                    t,
+                    uncond_feats,
+                    ref_latents=ref_latents,
+                )
+                v = v_uncond + guidance_scale * (v_cond - v_uncond)
+                if renorm:
+                    # Rescale the guided velocity per token back to the
+                    # conditional velocity's norm (reduces oversaturation).
+                    v = v * (
+                        torch.norm(v_cond, dim=1, keepdim=True)
+                        / (torch.norm(v, dim=1, keepdim=True) + 1e-6)
+                    )
+            else:
+                v = v_cond
+            latents = latents + (s_next - s_cur) * v.to(torch.float32)
+
+        images = model.decode_latents(latents, device=device, dtype=dtype)
+        images = images.float().clamp(-1.0, 1.0)
+        images = ((images + 1.0) * 127.5).round().to(torch.uint8)
+        images = images.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(arr) for arr in images]

extensions_built_in/diffusion_models/mageflow/src/text_encoder.py

@@ -0,0 +1,169 @@
+"""Qwen3-VL prompt encoding for Mage-Flow.
+
+The reference implementation (mage_flow/models/modules/text_encoder.py) packs
+many prompts into one varlen forward; per-sequence the math reduces to a plain
+causal forward with positions ``arange(L)`` replicated on every mrope axis and
+the leading system-prompt tokens dropped from the output. ai-toolkit encodes
+prompts one at a time (embeds are cached per caption), so that per-sequence
+form is implemented here directly on the stock HF Qwen3-VL module.
+
+The DiT consumes the final hidden states ``[L - drop_idx, 2560]``; the pooled
+"vec" of the reference is unused by the released checkpoints (vec_type null —
+the transformer adds a zero vector), so it is not computed.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+# Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx``
+# is the number of leading (system prompt) tokens dropped from the encoder
+# output before it conditions the DiT.
+PROMPT_TEMPLATE = {
+    "mage-flow": {
+        "template": (
+            "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, "
+            "text, spatial relationships of the objects and background:"
+            "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        ),
+        "start_idx": 34,
+    },
+    "mage-flow-edit": {
+        "template": (
+            "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture,"
+            " objects, background), then explain how the user's text instruction should alter or modify the image. "
+            "Generate a new image that meets the user's requirements while maintaining consistency with the original "
+            "input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        ),
+        "start_idx": 64,
+    },
+}
+
+# Fixed image placeholder used at edit training time (one per reference image).
+EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
+
+
+def edit_prompt_body(instruction: str, num_refs: int) -> str:
+    """Training-time multi-reference prompt body:
+    ``Image 1: <ph>Image 2: <ph>…{instruction}``."""
+    prefix = "".join(
+        f"Image {j}: {EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1)
+    )
+    return prefix + instruction
+
+
+def patch_qwen_vl_patch_embed(model):
+    """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain
+    linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
+    falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear
+    (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves.
+    Returns the number of patch_embed modules patched. (Same patch as the krea2
+    extension / Qwen3VLCaptioner.)"""
+    patched = 0
+    for module in model.modules():
+        proj = getattr(module, "proj", None)
+        if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
+            proj.stride
+        ):
+
+            def fast_forward(hidden_states, _proj=proj):
+                w = _proj.weight.reshape(_proj.weight.shape[0], -1)
+                x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
+                return F.linear(x, w, _proj.bias)
+
+            module.forward = fast_forward
+            patched += 1
+    return patched
+
+
+def resize_vl_images(
+    images: List[torch.Tensor], max_long_edge: int = 384
+) -> List["PIL.Image.Image"]:
+    """Prepare reference images for the Qwen3-VL conditioning pass.
+
+    Matches the reference's ``_resize_long_edge`` (cap the long edge at
+    ``max_long_edge``, preserving aspect ratio, BICUBIC, never upscale) — the
+    MLLM only needs a coarse view of the reference; full-resolution detail
+    flows through the VAE reference latents. Input tensors are ``(C, H, W)``
+    or ``(1, C, H, W)`` in [0, 1]; output is PIL for the Qwen processor.
+    """
+    from torchvision.transforms.functional import to_pil_image
+
+    out = []
+    for img in images:
+        if img.dim() == 4:
+            img = img[0]
+        img = img.float().clamp(0, 1).cpu()
+        pil = to_pil_image(img)
+        if max_long_edge and max_long_edge > 0:
+            w, h = pil.size
+            long_edge = max(w, h)
+            if long_edge > max_long_edge:
+                scale = max_long_edge / long_edge
+                new_w = max(1, int(round(w * scale)))
+                new_h = max(1, int(round(h * scale)))
+                pil = pil.resize((new_w, new_h), resample=3)  # PIL.Image.BICUBIC
+        out.append(pil)
+    return out
+
+
+@torch.no_grad()
+def encode_mageflow_prompt(
+    text_encoder,  # Qwen3VLForConditionalGeneration
+    tokenizer,
+    prompt: str,
+    template_name: str = "mage-flow",
+    max_length: int = 2048,
+    images: Optional[list] = None,  # list of PIL images (already VL-resized)
+    processor=None,  # AutoProcessor, required when images are given
+    dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Encode one prompt (optionally with reference images) to DiT conditioning.
+
+    Returns ``(L - start_idx, 2560)`` final hidden states with the system
+    prompt dropped, matching the reference ``TextEncoder.forward`` txt output.
+    """
+    info = PROMPT_TEMPLATE[template_name]
+    template = info["template"]
+    drop_idx = int(info["start_idx"])
+    device = next(text_encoder.parameters()).device
+
+    if images:
+        if processor is None:
+            raise ValueError(
+                "encoding reference images requires the Qwen3-VL AutoProcessor"
+            )
+        formatted = template.format(edit_prompt_body(prompt, len(images)))
+        vl = processor(text=[formatted], images=list(images), return_tensors="pt")
+        input_ids = vl["input_ids"].to(device)
+        extra = {}
+        for key in ("pixel_values", "image_grid_thw"):
+            if vl.get(key) is not None:
+                extra[key] = vl[key].to(device)
+    else:
+        formatted = template.format(prompt)
+        input_ids = tokenizer(
+            formatted,
+            max_length=max_length + drop_idx,
+            truncation=True,
+            return_tensors="pt",
+        ).input_ids.to(device)
+        extra = {}
+
+    # The reference packed encoder gives every sequence positions arange(L)
+    # replicated on all mrope axes (even for image tokens); a 2D position_ids
+    # is expanded to exactly that by the HF text model.
+    position_ids = torch.arange(input_ids.shape[1], device=device).unsqueeze(0)
+
+    out = text_encoder.model(
+        input_ids=input_ids,
+        position_ids=position_ids,
+        **extra,
+    )
+    hidden = out.last_hidden_state[0]  # (L, D)
+    return hidden[drop_idx:].to(dtype)

extensions_built_in/diffusion_models/mageflow/src/transformer.py

@@ -0,0 +1,809 @@
+"""Mage-Flow NR-MMDiT (dual-stream, native-resolution packed) for ai-toolkit.
+
+Vendored from the reference implementation
+(github.com/microsoft/Mage, mage_flow/models/mage_flow.py + modules/mage_layers.py)
+with these deviations:
+
+  - loguru / pydantic / torch._dynamo decorators dropped,
+  - gradient checkpointing is gated on ``self.gradient_checkpointing and
+    torch.is_grad_enabled()`` (instead of the reference's ``self.training and
+    self.checkpoint``) and exposed via ``enable_gradient_checkpointing`` /
+    ``disable_gradient_checkpointing`` so ai-toolkit's trainer can toggle it,
+  - flash-attn is imported through the local shim (``.attn``) which falls back
+    to a per-sequence SDPA loop when flash-attn is unavailable.
+
+State-dict keys are identical to the reference checkpoints
+(``transformer/diffusion_pytorch_model.safetensors`` in
+microsoft/Mage-Flow-Base and microsoft/Mage-Flow-Edit-Base).
+
+Sequence layout: images and text are packed as variable-length sequences of
+batch dim 1 (``[1, sum_len, C]``) with per-sample ``cu_seqlens``; joint text+image
+attention runs in a single varlen kernel. 2D multi-scale RoPE is applied to
+image tokens only (text tokens are not rotated). ``timesteps`` is the flow
+sigma in [0, 1] (1 = pure noise), one entry per packed sample.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.models.normalization import RMSNorm
+
+from .attn import flash_attn_varlen_func
+
+
+def apply_rotary_emb_mageflow(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex rotary embeddings to `x` ([S, H, D]) using `freqs_cis`
+    (the MageFlowEmbedRope 2D multi-scale RoPE, adjacent-pair complex convention)."""
+    x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.unsqueeze(1)
+    x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(-2)
+    return x_out.type_as(x)
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """Sinusoidal timestep embeddings (DDPM convention).
+
+    NOTE: kept vendored (not diffusers') because the frequency table is
+    downcast to ``timesteps.dtype`` (bf16) here — the model was trained with
+    this exact bf16 rounding, so diffusers' fp32 variant produces a slightly
+    different embedding and degrades outputs.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent).to(timesteps.dtype)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+class Timesteps(nn.Module):
+    def __init__(
+        self,
+        num_channels: int,
+        flip_sin_to_cos: bool,
+        downscale_freq_shift: float,
+        scale: int = 1,
+    ):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        return get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+        )
+
+
+class MageFlowTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000
+        )
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+
+    def forward(self, timestep, hidden_states):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(
+            timesteps_proj.to(dtype=hidden_states.dtype)
+        )  # (N, D)
+        return timesteps_emb
+
+
+class MageFlowEmbedRope(nn.Module):
+    def __init__(self, theta: int, axes_dim: list, scale_rope=False):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+        self.pos_freqs = torch.cat(
+            [
+                self.rope_params(pos_index, self.axes_dim[0], self.theta),
+                self.rope_params(pos_index, self.axes_dim[1], self.theta),
+                self.rope_params(pos_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.neg_freqs = torch.cat(
+            [
+                self.rope_params(neg_index, self.axes_dim[0], self.theta),
+                self.rope_params(neg_index, self.axes_dim[1], self.theta),
+                self.rope_params(neg_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+
+        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
+        self.scale_rope = scale_rope
+        self.video_freq_cache = {}
+
+    def rope_params(self, index, dim, theta=10000):
+        assert dim % 2 == 0
+        freqs = torch.outer(
+            index,
+            1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)),
+        )
+        freqs = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs
+
+    def forward(
+        self, video_fhw, device: torch.device, max_img_len: int = None
+    ) -> torch.Tensor:
+        """Compute the vision RoPE frequencies for the packed image tokens.
+        Text tokens are NOT rotated, so no text RoPE is computed.
+
+        ``video_fhw`` is a list of (frame, height, width) tuples, one per image
+        segment in the packed sequence; segment index doubles as the "frame"
+        RoPE coordinate (edit reference images land on later frame indices).
+        """
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        if isinstance(video_fhw, list):
+            video_fhw = video_fhw[0]
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+
+        vid_freqs = []
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            key = (frame, height, width, idx)
+            if key not in self.video_freq_cache:
+                self.video_freq_cache[key] = self._compute_video_freqs(
+                    frame, height, width, idx
+                )
+            vid_freqs.append(self.video_freq_cache[key].to(device))
+
+        vid_freqs = torch.cat(vid_freqs, dim=0)
+
+        if max_img_len is not None and vid_freqs.shape[0] < max_img_len:
+            pad_len = max_img_len - vid_freqs.shape[0]
+            vid_freqs = torch.nn.functional.pad(vid_freqs, (0, 0, 0, pad_len))
+
+        return vid_freqs
+
+    def _compute_video_freqs(
+        self, frame: int, height: int, width: int, idx: int = 0
+    ) -> torch.Tensor:
+        seq_lens = frame * height * width
+        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+        freqs_frame = (
+            freqs_pos[0][idx : idx + frame]
+            .view(frame, 1, 1, -1)
+            .expand(frame, height, width, -1)
+        )
+        if self.scale_rope:
+            freqs_height = torch.cat(
+                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
+                dim=0,
+            )
+            freqs_height = freqs_height.view(1, height, 1, -1).expand(
+                frame, height, width, -1
+            )
+            freqs_width = torch.cat(
+                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
+                dim=0,
+            )
+            freqs_width = freqs_width.view(1, 1, width, -1).expand(
+                frame, height, width, -1
+            )
+        else:
+            freqs_height = (
+                freqs_pos[1][:height]
+                .view(1, height, 1, -1)
+                .expand(frame, height, width, -1)
+            )
+            freqs_width = (
+                freqs_pos[2][:width]
+                .view(1, 1, width, -1)
+                .expand(frame, height, width, -1)
+            )
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(
+            seq_lens, -1
+        )
+        return freqs.clone().contiguous()
+
+
+class Attention(nn.Module):
+    """Joint text+image attention (reference subset: the double-stream config
+    MageFlow actually instantiates — qk rms_norm, added kv projections)."""
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        added_kv_proj_dim: int = None,
+        added_proj_bias: bool = True,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        processor=None,
+        out_dim: int = None,
+        out_context_dim: int = None,
+        elementwise_affine: bool = True,
+    ):
+        super().__init__()
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.cross_attention_dim = query_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.out_context_dim = (
+            out_context_dim if out_context_dim is not None else query_dim
+        )
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.added_kv_proj_dim = added_kv_proj_dim
+        self.added_proj_bias = added_proj_bias
+
+        # qk_norm is always "rms_norm" for MageFlow.
+        self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+        self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+
+        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+
+        self.add_k_proj = nn.Linear(
+            added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias
+        )
+        self.add_v_proj = nn.Linear(
+            added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias
+        )
+        self.add_q_proj = nn.Linear(
+            added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
+        )
+        self.norm_added_q = RMSNorm(dim_head, eps=eps)
+        self.norm_added_k = RMSNorm(dim_head, eps=eps)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(nn.Dropout(dropout))
+
+        self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+
+        self.processor = processor
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        txt_cu_lens: torch.Tensor = None,
+        img_cu_lens: torch.Tensor = None,
+        image_rotary_emb: torch.Tensor = None,
+        **attention_kwargs,
+    ) -> torch.Tensor:
+        return self.processor(
+            self,
+            hidden_states,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            image_rotary_emb=image_rotary_emb,
+            **attention_kwargs,
+        )
+
+
+class MageDoubleStreamAttnProcessor:
+    """Attention processor for the Mage double-stream architecture. Implements
+    joint attention where the packed text and image streams are concatenated
+    per sample (order: [text, image]) and run through one varlen kernel."""
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,  # Image stream [1, sum_img, D]
+        img_cu_lens: torch.LongTensor,
+        encoder_hidden_states: torch.FloatTensor = None,  # Text stream [1, sum_txt, D]
+        txt_cu_lens: torch.LongTensor = None,
+        image_rotary_emb: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if encoder_hidden_states is None:
+            raise ValueError(
+                "MageDoubleStreamAttnProcessor requires encoder_hidden_states (text stream)"
+            )
+
+        # Compute QKV for image stream (sample projections)
+        img_query = attn.to_q(hidden_states)
+        img_key = attn.to_k(hidden_states)
+        img_value = attn.to_v(hidden_states)
+
+        # Compute QKV for text stream (context projections)
+        txt_query = attn.add_q_proj(encoder_hidden_states)
+        txt_key = attn.add_k_proj(encoder_hidden_states)
+        txt_value = attn.add_v_proj(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        img_query = img_query.unflatten(-1, (attn.heads, -1))
+        img_key = img_key.unflatten(-1, (attn.heads, -1))
+        img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+        if img_query.ndim == 4:
+            img_query = img_query.flatten(0, 1)
+            img_key = img_key.flatten(0, 1)
+            img_value = img_value.flatten(0, 1)
+
+        if txt_query.ndim == 4:
+            txt_query = txt_query.flatten(0, 1)
+            txt_key = txt_key.flatten(0, 1)
+            txt_value = txt_value.flatten(0, 1)
+
+        # Apply QK normalization
+        img_query = attn.norm_q(img_query)
+        img_key = attn.norm_k(img_key)
+        txt_query = attn.norm_added_q(txt_query)
+        txt_key = attn.norm_added_k(txt_key)
+
+        # Apply 2D multi-scale RoPE (MageFlowEmbedRope) to image tokens
+        img_freqs = image_rotary_emb
+        img_query = apply_rotary_emb_mageflow(img_query, img_freqs)
+        img_key = apply_rotary_emb_mageflow(img_key, img_freqs)
+
+        # Calculate lengths
+        img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
+        txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
+
+        # Calculate joint cu_seqlens
+        joint_lens = txt_lens + img_lens
+        joint_cu_lens = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32, device=joint_lens.device),
+                torch.cumsum(joint_lens, dim=0, dtype=torch.int32),
+            ],
+            dim=0,
+        )
+
+        device = joint_lens.device
+        batch_size = len(txt_lens)
+        sample_indices = torch.arange(batch_size, device=device)
+
+        txt_sample_ids = torch.repeat_interleave(sample_indices, txt_lens)
+        img_sample_ids = torch.repeat_interleave(sample_indices, img_lens)
+
+        txt_intra_pos = (
+            torch.arange(txt_query.shape[0], device=device)
+            - txt_cu_lens[txt_sample_ids]
+        )
+        img_intra_pos = (
+            torch.arange(img_query.shape[0], device=device)
+            - img_cu_lens[img_sample_ids]
+        )
+
+        txt_dest_indices = joint_cu_lens[txt_sample_ids] + txt_intra_pos
+        img_dest_indices = (
+            joint_cu_lens[img_sample_ids] + txt_lens[img_sample_ids] + img_intra_pos
+        )
+
+        total_tokens = joint_cu_lens[-1]
+        joint_query = torch.empty(
+            (total_tokens, *txt_query.shape[1:]), dtype=txt_query.dtype, device=device
+        )
+        joint_key = torch.empty(
+            (total_tokens, *txt_key.shape[1:]), dtype=txt_key.dtype, device=device
+        )
+        joint_value = torch.empty(
+            (total_tokens, *txt_value.shape[1:]), dtype=txt_value.dtype, device=device
+        )
+
+        joint_query[txt_dest_indices] = txt_query
+        joint_query[img_dest_indices] = img_query
+
+        joint_key[txt_dest_indices] = txt_key
+        joint_key[img_dest_indices] = img_key
+
+        joint_value[txt_dest_indices] = txt_value
+        joint_value[img_dest_indices] = img_value
+
+        max_seqlen = joint_lens.max().item()
+        joint_attn_output = flash_attn_varlen_func(
+            joint_query,
+            joint_key,
+            joint_value,
+            cu_seqlens_q=joint_cu_lens,
+            cu_seqlens_k=joint_cu_lens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            dropout_p=0.0,
+            softmax_scale=None,
+            causal=False,
+        )
+
+        txt_attn_output = joint_attn_output[txt_dest_indices]
+        img_attn_output = joint_attn_output[img_dest_indices]
+
+        img_attn_output = img_attn_output.flatten(1, 2)  # (N, H, D) -> (N, H*D)
+        img_attn_output = img_attn_output.to(joint_query.dtype)
+
+        txt_attn_output = txt_attn_output.flatten(1, 2)  # (N, H, D) -> (N, H*D)
+        txt_attn_output = txt_attn_output.to(joint_query.dtype)
+
+        img_attn_output = attn.to_out[0](img_attn_output)
+        img_attn_output = attn.to_out[1](img_attn_output)  # dropout
+
+        txt_attn_output = attn.to_add_out(txt_attn_output)
+        txt_attn_output = txt_attn_output.view(
+            encoder_hidden_states.shape[0],
+            encoder_hidden_states.shape[1],
+            txt_attn_output.shape[-1],
+        )
+
+        return img_attn_output, txt_attn_output
+
+
+class MageFlowTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+
+        # Image processing modules
+        self.img_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                dim, 6 * dim, bias=True
+            ),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = Attention(
+            query_dim=dim,
+            added_kv_proj_dim=dim,  # Enable added KV projections for text stream
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=MageDoubleStreamAttnProcessor(),
+            eps=eps,
+        )
+        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = FeedForward(
+            dim=dim, dim_out=dim, activation_fn="gelu-approximate"
+        )
+
+        # Text processing modules
+        self.txt_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                dim, 6 * dim, bias=True
+            ),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        # Text doesn't need separate attention - it's handled by img attn joint computation
+        self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_mlp = FeedForward(
+            dim=dim, dim_out=dim, activation_fn="gelu-approximate"
+        )
+
+    def _modulate(self, x, mod_params, cu_lens=None):
+        """Apply modulation to input tensor"""
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        if cu_lens is not None:
+            assert x.shape[0] == 1, "x must be of shape (1, *) when cu_lens is not None"
+            x_flattened = x.view(-1, x.shape[-1])
+            lengths = cu_lens[1:] - cu_lens[:-1]
+            shift_t = shift.repeat_interleave(lengths, dim=0)
+            scale_t = scale.repeat_interleave(lengths, dim=0)
+            gate_t = gate.repeat_interleave(lengths, dim=0)
+
+            x_flattened = x_flattened * (1 + scale_t) + shift_t
+            x = x_flattened.view(x.shape)
+            return x, gate_t
+        else:
+            return x * (1 + scale) + shift, gate
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: torch.Tensor,
+        txt_cu_lens: torch.Tensor,
+        img_cu_lens: torch.Tensor,
+        joint_attention_kwargs: dict = None,
+    ):
+        # Get modulation parameters for both streams
+        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+        # Split modulation parameters for norm1 and norm2
+        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+
+        # Process image stream - norm1 + modulation
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = self._modulate(
+            img_normed, img_mod1, cu_lens=img_cu_lens
+        )
+
+        # Process text stream - norm1 + modulation
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = self._modulate(
+            txt_normed, txt_mod1, cu_lens=txt_cu_lens
+        )
+
+        # Joint attention: computes QKV for both streams, applies QK norm and
+        # RoPE, concatenates per sample and splits the results back.
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=img_modulated,
+            encoder_hidden_states=txt_modulated,
+            image_rotary_emb=image_rotary_emb,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            **joint_attention_kwargs,
+        )
+
+        img_attn_output, txt_attn_output = attn_output
+
+        # Apply attention gates and add residual
+        hidden_states = hidden_states + img_gate1 * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+        # Process image stream - norm2 + MLP
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = self._modulate(
+            img_normed2, img_mod2, cu_lens=img_cu_lens
+        )
+        img_mlp_output = self.img_mlp(img_modulated2)
+        hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+        # Process text stream - norm2 + MLP
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = self._modulate(
+            txt_normed2, txt_mod2, cu_lens=txt_cu_lens
+        )
+        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+        # Clip to prevent overflow for fp16
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class AdaLayerNormContinuous(nn.Module):
+    """Adaptive output norm; supports packed sequences via ``cu_seqlens``."""
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(
+            conditioning_embedding_dim, embedding_dim * 2, bias=bias
+        )
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        conditioning_embedding: torch.Tensor,
+        cu_seqlens: torch.Tensor = None,
+    ) -> torch.Tensor:
+        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
+        if cu_seqlens is None:
+            scale, shift = torch.chunk(emb, 2, dim=-1)
+            x = self.norm(x) * (1 + scale) + shift
+        else:
+            sample_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            flattened_x = x.view(-1, x.shape[-1])
+            scale, shift = torch.chunk(emb, 2, dim=-1)
+            scale_t = torch.repeat_interleave(scale, sample_lens, dim=0)
+            shift_t = torch.repeat_interleave(shift, sample_lens, dim=0)
+            flattened_x = self.norm(flattened_x) * (1 + scale_t) + shift_t
+            x = flattened_x.view(x.shape)
+        return x
+
+
+@dataclass
+class MageFlowParams:
+    in_channels: int
+    out_channels: int
+    context_in_dim: int
+    hidden_size: int
+    num_heads: int
+    depth: int
+    axes_dim: list
+    checkpoint: bool
+    patch_size: int = 1
+
+
+class MageFlow(nn.Module):
+    def __init__(self, params: MageFlowParams):
+        super().__init__()
+        self.params = params
+        self.gradient_checkpointing = bool(params.checkpoint)
+        self.in_channels = params.in_channels
+        self.out_channels = params.out_channels
+        self.inner_dim = params.hidden_size  # num_attention_heads * attention_head_dim
+        self.axes_dim = params.axes_dim
+        self.num_attention_heads = params.num_heads
+        self.attention_head_dim = self.inner_dim // self.num_attention_heads
+        self.patch_size = params.patch_size
+        assert sum(self.axes_dim) == self.attention_head_dim
+
+        self.pos_embed = MageFlowEmbedRope(
+            theta=10000, axes_dim=self.axes_dim, scale_rope=True
+        )
+        self.img_in = nn.Linear(self.in_channels, self.inner_dim)
+        self.txt_norm = RMSNorm(params.context_in_dim, eps=1e-6)
+        self.txt_in = nn.Linear(params.context_in_dim, self.inner_dim)
+
+        self.time_text_embed = MageFlowTimestepProjEmbeddings(
+            embedding_dim=self.inner_dim
+        )
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MageFlowTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
+        )
+        self.proj_out = nn.Linear(
+            self.inner_dim,
+            self.patch_size * self.patch_size * self.out_channels,
+            bias=True,
+        )
+
+    def enable_gradient_checkpointing(self, enable: bool = True):
+        self.gradient_checkpointing = enable
+
+    def disable_gradient_checkpointing(self):
+        self.gradient_checkpointing = False
+
+    def reset_rope(self):
+        """Rebuild the RoPE frequency tables. Needed after a meta-device init
+        (the tables are plain tensors, not buffers, so ``load_state_dict`` /
+        ``assign=True`` never materializes them)."""
+        self.pos_embed = MageFlowEmbedRope(
+            theta=10000, axes_dim=self.axes_dim, scale_rope=True
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def forward(
+        self,
+        img: Tensor,
+        txt: Tensor,
+        timesteps: Tensor,
+        img_shapes=None,
+        img_cu_seqlens: Tensor = None,
+        txt_cu_seqlens: Tensor = None,
+        attention_kwargs: dict = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # Prepare vision RoPE (msrope); text tokens are not rotated.
+        ms_pe = self.pos_embed(img_shapes, device=img.device)
+
+        img = self.img_in(img)
+        txt = self.txt_norm(txt)
+
+        timesteps = timesteps.to(img.dtype)
+        temb = self.time_text_embed(timesteps, img)
+
+        txt = self.txt_in(txt)
+        # The reference adds a zero "text vector" to temb (vec_type is null in
+        # the released config); kept for parity with the reference forward.
+        txt_vec = torch.zeros(
+            txt.shape[0], self.inner_dim, dtype=txt.dtype, device=txt.device
+        )
+        temb = temb + txt_vec
+
+        attention_kwargs = attention_kwargs or {}
+
+        for block in self.transformer_blocks:
+            if self.gradient_checkpointing and torch.is_grad_enabled():
+                txt, img = checkpoint(
+                    block,
+                    img,  # hidden_states
+                    txt,  # encoder_hidden_states
+                    temb,  # temb
+                    ms_pe,  # image_rotary_emb
+                    txt_cu_seqlens,  # txt_cu_lens
+                    img_cu_seqlens,  # img_cu_lens
+                    use_reentrant=False,
+                )
+            else:
+                txt, img = block(
+                    hidden_states=img,
+                    encoder_hidden_states=txt,
+                    txt_cu_lens=txt_cu_seqlens,
+                    img_cu_lens=img_cu_seqlens,
+                    temb=temb,
+                    image_rotary_emb=ms_pe,
+                    joint_attention_kwargs=attention_kwargs,
+                )
+
+        # Use only the image part (hidden_states) from the dual-stream blocks
+        img = self.norm_out(img, temb, cu_seqlens=img_cu_seqlens)
+        img = self.proj_out(img)
+        return img

extensions_built_in/diffusion_models/mageflow/src/vae.py

@@ -0,0 +1,703 @@
+"""
+MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) wrapper.
+
+Vendored from the reference implementation
+(github.com/microsoft/Mage, mage_flow/models/modules/mage_vae.py); loguru
+swapped for prints, otherwise unchanged. Loads the kl0.1 CoD ckpt layout used
+by ``vae/diffusion_pytorch_model.safetensors`` in the Mage-Flow HF repos:
+  encoder weights: 'student.dconv_encoder.*' (packed mean+logvar, out_ch_mult=2)
+  decoder weights: 'pipeline.*' (denoiser + y_embedder.decoder)
+
+Latent shape: [B, 128, H/16, W/16] — no patch packing, no normalization.
+Encode input / decode output are [-1, 1] pixel images with H, W multiples of 16.
+"""
+
+import math
+import os
+from functools import lru_cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Primitive layers (vendored from GenCodec, inference subset)
+# ---------------------------------------------------------------------------
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(
+        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+def modulate(x, shift, scale):
+    if x.dim() == 4:
+        b, c = x.shape[:2]
+        return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1)
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x):
+        # .contiguous() prevents a channels_last-strided NCHW view from
+        # propagating into downstream depthwise convs, which would otherwise
+        # hit a slow cuDNN path with a per-shape heuristic search.
+        x = x.permute(0, 2, 3, 1).contiguous()
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class _EncoderLayerNorm2d(LayerNorm2d):
+    pass
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        in_dtype = x.dtype
+        x = x.to(torch.float32)
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.variance_epsilon)
+        return self.weight * x.to(in_dtype)
+
+
+class TimestepEmbedder(nn.Module):
+    """DConv-style timestep MLP (max_period=10000, freq_size=256, hidden=384)."""
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
+        ).to(t.device)
+        args = t[:, None].float() * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+        return emb
+
+    def forward(self, t):
+        emb = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class BottleneckPatchEmbed(nn.Module):
+    """Image patch embed concatenated with a per-patch conditioning vector."""
+
+    def __init__(
+        self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True
+    ):
+        super().__init__()
+        self.proj1 = nn.Conv2d(
+            in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False
+        )
+        self.proj2 = nn.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
+
+    def forward(self, x, cond):
+        return self.proj2(torch.cat([self.proj1(x), cond], dim=1))
+
+
+class DiCoBlock(nn.Module):
+    """DConv block with adaLN modulation."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = nn.Conv2d(
+            hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True
+        )
+        self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
+
+        self.norm1 = LayerNorm2d(hidden_size, affine=False)
+        self.norm2 = LayerNorm2d(hidden_size, affine=False)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, inp, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaLN_modulation(c).chunk(6, dim=1)
+        )
+        x = modulate(self.norm1(inp), shift_msa, scale_msa)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + gate_msa[..., None, None] * x
+        x = x + gate_mlp[..., None, None] * self.conv5(
+            F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp)))
+        )
+        return x
+
+
+class _EncoderDiCoBlock(nn.Module):
+    """DiCoBlock without adaLN, for the encoder pathway."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = nn.Conv2d(
+            hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True
+        )
+        self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
+        self.norm1 = _EncoderLayerNorm2d(hidden_size)
+        self.norm2 = _EncoderLayerNorm2d(hidden_size)
+
+    def forward(self, inp):
+        x = self.norm1(inp)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + x
+        return x + self.conv5(F.gelu(self.conv4(self.norm2(x))))
+
+
+class NerfEmbedder(nn.Module):
+    """Patch-position embedder used by the DConv decoder x-pathway."""
+
+    def __init__(self, in_channels, hidden_size_input, max_freqs=8):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.embedder = nn.Sequential(
+            nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True),
+        )
+
+    @lru_cache
+    def fetch_pos(self, patch_size, device, dtype):
+        pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
+        pos_x = pos_x.reshape(-1, 1, 1)
+        pos_y = pos_y.reshape(-1, 1, 1)
+        freqs = torch.linspace(
+            0, self.max_freqs, self.max_freqs, dtype=dtype, device=device
+        )
+        fx = freqs[None, :, None]
+        fy = freqs[None, None, :]
+        coeffs = (1 + fx * fy) ** -1
+        dct_x = torch.cos(pos_x * fx * torch.pi)
+        dct_y = torch.cos(pos_y * fy * torch.pi)
+        return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2)
+
+    def forward(self, x):
+        B, P2, _ = x.shape
+        ps = int(P2**0.5)
+        dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1)
+        return self.embedder(torch.cat([x, dct], dim=-1))
+
+
+class NerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        return self.linear(self.norm(x))
+
+
+class SimpleMLPAdaLN(nn.Module):
+    """Final small MLP that maps NerfEmbedder features to per-patch RGB."""
+
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        out_channels,
+        z_channels,
+        num_res_blocks,
+        patch_size,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.patch_size = patch_size
+
+        self.cond_embed = nn.Linear(z_channels, patch_size**2 * model_channels)
+        self.input_proj = nn.Linear(in_channels, model_channels)
+
+        self.res_blocks = nn.ModuleList(
+            _MLPResBlock(model_channels) for _ in range(num_res_blocks)
+        )
+
+    def forward(self, x, c):
+        x = self.input_proj(x)
+        c = self.cond_embed(c).reshape(c.shape[0], self.patch_size**2, -1)
+        for block in self.res_blocks:
+            x = block(x, c)
+        return x
+
+
+class _MLPResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.in_ln = nn.LayerNorm(channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(channels, channels, bias=True),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(channels, 3 * channels, bias=True),
+        )
+
+    def forward(self, x, y):
+        shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = self.in_ln(x) * (1 + scale) + shift
+        return x + gate * self.mlp(h)
+
+
+class ResnetBlock(nn.Module):
+    """GroupNorm + Conv ResBlock used by the CoD Decoder."""
+
+    def __init__(self, *, in_channels, out_channels=None, dropout=0.0):
+        super().__init__()
+        out_channels = out_channels or in_channels
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1)
+
+    def forward(self, x):
+        h = self.conv1(nonlinearity(self.norm1(x)))
+        h = self.conv2(self.dropout(nonlinearity(self.norm2(h))))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Patched self-attention used at inference (eval mode of the original)."""
+
+    def __init__(self, in_channels, patch_size=32):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, 1)
+        self.k = nn.Conv2d(in_channels, in_channels, 1)
+        self.v = nn.Conv2d(in_channels, in_channels, 1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, 1)
+
+    def forward(self, x):
+        h_ = self.norm(x)
+        Q = self.q(h_)
+        K = self.k(h_)
+        V = self.v(h_)
+
+        d = self.patch_size
+        b, c, H, W = Q.shape
+        pad_h = (d - H % d) % d
+        pad_w = (d - W % d) % d
+        if pad_h or pad_w:
+            Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate")
+            K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate")
+            V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate")
+        _, _, H_pad, W_pad = Q.shape
+        nph, npw = H_pad // d, W_pad // d
+        np_ = nph * npw
+
+        def to_patches(t):
+            return (
+                t.reshape(b, c, nph, d, npw, d)
+                .permute(0, 2, 4, 1, 3, 5)
+                .reshape(b * np_, c, d * d)
+            )
+
+        Q = to_patches(Q)
+        K = to_patches(K)
+        V = to_patches(V)
+
+        w_ = torch.bmm(Q.permute(0, 2, 1), K) * (c**-0.5)
+        w_ = F.softmax(w_, dim=2).permute(0, 2, 1)
+        h_ = (
+            torch.bmm(V, w_)
+            .reshape(b, nph, npw, c, d, d)
+            .permute(0, 3, 1, 4, 2, 5)
+            .reshape(b, c, H_pad, W_pad)
+        )
+        if pad_h or pad_w:
+            h_ = h_[:, :, :H, :W]
+        return x + self.proj_out(h_)
+
+
+# ---------------------------------------------------------------------------
+# adaLN constant-folding: at fixed t=0, adaLN_modulation(c) is constant.
+# Replace the MLP with a buffer so DiCoBlock.forward stays unchanged.
+# ---------------------------------------------------------------------------
+class _ConstAdaLN(nn.Module):
+    def __init__(self, modulation: torch.Tensor):
+        super().__init__()
+        self.register_buffer("modulation", modulation.detach().clone())
+
+    def forward(self, c):
+        b = c.shape[0]
+        if self.modulation.shape[0] != b:
+            return self.modulation.expand(b, *self.modulation.shape[1:])
+        return self.modulation
+
+
+def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int:
+    # Only DiCoBlock is targeted: its adaLN is conditioned solely on t.
+    # Other adaLN_modulation submodules (e.g. _MLPResBlock in the decoder MLP)
+    # take a per-position latent and must not be folded.
+    n = 0
+    for child in module.modules():
+        if not isinstance(child, DiCoBlock):
+            continue
+        adaln = child.adaLN_modulation
+        if isinstance(adaln, _ConstAdaLN):
+            continue
+        with torch.no_grad():
+            mod = adaln(c)
+        child.adaLN_modulation = _ConstAdaLN(mod)
+        n += 1
+    return n
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder: latent → conditioning features for the denoiser
+# ---------------------------------------------------------------------------
+class _Decoder(nn.Module):
+    """ds=16, up2x=True, light=True only."""
+
+    def __init__(self, out_ch=384, z_ch=128):
+        super().__init__()
+        self.conv_in = nn.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.block = nn.Sequential(
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+        )
+        self.norm_out = Normalize(out_ch)
+        self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.ada = nn.Identity()
+
+    def forward(self, z):
+        h = self.block(self.conv_in(z))
+        h = self.conv_out(nonlinearity(self.norm_out(h)))
+        return self.ada(h)
+
+
+# ---------------------------------------------------------------------------
+# DConvEncoder: image → packed (mean, logvar) latent
+# ---------------------------------------------------------------------------
+class _DConvEncoder(nn.Module):
+    def __init__(
+        self,
+        z_ch=128,
+        hidden_size=384,
+        num_blocks=21,
+        patch_size=16,
+        mlp_ratio=4.0,
+        head_size=768,
+        num_head_blocks=2,
+        out_ch_mult=2,
+    ):
+        super().__init__()
+        self.z_ch = z_ch
+        self.patch_size = patch_size
+        self.patch_cond_embed = nn.Conv2d(
+            3, head_size, kernel_size=patch_size, stride=patch_size, bias=True
+        )
+        self.head_blocks = nn.ModuleList(
+            [
+                _EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio)
+                for _ in range(num_head_blocks)
+            ]
+        )
+        self.proj_down = nn.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
+        self.z_proj = nn.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True)
+        self.fuse_proj = nn.Conv2d(
+            hidden_size * 2, hidden_size, kernel_size=1, bias=True
+        )
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.blocks = nn.ModuleList(
+            [DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)]
+        )
+        self.norm_out = LayerNorm2d(hidden_size)
+        self.proj_out = nn.Conv2d(
+            hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True
+        )
+
+    def forward_pred(self, z_t, t, y):
+        cond = self.patch_cond_embed(y)
+        for block in self.head_blocks:
+            cond = block(cond)
+        cond = self.proj_down(cond)
+
+        s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1))
+        c = self.t_embedder(t.view(-1))
+        for block in self.blocks:
+            s = block(s, c)
+        return self.proj_out(self.norm_out(s))
+
+
+# ---------------------------------------------------------------------------
+# DConv denoiser: latent (via cond) + zero noise → reconstructed image
+# ---------------------------------------------------------------------------
+class _YEmbedder(nn.Module):
+    """Holds only the CoD decoder; the original Flux2 VAE encoder side is omitted."""
+
+    def __init__(self, ch=384, z_ch=128):
+        super().__init__()
+        self.decoder = _Decoder(out_ch=ch, z_ch=z_ch)
+
+
+class _DConvDenoiser(nn.Module):
+    def __init__(
+        self,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=384,
+        hidden_size_x=32,
+        mlp_ratio=4.0,
+        num_blocks=24,
+        num_cond_blocks=21,
+        bottleneck_dim=128,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.hidden_size = hidden_size
+        self.num_cond_blocks = num_cond_blocks
+
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder_x = nn.Conv2d(
+            hidden_size, hidden_size_x * patch_size**2, 1, 1, 0
+        )
+        self.x_embedder = NerfEmbedder(
+            in_channels + hidden_size_x, hidden_size_x, max_freqs=8
+        )
+        self.s_embedder = BottleneckPatchEmbed(
+            patch_size, in_channels, bottleneck_dim, hidden_size, bias=True
+        )
+        self.blocks = nn.ModuleList(
+            [
+                DiCoBlock(hidden_size, mlp_ratio=mlp_ratio)
+                for _ in range(num_cond_blocks)
+            ]
+        )
+        self.dec_net = SimpleMLPAdaLN(
+            in_channels=hidden_size_x,
+            model_channels=hidden_size_x,
+            out_channels=in_channels,
+            z_channels=hidden_size,
+            num_res_blocks=num_blocks - num_cond_blocks,
+            patch_size=patch_size,
+        )
+        self.final_layer = NerfFinalLayer(hidden_size_x, in_channels)
+        self.y_embedder = _YEmbedder(ch=hidden_size, z_ch=bottleneck_dim)
+
+    def forward(self, x, t, cond):
+        b, _, h, w = x.shape
+        c = self.t_embedder(t.view(-1))
+
+        s = self.s_embedder(x, cond)
+        for block in self.blocks:
+            s = block(s, c)
+
+        length = s.shape[-2] * s.shape[-1]
+        s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+        x = torch.nn.functional.unfold(
+            x, kernel_size=self.patch_size, stride=self.patch_size
+        )
+        x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1)
+        x = (
+            x.reshape(b, -1, self.patch_size**2, length)
+            .permute(0, 3, 2, 1)
+            .flatten(0, 1)
+        )
+        x = self.x_embedder(x)
+
+        x = self.dec_net(x, s)
+        x = self.final_layer(x)
+        x = x.transpose(1, 2).reshape(b, length, -1)
+        return torch.nn.functional.fold(
+            x.transpose(1, 2).contiguous(),
+            (h, w),
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Wrapper
+# ---------------------------------------------------------------------------
+def _load_state_dict(ckpt_path: str):
+    if ckpt_path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        return load_file(ckpt_path, device="cpu")
+    if os.path.exists(os.path.join(ckpt_path, "checkpoint-state_dict.pt")):
+        ckpt_path = os.path.join(ckpt_path, "checkpoint-state_dict.pt")
+    elif os.path.isdir(ckpt_path):
+        ckpt_path = os.path.join(ckpt_path, "checkpoint", "mp_rank_00_model_states.pt")
+    state = torch.load(ckpt_path, map_location="cpu")
+    if "module" in state:
+        return state["module"]
+    if "state_dict" in state:
+        return state["state_dict"]
+    return state
+
+
+class MageVAE(nn.Module):
+    """
+    Encode: DConvEncoder (one-step diffusion) → latent [B, 128, H/16, W/16]
+    Decode: DConvDenoiser + CoD Decoder       → image  [B, 3, H, W] in [-1, 1]
+    """
+
+    latent_channels = 128
+    downsample_factor = 16
+
+    def __init__(self, ckpt_path: str, sample_posterior: bool = True):
+        super().__init__()
+        self.sample_posterior = sample_posterior
+
+        self.dconv_encoder = _DConvEncoder()
+        self.decoder_model = _DConvDenoiser()
+
+        sd = _load_state_dict(ckpt_path)
+        self._load_encoder(sd, ckpt_path)
+        self._load_decoder(sd, ckpt_path)
+
+        # adaLN modulation depends only on t, and we always run at t=0.
+        # Precompute and drop the MLPs once at construction (~37M params saved).
+        self._freeze_adaln_cache()
+
+    def _load_encoder(self, sd, ckpt_path):
+        prefix = "student.dconv_encoder."
+        enc_sd = {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)}
+        if not enc_sd:
+            raise RuntimeError(f"MageVAE encoder: no '{prefix}*' keys in {ckpt_path}")
+        proj = enc_sd.get("proj_out.weight")
+        if proj is None or proj.shape[0] != 2 * self.latent_channels:
+            raise RuntimeError(
+                f"MageVAE encoder: expected packed mean+logvar (proj_out out_channels="
+                f"{2 * self.latent_channels}), got {None if proj is None else tuple(proj.shape)}"
+            )
+        missing, unexpected = self.dconv_encoder.load_state_dict(enc_sd, strict=False)
+        if missing:
+            print(f"MageVAE encoder missing keys: {missing[:10]}")
+
+    def _load_decoder(self, sd, ckpt_path):
+        prefix = "pipeline."
+        if not any(k.startswith(prefix) for k in sd):
+            raise RuntimeError(f"MageVAE decoder: no '{prefix}*' keys in {ckpt_path}")
+        model_dict = self.decoder_model.state_dict()
+        matched = {}
+        for k, v in sd.items():
+            if not k.startswith(prefix):
+                continue
+            new_k = k[len(prefix) :]
+            if new_k.startswith("y_embedder.encoder.") or new_k.startswith(
+                "y_embedder.bottleneck."
+            ):
+                continue
+            if new_k in model_dict and model_dict[new_k].shape == v.shape:
+                matched[new_k] = v
+        self.decoder_model.load_state_dict(matched, strict=False)
+        if not matched:
+            raise RuntimeError(f"MageVAE decoder: 0 params matched from {ckpt_path}")
+
+    @torch.no_grad()
+    def _moments(self, x: torch.Tensor):
+        B, _, H, W = x.shape
+        ps = self.dconv_encoder.patch_size
+        z_t = torch.zeros(
+            B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype
+        )
+        t = torch.zeros(B, device=x.device, dtype=x.dtype)
+        out = self.dconv_encoder.forward_pred(z_t, t, x)
+        mean = out[:, : self.latent_channels]
+        logvar = out[:, self.latent_channels :].clamp(min=-20.0, max=10.0)
+        return mean, logvar
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        ps = self.dconv_encoder.patch_size
+        H, W = x.shape[-2], x.shape[-1]
+        if H % ps or W % ps:
+            raise ValueError(f"H, W must be multiples of {ps}, got ({H}, {W})")
+        mean, logvar = self._moments(x)
+        if self.sample_posterior:
+            return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean)
+        return mean
+
+    @torch.no_grad()
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        cond = self.decoder_model.y_embedder.decoder(z)
+        B = z.shape[0]
+        H = z.shape[2] * self.downsample_factor
+        W = z.shape[3] * self.downsample_factor
+        noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype)
+        t = torch.zeros(B, device=z.device, dtype=z.dtype)
+        return self.decoder_model.forward(noise, t, cond)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def _freeze_adaln_cache(self):
+        """Constant-fold adaLN_modulation MLPs at t=0 (encoder + decoder)."""
+        device = next(self.parameters()).device
+        dtype = next(self.parameters()).dtype
+        t = torch.zeros(1, device=device, dtype=dtype)
+        c_enc = self.dconv_encoder.t_embedder(t)
+        _replace_adaln_with_const(self.dconv_encoder, c_enc)
+        c_dec = self.decoder_model.t_embedder(t)
+        _replace_adaln_with_const(self.decoder_model, c_dec)

ui/src/app/jobs/new/options.ts

@@ -1196,6 +1196,55 @@ export const modelArchs: ModelArch[] = [
       'model.model_kwargs.kv_cache',
     ],
   },
+  {
+    name: 'mageflow',
+    label: 'Mage-Flow',
+    group: 'image',
+    defaults: {
+      'config.process[0].model.name_or_path': ['microsoft/Mage-Flow-Base', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].network.conv': [undefined, 16],
+      'config.process[0].network.conv_alpha': [undefined, 16],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].sample.guidance_scale': [4, 4],
+      'config.process[0].sample.sample_steps': [25, 25],
+    },
+    disableSections: [
+      'network.conv',
+    ],
+    additionalSections: [
+      'model.low_vram',
+      'model.layer_offloading',
+    ],
+  },
+  {
+    name: 'mageflow_edit',
+    label: 'Mage-Flow Edit',
+    group: 'instruction',
+    defaults: {
+      'config.process[0].model.name_or_path': ['microsoft/Mage-Flow-Edit-Base', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].network.conv': [undefined, 16],
+      'config.process[0].network.conv_alpha': [undefined, 16],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].sample.guidance_scale': [4, 4],
+      'config.process[0].sample.sample_steps': [25, 25],
+      'config.process[0].train.unload_text_encoder': [false, false],
+    },
+    disableSections: [
+      'network.conv', 'train.unload_text_encoder',
+    ],
+    additionalSections: [
+      'datasets.multi_control_paths',
+      'sample.multi_ctrl_imgs',
+      'model.low_vram',
+      'model.layer_offloading',
+    ],
+  },
   {
     name: 'boogu_image',
     label: 'Boogu Image',