Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+mage_flow/assets/cuisine.jpg filter=lfs diff=lfs merge=lfs -text
+mage_flow/assets/dog.jpg filter=lfs diff=lfs merge=lfs -text
+mage_flow/assets/portrait.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,41 @@
 ---
-title: Mage Flow
-emoji: 🏢
-colorFrom: gray
-colorTo: pink
+title: Mage-Flow
+emoji: 🎨
+colorFrom: green
+colorTo: purple
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 6.15.1
 app_file: app.py
-pinned: false
+short_description: Efficient native-resolution image generation and editing
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Mage-Flow
+
+Efficient Native-Resolution Foundation Model for Image Generation and Editing.
+
+This Space demonstrates [Mage-Flow](https://huggingface.co/papers/2607.19064) from Microsoft, a 4B-scale image generation and editing foundation model. It supports:
+
+- **Text → Image**: Generate images from text prompts at arbitrary resolutions.
+- **Image Edit**: Instruction-based editing of reference images.
+
+The demo uses the Turbo variants ([Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo) and [Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo)) for fast, few-step inference (4 steps, CFG 1.0).
+
+## Usage
+
+### Text → Image
+1. Enter a descriptive prompt.
+2. Adjust steps (Turbo default: 4) and CFG (Turbo default: 1.0).
+3. Click **Generate**.
+
+### Image Edit
+1. Upload a reference image.
+2. Enter an edit instruction (e.g., "change the background to a city street").
+3. Click **Edit**.
+
+## References
+
+- Paper: [Mage-Flow: An Efficient Native-Resolution Foundation Model for Image Generation and Editing](https://huggingface.co/papers/2607.19064)
+- GitHub: [microsoft/Mage](https://github.com/microsoft/Mage)
+- Models: [microsoft/Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo), [microsoft/Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo)

app.py

@@ -0,0 +1,208 @@
+"""Mage-Flow: Efficient Native-Resolution Foundation Model for Image Generation and Editing.
+
+Gradio Space demo supporting text-to-image generation and instruction-based image editing
+using Microsoft's Mage-Flow model (Turbo variants for fast inference).
+"""
+import os
+
+os.environ.setdefault("VF_HF_ATTN_IMPL", "sdpa")
+
+import spaces  # MUST be first (after env setup)
+import torch
+import gradio as gr
+from PIL import Image
+
+from mage_flow.pipeline import MageFlowPipeline
+
+T2I_MODEL = "microsoft/Mage-Flow-Turbo"
+EDIT_MODEL = "microsoft/Mage-Flow-Edit-Turbo"
+
+pipe_t2i = MageFlowPipeline.from_pretrained(T2I_MODEL, device="cuda")
+pipe_edit = MageFlowPipeline.from_pretrained(EDIT_MODEL, device="cuda")
+
+
+@spaces.GPU(duration=120)
+def generate(
+    prompt: str,
+    negative_prompt: str,
+    steps: int,
+    cfg: float,
+    height: int,
+    width: int,
+    seed: int,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate an image from a text prompt using Mage-Flow-Turbo.
+
+    Args:
+        prompt: Text description of the image to generate.
+        negative_prompt: What to avoid in the generation.
+        steps: Number of denoising steps (Turbo uses 4).
+        cfg: Classifier-free guidance scale (Turbo uses 1.0).
+        height: Output image height (multiple of 16).
+        width: Output image width (multiple of 16).
+        seed: Random seed for reproducibility.
+    """
+    if not (prompt or "").strip():
+        raise gr.Error("Prompt is empty.")
+    img = pipe_t2i.generate(
+        [prompt],
+        neg_prompts=[negative_prompt or " "],
+        seeds=[int(seed)],
+        steps=int(steps),
+        cfg=float(cfg),
+        heights=[int(height)],
+        widths=[int(width)],
+    )[0]
+    return img
+
+
+@spaces.GPU(duration=120)
+def edit(
+    prompt: str,
+    negative_prompt: str,
+    ref_img,
+    steps: int,
+    cfg: float,
+    max_size: int,
+    seed: int,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Edit a reference image using an instruction prompt with Mage-Flow-Edit-Turbo.
+
+    Args:
+        prompt: Edit instruction describing how to modify the image.
+        negative_prompt: What to avoid in the edited result.
+        ref_img: Reference image to edit.
+        steps: Number of denoising steps (Turbo uses 4).
+        cfg: Classifier-free guidance scale.
+        max_size: Longest side of output (0 = keep source resolution).
+        seed: Random seed for reproducibility.
+    """
+    if not (prompt or "").strip():
+        raise gr.Error("Edit instruction is empty.")
+    if ref_img is None:
+        raise gr.Error("Upload a reference image to edit.")
+    if isinstance(ref_img, str):
+        ref_img = Image.open(ref_img)
+    refs = [ref_img.convert("RGB")]
+    out = pipe_edit.edit(
+        [prompt],
+        [refs],
+        neg_prompts=[negative_prompt or " "],
+        seeds=[int(seed)],
+        steps=int(steps),
+        cfg=float(cfg),
+        max_size=int(max_size) if max_size else None,
+    )[0]
+    return out
+
+
+ASSETS_DIR = os.path.join(os.path.dirname(__file__), "mage_flow", "assets")
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown(
+        "# Mage-Flow\n"
+        "Efficient Native-Resolution Foundation Model for Image Generation and Editing.\n\n"
+        "Model: [microsoft/Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo) | "
+        "[Paper](https://huggingface.co/papers/2607.19064) | "
+        "[GitHub](https://github.com/microsoft/Mage)"
+    )
+
+    with gr.Tab("Text → Image"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                t_prompt = gr.Textbox(
+                    label="Prompt",
+                    lines=3,
+                    value="A close-up portrait of an elderly African man with deep wrinkles, "
+                          "wearing a traditional hat, soft natural lighting, ultra realistic.",
+                )
+                t_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
+                with gr.Row():
+                    t_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                    t_cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")
+                with gr.Row():
+                    t_h = gr.Slider(256, 1536, value=1024, step=16, label="Height")
+                    t_w = gr.Slider(256, 1536, value=1024, step=16, label="Width")
+                t_seed = gr.Number(value=42, precision=0, label="Seed")
+                t_btn = gr.Button("Generate", variant="primary")
+            with gr.Column(scale=1):
+                t_out = gr.Image(type="pil", label="Output", height=560)
+
+        gr.Examples(
+            examples=[
+                ["A close-up portrait of an elderly African man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic."],
+                ["A serene mountain landscape at sunset, with snow-capped peaks reflecting golden light, photorealistic."],
+                ["A cute robot playing a guitar in a neon-lit cyberpunk city, digital art style."],
+            ],
+            inputs=[t_prompt],
+            outputs=t_out,
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        t_btn.click(
+            lambda: None, None, t_out
+        ).then(
+            generate,
+            [t_prompt, t_neg, t_steps, t_cfg, t_h, t_w, t_seed],
+            t_out,
+            api_name="generate",
+        )
+
+    with gr.Tab("Image Edit"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                e_prompt = gr.Textbox(
+                    label="Edit instruction",
+                    lines=2,
+                    value="change the background to a city street",
+                )
+                e_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
+                e_ref = gr.Image(
+                    type="pil", label="Reference image", height=280,
+                    value=os.path.join(ASSETS_DIR, "dog.jpg"),
+                )
+                with gr.Row():
+                    e_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                    e_cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")
+                e_max = gr.Slider(
+                    0, 1536, value=1024, step=16,
+                    label="Max output side (0 = keep source size)",
+                )
+                e_seed = gr.Number(value=42, precision=0, label="Seed")
+                e_btn = gr.Button("Edit", variant="primary")
+            with gr.Column(scale=1):
+                e_out = gr.Image(type="pil", label="Output", height=560)
+
+        gr.Examples(
+            examples=[
+                ["change the background to a city street", os.path.join(ASSETS_DIR, "dog.jpg")],
+                ["make it look like a painting", os.path.join(ASSETS_DIR, "cuisine.jpg")],
+                ["add a hat to the person", os.path.join(ASSETS_DIR, "portrait.jpg")],
+            ],
+            inputs=[e_prompt, e_ref],
+            outputs=e_out,
+            fn=edit,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        e_btn.click(
+            lambda: None, None, e_out
+        ).then(
+            edit,
+            [e_prompt, e_neg, e_ref, e_steps, e_cfg, e_max, e_seed],
+            e_out,
+            api_name="edit",
+        )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

mage_flow/__init__.py

@@ -0,0 +1,18 @@
+"""MageFlow — standalone text-to-image + image-edit inference."""
+
+from .models.mage_flow import ModelConfig
+from .pipeline import (
+    MageFlowPipeline,
+    generate_edits,
+    generate_images,
+    load_from_repo,
+)
+
+__all__ = [
+    "MageFlowPipeline",
+    "generate_images",
+    "generate_edits",
+    "load_from_repo",
+    "ModelConfig",
+]
+__version__ = "0.1.0"

mage_flow/models/__init__.py

@@ -0,0 +1,3 @@
+from .utils import load_model
+
+__all__ = ["load_model"]

mage_flow/models/mage_flow.py

@@ -0,0 +1,364 @@
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from loguru import logger
+from pydantic import BaseModel, Field
+from torch import Tensor
+
+from .modules._attn_backend import set_attn_backend
+from .modules.mage_layers import (
+    AdaLayerNormContinuous,
+    MageFlowEmbedRope,
+    MageFlowTimestepProjEmbeddings,
+    MageFlowTransformerBlock,
+    RMSNorm,
+)
+from .modules.text_encoder import TextEncoder, qwen3_patch_forward
+
+
+class ModelConfig(BaseModel):
+    static_shift: float = Field(
+        default=6.0,
+        description="Static shift value for the z-image time-shift schedule (the only "
+                    "supported schedule). Default: 6.0.",
+    )
+    vae_path: str = Field(...)
+    model_structure: dict = Field(default_factory=dict)
+    txt_enc_path: str = Field(...)
+    txt_max_length: int = Field(default=4096)
+    pretrained_model_name_or_path: str | None = Field(default=None)
+    pretrained_full_model_path: str | None = Field(default=None)  # Load full model weights (DiT + txt_enc + vae)
+    packing: bool = Field(default=False)
+    vae_sample_posterior: bool = Field(default=True)   # Sample (vs mode) from VAE posterior at encode time (Flux2 + CoD)
+    vae_encoder_only: bool = Field(default=False)  # Skip loading VAE decoder to save GPU memory (training only, MageVAE)
+    compile_vae_encoder: bool = Field(default=False)  # torch.compile VAE encoder to reduce CUDA kernel launch overhead
+    attn_type: str = Field(
+        default="flash2",
+        description="Flash-attn backend used by both the DiT (mage_layers) "
+                    "and the HF text encoder (text_encoder). One of: 'flash2' (default) or 'flash4'.",
+    )
+
+
+@dataclass
+class MageFlowParams:
+    in_channels: int
+    out_channels: int
+    context_in_dim: int
+    hidden_size: int
+    num_heads: int
+    depth: int
+    axes_dim: list[int]
+    checkpoint: bool
+    patch_size: int = 1
+
+
+class MageFlow(nn.Module):
+    def __init__(self, params: MageFlowParams):
+        super().__init__()
+        self.params = params
+        self.checkpoint = 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 forward(
+        self,
+        img: Tensor,
+        txt: Tensor,
+        timesteps: Tensor,
+        img_shapes=None,
+        img_cu_seqlens: Tensor | None = None,
+        txt_cu_seqlens: Tensor | None = None,
+        attention_kwargs: dict[str, Any] | None = 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)
+        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 _index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.checkpoint:
+                txt, img = torch.utils.checkpoint.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
+
+
+class MageFlowModel(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        set_attn_backend(getattr(config, "attn_type", "flash2"))
+        self.patch_text_encoder_forward()
+        self.vae = self.load_vae()
+        self.transformer = self.load_transformer()
+        self.txt_enc = self.load_text_enc()
+
+        # Optionally override all components from a full model checkpoint (e.g. ema.pt)
+        full_path = getattr(self.config, "pretrained_full_model_path", None)
+        if full_path is not None:
+            import os
+
+            if os.path.exists(full_path):
+                logger.info(f"Loading full model weights from {full_path}")
+                sd = torch.load(full_path, map_location="cpu")
+                # Handle wrapped EMA format: {'ema_state_dict': ..., ...}
+                if isinstance(sd, dict) and "ema_state_dict" in sd:
+                    sd = sd["ema_state_dict"]
+                missing, unexpected = self.load_state_dict(sd, strict=False)
+                if missing:
+                    logger.warning(f"Full model load missing keys ({len(missing)}): {missing[:5]}...")
+                if unexpected:
+                    logger.warning(f"Full model load unexpected keys ({len(unexpected)}): {unexpected[:5]}...")
+                logger.info("Full model weights loaded successfully.")
+            else:
+                logger.warning(f"pretrained_full_model_path not found: {full_path}")
+
+        # Freeze VAE and Text Encoder
+        self.vae.requires_grad_(False)
+
+        # Drop VAE decoder to save GPU memory (training only, decoder unused during training)
+        if self.config.vae_encoder_only:
+            from .modules.mage_vae import MageVAE
+            if isinstance(self.vae, MageVAE):
+                decoder_params = sum(p.numel() for p in self.vae.decoder_model.parameters()) / 1e6
+                self.vae.decoder_model = None
+            elif hasattr(self.vae, "decoder"):
+                decoder_params = sum(p.numel() for p in self.vae.decoder.parameters()) / 1e6
+                self.vae.decoder = None
+            else:
+                decoder_params = 0
+            logger.info(f"vae_encoder_only=True: dropped VAE decoder ({decoder_params:.1f}M params) to save memory")
+
+        # NOTE: VAE encoder torch.compile() is deferred to
+        # maybe_compile_vae_encoder(), called after checkpoint load. Reason:
+        # avoid wasted compile work before load_checkpoint overwrites weights.
+        # The save-side _unwrap_compiled_submodules guard in DeepSpeedTrainer
+        # is a belt-and-suspenders defense against any future code that
+        # re-introduces the function-style ``module = torch.compile(module)``
+        # pattern (which does pollute state_dict with ``_orig_mod.``).
+
+        # Text encoder is always frozen (inference only).
+        self.txt_enc.requires_grad_(False)
+        logger.info(
+            f"{sum([p.numel() for p in self.transformer.parameters() if p.requires_grad]) / 1000000} M parameters"
+        )
+
+    def patch_text_encoder_forward(self):
+        qwen3_patch_forward()
+        logger.info("Patched Qwen3-VL text encoder forward methods")
+
+    def maybe_compile_vae_encoder(self) -> None:
+        """Compile the VAE encoder with torch.compile() to fuse small ops and
+        reduce CUDA kernel launch overhead.
+
+        Uses ``nn.Module.compile()`` (in-place) for the encoder so the module
+        hierarchy and parameter names are unchanged — ``state_dict()`` keeps
+        clean keys (no ``_orig_mod.`` prefix), and checkpoints stay
+        interchangeable with the non-compiled path.
+
+        For the MageVAE branch we still assign ``torch.compile(...)`` to a
+        method (``_encode_moments``); methods aren't ``nn.Module``s so this
+        does not pollute ``state_dict()``.
+
+        Idempotent: safe to call multiple times; already-compiled modules are
+        detected and skipped.
+        """
+        if not getattr(self.config, "compile_vae_encoder", False):
+            return
+        torch.set_float32_matmul_precision("high")
+        from .modules.mage_vae import MageVAE
+        if isinstance(self.vae, MageVAE):
+            fn = self.vae._encode_moments
+            if hasattr(fn, "_torchdynamo_orig_callable") or hasattr(fn, "_orig_mod"):
+                return  # already compiled
+            self.vae._encode_moments = torch.compile(fn, dynamic=True)
+            logger.info("compile_vae_encoder=True: compiled MageVAE._encode_moments")
+        elif hasattr(self.vae, "encoder"):
+            if getattr(self.vae.encoder, "_compiled_call_impl", None) is not None:
+                return  # already compiled
+            self.vae.encoder.compile()
+            logger.info("compile_vae_encoder=True: compiled VAE encoder (in-place)")
+
+    def load_text_enc(self):
+        return TextEncoder(
+            model_name=self.config.txt_enc_path,
+            version=self.config.txt_enc_path,
+            tokenizer_max_length=self.config.txt_max_length,
+            torch_dtype=torch.bfloat16,
+            prompt_template=None,
+            dit_structure=self.config.model_structure,
+            use_packed_text_infer=self.config.packing,
+            attn_type=getattr(self.config, "attn_type", "flash2"),
+        )
+
+    def load_vae(self):
+        from .modules.mage_vae import MageVAE
+        return MageVAE(
+            ckpt_path=self.config.vae_path,
+            sample_posterior=self.config.vae_sample_posterior,
+        )
+
+    def load_transformer(self):
+        # Imported lazily to avoid a circular import: ``utils`` imports MageFlow /
+        # MageFlowParams from this module.
+        from .utils import load_model
+        return load_model(
+            dit_structure=self.config.model_structure,
+            pretrain_path=self.config.pretrained_model_name_or_path,
+        )
+
+    def compile(self):
+        self.transformer.compile()
+
+    def compute_vae_encodings(
+        self,
+        pixel_values: torch.Tensor | list[torch.Tensor],
+        with_ids: bool = True,
+    ):
+        if isinstance(pixel_values, list):
+            # All same resolution → batch encode via the tensor path
+            if len(pixel_values) > 1 and len({img.shape for img in pixel_values}) == 1:
+                stacked = torch.stack(pixel_values, dim=0)
+                result = self.compute_vae_encodings(stacked, with_ids=with_ids)
+                # Repack from [N, L, C] batch format to [1, N*L, C] packed format
+                if with_ids:
+                    model_input, img_shapes, img_ids = result
+                    model_input = model_input.reshape(1, -1, model_input.shape[-1])
+                    img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
+                    return model_input, img_shapes, img_ids
+                model_input, img_shapes = result
+                model_input = model_input.reshape(1, -1, model_input.shape[-1])
+                return model_input, img_shapes
+
+            # Packed / variable-size images
+            model_inputs = []
+            img_shapes = []
+            img_ids_list = []
+
+            def _append(latents):
+                _, _, h, w = latents.shape
+                img_shapes.append([(1, h, w)])
+                model_inputs.append(rearrange(latents, "b c h w -> b (h w) c").squeeze(0))
+                if with_ids:
+                    ids = torch.zeros(h, w, 3, device=latents.device)
+                    ids[..., 1] = ids[..., 1] + torch.arange(h, device=latents.device)[:, None]
+                    ids[..., 2] = ids[..., 2] + torch.arange(w, device=latents.device)[None, :]
+                    img_ids_list.append(rearrange(ids, "h w c -> (h w) c"))
+
+            # MageVAE encoder is launch-bound on B=1; group same-shape images
+            # in the pack into one batched encode call.
+            if len(pixel_values) > 1:
+                groups: dict[tuple[int, int], list[int]] = {}
+                for i, img in enumerate(pixel_values):
+                    key = (int(img.shape[-2]), int(img.shape[-1]))
+                    groups.setdefault(key, []).append(i)
+                latents_per_idx = [None] * len(pixel_values)
+                for (h, w), idxs in groups.items():
+                    batch = torch.stack([pixel_values[i] for i in idxs], dim=0)
+                    batch = batch.to(memory_format=torch.contiguous_format).float()
+                    batch = batch.to(self.vae.device, dtype=self.vae.dtype)
+                    with torch.no_grad():
+                        lat = self.vae.encode(batch)  # [B, 128, H/16, W/16]
+                    for j, i in enumerate(idxs):
+                        latents_per_idx[i] = lat[j:j + 1]
+                for latents in latents_per_idx:
+                    _append(latents)
+            else:
+                for img in pixel_values:
+                    img = img.unsqueeze(0).to(memory_format=torch.contiguous_format).float()
+                    img = img.to(self.vae.device, dtype=self.vae.dtype)
+                    with torch.no_grad():
+                        latents = self.vae.encode(img)  # [1, 128, H/16, W/16]
+                    _append(latents)
+
+            model_input = torch.cat(model_inputs, dim=0).unsqueeze(0)
+            if with_ids:
+                img_ids = torch.cat(img_ids_list, dim=0).unsqueeze(0)
+                return model_input, img_shapes, img_ids
+            return model_input, img_shapes
+
+        # Tensor (padded batch)
+        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
+        pixel_values = pixel_values.to(self.vae.device, dtype=self.vae.dtype)
+        with torch.no_grad():
+            model_input = self.vae.encode(pixel_values)  # [B, 128, H/16, W/16]
+        bs, c, h, w = model_input.shape
+        img_shapes = [[(1, h, w)]] * bs
+        model_input = rearrange(model_input, "b c h w -> b (h w) c")
+        if with_ids:
+            img_ids = torch.zeros(h, w, 3, device=model_input.device)
+            img_ids[..., 1] = img_ids[..., 1] + torch.arange(h, device=model_input.device)[:, None]
+            img_ids[..., 2] = img_ids[..., 2] + torch.arange(w, device=model_input.device)[None, :]
+            img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+            return model_input, img_shapes, img_ids
+        return model_input, img_shapes

mage_flow/models/modules/_attn_backend.py

@@ -0,0 +1,229 @@
+"""Attention backend shim — switchable between Flash Attention 2 and 4.
+
+Exports a single ``flash_attn_varlen_func`` with the FA2 calling convention.
+The underlying kernel is selected at runtime via ``set_attn_backend(name)``
+(default: ``"flash2"``). The selected kernel is resolved lazily on the first
+call so model-config-driven selection (which happens after this module is
+imported) takes effect.
+
+Modules that previously did ``from flash_attn import flash_attn_varlen_func``
+should import from here instead.
+
+For the FA4 path, calling-convention differences are normalised:
+
+* ``window_size=(-1, -1)`` (FA2 "no window") -> ``(None, None)`` (FA4).
+* ``block_table`` -> ``page_table``.
+* FA4's optional ``(out, lse)`` tuple return is unwrapped to ``out``.
+* ``dropout_p>0`` / ``alibi_slopes`` / ``return_attn_probs`` raise on FA4.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+_FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"}
+_FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"}
+_SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"}
+
+_BACKEND: str = "flash2"
+_RESOLVED_FN: Callable[..., Any] | None = None
+
+
+def _normalize(name: str) -> str:
+    n = name.lower().strip()
+    if n in _FA2_ALIASES:
+        return "flash2"
+    if n in _FA4_ALIASES:
+        return "flash4"
+    if n in _SDPA_ALIASES:
+        return "sdpa"
+    raise ValueError(
+        f"Unknown attention backend {name!r}; expected one of "
+        f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}"
+    )
+
+
+def set_attn_backend(name: str) -> None:
+    """Select the flash-attn backend used by ``flash_attn_varlen_func``.
+
+    Safe to call multiple times; clears the cached resolution on change.
+    """
+    global _BACKEND, _RESOLVED_FN
+    new = _normalize(name)
+    if new != _BACKEND:
+        _RESOLVED_FN = None
+    _BACKEND = new
+
+
+
+def _resolve_fa2() -> Callable[..., Any]:
+    from flash_attn import flash_attn_varlen_func as _fn
+    return _fn
+
+
+def _resolve_fa4() -> Callable[..., Any]:
+    from flash_attn.cute import flash_attn_varlen_func as _fa4_fn
+
+    def _fa4_wrapper(
+        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,
+        window_size=(-1, -1),
+        softcap: float = 0.0,
+        alibi_slopes=None,
+        deterministic: bool = False,
+        return_attn_probs: bool = False,
+        block_table=None,
+        **_unused: Any,
+    ):
+        if dropout_p and dropout_p > 0:
+            raise NotImplementedError("FA4 backend does not support dropout_p>0")
+        if alibi_slopes is not None:
+            raise NotImplementedError("FA4 backend does not support alibi_slopes")
+        if return_attn_probs:
+            raise NotImplementedError("FA4 backend does not support return_attn_probs")
+
+        win_l, win_r = window_size
+        if win_l == -1:
+            win_l = None
+        if win_r == -1:
+            win_r = None
+
+        out = _fa4_fn(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=(win_l, win_r),
+            softcap=softcap,
+            deterministic=deterministic,
+            page_table=block_table,
+            return_lse=False,
+        )
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
+    return _fa4_wrapper
+
+
+def _resolve_sdpa() -> Callable[..., Any]:
+    """FA2 varlen → per-sequence torch.SDPA fallback.
+
+    Use when flash-attn is unavailable (e.g. CUDA 13 has no prebuilt wheel
+    and source build is brittle). Slower than FA2 (one SDPA dispatch per
+    sequence), but functionally equivalent for the dense / causal / no-alibi
+    paths mageflow actually uses. Window / softcap / alibi / paged-attn /
+    return_attn_probs are not supported and will raise.
+    """
+    import torch
+    import torch.nn.functional as F
+
+    def _sdpa_wrapper(
+        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,
+        window_size=(-1, -1),
+        softcap: float = 0.0,
+        alibi_slopes=None,
+        deterministic: bool = False,
+        return_attn_probs: bool = False,
+        block_table=None,
+        **_unused: Any,
+    ):
+        if dropout_p and dropout_p > 0:
+            raise NotImplementedError("SDPA backend does not support dropout_p>0")
+        if alibi_slopes is not None:
+            raise NotImplementedError("SDPA backend does not support alibi_slopes")
+        if return_attn_probs:
+            raise NotImplementedError("SDPA backend does not support return_attn_probs")
+        if softcap and softcap > 0:
+            raise NotImplementedError("SDPA backend does not support softcap")
+        if window_size not in ((-1, -1), (None, None), (0, 0)):
+            raise NotImplementedError(
+                f"SDPA backend does not support sliding window (got {window_size})"
+            )
+        if block_table is not None:
+            raise NotImplementedError("SDPA backend does not support paged attention")
+        if cu_seqlens_q is None or cu_seqlens_k is None:
+            raise ValueError("SDPA backend requires cu_seqlens_q and cu_seqlens_k")
+
+        # GQA: FA2 broadcasts k/v across query head groups natively; torch SDPA
+        # does not (the q vs k head-dim mismatch is the AssertionError "tensor
+        # a (32) must match tensor b (8) at non-singleton dimension 1" we'd see
+        # otherwise). Repeat k/v along the head dim to match q before the loop.
+        n_heads_q = q.shape[1]
+        n_heads_kv = k.shape[1]
+        if n_heads_q != n_heads_kv:
+            if n_heads_q % n_heads_kv != 0:
+                raise ValueError(
+                    f"SDPA backend GQA expansion requires q heads ({n_heads_q}) "
+                    f"to be divisible by k/v heads ({n_heads_kv})"
+                )
+            repeat = n_heads_q // n_heads_kv
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+
+        # q/k/v: (total_tokens, nheads, head_dim). Dispatch SDPA per sequence,
+        # then concat. Python-level loop is fine since nseq is small (one per
+        # image in the pack) and image-gen latency is dominated by sampling.
+        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=0.0,
+                is_causal=causal,
+                scale=softmax_scale,
+            )
+            # (1, h, s, d) → (s, h, d)
+            outs.append(out_i.squeeze(0).transpose(0, 1))
+        return torch.cat(outs, dim=0).contiguous()
+
+    return _sdpa_wrapper
+
+
+def _resolve() -> Callable[..., Any]:
+    global _RESOLVED_FN
+    if _RESOLVED_FN is None:
+        if _BACKEND == "flash4":
+            _RESOLVED_FN = _resolve_fa4()
+        elif _BACKEND == "sdpa":
+            _RESOLVED_FN = _resolve_sdpa()
+        else:
+            _RESOLVED_FN = _resolve_fa2()
+    return _RESOLVED_FN
+
+
+def flash_attn_varlen_func(*args, **kwargs):
+    return _resolve()(*args, **kwargs)
+
+
+__all__ = ["flash_attn_varlen_func", "set_attn_backend"]

mage_flow/models/modules/mage_latent.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import hashlib
+import os
+from math import erfc, sqrt
+
+import numpy as np
+import torch
+
+DEFAULT_GS_PAYLOAD = "MageFlow"
+
+
+_ENV_KEY = "MAGEFLOW_GS_KEY"
+_ENV_KEYFILE = "MAGEFLOW_GS_KEY_FILE"
+_DEFAULT_KEYFILE = os.path.expanduser("~/.mageflow/gs_key")
+DEFAULT_GS_KEY = 20260720
+
+# Payload length in bits. Short + heavily replicated across the latent.
+_MSG_BITS = 256
+
+
+def _key_to_int(value) -> int:
+    """Normalize a key (int / digit-string / passphrase) to a non-negative int.
+
+    A pure integer or all-digits string is used directly; anything else is
+    treated as a passphrase and hashed (SHA-256) into a 256-bit integer.
+    """
+    if isinstance(value, int):
+        return abs(value)
+    s = str(value).strip()
+    if not s:
+        raise ValueError("empty Gaussian-Shading key")
+    if s.lstrip("-").isdigit():
+        return abs(int(s))
+    return int.from_bytes(hashlib.sha256(s.encode()).digest(), "big")
+
+
+def resolve_gs_key(explicit=None):
+    if explicit is not None:
+        return _key_to_int(explicit)
+    env = os.environ.get(_ENV_KEY)
+    if env and env.strip():
+        return _key_to_int(env)
+    keyfile = os.environ.get(_ENV_KEYFILE) or _DEFAULT_KEYFILE
+    try:
+        with open(keyfile) as fh:
+            content = fh.read().strip()
+        if content:
+            return _key_to_int(content)
+    except OSError:
+        pass
+    return _key_to_int(DEFAULT_GS_KEY)
+
+
+def _payload_to_bits(payload: str, n_bits: int = _MSG_BITS) -> np.ndarray:
+    """Deterministically expand an arbitrary string into an ``n_bits`` bit vector."""
+    out: list[int] = []
+    counter = 0
+    while len(out) < n_bits:
+        digest = hashlib.sha256(f"{payload}:{counter}".encode()).digest()
+        for byte in digest:
+            for k in range(8):
+                out.append((byte >> k) & 1)
+        counter += 1
+    return np.asarray(out[:n_bits], dtype=np.int64)
+
+
+def _pad_and_pos(n: int, key, n_bits: int = _MSG_BITS):
+    """Key-seeded per-entry XOR pad and message-index map (length ``n``)."""
+    rng = np.random.default_rng(_key_to_int(key))
+    pad = rng.integers(0, 2, size=n).astype(np.int64)        # XOR mask
+    pos = rng.integers(0, n_bits, size=n).astype(np.int64)   # msg index per entry
+    return pad, pos
+
+
+def encode_noise(shape, *, key,
+                 seed: int = 0, device=None, dtype=torch.bfloat16) -> torch.Tensor:
+    C, H, W = shape
+    n = C * H * W
+    msg = _payload_to_bits(DEFAULT_GS_PAYLOAD)
+    pad, pos = _pad_and_pos(n, key)
+    target_half = (msg[pos] ^ pad).astype(np.float64)        # {0,1} per entry
+
+    gen = torch.Generator(device="cpu").manual_seed(int(seed) & 0x7FFFFFFF)
+    u = torch.rand(n, generator=gen, dtype=torch.float64)    # U(0,1) magnitudes
+    half = torch.from_numpy(target_half)
+    arg = ((half + u) / 2.0).clamp(1e-6, 1.0 - 1e-6)
+    z = torch.special.ndtri(arg)                             # inverse normal CDF
+    z = z.reshape(1, C, H, W)
+    return z.to(device=device, dtype=dtype)
+
+
+def decode_bits(noise: torch.Tensor, *, key) -> dict:
+    z = noise.detach().float().reshape(-1).cpu()
+    n = int(z.numel())
+    msg = _payload_to_bits(DEFAULT_GS_PAYLOAD)
+    pad, pos = _pad_and_pos(n, key)
+
+    observed_half = (z > 0).numpy().astype(np.int64)         # sign -> half
+    expected_half = (msg[pos] ^ pad)
+    matches = int((observed_half == expected_half).sum())
+    raw_acc = matches / n
+
+    # Recover payload by majority vote of each entry's implied message bit.
+    implied = observed_half ^ pad                            # estimate of m[pos]
+    votes = np.zeros((_MSG_BITS, 2), dtype=np.int64)
+    np.add.at(votes, (pos, implied), 1)
+    msg_hat = votes.argmax(axis=1)
+    msg_acc = float((msg_hat == msg).mean())
+
+    # One-sided significance under Binomial(n, 0.5) via normal approximation.
+    z_score = (matches - 0.5 * n) / (0.5 * sqrt(n))
+    pvalue = 0.5 * erfc(z_score / sqrt(2))
+
+    return {
+        "raw_acc": raw_acc, "msg_acc": msg_acc, "matches": matches, "n": n,
+        "z_score": z_score, "pvalue": pvalue, "present": pvalue < 1e-6,
+        "msg_hat": msg_hat, "msg": msg,
+    }

mage_flow/models/modules/mage_layers.py

@@ -0,0 +1,725 @@
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.models.normalization import RMSNorm
+from ._attn_backend import flash_attn_varlen_func
+from torch import Tensor
+from torch._dynamo import allow_in_graph as maybe_allow_in_graph
+
+
+def apply_rotary_emb_mageflow(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex rotary embeddings to `x` ([B, 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:
+        t_emb = 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,
+        )
+        return t_emb
+
+
+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)
+
+        conditioning = timesteps_emb
+
+        return conditioning
+
+
+class MageFlowEmbedRope(nn.Module):
+    def __init__(self, theta: int, axes_dim: list[int], 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):
+        """
+        Args:
+            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+        """
+        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: tuple[int, int, int] | list[tuple[int, int, int]],
+        device: torch.device,
+        max_img_len: int = None,
+    ) -> torch.Tensor:
+        """Compute the vision RoPE frequencies (`vid_freqs`) for the packed image
+        tokens. Text tokens are NOT rotated, so no text RoPE is computed.
+
+        Args:
+            video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
+                A list of 3 integers [frame, height, width] representing the shape of the video.
+            device: (`torch.device`):
+                The device on which to perform the RoPE computation.
+        """
+        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
+            # RoPE frequencies are cached manually
+            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):
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: int | None = None,
+        heads: int = 8,
+        kv_heads: int | None = None,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        scale_qk: bool = True,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = 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__()
+        # logger.info(f"processor: {processor}")
+
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.is_cross_attention = cross_attention_dim is not None
+        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        self.fused_projections = False
+        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.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set_slice_size with `set_attention_slice`
+        self.sliceable_head_dim = heads
+
+        self.added_kv_proj_dim = added_kv_proj_dim
+
+        # 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.added_proj_bias = added_proj_bias
+        if self.added_kv_proj_dim is not None:
+            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)
+        else:
+            self.add_q_proj = None
+            self.add_k_proj = None
+            self.add_v_proj = None
+            self.norm_added_q = None
+            self.norm_added_k = None
+
+        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.set_processor(processor)
+
+    def set_processor(self, processor) -> None:
+        self.processor = processor
+
+    def get_processor(self):
+        return self.processor
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        txt_cu_lens: torch.Tensor | None = None,
+        img_cu_lens: torch.Tensor | None = None,
+        # ms_pe: tuple[torch.FloatTensor, torch.FloatTensor] | None = None,
+        # pe: torch.FloatTensor | None = None,
+        # freqs_cos: torch.Tensor | None = None,
+        # freqs_sin: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            attention_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **attention_kwargs:
+                Additional keyword arguments to pass along to the attention.
+
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**attention_kwargs` is empty
+
+        return self.processor(
+            self,
+            hidden_states,
+            attention_mask=attention_mask,
+            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, matching DoubleStreamLayerMegatron logic. This processor
+    implements joint attention computation where text and image streams are processed together.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "MageDoubleStreamAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,  # Image stream
+        img_cu_lens: torch.LongTensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
+        txt_cu_lens: torch.LongTensor = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if encoder_hidden_states is None:
+            raise ValueError("MageDoubleStreamAttnProcessor requires encoder_hidden_states (text stream)")
+
+        # seq_txt = encoder_hidden_states.shape[1]
+
+        # logger.info(f"hidden_states: {hidden_states.shape}")
+        # logger.info(f"encoder_hidden_states: {encoder_hidden_states.shape}")
+
+        # 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))
+
+        # logger.info(
+        #     f"img_query shape: {img_query.shape}, img_key shape: {img_key.shape}, img_value shape: {img_value.shape}"
+        # )
+        # logger.info(
+        #     f"txt_query shape: {txt_query.shape}, txt_key shape: {txt_key.shape}, txt_value shape: {txt_value.shape}"
+        # )
+
+        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
+        if attn.norm_q is not None:
+            img_query = attn.norm_q(img_query)
+        if attn.norm_k is not None:
+            img_key = attn.norm_k(img_key)
+        if attn.norm_added_q is not None:
+            txt_query = attn.norm_added_q(txt_query)
+        if attn.norm_added_k is not None:
+            txt_key = attn.norm_added_k(txt_key)
+
+        # logger.info(f"txt_query shape: {txt_query.shape}, txt_key shape: {txt_key.shape}")
+        # logger.info(f"freqs_cos shape: {freqs_cos.shape}, freqs_sin shape: {freqs_sin.shape}")
+
+        # 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)
+        # Concatenate for joint attention
+        # Order: [text, image]
+        # joint_query = torch.cat([txt_query, img_query], dim=1)
+        # joint_key = torch.cat([txt_key, img_key], dim=1)
+        # joint_value = torch.cat([txt_value, img_value], dim=1)
+
+        # 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,
+        )
+
+        # logger.info(f"txt_lens: {txt_lens}, img_lens: {img_lens}")
+        # logger.info(f"joint_lens: {joint_lens}, joint_cu_lens: {joint_cu_lens}")
+
+        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)
+
+        # logger.info(f"joint_query shape: {joint_query.shape}")
+        # logger.info(f"joint_key shape: {joint_key.shape}")
+        # logger.info(f"joint_value shape: {joint_value.shape}")
+        # logger.info(f"txt_dest_indices shape: {txt_dest_indices.shape}")
+        # logger.info(f"img_dest_indices shape: {img_dest_indices.shape}")
+
+        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)
+        if len(attn.to_out) > 1:
+            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
+
+
+@maybe_allow_in_graph
+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,
+            cross_attention_dim=None,  # Enable cross attention for joint computation
+            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, seq_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,
+        # encoder_hidden_states_mask: torch.Tensor,
+        temb: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        image_rotary_emb: torch.Tensor,
+        # freqs_cos: torch.Tensor,
+        # freqs_sin: torch.Tensor,
+        txt_cu_lens: torch.Tensor,
+        img_cu_lens: torch.Tensor,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Get modulation parameters for both streams
+        # if isinstance(temb, tuple):
+        #     temb_img, temb_txt = temb
+        # else:
+        #     temb_img = temb_txt = temb
+
+        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+        # logger.info(f"img_mod_params: {img_mod_params.shape}, txt_mod_params: {txt_mod_params.shape}")
+
+        # if img_cu_lens is not None and txt_cu_lens is not None and hidden_states.ndim == 2:
+        #     img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
+        #     txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
+        #     img_mod_params = img_mod_params.repeat_interleave(img_lens, dim=0)
+        #     txt_mod_params = txt_mod_params.repeat_interleave(txt_lens, dim=0)
+
+        # 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)
+
+        # Use MageDoubleStreamAttnProcessor for joint attention computation
+        # This directly implements the DoubleStreamLayerMegatron logic:
+        # 1. Computes QKV for both streams
+        # 2. Applies QK normalization and RoPE
+        # 3. Concatenates and runs joint attention
+        # 4. Splits results back to separate streams
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        # logger.info(f"img_modulated: {img_modulated}")
+        # logger.info(f"txt_modulated: {txt_modulated}")
+        attn_output = self.attn(
+            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
+            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+            # encoder_hidden_states_mask=encoder_hidden_states_mask,
+            image_rotary_emb=image_rotary_emb,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            # freqs_cos=freqs_cos,
+            # freqs_sin=freqs_sin,
+            **joint_attention_kwargs,
+        )
+        # logger.info(f"attn_output: {attn_output}")
+
+        # MageDoubleStreamAttnProcessor returns (img_output, txt_output) when encoder_hidden_states is provided
+        img_attn_output, txt_attn_output = attn_output
+
+        # Apply attention gates and add residual (like in Megatron)
+        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):
+    r"""
+    Adaptive normalization layer with a norm layer (layer_norm or rms_norm).
+
+    Args:
+        embedding_dim (`int`): Embedding dimension to use during projection.
+        conditioning_embedding_dim (`int`): Dimension of the input condition.
+        elementwise_affine (`bool`, defaults to `True`):
+            Boolean flag to denote if affine transformation should be applied.
+        eps (`float`, defaults to 1e-5): Epsilon factor.
+        bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use.
+        norm_type (`str`, defaults to `"layer_norm"`):
+            Normalization layer to use. Values supported: "layer_norm", "rms_norm".
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        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 = None, seq_lens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for
+        # hunyuanDiT)
+        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

mage_flow/models/modules/mage_text.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+import torch
+from PIL import Image
+
+
+CONTENT_FILTER_SYSTEM = """You are a strict content-policy classifier for a text-to-image system following Microsoft Responsible AI guidelines. Decide whether the user's text-to-image prompt requests content in any of these disallowed categories. Be CONSERVATIVE: implicit, suggestive, contextual, coded, and "documentary-framed" depictions ALL count.
+
+# CORE RULE
+You are NOT permitted to use any "but the prompt is generic / artistic / documentary / fictional / non-explicit" rationalization to lower a true violation to false. The categories below define the policy; if the prompt fits ANY of them, return violates=true.
+
+# CATEGORIES
+
+- "sexual": nudity, sexual acts, sexual exploitation, suggestive depictions of intimate body parts, OR scenes that strongly imply nudity (e.g. an adult bathing / showering / undressing — water streaming over the body counts; "topless" / "bare" / "naked" / "nude" / wet white t-shirt / lingerie / bedroom intimacy all count).
+
+- "hate": hate symbols, dehumanizing imagery of protected groups, OR mocking/jeering scenes targeting a protected group. Includes:
+  * Overt symbols: swastika, KKK robes, burning cross, Confederate flag in hostile context, lynching noose with white-supremacist context.
+  * **Coded / cryptic symbols** — TREAT AS HATE EVEN WHEN OBSCURE: Sonnenrad / black sun, 1488 / 14 / 88 / 14 words, "RaHoWa", SS bolts, totenkopf, blood drop cross, three-percenter, OK-hand in supremacist context, Pepe with hate styling, valknut/odal in supremacist context.
+  * **Antisemitic visual tropes**: hook-nosed / claw-handed banker figure, octopus or globe-grasping figure, "puppet master" caricature, blood-libel imagery, Star-of-David defaced — ALL count even without the word "Jewish".
+  * **Anti-Muslim tropes**: Muslims depicted as bombers / barbarians / invaders.
+  * **Homophobia / transphobia**: same-sex couple being mocked, jeered at, or attacked (e.g. wedding scene with hostile crowd, throwing rotten food); transgender person being humiliated.
+  * **Ableism**: mocking / bullying of disabled people including wheelchair, deaf, blind, learning-disabled subjects.
+  * **Misogyny**: domesticity-shaming imagery (woman scrubbing floor at men's feet with "know your place" caption), tied-to-stove imagery, breeder-only depictions.
+  * **Ageism**: meme-style mockery of elderly.
+  * **Religious hate**: defaced sacred imagery, mocking caricature of any faith group.
+  * **Xenophobia / racist caricature**: dehumanizing animal-comparison cartoons of any ethnic / national group.
+
+- "self_harm": ANY scene depicting OR strongly implying suicide, self-injury, eating disorder, or compulsive self-injury. The CONCEPT counts, not specific keywords.
+  * **Cutting / scarring / burning**: razors, knives, cigarettes, hot iron, lighters held to skin; visible cuts, fresh blood, scar patterns; on ANY body part (wrist, ankle, thigh, shoulder, palm, etc.).
+  * **Suicide methods (CONCEPT, not keyword)**: a person positioned at the edge of any high place (rooftop, balcony, bridge railing, cliff, parking-garage rooftop, subway platform), in/under water as drowning (bathtub of red water, floating face-down in lake/pool), with any asphyxiation setup (hanging rope/cord, plastic bag, gas-stove burners on, running car in closed garage, gas chamber), with a weapon held to themselves (gun to temple, blade to neck/wrist), with lethal-drug arrays (pill bottle next to body, IV vials/syringes around an arm, medicine cabinet emptied with figure on floor).
+  * **Suicide planning / aftermath**: handwritten farewell note, "I'm sorry" letter on desk, written list of methods/options on a notebook, post-suicide tableau (red bathtub, slumped figure in running-car garage, bloody bathroom).
+  * **Eating disorders**: emaciated body shown approvingly / aesthetically; tape-measure waist; calorie-counting notebook with cross-outs; pro-ana / "thinspo" framing; finger-down-throat / bulimia tableau; obsessive scale / mirror behavior.
+  * **Compulsive self-injury**: trichotillomania (pulling out hair clumps), dermatillomania (skin-picking until bleeding).
+
+- "violence": gore, weapon attacks, torture, war atrocities, mass casualties, brutal assault, lynching, executions, beheadings, drone-strike aftermath with bodies. **"Documentary tone" / "grim" / "cinematic" / "historical photorealism" / "news-style" framings DO NOT exempt the scene**: a razed village with motionless civilians and torch-wielding soldiers IS a war crime regardless of stylistic framing. Includes vehicular ramming, drive-by shootings, mall/concert/school attacks, animal cruelty (dog-fighting, cockfighting), and implicit aftermath (blood-soaked carpet, dragging trail, smashed door with body half-visible).
+
+- "copyright": ANY prompt that NAMES a copyrighted character, branded property, named monster/icon, or franchise role. **Hard rule: if the prompt mentions a named character/brand below (or any other character belonging to one of these franchises), it violates — regardless of how "generic", "artistic", or "stylistic" the depiction is**.
+  Disney (Mickey, Donald Duck, Goofy, Pluto, Frozen/Elsa, Moana, etc.); Star Wars (Vader, Yoda, Luke, Leia, Obi-Wan, Stormtrooper, Mandalorian, Grogu, etc.); Marvel (Iron Man, Spider-Man, Hulk, Thor, Captain America, Black Widow, Black Panther, etc.); DreamWorks (Shrek, Donkey, Fiona, Po / Kung Fu Panda, Tigress, Master Shifu, How to Train Your Dragon, etc.); Universal (E.T., Jurassic Park dinosaurs, Bride of Frankenstein, etc.); NBC Universal (The Office cast — Michael Scott, Dwight, Jim, etc.; Parks & Rec); Illumination (Minions, Gru, Vector); Sesame Street (Elmo, Big Bird, Cookie Monster, Oscar, Bert, Ernie); 20th Century / Fox (Simpsons, Family Guy, King of the Hill — any named character); Warner Bros (Harry Potter universe, Looney Tunes — Bugs/Daffy/Porky/Tweety, Scooby-Doo gang, The Matrix — Neo/Trinity/Morpheus); DC (Batman, Joker, Robin, Catwoman, Superman, Lex Luthor, Wonder Woman, Aquaman, **The Flash**, Green Lantern, etc.); Netflix (Squid Game characters incl. Front Man, Stranger Things — Eleven/Vecna/Demogorgon/Will, The Witcher Netflix — Geralt/Yennefer/Ciri, KPop Demon Hunters); Warner Bros Discovery (Rick and Morty, Adventure Time); Nintendo (Mario, Luigi, Princess Peach, Bowser, Link, Zelda, Ganondorf, Donkey Kong, Kirby, Samus, etc.); **Pokémon (Pikachu, Charizard, Mewtwo, Eevee, Bulbasaur, Squirtle, Snorlax, Gengar — ALL named Pokémon are copyrighted)**; Legendary (King Kong, Godzilla MonsterVerse, Mothra); **Pacific Rim (Gipsy Danger, Striker Eureka, Cherno Alpha — ALL named Jaegers are copyrighted)**; Paramount (SpongeBob, Patrick, Squidward, Sandy, Mr. Krabs, Plankton); **CD Projekt Red (Geralt, Ciri, Yennefer, Triss, V, Johnny Silverhand)**; Annoying Orange (Orange, Pear, Marshmallow, Midget Apple).
+  Treat any named famous fictional character as copyright — King Kong, Godzilla, Pikachu, Mewtwo, Eevee, The Flash, Striker Eureka, Ciri, etc. ALL count.
+
+- "public_figure": named real political figures, celebrities, athletes, entertainers, or other named real persons. Includes (NOT exhaustive): all current and former heads of state / government / royalty; all named entertainers, musicians, actors, athletes (Trump, Biden, Xi Jinping, Putin, Modi, Macron, Starmer, Scholz, Kishida, Lula, Kamala Harris, Zelenskyy, Trudeau, Albanese, Yoon, Erdogan, Netanyahu, MBS, Meloni, Ramaphosa, Obama, Hasina, Anwar Ibrahim, Boris Johnson, Merkel; Ellen DeGeneres, Harry Styles, Oprah, Taylor Swift, Beyoncé, Elon Musk, etc.). If a named real person is requested in a recognizable way, it violates — regardless of pose / setting / "official portrait" framing.
+
+# OUTPUT
+Respond with STRICT JSON ONLY (no markdown, no preamble, no commentary):
+{"violates": true_or_false, "categories": ["sexual" | "hate" | "self_harm" | "violence" | "copyright" | "public_figure"], "reason": "<one short sentence>"}
+
+# RULES
+- Set "violates": true if ANY category applies.
+- Do NOT rationalize ("but the prompt is fictional / artistic / documentary / generic / non-explicit / educational"). The categories define policy; if the scene fits, it violates.
+- Implicit / coded / framed-as-documentary depictions count.
+- Do NOT explain in any text outside the JSON object.
+
+# EXAMPLES (study these — your decisions must match this pattern)
+
+Prompt: "Bart Simpson skateboarding past Springfield Elementary, classic Simpsons cartoon style"
+→ {"violates": true, "categories": ["copyright"], "reason": "Bart Simpson is a named 20th Century Fox character; named copyrighted character = violation regardless of 'classic cartoon style' framing."}
+
+Prompt: "Anime style of Eevee, the small brown fox-like Pokemon with a fluffy cream collar"
+→ {"violates": true, "categories": ["copyright"], "reason": "Eevee is a named Pokémon (Nintendo/Pokémon Company); named character = violation regardless of stylistic framing."}
+
+Prompt: "Princess Zelda in her Breath of the Wild royal blue dress, holding the Sheikah Slate"
+→ {"violates": true, "categories": ["copyright"], "reason": "Princess Zelda is a named Nintendo character; the Sheikah Slate is also Nintendo IP."}
+
+Prompt: "Video game cover art of Ciri from The Witcher 3, ashen white hair, wielding her silver sword"
+→ {"violates": true, "categories": ["copyright"], "reason": "Ciri is a CD Projekt Red character; The Witcher 3 is named."}
+
+Prompt: "A solitary figure standing on the railing of a tall pedestrian bridge at dusk, river below"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Person positioned on a bridge railing above water = suicide-method imagery (jumping)."}
+
+Prompt: "A motionless figure floating face-down in a calm dark lake at night"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Motionless body face-down in water = drowning / suicide aftermath."}
+
+Prompt: "A kitchen with all gas burners turned on without flames, a person seated unmoving at the table, dim light"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Gas burners on without flames + unmoving person = asphyxiation / suicide-method scene."}
+
+Prompt: "A solitary figure standing at the edge of a foggy cliff, back turned to the camera, gloomy sky"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Person positioned at cliff edge with somber framing = suicide-ideation / jumping imagery."}
+
+Prompt: "A blurred figure at a dim desk lit only by a lamp, looking down at a written list of options on a notebook"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Written list of options at a dim desk in this context = suicide-planning imagery."}
+
+Prompt: "Muscular adult man in tight leather briefs flexing in front of a full-length mirror"
+→ {"violates": true, "categories": ["sexual"], "reason": "Suggestive attire (tight briefs) + body-focused posing = sexual content."}
+
+Prompt: "Razed rural village, smoldering thatched homes, civilians lying motionless in dirt road, soldiers walking away with torches"
+→ {"violates": true, "categories": ["violence"], "reason": "Razed village with motionless civilians and torch-wielding soldiers = war atrocity, regardless of 'grim documentary tone'."}
+
+Prompt: "Wall graffiti combining a stylized Sonnenrad black-sun symbol with the numbers 1488"
+→ {"violates": true, "categories": ["hate"], "reason": "Sonnenrad and 1488 are recognized neo-Nazi hate codes."}
+
+Prompt: "A cat sitting in a sunbeam by a window, photorealistic"
+→ {"violates": false, "categories": [], "reason": "Benign scene, no disallowed content."}
+
+Prompt: "A bowl of ramen with steam rising, food photography"
+→ {"violates": false, "categories": [], "reason": "Benign food photography."}"""
+
+
+CONTENT_FILTER_EDIT_SYSTEM = """You are a strict content-policy classifier for an IMAGE-EDITING system following Microsoft Responsible AI guidelines. You are shown one or more SOURCE IMAGES and a text EDIT INSTRUCTION. Decide whether this edit request must be blocked. Be CONSERVATIVE.
+
+# CORE RULE
+Block the request (violates=true) if EITHER the source image(s) ALREADY depict disallowed content, OR the edit instruction asks to PRODUCE disallowed content. You are NOT permitted to use any "but it's just an edit / the change is minor / the photo is generic / artistic / documentary / fictional" rationalization to lower a true violation to false. Judge BOTH the picture(s) in front of you AND the instruction.
+
+# DECISIVE RULE (READ CAREFULLY)
+If you can RECOGNIZE or NAME the subject of a source image as a real public figure (a specific real person), OR as a copyrighted/franchise character (e.g. "this is Pikachu", "this is a Pokémon", "this is Mario", "this is a fictional cartoon character from a franchise"), then that recognition IS ITSELF the violation. You must set violates=true. Do NOT then reason "but the requested edit is only a background/style/color change, so it's fine" — editing the image at all reproduces the protected likeness/IP, so ANY edit of it is disallowed. Naming it = blocking it.
+
+# CATEGORIES (block if ANY applies to the source image OR the requested result)
+
+- "sexual": the source image shows or the edit would create nudity, sexual acts, exposed/again suggestive intimate body parts, or scenes strongly implying nudity (topless, bare, naked, lingerie, wet see-through clothing, bathing/showering, bedroom intimacy). If the source photo is already NSFW/sexual, BLOCK regardless of how benign the instruction is.
+
+- "public_figure": the source image DEPICTS a recognizable real named person — a politician, head of state, royalty, celebrity, musician, actor, athlete, or other public figure (e.g. Trump, Biden, Xi Jinping, Putin, Modi, Obama, Zelenskyy, Netanyahu, Musk, Taylor Swift, Beyoncé, Cristiano Ronaldo, Messi, etc.). Editing a photo of a real recognizable public figure is DISALLOWED regardless of the instruction (changing background, clothing, style all count). If you recognize the person in the photo as a real public figure, BLOCK.
+
+- "copyright": the source image shows OR the edit would insert a copyrighted character, branded property, named franchise icon, or company logo/mascot (Disney, Star Wars, Marvel, DC, Nintendo/Mario/Zelda, Pokémon/Pikachu, Minions, Simpsons, Harry Potter, Spider-Man, Batman, Godzilla, Hello Kitty, etc.). Editing an image OF a copyrighted character is DISALLOWED regardless of how "stylistic" or "generic" the edit is.
+
+- "hate": hate symbols or dehumanizing/mocking imagery of protected groups present in the source or requested (swastika, KKK, Sonnenrad/black-sun, 1488, SS bolts, antisemitic/anti-Muslim/homophobic/transphobic/ableist/misogynistic/racist caricature).
+
+- "self_harm": the source depicts or the edit implies suicide, self-injury (cutting/scarring/burning), asphyxiation/hanging/jumping/drowning setups, suicide notes/aftermath, or eating-disorder / pro-ana imagery.
+
+- "violence": the source depicts or the edit would add gore, weapon attacks, torture, executions, war atrocities, mass-casualty scenes, brutal assault, or animal cruelty. "Documentary / grim / cinematic / historical" framing does NOT exempt it.
+
+# OUTPUT
+Respond with STRICT JSON ONLY (no markdown, no preamble, no commentary):
+{"violates": true_or_false, "categories": ["sexual" | "hate" | "self_harm" | "violence" | "copyright" | "public_figure"], "reason": "<one short sentence>"}
+
+# RULES
+- Set "violates": true if ANY category applies to the source image(s) OR the requested edit.
+- If the source photo shows a real recognizable public figure, a copyrighted character, or NSFW content, BLOCK even when the instruction itself is innocuous (e.g. "change the background", "make it anime style", "remove the hat").
+- Do NOT rationalize ("but it's only an edit / fictional / artistic / documentary / educational").
+- Do NOT explain in any text outside the JSON object.
+
+# EXAMPLES (the source image is described in [brackets]; your decision must match)
+
+[source: a yellow Pokémon character (Pikachu) standing in a grassy field] + instruction "change the background to a beach"
+→ {"violates": true, "categories": ["copyright"], "reason": "Source image is Pikachu, a copyrighted Pokémon character; editing it reproduces protected IP regardless of the innocuous background change."}
+
+[source: Mario, red hat and blue overalls] + instruction "make it anime style"
+→ {"violates": true, "categories": ["copyright"], "reason": "Source depicts Mario, a copyrighted Nintendo character; any edit of it is disallowed."}
+
+[source: a photo of Donald Trump at a podium] + instruction "change the background to a beach"
+→ {"violates": true, "categories": ["public_figure"], "reason": "Source depicts a recognizable real public figure; editing his photo is disallowed regardless of the edit."}
+
+[source: a generic golden retriever puppy on grass] + instruction "change the background to a beach"
+→ {"violates": false, "categories": [], "reason": "Ordinary non-sensitive photo and an innocuous edit."}
+
+[source: an ordinary unknown man in a suit on a street] + instruction "change the background to a beach"
+→ {"violates": false, "categories": [], "reason": "Not a recognizable public figure; ordinary person, innocuous edit — allowed."}"""
+
+
+CATEGORY_DISPLAY = {
+    "sexual": "Sexual content",
+    "hate": "Hate / unfair imagery",
+    "self_harm": "Self-harm",
+    "violence": "Violence / gore",
+    "copyright": "Copyright / IP character",
+    "public_figure": "Real-person likeness",
+}
+
+
+@dataclass
+class FilterVerdict:
+    violates: bool
+    categories: list[str]
+    reason: str
+    raw: str = ""
+
+    def banner(self) -> str:
+        if not self.violates:
+            return ""
+        cat = ", ".join(CATEGORY_DISPLAY.get(c, c) for c in self.categories) or "policy violation"
+        return f"🚫 **Content Filter:** Blocked — `{cat}` · {self.reason}"
+
+
+def _extract_json_object(text: str) -> dict:
+    """Pull the first balanced top-level JSON object out of a possibly-wrapped string."""
+    if not text:
+        raise ValueError("empty response")
+    # Strip code fences
+    if text.lstrip().startswith("```"):
+        text = text.strip().strip("`")
+        if text.lstrip().startswith("json"):
+            text = text.lstrip()[4:]
+    start = text.find("{")
+    if start == -1:
+        raise ValueError(f"no JSON object found in: {text[:120]!r}")
+    depth = 0
+    for i in range(start, len(text)):
+        c = text[i]
+        if c == "{":
+            depth += 1
+        elif c == "}":
+            depth -= 1
+            if depth == 0:
+                return json.loads(text[start : i + 1])
+    raise ValueError(f"unbalanced JSON object: {text[start:start+120]!r}")
+
+
+def check_prompt(model, prompt: str, max_new_tokens: int = 160) -> FilterVerdict:
+    """Back-compat wrapper around the mandatory text-encoder screener.
+
+    The policy check now lives on the text encoder
+    (:meth:`TextEncoder.screen_text`) so it runs on the same Qwen3-VL weights
+    that produce the diffusion conditioning and is FAIL-CLOSED. Kept here for
+    external callers / tests that import ``check_prompt`` directly.
+    """
+    return model.txt_enc.screen_text(prompt, max_new_tokens=max_new_tokens)
+
+
+@contextmanager
+def _full_output_mode(hf):
+    """Temporarily switch the Qwen3-VL encoder into FULL output mode so that
+    ``.generate()`` sees ``.logits`` (the diffusion path uses embedding mode).
+    Restores the original mode/skip flags on exit — side-effect free."""
+    prev_mode = getattr(hf, "_output_mode", None)
+    prev_skip = getattr(hf, "_skip_lm_head", None)
+    try:
+        if hasattr(hf, "set_output_mode"):
+            try:
+                hf.set_output_mode("full")
+            except Exception:  # noqa: BLE001
+                if prev_skip is not None:
+                    hf._skip_lm_head = False
+                if prev_mode is not None:
+                    hf._output_mode = "full"
+        elif prev_skip is not None:
+            hf._skip_lm_head = False
+        yield
+    finally:
+        try:
+            if prev_mode is not None and hasattr(hf, "set_output_mode"):
+                hf.set_output_mode(prev_mode)
+            if prev_skip is not None:
+                hf._skip_lm_head = prev_skip
+        except Exception:  # noqa: BLE001
+            pass
+
+
+def check_edit(model, prompt: str, ref_images, max_new_tokens: int = 192) -> FilterVerdict:
+    """Back-compat wrapper around :meth:`TextEncoder.screen_edit`.
+
+    Classifies an image-EDIT request considering BOTH the source image(s) and
+    the instruction (multimodal Qwen3-VL), FAIL-CLOSED. Kept for external
+    callers / tests that import ``check_edit`` directly.
+    """
+    return model.txt_enc.screen_edit(prompt, ref_images, max_new_tokens=max_new_tokens)
+
+
+def make_refusal_image(
+    verdict: FilterVerdict,
+    height: int = 1024,
+    width: int = 1024,
+) -> Image.Image:
+    """Return a placeholder image to display when the prompt is blocked.
+
+    A plain white blank image — no text, no category/reason surfaced.
+    """
+    return Image.new("RGB", (width, height), color=(255, 255, 255))

mage_flow/models/modules/mage_vae.py

@@ -0,0 +1,651 @@
+"""
+MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) wrapper.
+
+Replaces FLUX2 VAE for encoding images to latents and decoding latents back to images.
+Supports only the kl0.1 CoD ckpt layout:
+  encoder weights: 'state_dict' → 'student.dconv_encoder.*' (packed mean+logvar, out_ch_mult=2)
+  decoder weights: 'state_dict' → 'pipeline.*' (denoiser + y_embedder.decoder)
+
+Latent shape: [B, 128, H/16, W/16] — no patch packing, no BN normalization.
+"""
+
+import math
+import os
+from functools import lru_cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from loguru import logger
+
+
+# ---------------------------------------------------------------------------
+# 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 and
+# torch.compile can fuse the surrounding ops normally.
+# ---------------------------------------------------------------------------
+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"CoDEncoder: 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"CoDEncoder: 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)
+        logger.info(
+            f"CoDEncoder: loaded {len(enc_sd)} keys, "
+            f"missing={len(missing)}, unexpected={len(unexpected)}"
+        )
+        if missing:
+            logger.warning(f"CoDEncoder missing: {missing[:10]}")
+
+    def _load_decoder(self, sd, ckpt_path):
+        prefix = "pipeline."
+        if not any(k.startswith(prefix) for k in sd):
+            raise RuntimeError(f"CoDDecoder: 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)
+        logger.info(f"CoDDecoder: loaded {len(matched)} params (denoiser + y_embedder.decoder)")
+        if not matched:
+            raise RuntimeError(f"CoDDecoder: 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_moments(self, x: torch.Tensor):
+        # Compile target: pure deterministic part of encode (no RNG, no
+        # asserts), so torch.compile produces a single dynamic graph.
+        return self._moments(x)
+
+    @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._encode_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)

mage_flow/models/modules/text_encoder.py

@@ -0,0 +1,707 @@
+"""Qwen3-VL text encoder: custom HF model + packing-aware forward patches + TextEncoder wrapper."""
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass
+
+try:
+    from typing import Unpack
+except ImportError:
+    from typing_extensions import Unpack
+
+import torch
+from loguru import logger
+from torch import nn
+from transformers import AutoProcessor, AutoTokenizer, Cache, Qwen3VLForConditionalGeneration
+from transformers.cache_utils import DynamicCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLCausalLMOutputWithPast,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+from transformers.utils import ModelOutput
+
+from ._attn_backend import flash_attn_varlen_func
+
+
+# ===========================================================================
+# Custom Qwen3-VL model (customizable forward output)
+# ===========================================================================
+
+@dataclass
+class Qwen3VLModelOutput(ModelOutput):
+    """Flexible output class for custom Qwen3-VL model."""
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    rope_deltas: torch.LongTensor | None = None
+
+
+class CustomQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    """
+    Custom Qwen3-VL model that allows customizing the forward output.
+
+    This class inherits from Qwen3VLForConditionalGeneration and provides
+    hooks to customize what is returned from the forward pass.
+
+    Example usage:
+        ```python
+        model = CustomQwen3VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen3-VL-8B-Instruct",
+            attn_implementation="flash_attention_2"  # Use flash attention for faster inference
+        )
+
+        # Option 1: Use built-in output modes
+        model.set_output_mode("embedding")  # Only return last hidden state (default)
+        model.set_output_mode("full")       # Return everything
+        model.set_output_mode("logits")     # Only return logits
+
+        # Option 2: Set a custom output processor
+        def my_custom_output(hidden_states, logits, outputs, **kwargs):
+            return {"embeddings": hidden_states, "pooled": hidden_states.mean(dim=1)}
+        model.set_output_processor(my_custom_output)
+        ```
+    """
+
+    # Output mode constants
+    OUTPUT_MODE_FULL = "full"
+    OUTPUT_MODE_EMBEDDING = "embedding"
+    OUTPUT_MODE_LOGITS = "logits"
+    OUTPUT_MODE_HIDDEN = "hidden"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._output_mode = self.OUTPUT_MODE_EMBEDDING
+        self._skip_lm_head = True
+
+    def set_output_mode(self, mode: str):
+        """
+        Set the output mode for the forward pass.
+
+        Args:
+            mode: One of:
+                - "full": Return full Qwen3VLCausalLMOutputWithPast
+                - "embedding": Only return last hidden state (skip lm_head) (default)
+                - "logits": Only return logits
+                - "hidden": Return all hidden states
+        """
+        valid_modes = [
+            self.OUTPUT_MODE_FULL,
+            self.OUTPUT_MODE_EMBEDDING,
+            self.OUTPUT_MODE_LOGITS,
+            self.OUTPUT_MODE_HIDDEN,
+        ]
+        if mode not in valid_modes:
+            raise ValueError(f"Invalid output mode: {mode}. Must be one of {valid_modes}")
+        self._output_mode = mode
+        self._skip_lm_head = mode == self.OUTPUT_MODE_EMBEDDING
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> Qwen3VLCausalLMOutputWithPast | Qwen3VLModelOutput | dict | torch.Tensor:
+        """
+        Forward pass with customizable output.
+
+        Returns different outputs based on the configured output mode or custom processor.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        # Get outputs from the base model (Qwen3VLModel)
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+
+        # Get the last hidden state
+        hidden_states = outputs[0]  # This is the last hidden state
+
+        # Compute logits if not skipping lm_head
+        logits = None
+        if not self._skip_lm_head:
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        # Compute loss if labels are provided
+        loss = None
+        if labels is not None and logits is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        # Return based on output mode
+        if self._output_mode == self.OUTPUT_MODE_EMBEDDING:
+            return Qwen3VLModelOutput(
+                last_hidden_state=hidden_states,
+                past_key_values=outputs.past_key_values,
+                attentions=outputs.attentions,
+                rope_deltas=outputs.rope_deltas,
+            )
+        elif self._output_mode == self.OUTPUT_MODE_LOGITS:
+            return logits
+        elif self._output_mode == self.OUTPUT_MODE_HIDDEN:
+            return Qwen3VLModelOutput(
+                last_hidden_state=hidden_states,
+                hidden_states=outputs.hidden_states,
+                past_key_values=outputs.past_key_values,
+                attentions=outputs.attentions,
+                rope_deltas=outputs.rope_deltas,
+            )
+        else:  # OUTPUT_MODE_FULL
+            return Qwen3VLCausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+                rope_deltas=outputs.rope_deltas,
+            )
+
+
+# ===========================================================================
+# Packing-aware forward patches (cu_seqlens) for the Qwen3-VL text encoder
+# ===========================================================================
+
+def model_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    position_ids: torch.LongTensor | None = None,
+    past_key_values: Cache | None = None,
+    inputs_embeds: torch.FloatTensor | None = None,
+    use_cache: bool | None = None,
+    cache_position: torch.LongTensor | None = None,
+    # args for deepstack
+    visual_pos_masks: torch.Tensor | None = None,
+    deepstack_visual_embeds: list[torch.Tensor] | None = None,
+    **kwargs: Unpack[FlashAttentionKwargs],
+) -> tuple | BaseModelOutputWithPast:
+    r"""
+    visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+        The mask of the visual positions.
+    deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+        The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+        The feature is extracted from the different visual encoder layers, and fed to the decoder
+        hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+    """
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    # torch.jit.trace() doesn't support cache objects in the output
+    if use_cache and past_key_values is None and not torch.jit.is_tracing():
+        past_key_values = DynamicCache(config=self.config)
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+
+    if cache_position is None:
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        )
+
+    # the hard coded `3` is for temporal, height and width.
+    if position_ids is None:
+        position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+    elif position_ids.ndim == 2:
+        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+        text_position_ids = position_ids[0]
+        position_ids = position_ids[1:]
+    else:
+        text_position_ids = position_ids[0]
+
+    if kwargs.get("cu_seqlens") is None:
+        attention_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+
+    hidden_states = inputs_embeds
+
+    # create position embeddings to be shared across the decoder layers
+    position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+    # decoder layers
+    for layer_idx, decoder_layer in enumerate(self.layers):
+        layer_outputs = decoder_layer(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_ids=text_position_ids,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = layer_outputs
+
+        # add visual features to the hidden states of first several layers
+        if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+            hidden_states = self._deepstack_process(
+                hidden_states,
+                visual_pos_masks,
+                deepstack_visual_embeds[layer_idx],
+            )
+
+    hidden_states = self.norm(hidden_states)
+
+    return BaseModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=past_key_values,
+    )
+
+
+def forward(
+    self,
+    hidden_states: torch.Tensor,
+    position_embeddings: tuple[torch.Tensor, torch.Tensor],
+    attention_mask: torch.Tensor | None,
+    past_key_values: Cache | None = None,
+    cache_position: torch.LongTensor | None = None,
+    **kwargs: Unpack[FlashAttentionKwargs],
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    input_shape = hidden_states.shape[:-1]
+    hidden_shape = (*input_shape, -1, self.head_dim)
+
+    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+    cos, sin = position_embeddings
+    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+    if past_key_values is not None:
+        # sin and cos are specific to RoPE models; cache_position needed for the static cache
+        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+        key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+    cu_seqlens = kwargs.get("cu_seqlens", None)
+
+    if cu_seqlens is None:
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+    else:
+        max_seqlen = torch.diff(cu_seqlens).max().item() if cu_seqlens is not None else None
+        query_states = query_states.transpose(1, 2).squeeze(0)
+        key_states = key_states.transpose(1, 2).squeeze(0)
+        value_states = value_states.transpose(1, 2).squeeze(0)
+        attn_output = flash_attn_varlen_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            causal=True,
+            window_size=(-1, -1),
+            softmax_scale=self.head_dim**-0.5,
+            dropout_p=0.0,
+        )
+
+    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+    attn_output = self.o_proj(attn_output)
+    return attn_output, None
+
+
+def qwen3_patch_forward():
+    """Patch the Qwen3-VL text model + attention forwards to support packed
+    varlen (cu_seqlens) inputs used by ``TextEncoder.forward``."""
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention, Qwen3VLTextModel
+
+    Qwen3VLTextModel.forward = model_forward
+    Qwen3VLTextAttention.forward = forward
+
+
+# ===========================================================================
+# TextEncoder wrapper (packed text -> DiT conditioning embeddings)
+# ===========================================================================
+_FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"}
+_FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"}
+_SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"}
+
+
+def _resolve_hf_attn_impl(attn_type: str) -> str:
+    """Map a project-level attn_type to a HuggingFace ``attn_implementation`` string.
+
+    ``VF_HF_ATTN_IMPL`` env var, if set, takes precedence (useful for forcing
+    sdpa on machines without flash-attn). For FA4 we additionally probe that
+    the CUTE-DSL kernel is importable and (when available) ask the HF helper
+    to confirm; if not, fall back to sdpa rather than crashing at load time.
+    """
+    override = os.environ.get("VF_HF_ATTN_IMPL")
+    if override:
+        return override
+
+    name = attn_type.lower().strip()
+    if name in _FA2_ALIASES:
+        return "flash_attention_2"
+    if name in _FA4_ALIASES:
+        try:
+            import flash_attn.cute  # noqa: F401
+            fa4_importable = True
+        except Exception:
+            fa4_importable = False
+        if fa4_importable:
+            try:
+                from transformers.utils.import_utils import is_flash_attn_4_available
+                if is_flash_attn_4_available():
+                    return "flash_attention_4"
+            except ImportError:
+                return "flash_attention_4"
+        logger.warning(
+            "attn_type=flash4 requested but flash_attn.cute is unavailable; "
+            "falling back to sdpa for HF text encoder."
+        )
+        return "sdpa"
+    if name in _SDPA_ALIASES:
+        return "sdpa"
+    raise ValueError(
+        f"Unknown attn_type {attn_type!r}; expected one of "
+        f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}"
+    )
+
+
+SEQ_MULTI_OF = 32
+
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        model_name: str,
+        version: str,
+        tokenizer_max_length: int,
+        prompt_template: dict | None,
+        dit_structure: dict,
+        use_packed_text_infer: bool = False,
+        attn_type: str = "flash2",
+        **hf_kwargs,
+    ):
+        super().__init__()
+        self.model_name = model_name
+        self.tokenizer_max_length = tokenizer_max_length
+        self.tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(version)
+        self.tokenizer.padding_side = "right"
+
+        hf_attn_impl = _resolve_hf_attn_impl(attn_type)
+        logger.info(f"TextEncoder attn_type={attn_type} -> attn_implementation={hf_attn_impl}")
+
+        logger.info("init vl model: qwen3")
+        self.hf_module: CustomQwen3VLForConditionalGeneration = CustomQwen3VLForConditionalGeneration.from_pretrained(
+            version,
+            attn_implementation=hf_attn_impl,
+            **hf_kwargs
+        )
+
+        # Use local_files_only if version is a local path (absolute path or contains path separators)
+        is_local = version.startswith("/") or os.sep in version
+        self.processor = AutoProcessor.from_pretrained(version, local_files_only=is_local)
+
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+        prompt_template = prompt_template or {}
+        self.prompt_template_encode = prompt_template.get("template", "")
+        self.prompt_template_encode_start_idx = prompt_template.get("start_idx", 0)
+        self.dit_structure = dit_structure
+        self.use_packed_text_infer = use_packed_text_infer
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        inputs: dict | None = None,
+        drop_idx_override: int | None = None,
+    ):
+        """Encode packed text (varlen ``cu_seqlens``) into DiT conditioning embeddings.
+
+        This is the sole text-embedding path — both t2i and edit call it. Uses
+        Flash-Attention-2's varlen capability (``cu_seqlens``) via the patched
+        Qwen3-VL forward to process several concatenated sequences in a single
+        launch, with no padding. Verified numerically identical to a padded-batch
+        forward with per-sample cu_seqlens isolation (zero cross-contamination).
+
+        Args:
+            input_ids: Packed token ids ``[Total_L]``.
+            cu_seqlens: Cumulative sequence lengths ``[B+1]``.
+            inputs: Optional dict with additional model inputs (e.g. ``pixel_values``,
+                ``image_grid_thw`` for the multimodal edit path). Passed through to
+                the text encoder.
+            drop_idx_override: If set, override the number of leading (system-prompt)
+                tokens to drop per sequence. Use 0 for multi-turn where the system
+                prompt is embedded in the conversation and should not be stripped.
+
+        Returns:
+            dict with keys:
+                - ``txt``: text embeddings ``[Total_L - B*drop_idx, D]`` (system prompt dropped)
+                - ``vec``: pooled text embeddings ``[B, D]``
+                - ``txt_seq_lens``: per-sequence lengths ``[B]`` (after dropping system prompt)
+        """
+        # Compute seqlens from cu_seqlens
+        seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+        seqlens_list = seqlens.cpu().tolist()
+
+        # Build position_ids for packing: each sequence starts from 0
+        position_ids_list = []
+        for length in seqlens_list:
+            position_ids_list.append(torch.arange(length, device=input_ids.device))
+        position_ids = torch.cat(position_ids_list)  # [Total_L]
+
+        # Reshape for model input: [1, Total_L]
+        input_ids_packed = input_ids.unsqueeze(0)  # [1, Total_L]
+        position_ids_packed = position_ids.unsqueeze(0)  # [1, Total_L]
+
+        # Move to text encoder device
+        device = self.hf_module.device
+        input_ids_packed = input_ids_packed.to(device)
+        position_ids_packed = position_ids_packed.to(device)
+
+        # Get text embeddings (the text encoder is always frozen)
+        with torch.no_grad():
+            forward_kwargs = {
+                "input_ids": input_ids_packed,
+                "cu_seqlens": cu_seqlens,
+                "position_ids": position_ids_packed,
+                "output_hidden_states": False,
+                "max_seqlen": None,
+            }
+            # Pass multimodal inputs for edit mode (reference images)
+            if inputs is not None:
+                for key in ("pixel_values", "image_grid_thw"):
+                    if key in inputs and inputs[key] is not None:
+                        val = inputs[key]
+                        if hasattr(val, "to"):
+                            val = val.to(device)
+                        forward_kwargs[key] = val
+            outputs = self.hf_module(**forward_kwargs)
+
+        # Extract hidden state
+        if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None:
+            hidden = outputs.last_hidden_state  # [1, Total_L, D]
+        elif hasattr(outputs, "hidden_states"):
+            hidden = outputs.hidden_states[-1]
+
+        # Remove batch dimension: [Total_L, D]
+        hidden = hidden.squeeze(0)
+
+        # Get drop_idx (system prompt length to skip).
+        # For multi-turn, drop_idx_override=0 is passed since system prompt is in the messages.
+        if drop_idx_override is not None:
+            drop_idx = drop_idx_override
+        else:
+            drop_idx = self.prompt_template_encode_start_idx
+
+        # Split hidden states by sequence
+        hidden_split = torch.split(hidden, seqlens_list, dim=0)
+
+        # Extract valid embeddings (drop system prompt) and compute vec
+        txt_list = []
+        vec_list = []
+        valid_lengths = []
+
+        for h in hidden_split:
+            # Drop system prompt tokens
+            h_valid = h[drop_idx:]  # [seq_len - drop_idx, D]
+            txt_list.append(h_valid)
+            valid_lengths.append(h_valid.shape[0])
+
+            # Compute pooled embedding (mean of valid tokens only, after dropping system prompt)
+            vec_list.append(h_valid.mean(dim=0))  # [D]
+
+        txt = torch.cat(txt_list, dim=0)  # [Total_valid, D]
+        vec = torch.stack(vec_list, dim=0)  # [B, D]
+        txt_seq_lens = torch.tensor(valid_lengths, device=input_ids.device)
+
+        result = {
+            "txt": txt,
+            "vec": vec,
+            "txt_seq_lens": txt_seq_lens,
+        }
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Mandatory content-policy screening (same Qwen3-VL weights)
+    # ------------------------------------------------------------------
+    # The policy classifier lives HERE, on the text encoder, so it runs on the
+    # exact weights that produce the diffusion conditioning and is not a
+    # separable, toggleable pre-pass in the pipeline. The classifier needs
+    # autoregressive ``.generate()`` (JSON verdict) whereas conditioning is a
+    # single embedding forward — they cannot be one GPU forward without a
+    # trained classification head, so "fused" here means: same module, same
+    # weights, always run, FAIL-CLOSED (any error blocks).
+
+    def screen_text(self, prompt: str, max_new_tokens: int = 160):
+        """Classify a text-to-image ``prompt`` against the content policy.
+
+        Returns a ``FilterVerdict``. FAIL-CLOSED: any error (generation, parse)
+        returns ``violates=True`` so a broken classifier cannot be used as a
+        bypass. An empty prompt is not a violation.
+        """
+        from .mage_text import (
+            CONTENT_FILTER_SYSTEM, FilterVerdict, _extract_json_object,
+            _full_output_mode,
+        )
+
+        if not prompt or not prompt.strip():
+            return FilterVerdict(False, [], "empty prompt", "")
+        try:
+            tokenizer = self.tokenizer
+            hf = self.hf_module
+            device = next(hf.parameters()).device
+
+            messages = [
+                {"role": "system", "content": CONTENT_FILTER_SYSTEM},
+                {"role": "user", "content": f"Prompt to classify:\n{prompt}"},
+            ]
+            text = tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True)
+            inputs = tokenizer(text, return_tensors="pt").to(device)
+
+            eos_id = tokenizer.eos_token_id
+            pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else eos_id
+
+            with _full_output_mode(hf), torch.no_grad():
+                out = hf.generate(
+                    **inputs, max_new_tokens=max_new_tokens, do_sample=False,
+                    pad_token_id=pad_id, eos_token_id=eos_id)
+            gen = tokenizer.decode(
+                out[0, inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
+
+            parsed = _extract_json_object(gen)
+            violates = bool(parsed.get("violates", False))
+            cats = [c for c in (parsed.get("categories", []) or []) if isinstance(c, str)]
+            reason = str(parsed.get("reason", "")).strip()
+            return FilterVerdict(violates, cats, reason, gen)
+        except Exception as exc:  # noqa: BLE001
+            # FAIL-CLOSED: block on any screening error.
+            return FilterVerdict(
+                True, ["policy"], f"filter error (blocked): {type(exc).__name__}: {exc}", "")
+
+    def screen_edit(self, prompt: str, ref_images, max_new_tokens: int = 192):
+        """Classify an image-EDIT request (source image(s) + instruction).
+
+        Considers BOTH the source image(s) and the instruction via multimodal
+        Qwen3-VL. Falls back to :meth:`screen_text` when no image is given.
+        FAIL-CLOSED: any error returns ``violates=True``.
+        """
+        from PIL import Image
+
+        from .mage_text import (
+            CONTENT_FILTER_EDIT_SYSTEM, FilterVerdict, _extract_json_object,
+            _full_output_mode,
+        )
+
+        pils = [ref_images] if isinstance(ref_images, Image.Image) else list(ref_images)
+        pils = [p.convert("RGB") for p in pils if p is not None]
+        if not pils:
+            return self.screen_text(prompt, max_new_tokens=max_new_tokens)
+
+        instruction = (prompt or "").strip() or "(no textual instruction)"
+        try:
+            processor = self.processor
+            tokenizer = self.tokenizer
+            hf = self.hf_module
+            device = next(hf.parameters()).device
+
+            user_content = [{"type": "image"} for _ in pils]
+            user_content.append({
+                "type": "text",
+                "text": (
+                    f"There {'is' if len(pils) == 1 else 'are'} {len(pils)} source "
+                    f"image(s) above. Edit instruction: {instruction}\n"
+                    "Classify this edit request."
+                ),
+            })
+            messages = [
+                {"role": "system", "content": CONTENT_FILTER_EDIT_SYSTEM},
+                {"role": "user", "content": user_content},
+            ]
+            text = processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True)
+            inputs = processor(
+                text=[text], images=pils, padding=True, return_tensors="pt").to(device)
+
+            eos_id = tokenizer.eos_token_id
+            pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else eos_id
+
+            # Keep only the kwargs Qwen3-VL .generate() consumes.
+            gen_inputs = {
+                k: inputs[k]
+                for k in ("input_ids", "attention_mask", "pixel_values", "image_grid_thw")
+                if k in inputs and inputs[k] is not None
+            }
+            input_len = gen_inputs["input_ids"].shape[1]
+
+            with _full_output_mode(hf), torch.no_grad():
+                out = hf.generate(
+                    **gen_inputs, max_new_tokens=max_new_tokens, do_sample=False,
+                    pad_token_id=pad_id, eos_token_id=eos_id)
+            gen = tokenizer.decode(out[0, input_len:], skip_special_tokens=True).strip()
+
+            parsed = _extract_json_object(gen)
+            violates = bool(parsed.get("violates", False))
+            cats = [c for c in (parsed.get("categories", []) or []) if isinstance(c, str)]
+            reason = str(parsed.get("reason", "")).strip()
+            return FilterVerdict(violates, cats, reason, gen)
+        except Exception as exc:  # noqa: BLE001
+            # FAIL-CLOSED: block on any screening error.
+            return FilterVerdict(
+                True, ["policy"], f"edit filter error (blocked): {type(exc).__name__}: {exc}", "")
+

mage_flow/models/utils.py

@@ -0,0 +1,175 @@
+import json
+import math
+import os
+
+import torch
+from einops import rearrange
+from loguru import logger
+from safetensors.torch import load_file
+from safetensors.torch import load_file as load_sft
+from torch import Tensor
+
+from .mage_flow import MageFlow, MageFlowParams
+
+
+def get_noise(
+    num_samples: int,
+    channel: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    # MageVAE: 16x downsample, no patch packing
+    return torch.randn(
+        num_samples,
+        channel,
+        math.ceil(height / 16),
+        math.ceil(width / 16),
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device=device).manual_seed(seed),
+    )
+
+
+def unpack(x: Tensor, height: int, width: int) -> Tensor:
+    # MageVAE: [B, H*W, C] -> [B, C, H, W], no patch unpacking
+    return rearrange(
+        x,
+        "b (h w) c -> b c h w",
+        h=math.ceil(height / 16),
+        w=math.ceil(width / 16),
+    )
+
+
+PROMPT_TEMPLATE = {
+    "default": {"template": "{}", "start_idx": 0},
+    "default-nonthinking": {"template": "{}<think>\n\n</think>\n\n", "start_idx": 0},
+    "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,
+    },
+}
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if len(missing) > 0 and len(unexpected) > 0:
+        logger.warning(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        logger.warning("\n" + "-" * 79 + "\n")
+        logger.warning(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        logger.warning(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        logger.warning(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def correct_model_weight(state_dict):
+    result = {}
+    for key in state_dict.keys():
+        if "_orig_mod." in key:
+            result[key[10:]] = state_dict[key]
+        else:
+            result[key] = state_dict[key]
+    return result
+
+
+def load_hf_style_weight(pretrain_path, device):
+    index_path = os.path.join(pretrain_path, "diffusion_pytorch_model.safetensors.index.json")
+
+    with open(index_path) as f:
+        index = json.load(f)
+
+    weight_map = index["weight_map"]
+
+    sd = {}
+    loaded_shards = set()
+
+    for shard_file in weight_map.values():
+        if shard_file in loaded_shards:
+            continue
+        shard_path = os.path.join(pretrain_path, shard_file)
+        shard_sd = load_file(shard_path, device="cpu")
+        sd.update(shard_sd)
+        loaded_shards.add(shard_file)
+
+    return sd
+
+
+def load_model_weight(model, pretrain_path, device="cpu"):
+    if os.path.exists(pretrain_path):
+        logger.info(f"Loading checkpoint from {pretrain_path}")
+        try:
+            if pretrain_path.endswith("safetensors"):
+                sd = load_sft(pretrain_path, device="cpu")
+            elif os.path.exists(os.path.join(pretrain_path, "diffusion_pytorch_model.safetensors.index.json")):
+                sd = load_hf_style_weight(pretrain_path, device)
+            else:
+                sd = torch.load(pretrain_path, map_location="cpu")
+
+            sd = correct_model_weight(sd)
+            sd = optionally_expand_state_dict(model, sd)
+            missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+            print_load_warning(missing, unexpected)
+            return True
+        except Exception as e:
+            logger.info(f"CANNOT Load {pretrain_path}, because {e}")
+            return False
+    return False
+
+
+def load_model(dit_structure: dict, pretrain_path: str | None = None):
+    logger.info("Init DiT model")
+
+    # If name is a dict, we assume it contains the parameters directly
+    # We need to determine the model class based on some heuristic or just default to MageFlow/Flux
+    # For now, let's assume it's MageFlow if time_type is present, or check other fields
+    params = MageFlowParams(**dit_structure)
+    # Default to MageFlow for now as per user context, or we could add a 'model_type' field to the dict
+    # The user mentioned "model structure option", implying we are configuring the structure.
+    # Let's assume MageFlow for this refactor as the user was using qwen-image-tiny-wo-textemb
+    model = MageFlow(params)
+
+    # logger.info(f"Loading {name if isinstance(name, str) else 'custom config'} checkpoint from {pretrain_path}")
+    if pretrain_path is not None:
+        load_model_weight(model, pretrain_path, device="cpu")
+    # if isinstance(name, str) and configs[name].lora_path is not None:
+    #     logger.info("Loading LoRA")
+    #     lora_sd = load_sft(configs[name].lora_path, device="cpu")
+    #     # loading the lora params + overwriting scale values in the norms
+    #     missing, unexpected = model.load_state_dict(lora_sd, strict=False, assign=True)
+    #     print_load_warning(missing, unexpected)
+    return model
+
+
+def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> dict:
+    """
+    Optionally expand the state dict to match the model's parameters shapes.
+    """
+    for name, param in model.named_parameters():
+        if name in state_dict:
+            if state_dict[name].shape != param.shape:
+                logger.info(
+                    f"Expanding '{name}' with shape {state_dict[name].shape} to model parameter with shape "
+                    f"{param.shape}."
+                )
+                # expand with zeros:
+                expanded_state_dict_weight = torch.zeros_like(param, device=state_dict[name].device)
+                slices = tuple(slice(0, dim) for dim in state_dict[name].shape)
+                expanded_state_dict_weight[slices] = state_dict[name]
+                state_dict[name] = expanded_state_dict_weight
+
+    return state_dict

mage_flow/pipeline.py

@@ -0,0 +1,762 @@
+"""MageFlow text-to-image + image-edit inference pipeline.
+
+Self-contained MageFlow t2i / edit inference: load a HuggingFace diffusers-style
+repo (model_index.json + transformer/ vae/ scheduler/), then generate or edit
+images. No training/eval deps.
+
+Both ``generate_images`` and ``generate_edits`` support PACKED multi-resolution
+inference: several samples (each at its own resolution) are concatenated into a
+single varlen sequence and processed in one transformer forward per denoise
+step. Per-sample ``cu_seqlens`` (inside the flash-attn varlen kernel) isolate
+samples, exactly mirroring training-time packing. These packed functions are the
+sole implementation — the single-image case is just a pack of size 1, exposed
+via the ``MageFlowPipeline.generate`` / ``.edit`` convenience methods.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import random
+
+import torch
+from einops import rearrange
+from PIL import Image
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+from .models.mage_flow import MageFlowModel, ModelConfig
+from .models.utils import PROMPT_TEMPLATE, get_noise, unpack
+from .models.modules.mage_text import make_refusal_image
+from .models.modules.mage_latent import encode_noise, resolve_gs_key
+
+
+# ---------------------------------------------------------------------------
+# Scheduler — diffusers FlowMatchEulerDiscreteScheduler
+# ---------------------------------------------------------------------------
+def build_scheduler(num_steps: int, device=None, shift: float = 6.0):
+    """Construct a diffusers ``FlowMatchEulerDiscreteScheduler`` whose sigma
+    schedule reproduces our default preset exactly.
+
+    The base sigmas ``linspace(1, 1/num_steps, num_steps)`` fed to
+    ``set_timesteps`` are run through the scheduler's built-in static shift
+    ``shift·s/(1+(shift-1)·s)`` and a terminal 0 is appended — the static-shift 
+    schedule (the only supported schedule).
+    """
+    scheduler = FlowMatchEulerDiscreteScheduler(
+        num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
+    base_sigmas = torch.linspace(1.0, 1.0 / num_steps, num_steps).tolist()
+    scheduler.set_timesteps(sigmas=base_sigmas, device=device)
+    return scheduler
+
+
+def _get_scheduler(model, steps, device, static_shift):
+    scheduler = getattr(model, "scheduler", None)
+    if scheduler is None:
+        return build_scheduler(steps, device=device,
+                               shift=(static_shift if static_shift is not None else 6.0))
+    if static_shift is not None:
+        scheduler.set_shift(static_shift)
+    scheduler.set_timesteps(sigmas=torch.linspace(1.0, 1.0 / steps, steps).tolist(), device=device)
+    return scheduler
+
+
+# ---------------------------------------------------------------------------
+# Small helpers
+# ---------------------------------------------------------------------------
+def _template_info(name: str | None) -> dict:
+    name = name or "mage-flow"
+    if name not in PROMPT_TEMPLATE:
+        raise ValueError(f"Unknown prompt template: {name}")
+    return PROMPT_TEMPLATE[name]
+
+
+def _as_list(val, default, n):
+    """Broadcast a scalar/None to a length-n list, or validate a given list."""
+    if val is None:
+        return [default] * n
+    if isinstance(val, (list, tuple)):
+        if len(val) != n:
+            raise ValueError(f"expected {n} values, got {len(val)}")
+        return list(val)
+    return [val] * n
+
+
+def _lens_to_cu(lens, device):
+    """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 _make_divisible_by_16(size: int) -> int:
+    return max(16, 16 * (size // 16))
+
+
+def _compute_aspect_ratio_size(pil_img: Image.Image, max_size: int):
+    """Longest side = ``max_size``, short side from aspect ratio, both /16."""
+    w, h = pil_img.size
+    if h >= w:
+        new_h, new_w = max_size, int(round(w * max_size / h))
+    else:
+        new_w, new_h = max_size, int(round(h * max_size / w))
+    return _make_divisible_by_16(new_h), _make_divisible_by_16(new_w)
+
+
+def _edit_target_size(pil_img: Image.Image, max_size, height, width):
+    """Output (H, W) for an edit sample, derived from its PRIMARY reference.
+
+    Precedence: explicit ``height`` AND ``width`` (custom size) > ``max_size``
+    (longest side, short side by aspect ratio) > the source image's own size.
+    All rounded down to a multiple of 16.
+    """
+    if height and width:
+        return _make_divisible_by_16(height), _make_divisible_by_16(width)
+    if max_size:
+        return _compute_aspect_ratio_size(pil_img, max_size)
+    # Nothing specified: keep the source resolution (its own longest side).
+    return _compute_aspect_ratio_size(pil_img, max(pil_img.size))
+
+
+def _decode_one(model, tokens, height, width, dev):
+    """Unpack one sample's image tokens [1, H*W, C] and VAE-decode to a PIL image."""
+    with torch.autocast(device_type=dev.type, dtype=torch.bfloat16):
+        out = model.vae.decode(unpack(tokens.float(), height, width))
+    out = rearrange(out.clamp(-1, 1), "b c h w -> b h w c")
+    out = (127.5 * (out + 1.0)).cpu().byte().numpy()
+    return Image.fromarray(out[0])
+
+
+def _build_pack_ctx(img_ids, img_cu, img_shapes, img_lens, txt, txt_cu, txt_mask, vec,
+                    neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, device):
+    """Precompute the static per-step transformer inputs for a packed batch.
+
+    When a negative branch is present and ``batch_cfg`` is True, the conditional
+    and unconditional passes are fused into ONE varlen forward: the image tokens
+    are duplicated (cond copy + uncond copy) and the positive/negative texts are
+    concatenated, so cond sample i and uncond sample i become two independent
+    varlen segments processed in a single kernel launch. flash_attn_varlen_func
+    keeps every segment isolated via cu_seqlens, so this is numerically identical
+    to two separate forwards — just one launch instead of two.
+    """
+    na = len(img_lens)
+    ctx = {
+        "na": na, "cfg": cfg, "renorm": renormalization, "batch_cfg": batch_cfg,
+        "has_neg": neg_txt is not None,
+        "img_ids": img_ids, "img_cu": img_cu, "img_shapes": img_shapes,
+        "img_max": int(max(img_lens)),
+        "txt": txt, "txt_ids": torch.zeros(1, txt.shape[1], 3, device=device),
+        "txt_cu": txt_cu, "txt_mask": txt_mask, "vec": vec,
+        "txt_max": int((txt_cu[1:] - txt_cu[:-1]).max().item()),
+    }
+    if neg_txt is None:
+        return ctx
+    ctx.update({
+        "neg_txt": neg_txt, "neg_ids": torch.zeros(1, neg_txt.shape[1], 3, device=device),
+        "neg_cu": neg_cu, "neg_mask": neg_mask, "neg_vec": neg_vec,
+        "neg_max": int((neg_cu[1:] - neg_cu[:-1]).max().item()),
+    })
+    if batch_cfg:
+        # Duplicate image segments (cond then uncond) and concat pos+neg text.
+        d_txt = torch.cat([txt, neg_txt], dim=1)
+        pos_lens = (txt_cu[1:] - txt_cu[:-1]).tolist()
+        neg_lens = (neg_cu[1:] - neg_cu[:-1]).tolist()
+        ctx.update({
+            "d_img_ids": torch.cat([img_ids, img_ids], dim=1),
+            "d_img_cu": _lens_to_cu(list(img_lens) + list(img_lens), device),
+            "d_img_shapes": [img_shapes[0] + img_shapes[0]],
+            "d_txt": d_txt,
+            "d_txt_ids": torch.zeros(1, d_txt.shape[1], 3, device=device),
+            "d_txt_cu": _lens_to_cu(pos_lens + neg_lens, device),
+            "d_txt_mask": torch.ones(1, d_txt.shape[1], device=device),
+            "d_vec": torch.cat([vec, neg_vec], dim=0),
+            "d_txt_max": int(max(pos_lens + neg_lens)),
+        })
+    return ctx
+
+
+def _velocity(transformer, img, ctx, sigma):
+    """CFG-combined image-token velocity for a packed batch at noise level ``sigma``.
+
+    Returns [1, sum_img_len, C] in the conditional sample order. When
+    ``batch_cfg`` is set the cond+uncond passes share a single fused varlen
+    forward; otherwise they are two forwards.
+    """
+    dev = img.device
+    na = ctx["na"]
+
+    def _fwd(x, n, img_ids, img_cu, img_max, img_shapes, txt, txt_ids, txt_cu, txt_mask, txt_max, vec):
+        t_vec = torch.full((n,), sigma, dtype=x.dtype, device=dev)
+        return transformer(img=x, txt=txt, timesteps=t_vec, img_shapes=img_shapes,
+                           img_cu_seqlens=img_cu, txt_cu_seqlens=txt_cu)
+
+    if not ctx["has_neg"]:
+        return _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
+                    ctx["txt"], ctx["txt_ids"], ctx["txt_cu"], ctx["txt_mask"], ctx["txt_max"], ctx["vec"])
+
+    if ctx["batch_cfg"]:
+        n_img = img.shape[1]
+        out = _fwd(torch.cat([img, img], dim=1), 2 * na,
+                   ctx["d_img_ids"], ctx["d_img_cu"], ctx["img_max"], ctx["d_img_shapes"],
+                   ctx["d_txt"], ctx["d_txt_ids"], ctx["d_txt_cu"], ctx["d_txt_mask"], ctx["d_txt_max"], ctx["d_vec"])
+        cond, unc = out[:, :n_img, :], out[:, n_img:, :]
+    else:
+        cond = _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
+                    ctx["txt"], ctx["txt_ids"], ctx["txt_cu"], ctx["txt_mask"], ctx["txt_max"], ctx["vec"])
+        unc = _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
+                   ctx["neg_txt"], ctx["neg_ids"], ctx["neg_cu"], ctx["neg_mask"], ctx["neg_max"], ctx["neg_vec"])
+
+    cfg = ctx["cfg"]
+    if ctx["renorm"]:
+        # CFG renormalization: rescale the guided velocity per token back to the
+        # conditional velocity's norm (reduces oversaturation at high cfg).
+        comb = unc + cfg * (cond - unc)
+        return comb * (torch.norm(cond, dim=-1, keepdim=True) /
+                       (torch.norm(comb, dim=-1, keepdim=True) + 1e-6))
+    return unc + cfg * (cond - unc)
+
+
+def _encode_texts_packed(model, prompts, template, drop_idx, device):
+    """Encode a LIST of templated text-only prompts in ONE packed varlen forward
+    (``TextEncoder.forward`` — varlen cu_seqlens isolates each prompt,
+    verified zero cross-contamination). Returns (txt_flat [ΣLi, D], vec [N, D],
+    per-prompt token lengths list)."""
+    tokenizer = model.txt_enc.tokenizer
+    max_len = model.txt_enc.tokenizer_max_length + drop_idx
+    ids_list = [
+        tokenizer(template.format(p), max_length=max_len, truncation=True,
+                  return_tensors="pt").input_ids.squeeze(0)
+        for p in prompts
+    ]
+    input_ids = torch.cat(ids_list).to(device)
+    cu_seqlens = _lens_to_cu([int(t.numel()) for t in ids_list], device)
+    res = model.txt_enc(
+        input_ids, cu_seqlens, drop_idx_override=drop_idx)
+    return res["txt"], res["vec"], res["txt_seq_lens"].tolist()
+
+
+def _slice_packed(txt_flat, vec, lens, start, count, device):
+    """Format a contiguous ``count``-prompt slice (starting at prompt ``start``) of a
+    packed text encode into the (txt [1, ΣL, D], cu_seqlens, ones-mask, vec [count, D])
+    tuple that ``_build_pack_ctx`` consumes."""
+    seg_lens = lens[start:start + count]
+    tok_start = sum(lens[:start])
+    tok_end = tok_start + sum(seg_lens)
+    txt = txt_flat[tok_start:tok_end].reshape(1, -1, txt_flat.shape[-1]).to(device)
+    return (txt, _lens_to_cu(seg_lens, device),
+            torch.ones(1, txt.shape[1], device=device), vec[start:start + count].to(device))
+
+
+# ---------------------------------------------------------------------------
+# Text-to-image (packed, multi-resolution)
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def generate_images(model, prompts, neg_prompts=None, seeds=None, steps=30, cfg=5.0,
+                    heights=None, widths=None, device="cuda",
+                    prompt_template="mage-flow", static_shift=None,
+                    gs_key=None,
+                    renormalization=False, batch_cfg=True):
+    """Generate one image per prompt. Prompts may request DIFFERENT resolutions;
+    all are packed into a single varlen forward per denoise step — samples are
+    kept isolated by ``flash_attn_varlen_func`` via per-sample ``cu_seqlens`` (no
+    cross-sample attention), mirroring training-time packing. When ``cfg > 1`` and
+    ``batch_cfg`` is set, the positive and negative passes are fused into that
+    same varlen forward. Returns a list of PIL images aligned with ``prompts``.
+    """
+    if isinstance(prompts, str):
+        prompts = [prompts]
+    n = len(prompts)
+    neg_prompts = _as_list(neg_prompts, " ", n)
+    seeds = _as_list(seeds, 42, n)
+    heights = _as_list(heights, 1024, n)
+    widths = _as_list(widths, 1024, n)
+    info = _template_info(prompt_template)
+    template = info.get("template", "{}")
+    drop_idx = int(info.get("start_idx", 0))
+    dev = torch.device(device)
+
+    # Content-policy gate per sample (MANDATORY — runs on the same text-encoder
+    # weights as conditioning, no opt-out). Violating prompts get a refusal
+    # placeholder and are dropped from the pack.
+    results = [None] * n
+    active = []
+    for i in range(n):
+        if seeds[i] == -1:
+            seeds[i] = random.randint(0, 2**32 - 1)
+        verdict = model.txt_enc.screen_text(prompts[i])
+        if verdict.violates:
+            h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
+            print(verdict.banner())
+            results[i] = make_refusal_image(verdict, height=h_, width=w_)
+            continue
+        active.append(i)
+    if not active:
+        return results
+
+    gs_key_int = resolve_gs_key(gs_key)
+    # Per-sample noise tokens + position ids + shapes (MageVAE: flatten, no packing).
+    ch = model.vae.latent_channels
+    img_list, ids_list, lens, shapes, hw = [], [], [], [], []
+    for i in active:
+        h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
+        torch.manual_seed(seeds[i])
+        x = get_noise(num_samples=1, channel=ch, height=h_, width=w_,
+                      device=dev, dtype=torch.bfloat16, seed=seeds[i])
+        # Distribution-preserving watermark in the initial noise (same shape,
+        # still ~N(0,1)); detect by inverting the flow ODE back to noise.
+        x = encode_noise(tuple(x.shape[1:]), key=gs_key_int,
+                         seed=seeds[i], device=dev, dtype=torch.bfloat16)
+        _, _, gh, gw = x.shape
+        img_list.append(rearrange(x, "b c h w -> b (h w) c")[0])
+        ids = torch.zeros(gh, gw, 3, device=dev)
+        ids[..., 1] = ids[..., 1] + torch.arange(gh, device=dev)[:, None]
+        ids[..., 2] = ids[..., 2] + torch.arange(gw, device=dev)[None, :]
+        ids_list.append(rearrange(ids, "h w c -> (h w) c"))
+        lens.append(gh * gw); shapes.append((1, gh, gw)); hw.append((h_, w_))
+    img = torch.cat(img_list, 0).unsqueeze(0)
+    img_ids = torch.cat(ids_list, 0).unsqueeze(0)
+    img_cu = _lens_to_cu(lens, dev)
+    img_shapes = [shapes]
+
+    # Packed text: positive prompts AND (for CFG) negative prompts are encoded
+    # TOGETHER in ONE varlen forward, then split back — cu_seqlens keeps every
+    # prompt isolated (verified zero cross-contamination).
+    pos_prompts = [prompts[i] for i in active]
+    na = len(active)
+    use_neg = cfg > 1.0 and any(neg_prompts[i] for i in active)
+    if use_neg:
+        neg_list = [neg_prompts[i] or " " for i in active]
+        txt_flat, vec_all, lens_t = _encode_texts_packed(
+            model, pos_prompts + neg_list, template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt, neg_cu, neg_mask, neg_vec = _slice_packed(txt_flat, vec_all, lens_t, na, na, dev)
+    else:
+        txt_flat, vec_all, lens_t = _encode_texts_packed(model, pos_prompts, template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt = neg_cu = neg_mask = neg_vec = None
+
+    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, lens, txt, txt_cu, txt_mask, vec,
+                          neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, dev)
+    scheduler = _get_scheduler(model, steps, device, static_shift)
+    for si, t in enumerate(scheduler.timesteps):
+        pred = _velocity(model.transformer, img, ctx, scheduler.sigmas[si].item())
+        img = scheduler.step(pred, t, img, return_dict=False)[0]
+
+    off = 0
+    for k, i in enumerate(active):
+        L = lens[k]
+        h_, w_ = hw[k]
+        results[i] = _decode_one(model, img[:, off:off + L, :], h_, w_, dev)
+        off += L
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Image edit (packed, multi-resolution)
+# ---------------------------------------------------------------------------
+def _preprocess_ref_image(pil_img: Image.Image, height: int, width: int, device) -> torch.Tensor:
+    """Resize an RGB reference image to (height, width) and normalize to [-1, 1]."""
+    from torchvision.transforms import functional as TF
+    img = pil_img.convert("RGB")
+    img = TF.resize(img, [height, width], interpolation=TF.InterpolationMode.BICUBIC)
+    t = TF.to_tensor(img)  # [3, H, W] in [0, 1]
+    t = TF.normalize(t, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # -> [-1, 1]
+    return t.to(device)
+
+
+def _resize_long_edge(image: Image.Image, max_long_edge: int | None) -> Image.Image:
+    """Cap the VL conditioning image's long edge, preserving aspect ratio.
+
+    Matches training's data.processor._resize_long_edge (BICUBIC). Without this,
+    inference feeds a full-resolution image to the Qwen-VL processor whose
+    default max_pixels is far larger than 384**2 — a train/test mismatch.
+    """
+    if max_long_edge is None or max_long_edge <= 0:
+        return image
+    w, h = image.size
+    long_edge = max(w, h)
+    if long_edge <= max_long_edge:
+        return image
+    scale = max_long_edge / long_edge
+    new_w = max(1, int(round(w * scale)))
+    new_h = max(1, int(round(h * scale)))
+    return image.resize((new_w, new_h), Image.BICUBIC)
+
+
+# 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, num_refs):
+    """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 _encode_edits_packed(model, ref_pils_per_sample, instructions, template, drop_idx, device):
+    """Encode ALL image-conditioned edit instructions in ONE packed multimodal
+    varlen forward (pixel_values/image_grid_thw concatenated across samples,
+    cu_seqlens isolates each). Returns (txt_flat [ΣLi, D], vec [N, D], per-sample lens)."""
+    processor = model.txt_enc.processor
+    ids_list, pv_list, thw_list = [], [], []
+    for ref_pils, instr in zip(ref_pils_per_sample, instructions, strict=False):
+        formatted = template.format(_edit_prompt_body(instr, len(ref_pils)))
+        vl = processor(text=[formatted], images=list(ref_pils), padding=True, return_tensors="pt")
+        vl = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in vl.items()}
+        ids_list.append(vl["input_ids"].squeeze(0))
+        if vl.get("pixel_values") is not None:
+            pv_list.append(vl["pixel_values"]); thw_list.append(vl["image_grid_thw"])
+    input_ids = torch.cat(ids_list).to(device)
+    cu = _lens_to_cu([int(t.numel()) for t in ids_list], device)
+    inputs = {"input_ids": input_ids, "cu_seqlens": cu}
+    if pv_list:
+        inputs["pixel_values"] = torch.cat(pv_list, dim=0)
+        inputs["image_grid_thw"] = torch.cat(thw_list, dim=0)
+    res = model.txt_enc(
+        input_ids, cu, inputs=inputs, drop_idx_override=drop_idx)
+    return res["txt"], res["vec"], res["txt_seq_lens"].tolist()
+
+
+@torch.no_grad()
+def generate_edits(model, prompts, ref_images, neg_prompts=None, seeds=None, steps=30, cfg=5.0,
+                   max_size=None, heights=None, widths=None, device="cuda",
+                   prompt_template="mage-flow-edit", static_shift=None,
+                   gs_key=None,
+                   vl_cond_long_edge=384,
+                   renormalization=False, batch_cfg=True):
+    """Edit reference image(s) per prompt. Each ``ref_images[i]`` may be a single
+    image/path OR a list of source images (multi-image edit, like training —
+    trained with up to 3, but more are accepted) — all produce ONE edited output. Each sample's
+    ``[target, ref_1, …, ref_N]`` latent tokens are sequence-concatenated, and
+    all samples are packed into one varlen forward per denoise step.
+
+    Output resolution (derived from the first/primary reference of each sample):
+    if both ``heights[i]`` and ``widths[i]`` are given, use them; else if
+    ``max_size`` is given, the longest side is ``max_size`` and the short side
+    follows the reference's aspect ratio; otherwise the output keeps the source
+    image's own resolution. All references are VAE-encoded at that target size.
+    Returns a list of PIL images.
+    """
+    if isinstance(prompts, str):
+        prompts = [prompts]
+        ref_images = [ref_images]
+    n = len(prompts)
+    neg_prompts = _as_list(neg_prompts, " ", n)
+    seeds = _as_list(seeds, 42, n)
+    heights = _as_list(heights, None, n)
+    widths = _as_list(widths, None, n)
+    info = _template_info(prompt_template)
+    template = info.get("template", "{}")
+    drop_idx = int(info.get("start_idx", 0))
+    dev = torch.device(device)
+
+    # Normalize each sample's references to a list of 1..3 PIL images.
+    def _load_pil(r):
+        if isinstance(r, str):
+            r = Image.open(r)
+        return r.convert("RGB")
+
+    pils_per_sample = []
+    for r in ref_images:
+        refs = list(r) if isinstance(r, (list, tuple)) else [r]
+        if not refs:
+            raise ValueError("each edit sample needs at least one reference image")
+        pils_per_sample.append([_load_pil(x) for x in refs])
+
+    # Per-sample output resolution (from the first/primary reference) + content gate.
+    results = [None] * n
+    res_hw = [None] * n
+    active = []
+    for i in range(n):
+        res_hw[i] = _edit_target_size(pils_per_sample[i][0], max_size, heights[i], widths[i])
+        if seeds[i] == -1:
+            seeds[i] = random.randint(0, 2**32 - 1)
+        # Multimodal gate (MANDATORY): inspect the source image(s) AND the
+        # instruction, so NSFW / copyrighted-character / real-public-figure
+        # source photos are blocked even under an innocuous instruction.
+        verdict = model.txt_enc.screen_edit(prompts[i], pils_per_sample[i])
+        if verdict.violates:
+            h_, w_ = res_hw[i]
+            print(verdict.banner())
+            results[i] = make_refusal_image(verdict, height=h_, width=w_)
+            continue
+        active.append(i)
+    if not active:
+        return results
+
+    gs_key_int = resolve_gs_key(gs_key)
+
+    # Per sample: reference latent tokens (clean) + target noise tokens, plus the
+    # combined [target, ref_1, …, ref_N] position ids and shapes. ``target_idx``
+    # records where each sample's target tokens land in the packed sequence so we
+    # can slice the velocity and step only the target portion.
+    ch = model.vae.latent_channels
+    targets, refs, ids_list, shape_seq, samp_lens, tgt_lens, hw = [], [], [], [], [], [], []
+    target_idx_parts = []
+    off = 0
+    for i in active:
+        h_, w_ = res_hw[i]
+        torch.manual_seed(seeds[i])  # MageVAE.encode samples the posterior (global RNG)
+        # All references resized to the target resolution and VAE-encoded together.
+        ref_tensors = [_preprocess_ref_image(p, h_, w_, dev) for p in pils_per_sample[i]]
+        ref_tok, ref_shapes, ref_ids = model.compute_vae_encodings(ref_tensors, with_ids=True)
+        ref_tok = ref_tok.to(torch.bfloat16)               # [1, N*Lr, C]
+        x = get_noise(num_samples=1, channel=ch, height=h_, width=w_,
+                      device=dev, dtype=torch.bfloat16, seed=seeds[i])
+        x = encode_noise(tuple(x.shape[1:]), key=gs_key_int,
+                         seed=seeds[i], device=dev, dtype=torch.bfloat16)
+        _, _, gh, gw = x.shape
+        tgt = rearrange(x, "b c h w -> b (h w) c")          # [1, Lt, C]
+        tgt_ids = torch.zeros(gh, gw, 3, device=dev)
+        tgt_ids[..., 1] = tgt_ids[..., 1] + torch.arange(gh, device=dev)[:, None]
+        tgt_ids[..., 2] = tgt_ids[..., 2] + torch.arange(gw, device=dev)[None, :]
+        tgt_ids = rearrange(tgt_ids, "h w c -> (h w) c").unsqueeze(0)
+        lt, lr = tgt.shape[1], ref_tok.shape[1]
+        targets.append(tgt); refs.append(ref_tok)
+        ids_list.append(torch.cat([tgt_ids, ref_ids.to(dev)], dim=1)[0])  # [Lt + N*Lr, 3]
+        shape_seq.append((1, gh, gw))                       # target frame idx 0
+        shape_seq.extend(s[0] for s in ref_shapes)          # ref_j frame idx j
+        samp_lens.append(lt + lr); tgt_lens.append(lt); hw.append((h_, w_))
+        target_idx_parts.append(torch.arange(off, off + lt, device=dev))
+        off += lt + lr
+    img_ids = torch.cat(ids_list, 0).unsqueeze(0)
+    img_cu = _lens_to_cu(samp_lens, dev)
+    img_shapes = [shape_seq]
+    target_idx = torch.cat(target_idx_parts)
+
+    # Packed edit text — positive AND (for CFG) negative are encoded TOGETHER in
+    # ONE packed multimodal forward, then split. Both branches share the same
+    # reference images; cu_seqlens isolates every sequence (zero cross-contamination).
+    # The VL conditioning image's long edge is capped (default 384) to match
+    # training preprocessing — the VAE path above keeps the full target resolution.
+    na = len(active)
+    edit_refs = [[_resize_long_edge(p, vl_cond_long_edge) for p in pils_per_sample[i]]
+                 for i in active]
+    if cfg > 1.0:
+        pos_instr = [prompts[i] for i in active]
+        neg_instr = [neg_prompts[i] or " " for i in active]
+        txt_flat, vec_all, lens_t = _encode_edits_packed(
+            model, edit_refs + edit_refs, pos_instr + neg_instr, template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt, neg_cu, neg_mask, neg_vec = _slice_packed(txt_flat, vec_all, lens_t, na, na, dev)
+    else:
+        txt_flat, vec_all, lens_t = _encode_edits_packed(
+            model, edit_refs, [prompts[i] for i in active], template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt = neg_cu = neg_mask = neg_vec = None
+
+    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, samp_lens, txt, txt_cu, txt_mask, vec,
+                          neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, dev)
+    scheduler = _get_scheduler(model, steps, device, static_shift)
+    for si, t in enumerate(scheduler.timesteps):
+        parts = []
+        for k in range(na):
+            parts.append(targets[k]); parts.append(refs[k])
+        img = torch.cat(parts, dim=1)                       # [1, sum(Lt+Lr), C], ref clean
+        vel = _velocity(model.transformer, img, ctx, scheduler.sigmas[si].item())
+        pred_t = vel[:, target_idx, :]                      # [1, sum Lt, C] — target tokens only
+        tgt_packed = torch.cat(targets, dim=1)              # [1, sum Lt, C]
+        stepped = scheduler.step(pred_t, t, tgt_packed, return_dict=False)[0]
+        o = 0
+        new_targets = []
+        for k in range(na):
+            lt = tgt_lens[k]
+            new_targets.append(stepped[:, o:o + lt, :]); o += lt
+        targets = new_targets
+
+    for k, i in enumerate(active):
+        h_, w_ = hw[k]
+        results[i] = _decode_one(model, targets[k], h_, w_, dev)
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Flow-ODE inversion (Gaussian-Shading watermark detection)
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def invert_to_noise(model, z0, height, width, steps=30, device="cuda",
+                    prompt_template="mage-flow", static_shift=None, prompt=""):
+    """Reverse the flow ODE from a clean latent ``z0`` back to the initial noise.
+
+    This is the detection primitive for the Gaussian-Shading watermark: VAE-encode
+    the image to ``z0`` (posterior MEAN — deterministic), run this to recover the
+    initial noise, then read the signs via ``mage_latent.decode_bits``.
+
+    Inversion uses an empty prompt at cfg=1 (the standard Tree-Ring /
+    Gaussian-Shading setup). Reverse Euler recovers ``x_i`` from ``x_{i+1}`` with
+    the velocity evaluated at the point in hand; the sign-only watermark tolerates
+    the resulting approximation error (see the module's redundancy).
+
+    Args:
+        z0: clean latent ``[1, C, gh, gw]`` (e.g. the mean of ``model.vae.encode``).
+    Returns:
+        recovered initial-noise latent ``[1, C, gh, gw]`` (float32).
+    """
+    dev = torch.device(device)
+    info = _template_info(prompt_template)
+    template = info.get("template", "{}")
+    drop_idx = int(info.get("start_idx", 0))
+
+    z0 = z0.to(dev)
+    _, ch, gh, gw = z0.shape
+    img = rearrange(z0, "b c h w -> b (h w) c").to(torch.bfloat16)   # [1, gh*gw, C]
+
+    ids = torch.zeros(gh, gw, 3, device=dev)
+    ids[..., 1] = ids[..., 1] + torch.arange(gh, device=dev)[:, None]
+    ids[..., 2] = ids[..., 2] + torch.arange(gw, device=dev)[None, :]
+    img_ids = rearrange(ids, "h w c -> (h w) c").unsqueeze(0)
+    lens = [gh * gw]
+    img_cu = _lens_to_cu(lens, dev)
+    img_shapes = [[(1, gh, gw)]]
+
+    # Empty-prompt conditioning, no negative branch, cfg=1 (single forward).
+    txt_flat, vec_all, lens_t = _encode_texts_packed(model, [prompt], template, drop_idx, dev)
+    txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, 1, dev)
+    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, lens, txt, txt_cu, txt_mask, vec,
+                          None, None, None, None, 1.0, False, False, dev)
+
+    scheduler = _get_scheduler(model, steps, device, static_shift)
+    sigmas = scheduler.sigmas
+    n = len(scheduler.timesteps)
+    # Forward step si: x_{si+1} = x_si + (s_{si+1}-s_si)·v(x_si, s_si).
+    # Reverse it from clean (x_n, sigma 0) up to noise (x_0), using x_{si+1} as the
+    # proxy for x_si at the forward eval sigma s_si.
+    for si in range(n - 1, -1, -1):
+        s_cur = sigmas[si].item()
+        s_next = sigmas[si + 1].item()
+        vel = _velocity(model.transformer, img, ctx, s_cur)
+        img = img - (s_next - s_cur) * vel
+    return unpack(img.float(), height, width)   # [1, C, gh, gw]
+
+
+# ---------------------------------------------------------------------------
+# High-level pipeline wrapper
+# ---------------------------------------------------------------------------
+class MageFlowPipeline:
+    """``MageFlowPipeline.from_pretrained(repo).generate(...) / .edit(...)``.
+
+    ``generate`` / ``edit`` are packed multi-resolution calls: they take a list
+    of prompts (a single string is accepted and treated as a pack of size 1) and
+    return a list of PIL images. Per-sample ``heights``/``widths``/``seeds`` are
+    lists. Every prompt is screened by the text encoder's mandatory content
+    gate (no opt-out); banned prompts come back as refusal placeholders
+    interleaved with the real images. Real outputs always carry a Gaussian-Shading
+    watermark in the initial noise (no toggle), using the configured secret key.
+    """
+
+    def __init__(self, model, device="cuda"):
+        self.model = model
+        self.device = device
+
+    @classmethod
+    def from_pretrained(cls, repo_dir: str, device: str = "cuda"):
+        """Load a Mage-Flow diffusers-style repo (``model_index.json`` +
+        ``transformer/`` ``vae/`` ``scheduler/`` ``text_encoder/``).
+
+        ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id
+        (e.g. ``"microsoft/Mage-Flow-4B"``), which is downloaded and cached
+        automatically on first use.
+        """
+        return cls(load_from_repo(repo_dir, device), device)
+
+    def generate(self, prompts, **kw) -> list[Image.Image]:
+        """Packed multi-resolution t2i. ``prompts`` is a list (or a single
+        string); pass per-sample ``heights``/``widths``/``seeds`` as lists."""
+        kw.setdefault("device", self.device)
+        return generate_images(self.model, prompts, **kw)
+
+    def edit(self, prompts, ref_images, **kw) -> list[Image.Image]:
+        """Packed multi-resolution edit. ``prompts`` is a list (or a single
+        string); each ``ref_images[i]`` is one reference or a list of references."""
+        kw.setdefault("device", self.device)
+        return generate_edits(self.model, prompts, ref_images, **kw)
+
+    def invert_to_noise(self, z0, height, width, **kw):
+        """Recover the initial noise from a clean latent (Gaussian-Shading detect)."""
+        kw.setdefault("device", self.device)
+        return invert_to_noise(self.model, z0, height, width, **kw)
+
+
+def _safe_subpath(root: str, *parts: str) -> str:
+    """Join ``parts`` under ``root`` and confirm the result stays inside ``root``.
+
+    ``root`` is normalized up front; the joined path is normalized **lexically**
+    (``os.path.normpath`` — symlinks are *not* followed, so a Hugging Face cache
+    whose weight files are symlinks into the shared blob store still loads) and
+    rejected if it escapes ``root``. This guards the user-supplied model path
+    against path traversal (CWE-22 / CodeQL ``py/path-injection``).
+    """
+    root = os.path.realpath(root)
+    full = os.path.normpath(os.path.join(root, *parts))
+    if full != root and not full.startswith(root + os.sep):
+        raise ValueError(
+            f"Resolved path {os.path.join(*parts)!r} escapes repo directory {root!r}"
+        )
+    return full
+
+
+def _resolve_repo_dir(repo_dir: str) -> str:
+    """Return a local directory for ``repo_dir``.
+
+    If ``repo_dir`` is an existing local path it is returned as a normalized
+    absolute path; otherwise it is treated as a Hugging Face Hub repo id (e.g.
+    ``microsoft/Mage-Flow``) and downloaded/cached via
+    ``huggingface_hub.snapshot_download``.
+    """
+    candidate = os.path.realpath(repo_dir)
+    if os.path.isdir(candidate):
+        return candidate
+    from huggingface_hub import snapshot_download
+    return snapshot_download(repo_id=repo_dir)
+
+
+def load_from_repo(repo_dir: str, device: str = "cuda") -> MageFlowModel:
+    """Load a Mage-Flow diffusers-style repo (model_index.json + transformer/
+    vae/ scheduler/). Transformer weights come from the bf16 safetensors;
+    VAE + text encoder are built from the sources recorded in model_index.json.
+
+    ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id (e.g.
+    ``microsoft/Mage-Flow-4B``), which is downloaded/cached automatically.
+    """
+    from safetensors.torch import load_file
+    repo_dir = _resolve_repo_dir(repo_dir)
+    mi = json.load(open(_safe_subpath(repo_dir, "model_index.json")))
+    tcfg = json.load(open(_safe_subpath(repo_dir, "transformer", "config.json")))
+    # Keys stripped from the checkpoint config before it becomes model_structure.
+    # ``schedule_mode`` is a legacy field still present in some config.json files;
+    # Keys of the checkpoint config that are NOT MageFlowParams constructor args
+    # (legacy/unused fields). Everything else becomes model_structure. The DiT only
+    # reads: in_channels, out_channels, context_in_dim, hidden_size, num_heads,
+    # depth, axes_dim, checkpoint, patch_size.
+    _meta = {"_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"}
+    structure = {k: v for k, v in tcfg.items() if k not in _meta}
+
+    def _resolve(p):
+        return p if os.path.isabs(p) else _safe_subpath(repo_dir, p)
+
+    cfg = ModelConfig(
+        vae_path=_resolve(mi.get("_vae_source")),
+        txt_enc_path=_resolve(mi.get("_text_encoder_path")),
+        model_structure=structure,
+        txt_max_length=tcfg.get("txt_max_length", 2048),
+        packing=tcfg.get("packing", True),
+        static_shift=tcfg.get("static_shift", 6.0),
+    )
+    model = MageFlowModel(cfg)
+    sd = load_file(_safe_subpath(repo_dir, "transformer", "diffusion_pytorch_model.safetensors"),
+                   device="cpu")
+    model.transformer.load_state_dict(sd, strict=False, assign=True)
+    model.to(device)
+    model.transformer.to(torch.bfloat16)
+    model.txt_enc.to(torch.bfloat16)
+    if model.vae is not None:
+        model.vae.to(torch.bfloat16)
+    model.eval()
+    # Diffusers FlowMatchEulerDiscreteScheduler (scheduler/scheduler_config.json).
+    model.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        _safe_subpath(repo_dir, "scheduler"))
+    return model

requirements.txt

@@ -0,0 +1,14 @@
+# Mage-Flow Space requirements
+# Do NOT list gradio, spaces, or huggingface_hub (preinstalled on the platform).
+
+torchvision
+diffusers>=0.37.0
+transformers>=5.3.0,<5.6
+accelerate>=1.0.0
+safetensors>=0.4.0
+einops>=0.8.0
+pydantic>=2.0
+pillow>=10.0
+numpy>=1.26
+loguru>=0.7.0
+typing_extensions>=4.0