add mage-flow

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -1,7 +1,11 @@
 # Change Log for SD.Next
 
-## Update for 2026-07-30
+## Update for 2026-07-31
 
+- **Models**
+  - [Microsoft Mage-Flow](https://huggingface.co/mage-flow-community/Mage-Flow) in *Base* and *Turbo* (distilled) variants  
+    Mage-Flow is a 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing  
+    *note*: Microsoft released and then unpublished the model, but we still have a mirror available for download  
 - **Features**
   - startup: optimized server startup
   - process: preserve audio when processing video

TODO.md

@@ -6,10 +6,6 @@
 
 ## Features
 
-- [MageFlow](https://github.com/huggingface/diffusers/pull/14295)
-- [SeFi](https://github.com/huggingface/diffusers/pull/14084)
-- [Boogu](https://github.com/huggingface/diffusers/pull/14040)
-
 ### Assigned
 
 - Chat-based interface, @vladmandic

data/reference-base.json

@@ -918,5 +918,13 @@
     "size": 47.98,
     "extras": "sampler: Default",
     "date": "2026 July"
+  },
+  "Microsoft Mage-Flow": {
+    "path": "vladmandic/Mage-Flow-4B",
+    "preview": "vladmandic--Mage-Flow-4B.jpg",
+    "desc": "Mage-Flow is a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing.",
+    "extras": "sampler: Default",
+    "size": 16.19,
+    "date": "2026 July"
   }
 }

data/reference-distilled.json

@@ -225,5 +225,13 @@
     "desc": "LongCat-Image-Edit-Turbo, the distilled version of LongCat-Image-Edit. It achieves high-quality image editing with only 8 NFEs (Number of Function Evaluations) , offering extremely low inference latency.",
     "size": 27.28,
     "date": "2026 February"
+  },
+  "Microsoft Mage-Flow Turbo": {
+    "path": "vladmandic/Mage-Flow-4B-Turbo",
+    "preview": "vladmandic--Mage-Flow-Turbo-4B.jpg",
+    "desc": "Mage-Flow is a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing.",
+    "extras": "sampler: Default",
+    "size": 16.19,
+    "date": "2026 July"
   }
 }

pipelines/generic_shared.py

@@ -107,7 +107,7 @@ shared_te_map = {
     'Qwen3-VL 4B Conditional': {
         'cls': transformers.Qwen3VLForConditionalGeneration,
         'target_repo': 'SeFi-Image/SeFi-Image-5B-Base',
-        'identifier': ['5b'],
+        'identifier': ['4b','5b'],
         'target_subfolder': 'Qwen3-VL-4B-Instruct',
     },
     'Qwen3-VL 8B Conditional': {

pipelines/ideogram/ideogram4.py

@@ -653,7 +653,7 @@ class Ideogram4Pipeline(DiffusionPipeline):
         # 4. Set up the resolution-aware logit-normal schedule on the scheduler.
         schedule_mu = _resolution_aware_mu(height=height, width=width, base_mu=mu)
         sigmas = _logit_normal_sigmas(num_inference_steps, schedule_mu, std=std, device=device)
-        self.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device)
+        self.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device) # pylint: disable=unexpected-keyword-arg
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps) # pylint: disable=attribute-defined-outside-init
 

pipelines/mageflow/__init__.py

@@ -1 +1,10 @@
-from .pipeline_mage import MageFlowPipeline
+import diffusers
+from .pipeline_mage_flow import MageFlowPipeline
+from .pipeline_output import MageFlowPipelineOutput
+from .autoencoder_mage_vae import AutoencoderMageVAE
+from .transformer_mage_flow import MageFlowTransformer2DModel
+
+diffusers.MageFlowPipeline = MageFlowPipeline
+diffusers.MageFlowPipelineOutput = MageFlowPipelineOutput
+diffusers.AutoencoderMageVAE = AutoencoderMageVAE
+diffusers.MageFlowTransformer2DModel = MageFlowTransformer2DModel

pipelines/mageflow/autoencoder_mage_vae.py

@@ -0,0 +1,768 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) autoencoder.
+
+Encodes images to 128-channel latents at 16x spatial downsampling using a one-step
+diffusion encoder, and decodes latents back to images using a DConv denoiser conditioned
+on a CoD (Cascaded-of-Decoders) decoder.
+
+Latent shape: [B, 128, H/16, W/16].
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.autoencoders.vae import DecoderOutput
+
+
+logger = logging.get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helper
+# ---------------------------------------------------------------------------
+def _mage_vae_modulate(x, shift, scale):
+    if x.dim() == 4:
+        batch_size, channels = x.shape[:2]
+        return x * (1 + scale.view(batch_size, channels, 1, 1)) + shift.view(batch_size, channels, 1, 1)
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+# ---------------------------------------------------------------------------
+# Primitive layers
+# ---------------------------------------------------------------------------
+class MageVAELayerNorm2d(nn.LayerNorm):
+    """Channel-last LayerNorm for NCHW tensors."""
+
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x):
+        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 MageVAERMSNorm(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):
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * x.to(input_dtype)
+
+
+class MageVAETimestepEmbedder(nn.Module):
+    """Timestep MLP (max_period=10000, freq_size=256)."""
+
+    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]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        embedding = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(embedding.to(self.mlp[0].weight.dtype))
+
+
+# ---------------------------------------------------------------------------
+# DConv blocks
+# ---------------------------------------------------------------------------
+class MageVAEDiCoBlock(nn.Module):
+    """DConv block with adaLN modulation, used in encoder and decoder."""
+
+    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_channels = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn_channels, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn_channels, hidden_size, 1, bias=True)
+
+        self.norm1 = MageVAELayerNorm2d(hidden_size, affine=False)
+        self.norm2 = MageVAELayerNorm2d(hidden_size, affine=False)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, hidden_states, conditioning):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(conditioning).chunk(
+            6, dim=1
+        )
+        residual = hidden_states
+        hidden_states = _mage_vae_modulate(self.norm1(residual), shift_msa, scale_msa)
+        hidden_states = F.gelu(self.conv2(self.conv1(hidden_states)))
+        hidden_states = hidden_states * self.ca(hidden_states)
+        hidden_states = self.conv3(hidden_states)
+        hidden_states = residual + gate_msa[..., None, None] * hidden_states
+
+        hidden_states = hidden_states + gate_mlp[..., None, None] * self.conv5(
+            F.gelu(self.conv4(_mage_vae_modulate(self.norm2(hidden_states), shift_mlp, scale_mlp)))
+        )
+        return hidden_states
+
+
+class MageVAEEncoderDiCoBlock(nn.Module):
+    """DConv block without adaLN modulation, for the encoder head 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_channels = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn_channels, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn_channels, hidden_size, 1, bias=True)
+
+        self.norm1 = MageVAELayerNorm2d(hidden_size, affine=True)
+        self.norm2 = MageVAELayerNorm2d(hidden_size, affine=True)
+
+    def forward(self, hidden_states):
+        residual = hidden_states
+        hidden_states = self.norm1(residual)
+        hidden_states = F.gelu(self.conv2(self.conv1(hidden_states)))
+        hidden_states = hidden_states * self.ca(hidden_states)
+        hidden_states = self.conv3(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states + self.conv5(F.gelu(self.conv4(self.norm2(hidden_states))))
+
+
+# ---------------------------------------------------------------------------
+# Nerf-style patch embedder and final layer
+# ---------------------------------------------------------------------------
+class MageVAENerfEmbedder(nn.Module):
+    """Patch-position embedder for 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),
+        )
+        self._pos_cache = {}
+
+    def _compute_pos(self, patch_size, device, dtype):
+        key = (patch_size, device, dtype)
+        if key in self._pos_cache:
+            return self._pos_cache[key]
+
+        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)
+
+        result = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2)
+        self._pos_cache[key] = result
+        return result
+
+    def forward(self, x):
+        batch_size, num_patches, _ = x.shape
+        patch_size = int(num_patches**0.5)
+        dct = self._compute_pos(patch_size, x.device, x.dtype).expand(batch_size, -1, -1)
+        return self.embedder(torch.cat([x, dct], dim=-1))
+
+
+class MageVAENerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = MageVAERMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        return self.linear(self.norm(x))
+
+
+# ---------------------------------------------------------------------------
+# MLP decoder (SimpleMLPAdaLN + MLPResBlock)
+# ---------------------------------------------------------------------------
+class _MageVAEMLPResBlock(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 MageVAESimpleMLPAdaLN(nn.Module):
+    """Small MLP that maps NerfEmbedder features to per-patch output, conditioned on spatial features."""
+
+    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([_MageVAEMLPResBlock(model_channels) for _ in range(num_res_blocks)])
+
+    def forward(self, x, conditioning):
+        x = self.input_proj(x)
+        conditioning = self.cond_embed(conditioning).reshape(conditioning.shape[0], self.patch_size**2, -1)
+        for block in self.res_blocks:
+            x = block(x, conditioning)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder building blocks (ResNet + Attention)
+# ---------------------------------------------------------------------------
+class MageVAEResnetBlock(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 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        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):
+        hidden_states = self.conv1(F.silu(self.norm1(x)))
+        hidden_states = self.conv2(self.dropout(F.silu(self.norm2(hidden_states))))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + hidden_states
+
+
+class MageVAEAttnBlock(nn.Module):
+    """Patched self-attention for the CoD Decoder."""
+
+    def __init__(self, in_channels, patch_size=32):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        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):
+        normalized = self.norm(x)
+        query = self.q(normalized)
+        key = self.k(normalized)
+        value = self.v(normalized)
+
+        d = self.patch_size
+        batch_size, channels, height, width = query.shape
+        pad_h = (d - height % d) % d
+        pad_w = (d - width % d) % d
+        if pad_h or pad_w:
+            query = F.pad(query, (0, pad_w, 0, pad_h), mode="replicate")
+            key = F.pad(key, (0, pad_w, 0, pad_h), mode="replicate")
+            value = F.pad(value, (0, pad_w, 0, pad_h), mode="replicate")
+
+        _, _, height_padded, width_padded = query.shape
+        num_patches_h = height_padded // d
+        num_patches_w = width_padded // d
+        num_patches = num_patches_h * num_patches_w
+
+        # Reshape to patches: [B*num_patches, C, d*d]
+        query = (
+            query.reshape(batch_size, channels, num_patches_h, d, num_patches_w, d)
+            .permute(0, 2, 4, 1, 3, 5)
+            .reshape(batch_size * num_patches, channels, d * d)
+        )
+        key = (
+            key.reshape(batch_size, channels, num_patches_h, d, num_patches_w, d)
+            .permute(0, 2, 4, 1, 3, 5)
+            .reshape(batch_size * num_patches, channels, d * d)
+        )
+        value = (
+            value.reshape(batch_size, channels, num_patches_h, d, num_patches_w, d)
+            .permute(0, 2, 4, 1, 3, 5)
+            .reshape(batch_size * num_patches, channels, d * d)
+        )
+
+        # Attention via F.scaled_dot_product_attention
+        # query/key/value: [B*np, C, d*d] -> [B*np, 1, d*d, C] for SDPA (batch, heads, seq, head_dim)
+        q = query.permute(0, 2, 1).unsqueeze(1)
+        k = key.permute(0, 2, 1).unsqueeze(1)
+        v = value.permute(0, 2, 1).unsqueeze(1)
+        h_ = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+        h_ = h_.squeeze(1).permute(0, 2, 1)  # back to [B*np, C, d*d]
+
+        # Reconstruct
+        hidden_states = (
+            h_
+            .reshape(batch_size, num_patches_h, num_patches_w, channels, d, d)
+            .permute(0, 3, 1, 4, 2, 5)
+            .reshape(batch_size, channels, height_padded, width_padded)
+        )
+        if pad_h or pad_w:
+            hidden_states = hidden_states[:, :, :height, :width]
+
+        return x + self.proj_out(hidden_states)
+
+
+# ---------------------------------------------------------------------------
+# Patch embedding
+# ---------------------------------------------------------------------------
+class MageVAEBottleneckPatchEmbed(nn.Module):
+    """Image patch embed concatenated with a per-patch conditioning vector."""
+
+    def __init__(self, patch_size=16, in_channels=3, bottleneck_dim=128, embed_dim=384, bias=True):
+        super().__init__()
+        self.proj1 = nn.Conv2d(in_channels, bottleneck_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(bottleneck_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
+
+    def forward(self, x, conditioning):
+        return self.proj2(torch.cat([self.proj1(x), conditioning], dim=1))
+
+
+# ---------------------------------------------------------------------------
+# adaLN constant folding
+# ---------------------------------------------------------------------------
+class _MageVAEConstAdaLN(nn.Module):
+    """Replaces an adaLN_modulation MLP with a precomputed constant buffer."""
+
+    def __init__(self, modulation: torch.Tensor):
+        super().__init__()
+        self.register_buffer("modulation", modulation.detach().clone())
+
+    def forward(self, conditioning):
+        batch_size = conditioning.shape[0]
+        if self.modulation.shape[0] != batch_size:
+            return self.modulation.expand(batch_size, *self.modulation.shape[1:])
+        return self.modulation
+
+
+# ---------------------------------------------------------------------------
+# DConv Encoder
+# ---------------------------------------------------------------------------
+class MageVAEDConvEncoder(nn.Module):
+    """One-step diffusion encoder: image -> packed (mean, logvar) latent."""
+
+    def __init__(
+        self,
+        latent_channels=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.latent_channels = latent_channels
+        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(
+            [MageVAEEncoderDiCoBlock(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(latent_channels, 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 = MageVAETimestepEmbedder(hidden_size)
+        self.blocks = nn.ModuleList([MageVAEDiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)])
+
+        self.norm_out = MageVAELayerNorm2d(hidden_size, affine=True)
+        self.proj_out = nn.Conv2d(hidden_size, latent_channels * out_ch_mult, kernel_size=1, bias=True)
+
+    def forward(self, z_t, t, image):
+        conditioning = self.patch_cond_embed(image)
+        for block in self.head_blocks:
+            conditioning = block(conditioning)
+        conditioning = self.proj_down(conditioning)
+
+        hidden_states = self.fuse_proj(torch.cat([conditioning, self.z_proj(z_t)], dim=1))
+        timestep_embedding = self.t_embedder(t.view(-1))
+        for block in self.blocks:
+            hidden_states = block(hidden_states, timestep_embedding)
+        return self.proj_out(self.norm_out(hidden_states))
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder: latent -> conditioning features for the denoiser
+# ---------------------------------------------------------------------------
+class MageVAEDecoder(nn.Module):
+    """CoD (Cascaded-of-Decoders) decoder: latent -> spatial conditioning features."""
+
+    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(
+            MageVAEResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            MageVAEAttnBlock(out_ch, patch_size=32),
+            MageVAEResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            MageVAEAttnBlock(out_ch, patch_size=32),
+            MageVAEResnetBlock(in_channels=out_ch, out_channels=out_ch),
+        )
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=out_ch, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z):
+        hidden_states = self.block(self.conv_in(z))
+        hidden_states = self.conv_out(F.silu(self.norm_out(hidden_states)))
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Y-Embedder wrapper (holds the CoD decoder)
+# ---------------------------------------------------------------------------
+class _MageVAEYEmbedder(nn.Module):
+    """Namespace wrapper for the CoD decoder, matching the original checkpoint's
+    ``pipeline.y_embedder.decoder.*`` weight key hierarchy."""
+
+    def __init__(self, hidden_size=384, latent_channels=128):
+        super().__init__()
+        self.decoder = MageVAEDecoder(out_ch=hidden_size, z_ch=latent_channels)
+
+
+# ---------------------------------------------------------------------------
+# DConv Denoiser: conditioning + zero noise -> reconstructed image
+# ---------------------------------------------------------------------------
+class MageVAEDConvDenoiser(nn.Module):
+    """One-step denoiser: takes conditioning from CoD decoder and produces the output image."""
+
+    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 = MageVAETimestepEmbedder(hidden_size)
+        self.y_embedder_x = nn.Conv2d(hidden_size, hidden_size_x * patch_size**2, 1, 1, 0)
+        self.x_embedder = MageVAENerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8)
+        self.s_embedder = MageVAEBottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+
+        self.blocks = nn.ModuleList(
+            [MageVAEDiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks)]
+        )
+
+        self.dec_net = MageVAESimpleMLPAdaLN(
+            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 = MageVAENerfFinalLayer(hidden_size_x, in_channels)
+        self.y_embedder = _MageVAEYEmbedder(hidden_size=hidden_size, latent_channels=bottleneck_dim)
+
+    def forward(self, x, t, conditioning, is_latent=False):
+        if is_latent:
+            conditioning = self.y_embedder.decoder(conditioning)
+
+        batch_size, _, height, width = x.shape
+        timestep_embedding = self.t_embedder(t.view(-1))
+
+        # Spatial conditioning path
+        spatial = self.s_embedder(x, conditioning)
+        for block in self.blocks:
+            spatial = block(spatial, timestep_embedding)
+
+        num_spatial = spatial.shape[-2] * spatial.shape[-1]
+        spatial_flat = spatial.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+        # Per-patch x-pathway
+        x_unfolded = F.unfold(x, kernel_size=self.patch_size, stride=self.patch_size)
+        y_x = self.y_embedder_x(conditioning).flatten(2)
+        x_combined = torch.cat([x_unfolded, y_x], dim=1)
+
+        # Reshape: [B, (in_ch + hidden_x), ps^2, num_spatial] -> [B*num_spatial, ps^2, (in_ch + hidden_x)]
+        x_combined = (
+            x_combined.reshape(batch_size, -1, self.patch_size**2, num_spatial).permute(0, 3, 2, 1).flatten(0, 1)
+        )
+
+        x_embedded = self.x_embedder(x_combined)
+        x_decoded = self.dec_net(x_embedded, spatial_flat)
+        x_final = self.final_layer(x_decoded)
+
+        # Fold patches back to image: [B*num_spatial, ps^2, in_ch] -> [B, in_ch, H, W]
+        x_final = x_final.transpose(1, 2).reshape(batch_size, num_spatial, -1)
+        return F.fold(
+            x_final.transpose(1, 2).contiguous(),
+            (height, width),
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Main autoencoder
+# ---------------------------------------------------------------------------
+class AutoencoderMageVAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    r"""
+    MageVAE autoencoder model using a one-step diffusion encoder and a DConv denoiser
+    with a CoD (Cascaded-of-Decoders) decoder.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods
+    implemented for all models (such as downloading or saving).
+
+    Encoder: DConvEncoder takes an image [B, 3, H, W] and produces a latent [B, 128, H/16, W/16].
+    Decoder: CoD Decoder + DConvDenoiser takes a latent and reconstructs the image [B, 3, H, W].
+
+    Args:
+        latent_channels (`int`, defaults to `128`):
+            Number of channels in the latent space.
+        downsample_factor (`int`, defaults to `16`):
+            Spatial downsampling factor from image to latent.
+        encoder_hidden_size (`int`, defaults to `384`):
+            Hidden dimension of the encoder DConv blocks.
+        encoder_num_blocks (`int`, defaults to `21`):
+            Number of adaLN-modulated DConv blocks in the encoder.
+        encoder_patch_size (`int`, defaults to `16`):
+            Patch size for the encoder's image tokenization.
+        encoder_head_size (`int`, defaults to `768`):
+            Channel dimension of the encoder's head blocks.
+        encoder_num_head_blocks (`int`, defaults to `2`):
+            Number of head blocks in the encoder (without adaLN).
+        decoder_hidden_size (`int`, defaults to `384`):
+            Hidden dimension of the decoder DConv blocks.
+        decoder_hidden_size_x (`int`, defaults to `32`):
+            Hidden dimension of the decoder's per-patch x-pathway.
+        decoder_num_blocks (`int`, defaults to `24`):
+            Total number of blocks in the decoder (cond blocks + MLP res blocks).
+        decoder_num_cond_blocks (`int`, defaults to `21`):
+            Number of adaLN-modulated DConv blocks in the decoder.
+        decoder_bottleneck_dim (`int`, defaults to `128`):
+            Bottleneck dimension for the patch embedding and CoD decoder input.
+        decoder_patch_size (`int`, defaults to `16`):
+            Patch size for the decoder.
+        sample_posterior (`bool`, defaults to `True`):
+            Whether to sample from the posterior (mean + noise * std) or use the mean directly.
+    """
+
+    _no_split_modules = ["MageVAEDiCoBlock", "MageVAEResnetBlock", "MageVAEAttnBlock"]
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels: int = 128,
+        downsample_factor: int = 16,
+        encoder_hidden_size: int = 384,
+        encoder_num_blocks: int = 21,
+        encoder_patch_size: int = 16,
+        encoder_head_size: int = 768,
+        encoder_num_head_blocks: int = 2,
+        decoder_hidden_size: int = 384,
+        decoder_hidden_size_x: int = 32,
+        decoder_num_blocks: int = 24,
+        decoder_num_cond_blocks: int = 21,
+        decoder_bottleneck_dim: int = 128,
+        decoder_patch_size: int = 16,
+        sample_posterior: bool = True,
+    ):
+        super().__init__()
+
+        self.encoder = MageVAEDConvEncoder(
+            latent_channels=latent_channels,
+            hidden_size=encoder_hidden_size,
+            num_blocks=encoder_num_blocks,
+            patch_size=encoder_patch_size,
+            head_size=encoder_head_size,
+            num_head_blocks=encoder_num_head_blocks,
+        )
+
+        self.decoder = MageVAEDConvDenoiser(
+            patch_size=decoder_patch_size,
+            in_channels=3,
+            hidden_size=decoder_hidden_size,
+            hidden_size_x=decoder_hidden_size_x,
+            num_blocks=decoder_num_blocks,
+            num_cond_blocks=decoder_num_cond_blocks,
+            bottleneck_dim=decoder_bottleneck_dim,
+        )
+
+    def encode(self, x: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        """
+        Encode images to latents.
+
+        Args:
+            x (`torch.Tensor`): Input images of shape `[B, 3, H, W]`. H and W must be
+                multiples of `encoder_patch_size`.
+            generator (`torch.Generator`, *optional*):
+                A torch generator for reproducible sampling.
+
+        Returns:
+            `torch.Tensor`: Latent of shape `[B, 128, H/16, W/16]`.
+        """
+        batch_size, _, height, width = x.shape
+        patch_size = self.config.encoder_patch_size
+        latent_channels = self.config.latent_channels
+
+        z_t = torch.zeros(
+            batch_size,
+            latent_channels,
+            height // patch_size,
+            width // patch_size,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        t = torch.zeros(batch_size, device=x.device, dtype=x.dtype)
+
+        out = self.encoder(z_t, t, x)
+        mean = out[:, :latent_channels]
+        logvar = out[:, latent_channels:].clamp(min=-20.0, max=10.0)
+
+        if self.config.sample_posterior:
+            noise = randn_tensor(mean.shape, generator=generator, device=mean.device, dtype=mean.dtype)
+            return mean + torch.exp(0.5 * logvar) * noise
+        return mean
+
+    def forward(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+        return self.decode(z, return_dict=return_dict)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+        """
+        Decode latents to images.
+
+        Args:
+            z (`torch.Tensor`): Latent of shape `[B, 128, H/16, W/16]`.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return a [`~models.autoencoders.vae.DecoderOutput`] or a plain tuple.
+
+        Returns:
+            [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
+                Decoded images of shape `[B, 3, H, W]`.
+        """
+        batch_size = z.shape[0]
+        height = z.shape[2] * self.config.downsample_factor
+        width = z.shape[3] * self.config.downsample_factor
+        noise = torch.zeros(batch_size, 3, height, width, device=z.device, dtype=z.dtype)
+        t = torch.zeros(batch_size, device=z.device, dtype=z.dtype)
+        sample = self.decoder(noise, t, z, is_latent=True)
+
+        if not return_dict:
+            return (sample,)
+        return DecoderOutput(sample=sample)
+
+    def freeze_adaln(self):
+        """Constant-fold adaLN_modulation MLPs at t=0 for both encoder and decoder.
+
+        At t=0 the adaLN modulation outputs are constant (they only depend on the
+        timestep embedding). This method precomputes those constants and replaces the
+        MLP modules with lightweight buffer wrappers, saving compute and parameters.
+        """
+        device = next(self.parameters()).device
+        dtype = next(self.parameters()).dtype
+        t = torch.zeros(1, device=device, dtype=dtype)
+
+        c_enc = self.encoder.t_embedder(t)
+        count_enc = self._replace_adaln_with_const(self.encoder, c_enc)
+
+        c_dec = self.decoder.t_embedder(t)
+        count_dec = self._replace_adaln_with_const(self.decoder, c_dec)
+
+        logger.info(f"MageVAE: folded {count_enc} encoder + {count_dec} decoder adaLN blocks")
+
+    @staticmethod
+    def _replace_adaln_with_const(module: nn.Module, conditioning: torch.Tensor) -> int:
+        """Replace adaLN_modulation MLPs in MageVAEDiCoBlock instances with constant buffers."""
+        count = 0
+        for child in module.modules():
+            if not isinstance(child, MageVAEDiCoBlock):
+                continue
+            adaln = child.adaLN_modulation
+            if isinstance(adaln, _MageVAEConstAdaLN):
+                continue
+            with torch.no_grad():
+                modulation = adaln(conditioning)
+            child.adaLN_modulation = _MageVAEConstAdaLN(modulation)
+            count += 1
+        return count

pipelines/mageflow/convert_mage_flow_to_diffusers.py

@@ -0,0 +1,269 @@
+"""Convert an original Mage-Flow HF repo layout into diffusers-format weights.
+
+Example
+-------
+
+    python scripts/convert_mage_flow_to_diffusers.py \
+        --input_dir  path/to/Mage-Flow-Base \
+        --output_dir path/to/Mage-Flow-Base-diffusers \
+        --dtype bfloat16
+"""
+
+import argparse
+import json
+import os
+import shutil
+from typing import Dict
+
+import safetensors.torch
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Transformer conversion
+# ---------------------------------------------------------------------------
+
+# Top-level 1:1 renames. Anything not matched here (transformer_blocks.*, the
+# time_text_embed / norm_out / proj_out subtrees) is passed through unchanged.
+TRANSFORMER_TOP_LEVEL_RENAMES: Dict[str, str] = {
+    "img_in.weight": "x_embedder.weight",
+    "img_in.bias": "x_embedder.bias",
+    "txt_norm.weight": "context_embedder_norm.weight",
+    "txt_in.weight": "context_embedder.weight",
+    "txt_in.bias": "context_embedder.bias",
+}
+
+
+TRANSFORMER_DIFFUSERS_CONFIG = {
+    "_class_name": "MageFlowTransformer2DModel",
+    "_diffusers_version": "0.37.0",
+    "in_channels": 128,
+    "out_channels": 128,
+    "context_in_dim": 2560,
+    "hidden_size": 3072,
+    "num_attention_heads": 24,
+    "num_layers": 12,
+    "axes_dim": [16, 56, 56],
+    "patch_size": 1,
+}
+
+
+def convert_transformer_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    new_state_dict: Dict[str, torch.Tensor] = {}
+    for key, tensor in state_dict.items():
+        if key in TRANSFORMER_TOP_LEVEL_RENAMES:
+            new_key = TRANSFORMER_TOP_LEVEL_RENAMES[key]
+        else:
+            # Pass-through: transformer_blocks.*, time_text_embed.*, norm_out.*, proj_out.*
+            new_key = key
+        if new_key in new_state_dict:
+            raise ValueError(f"Duplicate destination key while converting transformer: {new_key}")
+        new_state_dict[new_key] = tensor
+    return new_state_dict
+
+
+# ---------------------------------------------------------------------------
+# VAE conversion
+# ---------------------------------------------------------------------------
+
+VAE_ENCODER_PREFIX = "student.dconv_encoder."
+VAE_DECODER_PREFIX = "pipeline."
+
+# Sub-trees of the original decoder ("pipeline.*") that belong to the Flux2
+# encoder and must be dropped instead of being mapped into the diffusers
+# decoder namespace.
+VAE_DECODER_EXCLUDE_PREFIXES = (
+    "pipeline.y_embedder.encoder.",
+    "pipeline.y_embedder.bottleneck.",
+)
+
+
+VAE_DIFFUSERS_CONFIG = {
+    "_class_name": "AutoencoderMageVAE",
+    "_diffusers_version": "0.37.0",
+    "latent_channels": 128,
+    "downsample_factor": 16,
+    "encoder_hidden_size": 384,
+    "encoder_num_blocks": 21,
+    "encoder_patch_size": 16,
+    "encoder_head_size": 768,
+    "encoder_num_head_blocks": 2,
+    "decoder_hidden_size": 384,
+    "decoder_hidden_size_x": 32,
+    "decoder_num_blocks": 24,
+    "decoder_num_cond_blocks": 21,
+    "decoder_bottleneck_dim": 128,
+    "decoder_patch_size": 16,
+    "sample_posterior": False,
+}
+
+
+def convert_vae_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    new_state_dict: Dict[str, torch.Tensor] = {}
+    for key, tensor in state_dict.items():
+        if key.startswith(VAE_ENCODER_PREFIX):
+            new_key = "encoder." + key[len(VAE_ENCODER_PREFIX):]
+        elif key.startswith(VAE_DECODER_PREFIX):
+            if any(key.startswith(p) for p in VAE_DECODER_EXCLUDE_PREFIXES):
+                continue
+            # pipeline.y_embedder.decoder.* naturally maps to
+            # decoder.y_embedder.decoder.* under this rule, matching the spec.
+            new_key = "decoder." + key[len(VAE_DECODER_PREFIX):]
+        else:
+            raise ValueError(f"Unexpected VAE key with no known prefix: {key}")
+        if new_key in new_state_dict:
+            raise ValueError(f"Duplicate destination key while converting VAE: {new_key}")
+        new_state_dict[new_key] = tensor
+    return new_state_dict
+
+
+# ---------------------------------------------------------------------------
+# I/O helpers
+# ---------------------------------------------------------------------------
+
+DTYPE_MAP = {
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "float16": torch.float16,
+    "fp16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "bf16": torch.bfloat16,
+}
+
+
+def cast_state_dict(state_dict: Dict[str, torch.Tensor], dtype: torch.dtype) -> Dict[str, torch.Tensor]:
+    out: Dict[str, torch.Tensor] = {}
+    for key, tensor in state_dict.items():
+        # Leave integer / bool buffers alone (e.g. num_batches_tracked).
+        if tensor.is_floating_point():
+            out[key] = tensor.to(dtype)
+        else:
+            out[key] = tensor
+    return out
+
+
+def save_component(
+    state_dict: Dict[str, torch.Tensor],
+    config: Dict,
+    output_component_dir: str,
+) -> None:
+    os.makedirs(output_component_dir, exist_ok=True)
+    safetensors.torch.save_file(
+        state_dict,
+        os.path.join(output_component_dir, "diffusion_pytorch_model.safetensors"),
+    )
+    with open(os.path.join(output_component_dir, "config.json"), "w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2)
+        f.write("\n")
+
+
+def copy_scheduler(input_dir: str, output_dir: str) -> None:
+    src = os.path.join(input_dir, "scheduler", "scheduler_config.json")
+    dst_dir = os.path.join(output_dir, "scheduler")
+    os.makedirs(dst_dir, exist_ok=True)
+    shutil.copyfile(src, os.path.join(dst_dir, "scheduler_config.json"))
+
+
+def copy_text_encoder(input_dir: str, output_dir: str, symlink: bool) -> None:
+    src = os.path.join(input_dir, "text_encoder")
+    dst = os.path.join(output_dir, "text_encoder")
+    if os.path.lexists(dst):
+        if os.path.islink(dst) or os.path.isfile(dst):
+            os.remove(dst)
+        else:
+            shutil.rmtree(dst)
+    if symlink:
+        os.symlink(os.path.abspath(src), dst)
+    else:
+        shutil.copytree(src, dst)
+
+
+MODEL_INDEX = {
+    "_class_name": "MageFlowPipeline",
+    "_diffusers_version": "0.37.0",
+    "transformer": ["diffusers", "MageFlowTransformer2DModel"],
+    "vae": ["diffusers", "AutoencoderMageVAE"],
+    "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+    "text_encoder": ["transformers", "Qwen3VLForConditionalGeneration"],
+    "tokenizer": ["transformers", "AutoTokenizer"],
+}
+
+
+def write_model_index(output_dir: str) -> None:
+    with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
+        json.dump(MODEL_INDEX, f, indent=2)
+        f.write("\n")
+
+
+# ---------------------------------------------------------------------------
+# Driver
+# ---------------------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--input_dir", required=True, help="Original Mage-Flow HF repo path.")
+    parser.add_argument("--output_dir", required=True, help="Destination diffusers-format directory.")
+    parser.add_argument("--dtype", default="bfloat16", choices=sorted(DTYPE_MAP.keys()), help="Output tensor dtype.")
+    parser.add_argument(
+        "--text_encoder_mode",
+        default="symlink",
+        choices=["symlink", "copy"],
+        help="How to include the text_encoder directory in the output.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    dtype = DTYPE_MAP[args.dtype]
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # --- Transformer ---
+    print("[transformer] loading original weights ...")
+    transformer_sd = safetensors.torch.load_file(
+        os.path.join(args.input_dir, "transformer", "diffusion_pytorch_model.safetensors")
+    )
+    print(f"[transformer] converting {len(transformer_sd)} tensors ...")
+    transformer_sd = convert_transformer_state_dict(transformer_sd)
+    transformer_sd = cast_state_dict(transformer_sd, dtype)
+    save_component(
+        transformer_sd,
+        TRANSFORMER_DIFFUSERS_CONFIG,
+        os.path.join(args.output_dir, "transformer"),
+    )
+    print(f"[transformer] wrote {len(transformer_sd)} tensors to {args.output_dir}/transformer")
+    del transformer_sd
+
+    # --- VAE ---
+    print("[vae] loading original weights ...")
+    vae_sd = safetensors.torch.load_file(
+        os.path.join(args.input_dir, "vae", "diffusion_pytorch_model.safetensors")
+    )
+    print(f"[vae] converting {len(vae_sd)} tensors ...")
+    vae_sd = convert_vae_state_dict(vae_sd)
+    vae_sd = cast_state_dict(vae_sd, dtype)
+    save_component(
+        vae_sd,
+        VAE_DIFFUSERS_CONFIG,
+        os.path.join(args.output_dir, "vae"),
+    )
+    print(f"[vae] wrote {len(vae_sd)} tensors to {args.output_dir}/vae")
+    del vae_sd
+
+    # --- Scheduler ---
+    print("[scheduler] copying config ...")
+    copy_scheduler(args.input_dir, args.output_dir)
+
+    # --- Text encoder ---
+    print(f"[text_encoder] {args.text_encoder_mode} ...")
+    copy_text_encoder(args.input_dir, args.output_dir, symlink=(args.text_encoder_mode == "symlink"))
+
+    # --- model_index.json ---
+    write_model_index(args.output_dir)
+    print(f"Done. Diffusers-format repo written to: {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

pipelines/mageflow/pipeline_mage.py

@@ -1,13 +0,0 @@
-import diffusers
-
-
-class MageFlowPipeline(diffusers.DiffusionPipeline):
-    def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler):
-        super().__init__()
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-        )

pipelines/mageflow/pipeline_mage_flow.py

@@ -0,0 +1,665 @@
+# Copyright 2025 Microsoft and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from transformers import Qwen2Tokenizer, Qwen3VLForConditionalGeneration
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from .transformer_mage_flow import MageFlowTransformer2DModel
+from .pipeline_output import MageFlowPipelineOutput
+from .autoencoder_mage_vae import AutoencoderMageVAE
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import MageFlowPipeline
+
+        >>> pipe = MageFlowPipeline.from_pretrained("microsoft/Mage-Flow-4B", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+        >>> prompt = "A cat holding a sign that says hello world"
+        >>> image = pipe(prompt, num_inference_steps=30, guidance_scale=5.0).images[0]
+        >>> image.save("mage_flow.png")
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`list[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`list[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class MageFlowPipeline(DiffusionPipeline):
+    r"""
+    The Mage-Flow pipeline for text-to-image generation.
+
+    Args:
+        transformer ([`MageFlowTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderMageVAE`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`Qwen3VLForConditionalGeneration`]):
+            Qwen3-VL text encoder for producing text conditioning embeddings.
+        tokenizer (`AutoTokenizer`):
+            Tokenizer for the Qwen3-VL text encoder.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderMageVAE,
+        text_encoder: Qwen3VLForConditionalGeneration,
+        tokenizer: Qwen2Tokenizer,
+        transformer: MageFlowTransformer2DModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = 16  # MageVAE downsample factor
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = 2048
+        self.default_sample_size = 64  # 1024 / 16 = 64
+        # ChatML prompt template (same as QwenImage)
+        self.prompt_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"
+        )
+        self.prompt_template_start_idx = 34  # number of system-prompt tokens to skip
+
+    def _get_prompt_embeds(
+        self,
+        prompt: str | list[str] | None = None,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        device = device or self._execution_device
+        if self.text_encoder is None:
+            raise ValueError(
+                "Text encoder is not available. Please provide `prompt_embeds` directly "
+                "when the pipeline is initialized without a text encoder."
+            )
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        template = self.prompt_template
+        drop_idx = self.prompt_template_start_idx
+        txt = [template.format(e) for e in prompt]
+        txt_tokens = self.tokenizer(
+            txt, max_length=self.tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt"
+        ).to(device)
+        encoder_out = self.text_encoder(
+            input_ids=txt_tokens.input_ids,
+            attention_mask=txt_tokens.attention_mask,
+            output_hidden_states=True,
+        )
+        hidden_states = encoder_out.hidden_states[-1]
+
+        # Extract valid tokens per sample, drop system prompt prefix, then re-pad to uniform length.
+        bool_mask = txt_tokens.attention_mask.bool()
+        valid_lengths = bool_mask.sum(dim=1)
+        selected = hidden_states[bool_mask]
+        split_hidden_states = torch.split(selected, valid_lengths.tolist(), dim=0)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        prompt_embeds_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device | None = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 2048,
+    ):
+        r"""
+        Encode the text prompt into embeddings for the transformer.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Prompt to be encoded.
+            device (`torch.device`):
+                Torch device.
+            num_images_per_prompt (`int`):
+                Number of images that should be generated per prompt.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_embeds_mask (`torch.Tensor`, *optional*):
+                Attention mask for `prompt_embeds`.
+            max_sequence_length (`int`):
+                Maximum sequence length for the text embeddings.
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_prompt_embeds(prompt, device)
+
+        prompt_embeds = prompt_embeds[:, :max_sequence_length]
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        if prompt_embeds_mask is not None:
+            prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
+            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+            if prompt_embeds_mask.all():
+                prompt_embeds_mask = None
+
+        return prompt_embeds, prompt_embeds_mask
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_embeds_mask=None,
+        negative_prompt_embeds_mask=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor} but are {height} and {width}. "
+                "Dimensions will be resized accordingly"
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found "
+                f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_embeds_mask is None:
+            logger.warning(
+                "`prompt_embeds` is provided and `prompt_embeds_mask` is not provided, so the model will treat all"
+                " prompt tokens as valid. If `prompt_embeds` contains padding, you should provide the padding mask as"
+                " `prompt_embeds_mask`. Make sure to generate `prompt_embeds_mask` from the same text encoder that was"
+                " used to generate `prompt_embeds`."
+            )
+
+        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+            logger.warning(
+                "`negative_prompt_embeds` is provided and `negative_prompt_embeds_mask` is not provided, so the model"
+                " will treat all negative prompt tokens as valid. If `negative_prompt_embeds` contains padding, you"
+                " should provide the padding mask as `negative_prompt_embeds_mask`. Make sure to generate"
+                " `negative_prompt_embeds_mask` from the same text encoder that was used to generate"
+                " `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 2048:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 2048 but is {max_sequence_length}")
+
+    @staticmethod
+    def _prepare_latent_image_ids(height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3, device=device, dtype=dtype)
+        latent_image_ids[..., 1] = torch.arange(height, device=device, dtype=dtype)[:, None]
+        latent_image_ids[..., 2] = torch.arange(width, device=device, dtype=dtype)[None, :]
+        latent_image_ids = latent_image_ids.reshape(height * width, 3)
+        return latent_image_ids
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # MageVAE: 16x downsample, no patch packing
+        height = height // self.vae_scale_factor
+        width = width // self.vae_scale_factor
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        # Flatten to sequence: [B, C, H, W] -> [B, H*W, C]
+        latents = latents.permute(0, 2, 3, 1).reshape(batch_size, height * width, num_channels_latents)
+
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 5.0,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds_mask: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end: Callable[[int, int], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 2048,
+        sigmas: list[float] | None = None,
+    ) -> MageFlowPipelineOutput | tuple:
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `list[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                not greater than `1`).
+            height (`int`, *optional*, defaults to `self.default_sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.default_sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 30):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Classifier-free guidance scale. Enabled by setting `guidance_scale > 1`. Higher guidance scale
+                encourages images closely linked to the text `prompt`, usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_embeds_mask (`torch.Tensor`, *optional*):
+                Attention mask for `prompt_embeds`.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_prompt_embeds_mask (`torch.Tensor`, *optional*):
+                Attention mask for `negative_prompt_embeds`.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.mage_flow.MageFlowPipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int`, defaults to 2048):
+                Maximum sequence length to use with the `prompt`.
+            sigmas (`list[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.mage_flow.MageFlowPipelineOutput`] or `tuple`:
+            [`~pipelines.mage_flow.MageFlowPipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is a list with the generated images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode prompt
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_classifier_free_guidance:
+            if negative_prompt_embeds is None and self.text_encoder is None:
+                # text_encoder unavailable and no negative_prompt_embeds provided, skip CFG
+                do_classifier_free_guidance = False
+            else:
+                negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                    prompt=negative_prompt if negative_prompt is not None else [""] * batch_size,
+                    prompt_embeds=negative_prompt_embeds,
+                    prompt_embeds_mask=negative_prompt_embeds_mask,
+                    device=device,
+                    num_images_per_prompt=num_images_per_prompt,
+                    max_sequence_length=max_sequence_length,
+                )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare image position ids for RoPE
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        img_ids = self._prepare_latent_image_ids(latent_h, latent_w, device, prompt_embeds.dtype)
+
+        # 6. Prepare timesteps
+        # Mage-Flow uses base_sigmas = linspace(1, 1/N, N) which differs from the scheduler's
+        # default sigma computation. The scheduler's shift=6.0 is applied on top of these.
+        if sigmas is None:
+            sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        # 7. Denoising loop
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                if do_classifier_free_guidance:
+                    noise_pred_cond = self.transformer(
+                        hidden_states=latents,
+                        encoder_hidden_states=prompt_embeds,
+                        timestep=timestep / 1000,
+                        img_ids=img_ids,
+                        joint_attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred_uncond = self.transformer(
+                        hidden_states=latents,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        timestep=timestep / 1000,
+                        img_ids=img_ids,
+                        joint_attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                else:
+                    noise_pred = self.transformer(
+                        hidden_states=latents,
+                        encoder_hidden_states=prompt_embeds,
+                        timestep=timestep / 1000,
+                        img_ids=img_ids,
+                        joint_attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        # 8. VAE decode
+        if output_type == "latent":
+            image = latents
+        else:
+            # Unflatten: [B, H*W, C] -> [B, C, H, W]
+            latents = latents.reshape(batch_size * num_images_per_prompt, latent_h, latent_w, num_channels_latents)
+            latents = latents.permute(0, 3, 1, 2)
+            latents = latents.to(self.vae.dtype)
+            image = self.vae(latents, return_dict=False)[0]
+            image = image.clamp(-1, 1)
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return MageFlowPipelineOutput(images=image)

pipelines/mageflow/pipeline_output.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+import numpy as np
+import PIL.Image
+
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class MageFlowPipelineOutput(BaseOutput):
+    """
+    Output class for Mage-Flow pipelines.
+
+    Args:
+        images (`list[PIL.Image.Image]` or `np.ndarray`)
+            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+    """
+
+    images: list[PIL.Image.Image] | np.ndarray

pipelines/mageflow/transformer_mage_flow.py

@@ -0,0 +1,630 @@
+# Copyright 2025 The Mage Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.utils import logging
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin, FeedForward
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def _apply_rotary_emb_complex(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex rotary embeddings to ``x`` using MageFlow's adjacent-pair convention.
+
+    Args:
+        x: Query or key tensor of shape ``[B, S, H, D]``.
+        freqs_cis: Complex frequency tensor of shape ``[S, D_rope // 2]`` where
+            ``D_rope = sum(axes_dim)``. When ``D_rope < D`` only the first
+            ``D_rope`` dimensions are rotated; the rest pass through unchanged.
+
+    Returns:
+        Tensor of same shape and dtype as *x* with rotary embeddings applied.
+    """
+    rope_dim = freqs_cis.shape[-1] * 2  # complex dim -> real dim
+    head_dim = x.shape[-1]
+
+    if rope_dim < head_dim:
+        x_rope = x[..., :rope_dim]
+        x_pass = x[..., rope_dim:]
+    else:
+        x_rope = x
+        x_pass = None
+
+    # [B, S, H, rope_dim] -> [B, S, H, rope_dim/2] complex
+    x_complex = torch.view_as_complex(x_rope.float().reshape(*x_rope.shape[:-1], -1, 2))
+    # freqs_cis: [S, D_rope/2] -> [1, S, 1, D_rope/2] for broadcasting
+    freqs = freqs_cis.unsqueeze(0).unsqueeze(2)
+    x_rotated = torch.view_as_real(x_complex * freqs).flatten(-2)
+    x_rotated = x_rotated.to(x.dtype)
+
+    if x_pass is not None:
+        return torch.cat([x_rotated, x_pass], dim=-1)
+    return x_rotated
+
+
+class MageFlowPosEmbed(nn.Module):
+    """Complex RoPE with symmetric positive/negative frequency scaling for MageFlow.
+
+    Computes multi-scale rotary positional embeddings for video/image tokens using
+    three axes (frame, height, width). Height and width axes use symmetric
+    positive/negative frequency indices centered around the spatial midpoint.
+    """
+
+    def __init__(self, theta: int = 10000, axes_dim: list[int] | None = None):
+        super().__init__()
+        if axes_dim is None:
+            axes_dim = [16, 48, 48]
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+
+        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,
+        )
+        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,
+        )
+        self.register_buffer("pos_freqs_real", pos_freqs.real.contiguous(), persistent=False)
+        self.register_buffer("pos_freqs_imag", pos_freqs.imag.contiguous(), persistent=False)
+        self.register_buffer("neg_freqs_real", neg_freqs.real.contiguous(), persistent=False)
+        self.register_buffer("neg_freqs_imag", neg_freqs.imag.contiguous(), persistent=False)
+
+    @staticmethod
+    def _rope_params(index: torch.Tensor, dim: int, theta: float = 10000.0) -> torch.Tensor:
+        """Compute complex RoPE frequencies for a 1-D position index."""
+        freqs = torch.outer(
+            index.float(),
+            1.0 / torch.pow(theta, torch.arange(0, dim, 2, dtype=torch.float32).div(dim)),
+        )
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor:
+        seq_len = frame * height * width
+        pos_freqs = torch.complex(self.pos_freqs_real.float(), self.pos_freqs_imag.float())
+        neg_freqs = torch.complex(self.neg_freqs_real.float(), self.neg_freqs_imag.float())
+        freqs_pos = pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+        freqs_neg = 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)
+        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)
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_len, -1)
+        return freqs.clone().contiguous()
+
+    def forward(self, img_ids: torch.Tensor, height: int | None = None, width: int | None = None) -> torch.Tensor:
+        """Compute RoPE frequencies from image position ids.
+
+        Args:
+            img_ids: ``[seq_len, 3]`` tensor with (frame, height, width) position
+                indices for each image token.
+            height: Latent spatial height.
+            width: Latent spatial width.
+
+        Returns:
+            Complex frequency tensor of shape ``[seq_len, head_dim // 2]``.
+        """
+        frame = 1
+        freqs = self._compute_video_freqs(frame, height, width, idx=0)
+        return freqs.to(img_ids.device)
+
+    @staticmethod
+    def _infer_grid_size(img_ids: torch.Tensor) -> tuple[int, int]:
+        height = int(img_ids[:, 1].max().item()) + 1
+        width = int(img_ids[:, 2].max().item()) + 1
+        return height, width
+
+
+class MageFlowTimestepProjEmbeddings(nn.Module):
+    """Timestep projection embeddings for MageFlow.
+
+    Uses a custom sinusoidal embedding that downcasts the frequency table to the
+    input dtype before computing the embedding. The model was trained with this
+    exact bf16 rounding, so using diffusers' standard float32 variant degrades
+    output quality.
+    """
+
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.num_channels = 256
+        self.scale = 1000
+
+    @staticmethod
+    def _sinusoidal_embedding(
+        timesteps: torch.Tensor,
+        embedding_dim: int,
+        scale: float = 1.0,
+        max_period: int = 10000,
+    ) -> torch.Tensor:
+        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
+
+        # Downcast frequency table to input dtype (bf16) before multiplying —
+        # the model was trained with this exact rounding.
+        emb = torch.exp(exponent).to(timesteps.dtype)
+        emb = timesteps[:, None].float() * emb[None, :]
+        emb = scale * emb
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        # flip sin to cos
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+        return emb
+
+    def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+        timesteps_proj = self._sinusoidal_embedding(timestep, self.num_channels, scale=self.scale)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))
+        return timesteps_emb
+
+
+class MageFlowAttnProcessor:
+    """Attention processor for MageFlow double-stream (MMDiT) architecture.
+
+    Implements joint attention over concatenated ``[text, image]`` tokens. RoPE is
+    applied only to image query/key, not text.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.")
+
+    def __call__(
+        self,
+        attn: "MageFlowAttention",
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Compute QKV for image stream
+        img_query = attn.to_q(hidden_states)
+        img_key = attn.to_k(hidden_states)
+        img_value = attn.to_v(hidden_states)
+
+        # Reshape to multi-head: [B, S, inner_dim] -> [B, S, H, D]
+        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))
+
+        # Apply QK normalization
+        img_query = attn.norm_q(img_query)
+        img_key = attn.norm_k(img_key)
+
+        # Apply RoPE to image Q/K only (not text)
+        if image_rotary_emb is not None:
+            img_query = _apply_rotary_emb_complex(img_query, image_rotary_emb)
+            img_key = _apply_rotary_emb_complex(img_key, image_rotary_emb)
+
+        if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
+            # Compute QKV for text stream
+            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)
+
+            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))
+
+            txt_query = attn.norm_added_q(txt_query)
+            txt_key = attn.norm_added_k(txt_key)
+
+            # No RoPE on text — concatenate [text, image] for joint attention
+            query = torch.cat([txt_query, img_query], dim=1)
+            key = torch.cat([txt_key, img_key], dim=1)
+            value = torch.cat([txt_value, img_value], dim=1)
+        else:
+            query = img_query
+            key = img_key
+            value = img_value
+
+        # Joint attention via dispatch
+        attn_output = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
+        )
+        attn_output = attn_output.flatten(2, 3)
+        attn_output = attn_output.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            # Split back into text and image parts
+            txt_seq_len = encoder_hidden_states.shape[1]
+            txt_attn_output, img_attn_output = attn_output.split_with_sizes(
+                [txt_seq_len, attn_output.shape[1] - txt_seq_len], dim=1
+            )
+            img_attn_output = attn.to_out[0](img_attn_output)
+            img_attn_output = attn.to_out[1](img_attn_output)
+            txt_attn_output = attn.to_add_out(txt_attn_output)
+            return img_attn_output, txt_attn_output
+
+        return attn_output
+
+
+class MageFlowAttention(nn.Module, AttentionModuleMixin):
+    """Multi-head attention module for MageFlow with support for dual-stream (MMDiT) attention.
+
+    Follows the diffusers attention pattern with ``_default_processor_cls`` and
+    ``_available_processors`` for backend dispatch.
+    """
+
+    _default_processor_cls = MageFlowAttnProcessor
+    _available_processors = [MageFlowAttnProcessor]
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = True,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-6,
+        out_dim: int | None = None,
+        elementwise_affine: bool = True,
+        processor: "MageFlowAttnProcessor | None" = None,
+    ):
+        super().__init__()
+        self.head_dim = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.dropout = dropout
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.added_kv_proj_dim = added_kv_proj_dim
+        self.added_proj_bias = added_proj_bias
+
+        self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+        self.norm_k = nn.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(query_dim, self.inner_dim, bias=bias)
+        self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+        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))
+
+        if added_kv_proj_dim is not None:
+            self.norm_added_q = nn.RMSNorm(dim_head, eps=eps)
+            self.norm_added_k = nn.RMSNorm(dim_head, eps=eps)
+            self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+            self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+            self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+            self.to_add_out = nn.Linear(self.inner_dim, query_dim, bias=out_bias)
+
+        if processor is None:
+            processor = self._default_processor_cls()
+        self.set_processor(processor)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
+
+
+class MageFlowTransformerBlock(nn.Module):
+    """Double-stream MMDiT transformer block for MageFlow.
+
+    Each block processes image and text streams with separate modulation (AdaLN),
+    joint attention, and separate feed-forward networks.
+    """
+
+    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 stream modulation and layers
+        self.img_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),
+        )
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = MageFlowAttention(
+            query_dim=dim,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=MageFlowAttnProcessor(),
+            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 stream modulation and layers
+        self.txt_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),
+        )
+        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        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 forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Compute modulation parameters for both streams
+        img_mod_params = self.img_mod(temb)
+        txt_mod_params = self.txt_mod(temb)
+
+        # Split into norm1 and norm2 modulation parameters (each has shift, scale, gate)
+        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)
+        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)
+
+        # Image stream: norm1 + modulation
+        img_shift1, img_scale1, img_gate1 = img_mod1.chunk(3, dim=-1)
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated = img_normed * (1 + img_scale1.unsqueeze(1)) + img_shift1.unsqueeze(1)
+
+        # Text stream: norm1 + modulation
+        txt_shift1, txt_scale1, txt_gate1 = txt_mod1.chunk(3, dim=-1)
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated = txt_normed * (1 + txt_scale1.unsqueeze(1)) + txt_shift1.unsqueeze(1)
+
+        # Joint attention
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        img_attn_output, txt_attn_output = self.attn(
+            hidden_states=img_modulated,
+            encoder_hidden_states=txt_modulated,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        # Apply gates and residuals
+        hidden_states = hidden_states + img_gate1.unsqueeze(1) * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1.unsqueeze(1) * txt_attn_output
+
+        # Image stream: norm2 + MLP
+        img_shift2, img_scale2, img_gate2 = img_mod2.chunk(3, dim=-1)
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2 = img_normed2 * (1 + img_scale2.unsqueeze(1)) + img_shift2.unsqueeze(1)
+        img_mlp_output = self.img_mlp(img_modulated2)
+        hidden_states = hidden_states + img_gate2.unsqueeze(1) * img_mlp_output
+
+        # Text stream: norm2 + MLP
+        txt_shift2, txt_scale2, txt_gate2 = txt_mod2.chunk(3, dim=-1)
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2 = txt_normed2 * (1 + txt_scale2.unsqueeze(1)) + txt_shift2.unsqueeze(1)
+        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2.unsqueeze(1) * 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 MageFlowTransformer2DModel(
+    ModelMixin,
+    ConfigMixin,
+    PeftAdapterMixin,
+    FromOriginalModelMixin,
+    CacheMixin,
+    AttentionMixin,
+):
+    """Transformer model for MageFlow image generation.
+
+    A dual-stream (MMDiT) Transformer that processes image and text tokens jointly.
+    Uses complex multi-scale RoPE for image positional encoding and Qwen3-VL text
+    embeddings as conditioning.
+
+    Args:
+        in_channels (`int`, defaults to ``128``):
+            Number of channels in the input latent (MageVAE latent channels).
+        out_channels (`int`, defaults to ``128``):
+            Number of channels in the output.
+        context_in_dim (`int`, defaults to ``3584``):
+            Dimension of the text encoder hidden states (Qwen3-VL hidden size).
+        hidden_size (`int`, defaults to ``3072``):
+            Inner dimension of the transformer (num_attention_heads * attention_head_dim).
+        num_attention_heads (`int`, defaults to ``24``):
+            Number of attention heads.
+        num_layers (`int`, defaults to ``32``):
+            Number of dual-stream transformer blocks.
+        axes_dim (`list[int]``, defaults to ``[16, 48, 48]``):
+            RoPE dimension split across axes (frame, height, width). Must sum to
+            ``hidden_size // num_attention_heads``.
+        patch_size (`int`, defaults to ``1``):
+            Patch size for the output projection.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["MageFlowTransformerBlock"]
+    _repeated_blocks = ["MageFlowTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    main_input_name = "hidden_states"
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        context_in_dim: int = 3584,
+        hidden_size: int = 3072,
+        num_attention_heads: int = 24,
+        num_layers: int = 32,
+        axes_dim: list[int] = [16, 48, 48],
+        patch_size: int = 1,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.inner_dim = hidden_size
+        self.num_attention_heads = num_attention_heads
+        attention_head_dim = hidden_size // num_attention_heads
+
+        self.pos_embed = MageFlowPosEmbed(theta=10000, axes_dim=axes_dim)
+
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+        self.context_embedder_norm = nn.RMSNorm(context_in_dim, eps=1e-6)
+        self.context_embedder = nn.Linear(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=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        timestep: torch.Tensor = None,
+        img_ids: torch.Tensor = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        """
+        The [`MageFlowTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, img_seq_len, in_channels)`):
+                Flattened image latent tokens.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, txt_seq_len, context_in_dim)`):
+                Text encoder hidden states (Qwen3-VL embeddings).
+            timestep (`torch.Tensor`):
+                Raw sigma value in ``[0, 1]``.
+            img_ids (`torch.Tensor` of shape `(img_seq_len, 3)`):
+                Image position ids ``(frame, height, width)`` for RoPE computation.
+            joint_attention_kwargs (`dict`, *optional*):
+                Additional keyword arguments passed to the attention processor.
+            return_dict (`bool`, defaults to ``True``):
+                Whether to return a :class:`Transformer2DModelOutput` or a plain tuple.
+
+        Returns:
+            :class:`Transformer2DModelOutput` or ``tuple``.
+        """
+        # Embed image tokens
+        hidden_states = self.x_embedder(hidden_states)
+
+        # Embed text tokens: RMSNorm then linear projection
+        encoder_hidden_states = self.context_embedder_norm(encoder_hidden_states)
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        # Timestep embedding (Timesteps module handles the 1000x scaling internally via scale=1000)
+        timestep = timestep.to(hidden_states.dtype)
+        temb = self.time_text_embed(timestep, hidden_states)
+
+        # Compute image RoPE (text tokens are not rotated)
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        image_rotary_emb = self.pos_embed(img_ids, *MageFlowPosEmbed._infer_grid_size(img_ids))
+
+        # Transformer blocks
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        # Final norm and projection (image stream only)
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

pipelines/model_mageflow.py

@@ -1,4 +1,5 @@
 import diffusers
+import transformers
 from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
 from modules.logger import log
 from pipelines import generic
@@ -14,7 +15,12 @@ def load_mageflow(checkpoint_info, diffusers_load_config=None):
     load_args, _ = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
     log.debug(f'Load model: type=MageFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
 
-    from pipelines.mageflow import MageFlowPipeline
+    from pipelines.mageflow import MageFlowPipeline, MageFlowTransformer2DModel
+
+    log.debug(f'Load model: type=MageFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}')
+    transformer = generic.load_transformer(repo_id, cls_name=MageFlowTransformer2DModel, load_config=diffusers_load_config)
+    text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLForConditionalGeneration, load_config=diffusers_load_config)
+    tokenizer = transformers.Qwen2Tokenizer.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=shared.opts.diffusers_dir)
 
     if repo_id is None or repo_id.lower() == 'none':
         return None
@@ -25,16 +31,21 @@ def load_mageflow(checkpoint_info, diffusers_load_config=None):
 
     pipe = MageFlowPipeline.from_pretrained(
         repo_id,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
         cache_dir=shared.opts.diffusers_dir,
         **load_args,
     )
 
     pipe.task_args = {
         'output_type': 'pil',
-        'device': devices.device,
     }
 
     generic.load_vae_override(pipe, diffusers_load_config)
+    del transformer
+    del text_encoder
+    del tokenizer
 
     sd_hijack_te.init_hijack(pipe)
     sd_hijack_vae.init_hijack(pipe)

pipelines/sefi/convert_sefi_to_diffusers.py

@@ -0,0 +1,307 @@
+#!/usr/bin/env python
+# Copyright 2026 SeFi-Image Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+
+import torch
+import yaml
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+from diffusers import __version__
+from diffusers.models import SeFiTransformer2DModel
+
+
+SEFI_SCALE_PRESETS = {
+    "0p5b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 12,
+        "num_layers": 3,
+        "num_single_layers": 10,
+        "joint_attention_dim": 6144,
+    },
+    "1b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 16,
+        "num_layers": 4,
+        "num_single_layers": 12,
+        "joint_attention_dim": 6144,
+    },
+    "2b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 20,
+        "num_layers": 4,
+        "num_single_layers": 16,
+        "joint_attention_dim": 6144,
+    },
+    "3b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 22,
+        "num_layers": 5,
+        "num_single_layers": 18,
+        "joint_attention_dim": 7680,
+    },
+    "4b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 24,
+        "num_layers": 5,
+        "num_single_layers": 20,
+        "joint_attention_dim": 7680,
+    },
+    "5b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 26,
+        "num_layers": 6,
+        "num_single_layers": 21,
+        "joint_attention_dim": 7680,
+    },
+    "6b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 28,
+        "num_layers": 6,
+        "num_single_layers": 22,
+        "joint_attention_dim": 7680,
+    },
+    "8b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 30,
+        "num_layers": 7,
+        "num_single_layers": 24,
+        "joint_attention_dim": 7680,
+    },
+    "9b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 32,
+        "num_layers": 8,
+        "num_single_layers": 24,
+        "joint_attention_dim": 12288,
+    },
+}
+
+QWEN3VL_TEXT_HIDDEN_DIMS = {
+    "qwen3vl_2b": 2048,
+    "qwen3vl_4b": 2560,
+    "qwen3vl_8b": 4096,
+}
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Convert a SeFi-Image checkpoint to Diffusers format.")
+    parser.add_argument("--checkpoint", required=True, help="Local checkpoint folder or Hugging Face repo id.")
+    parser.add_argument("--output", required=True, help="Output Diffusers checkpoint folder.")
+    parser.add_argument("--cache-dir", default=None, help="Optional Hugging Face cache directory.")
+    parser.add_argument("--token", default=None, help="Optional Hugging Face token for gated checkpoints.")
+    parser.add_argument(
+        "--variant",
+        choices=["base", "rl", "turbo"],
+        default=None,
+        help="Model family. Inferred from checkpoint name if omitted.",
+    )
+    return parser.parse_args()
+
+
+def resolve_checkpoint(checkpoint: str, cache_dir: str | None, token: str | None) -> Path:
+    path = Path(checkpoint).expanduser()
+    if path.exists():
+        return path
+    return Path(snapshot_download(checkpoint, cache_dir=cache_dir, token=token))
+
+
+def copytree(src: Path, dst: Path, ignore=None):
+    if dst.exists():
+        shutil.rmtree(dst)
+    shutil.copytree(src, dst, ignore=ignore)
+
+
+def load_json(path: Path):
+    with open(path, "r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def load_yaml(path: Path):
+    with open(path, "r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def save_json(path: Path, payload: dict):
+    with open(path, "w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+
+
+def infer_variant(checkpoint: str, config: dict, explicit_variant: str | None) -> str:
+    if explicit_variant is not None:
+        return explicit_variant
+    configured = str(config.get("inference", {}).get("family", "") or config.get("model", {}).get("variant", ""))
+    text = f"{checkpoint} {configured}".lower()
+    if "turbo" in text or "distill" in text:
+        return "turbo"
+    if "rl" in text:
+        return "rl"
+    return "base"
+
+
+def default_steps(variant: str) -> int:
+    return 4 if variant == "turbo" else 50
+
+
+def default_guidance_scale(variant: str) -> float:
+    return 1.0 if variant == "turbo" else 4.0
+
+
+def texture_vae_config_path(root: Path, texture_vae_name: str) -> Path:
+    if texture_vae_name == "sd1.5":
+        return root / "vae" / "config.json"
+    if texture_vae_name in {"flux1", "flux2"}:
+        return root / "vae" / "config.json"
+    raise ValueError(f"Unsupported texture VAE: {texture_vae_name}")
+
+
+def build_transformer_config(root: Path, sefi_config: dict) -> dict:
+    model_config = sefi_config["model"]
+    transformer_config = load_json(root / "transformer" / "config.json")
+    transformer_config.pop("_class_name", None)
+    transformer_config.pop("_diffusers_version", None)
+    transformer_config.pop("_name_or_path", None)
+    transformer_config.pop("guidance_embeds", None)
+
+    scale = str(model_config.get("transformer_scale", "")).lower()
+    if scale and scale != "custom":
+        transformer_config.update(SEFI_SCALE_PRESETS[scale])
+    elif scale == "custom":
+        transformer_config.update(model_config.get("transformer_overrides", {}))
+
+    semantic_channels = int(model_config["semantic_channels"])
+    texture_vae_name = str(model_config["texture_vae"]["name"]).lower()
+    vae_config = load_json(texture_vae_config_path(root, texture_vae_name))
+    texture_channels = int(vae_config["latent_channels"]) * 4
+    total_channels = semantic_channels + texture_channels
+
+    text_config = model_config["text_encoder"]
+    hidden_layers = tuple(int(layer) for layer in text_config["hidden_layers"])
+    text_dim = int(QWEN3VL_TEXT_HIDDEN_DIMS[text_config["model_name"]]) * len(hidden_layers)
+
+    transformer_config["in_channels"] = total_channels
+    transformer_config["out_channels"] = total_channels
+    transformer_config["text_input_dim"] = text_dim
+    if int(transformer_config["joint_attention_dim"]) != text_dim:
+        raise ValueError(
+            "Text dimension mismatch: "
+            f"transformer joint_attention_dim={transformer_config['joint_attention_dim']} vs text_dim={text_dim}."
+        )
+    return transformer_config
+
+
+def load_transformer_state_dict(transformer_dir: Path) -> dict[str, torch.Tensor]:
+    index_path = transformer_dir / "diffusion_pytorch_model.safetensors.index.json"
+    single_path = transformer_dir / "diffusion_pytorch_model.safetensors"
+    bin_path = transformer_dir / "diffusion_pytorch_model.bin"
+
+    if index_path.exists():
+        index = load_json(index_path)
+        state_dict = {}
+        for shard in sorted(set(index["weight_map"].values())):
+            state_dict.update(load_file(transformer_dir / shard))
+        return state_dict
+    if single_path.exists():
+        return load_file(single_path)
+    if bin_path.exists():
+        return torch.load(bin_path, map_location="cpu")
+    raise FileNotFoundError(f"No supported transformer weights found under {transformer_dir}.")
+
+
+def copy_tokenizer_files(src: Path, dst: Path):
+    weight_patterns = {
+        "model*.safetensors",
+        "pytorch_model*.bin",
+        "*.index.json",
+    }
+
+    def ignore(_dir, names):
+        ignored = set()
+        for name in names:
+            for pattern in weight_patterns:
+                if Path(name).match(pattern):
+                    ignored.add(name)
+        return ignored
+
+    copytree(src, dst, ignore=ignore)
+
+
+def main():
+    args = parse_args()
+    root = resolve_checkpoint(args.checkpoint, args.cache_dir, args.token)
+    output = Path(args.output).expanduser()
+    output.mkdir(parents=True, exist_ok=True)
+
+    sefi_config = load_yaml(root / "sefi_config.yaml")
+    variant = infer_variant(args.checkpoint, sefi_config, args.variant)
+    transformer_config = build_transformer_config(root, sefi_config)
+
+    transformer = SeFiTransformer2DModel(**transformer_config)
+    state_dict = load_transformer_state_dict(root / "transformer")
+    missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
+    if missing or unexpected:
+        raise ValueError(f"Transformer state dict mismatch. Missing={missing[:20]}, unexpected={unexpected[:20]}")
+    transformer.save_pretrained(output / "transformer", safe_serialization=True)
+
+    copytree(root / "scheduler", output / "scheduler")
+    copytree(root / "vae", output / "vae")
+
+    text_encoder_name = sefi_config["model"]["text_encoder"]["model_name"]
+    qwen_dir_name = {
+        "qwen3vl_2b": "Qwen3-VL-2B-Instruct",
+        "qwen3vl_4b": "Qwen3-VL-4B-Instruct",
+        "qwen3vl_8b": "Qwen3-VL-8B-Instruct",
+    }[text_encoder_name]
+    qwen_dir = root / qwen_dir_name
+    copytree(qwen_dir, output / "text_encoder")
+    copy_tokenizer_files(qwen_dir, output / "tokenizer")
+
+    model_config = sefi_config["model"]
+    inference_config = sefi_config.get("inference", {})
+    training_sefi_config = sefi_config.get("training", {}).get("sefi", {})
+    texture_vae_name = str(model_config["texture_vae"]["name"]).lower()
+    vae_class = "AutoencoderKLFlux2" if texture_vae_name == "flux2" else "AutoencoderKL"
+    model_index = {
+        "_class_name": "SeFiPipeline",
+        "_diffusers_version": __version__,
+        "transformer": ["diffusers", "SeFiTransformer2DModel"],
+        "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+        "vae": ["diffusers", vae_class],
+        "text_encoder": ["transformers", "Qwen3VLForConditionalGeneration"],
+        "tokenizer": ["transformers", "Qwen2Tokenizer"],
+        "semantic_channels": int(model_config["semantic_channels"]),
+        "texture_vae_name": texture_vae_name,
+        "is_turbo": variant == "turbo",
+        "default_guidance_scale": float(inference_config.get("guidance_scale", default_guidance_scale(variant))),
+        "default_num_inference_steps": int(inference_config.get("steps", default_steps(variant))),
+        "delta_t": float(inference_config.get("delta_t", training_sefi_config.get("delta_t_max", 0.1))),
+        "timestep_shift_alpha": float(
+            inference_config.get("timestep_shift_alpha", 1.0 if variant == "turbo" else 0.3)
+        ),
+        "text_encoder_hidden_layers": [int(layer) for layer in model_config["text_encoder"]["hidden_layers"]],
+        "max_sequence_length": int(model_config["text_encoder"].get("max_length", 1024)),
+    }
+    save_json(output / "model_index.json", model_index)
+    shutil.copy2(root / "sefi_config.yaml", output / "sefi_config.yaml")
+    print(f"Saved SeFi-Image Diffusers checkpoint to {output}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -138,6 +138,7 @@ main.ignore-paths=[
     "pipelines/hdm",
     "pipelines/hidream",
     "pipelines/lumina_dimmo",
+    "pipelines/mageflow",
     "pipelines/meissonic",
     "pipelines/omnigen2",
     "pipelines/segmoe",
@@ -410,16 +411,16 @@ exclude = [
     "pipelines/flex2",
     "pipelines/hidream",
     "pipelines/lumina_dimmo",
+    "pipelines/mageflow",
     "pipelines/meissonic",
-    "pipelines/meissonic/",
     "pipelines/model_stablecascade.py",
-    "pipelines/omnigen2/",
-    "pipelines/sefi/",
-    "pipelines/step1x/",
-    "pipelines/ultraflux/",
-    "pipelines/vibe/",
-    "pipelines/xomni/",
-    "pipelines/zetachroma/",
+    "pipelines/omnigen2",
+    "pipelines/sefi",
+    "pipelines/step1x",
+    "pipelines/ultraflux",
+    "pipelines/vibe",
+    "pipelines/xomni",
+    "pipelines/zetachroma",
     "extensions-builtin/sd-extension-chainner/nodes",
 ]