Mage flow pipeline

[chenyangzhu1]

What does this PR do?

Add Mage-Flow text-to-image pipeline to diffusers. Mage-Flow is a 4B parameter model from Microsoft's Mage team, based on the NR-MMDiT (Native-Resolution Multimodal DiT) dual-stream architecture with Qwen3-VL text encoder and a custom MageVAE (128-channel latents, 16x downsampling). Supports any resolution from 512 to 2048 at arbitrary aspect ratios.

This PR adds:

• MageFlowTransformer2DModel — dual-stream MMDiT with complex RoPE and joint attention
• AutoencoderMageVAE — DConv encoder + CoD decoder with adaLN constant folding
• MageFlowPipeline — T2I pipeline with Qwen3-VL text encoding and FlowMatchEuler scheduling
• Weight conversion script, tests (70 passed / 0 failed), and API documentation

Usage

import torch
from diffusers import MageFlowPipeline, MageFlowTransformer2DModel, FlowMatchEulerDiscreteScheduler
from diffusers.models.autoencoders.autoencoder_mage_vae import AutoencoderMageVAE
from transformers import Qwen3VLForConditionalGeneration, AutoTokenizer
from safetensors.torch import load_file

model_path = "path/to/converted/model"

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(f"{model_path}/scheduler")
transformer = MageFlowTransformer2DModel(
    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,
)
transformer.load_state_dict(load_file(f"{model_path}/transformer/diffusion_pytorch_model.safetensors"))
transformer = transformer.to(torch.bfloat16).cuda()

vae = AutoencoderMageVAE()
vae.load_state_dict(load_file(f"{model_path}/vae/diffusion_pytorch_model.safetensors"))
vae = vae.to(torch.bfloat16).cuda()
vae.freeze_adaln()

text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
    f"{model_path}/text_encoder", torch_dtype=torch.bfloat16
).cuda()
tokenizer = AutoTokenizer.from_pretrained(f"{model_path}/text_encoder")

pipe = MageFlowPipeline(
    scheduler=scheduler, vae=vae, text_encoder=text_encoder,
    tokenizer=tokenizer, transformer=transformer,
)

image = pipe(
    prompt="A cat sitting on a windowsill, sunlight streaming in",
    height=1024, width=1024,
    num_inference_steps=30,
    guidance_scale=5.0,
).images[0]
image.save("mage_flow_output.png")

Quality Verification

All images 1024x1024, 30 steps, CFG=5.0. The original pipeline uses Gaussian-Shading watermark noise and packed varlen text encoding, so exact pixel-level reproduction is not expected — but visual quality and prompt adherence are equivalent.

Mountain Landscape (seed=88)

Original Pipeline	Diffusers Pipeline

Cat on Windowsill (seed=123)

Original Pipeline	Diffusers Pipeline

City Skyline (seed=77)

Original Pipeline	Diffusers Pipeline

Numerical Parity

Metric	Value
Transformer forward max_diff (bf16, GPU)	0.000000
Transformer forward cos_sim (bf16, GPU)	0.99999994
Weight key match (transformer)	397 / 397
Weight key match (VAE)	728 / 728

Note on image differences: Images are not pixel-identical because (1) the original pipeline uses Gaussian-Shading watermark noise while diffusers uses standard randn, (2) the original uses packed varlen text encoding while diffusers uses standard padded-batch encoding. Both produce equivalent visual quality.

Implementation Details

Transformer (MageFlowTransformer2DModel)

• NR-MMDiT dual-stream architecture with 12 dual-stream blocks
• Complex RoPE with symmetric positive/negative frequency scaling, stored as real-valued registered buffers for correct model.to(device) behavior
• Joint attention via dispatch_attention_fn (non-packed [txt, img] concatenation)
• Custom sinusoidal timestep embedding that downcasts the frequency table to input dtype (bf16) before computing — the model was trained with this exact rounding, using float32 frequencies produces visibly degraded output
• latent_height/latent_width forward params to avoid GPU-CPU sync in RoPE (torch.compile compatible)

VAE (AutoencoderMageVAE)

• DConv encoder (one-step diffusion) + CoD decoder with DConv denoiser
• 128-channel latents, 16x spatial downsampling
• adaLN constant folding via freeze_adaln() for inference
• F.scaled_dot_product_attention in CoD decoder attention blocks
• forward() delegates to decode() for CPU offload hook compatibility

Pipeline (MageFlowPipeline)

• Qwen3-VL (Qwen3VLForConditionalGeneration) text encoder with ChatML prompt template
• FlowMatchEulerDiscreteScheduler(shift=6.0) with custom base sigmas linspace(1, 1/N, N)
• Standard two-pass CFG with graceful fallback when text_encoder unavailable
• model_cpu_offload_seq = "text_encoder->transformer->vae"

Weight Conversion

• scripts/convert_mage_flow_to_diffusers.py
• Transformer: img_in → x_embedder, txt_norm → context_embedder_norm, txt_in → context_embedder
• VAE: student.dconv_encoder.* → encoder.*, pipeline.* → decoder.* (excluding y_embedder.encoder.*)

Test Results

73 passed, 0 failed, 3 skipped (torch 2.7.1+cu126)

The 3 skipped tests:

• 1x test_save_load_optional_components: Pipeline has no optional components (correct skip)
• 2x test_keep_in_fp32_modules: Mage-Flow does not need _keep_in_fp32_modules (correct skip)

Self-Review

Ran .ai/skills/self-review against .ai/review-rules.md, .ai/models.md, .ai/pipelines.md, .ai/testing.md, and .ai/skills/model-integration/pitfalls.md.

Blocking issues found and resolved

#	Issue	Resolution
B1	Transformer declared encoder_hidden_states_mask param but never used it	Removed from signature and pipeline calls
B2	Transformer declared unused txt_ids param	Removed
B3	Pipeline inlined custom sigma schedule	Added comment explaining why it differs from scheduler default
B4	Dead temb + torch.zeros(...) no-op	Removed
B5	.item() calls in RoPE break torch.compile	Added latent_height/latent_width params
B6	Hardcoded developer paths in conversion script	Replaced with placeholders
B7	Pipeline docstring referenced AutoencoderKL	Fixed to AutoencoderMageVAE
B8	Tests imported via internal module paths	Changed to public API

Non-blocking issues found and resolved

#	Issue	Resolution
NB1	PosEmbed stored complex tensors as plain attributes	Split into real/imag registered buffers
NB2	_MageVAEYEmbedder wrapper undocumented	Added docstring for weight-key compatibility
NB3	Denoiser dispatched on channel count	Changed to explicit is_latent parameter
NB4	VAE missing _no_split_modules	Added class attribute
NB5	encode() lacked generator for reproducibility	Added generator param with randn_tensor
NB6	Inline CFG	Consistent with Flux/QwenImage, kept as-is
NB9	MageVAEAttnBlock used torch.bmm	Replaced with F.scaled_dot_product_attention

Dead code removed

Location	What	Why dead
Transformer	temb + zeros block	No-op (adding zero)
Transformer	scale_rope=False branch	scale_rope always True
Transformer	encoder_hidden_states_mask / txt_ids params	Never read
VAE	MageVAEDecoder.ada = nn.Identity()	No-op identity layer
VAE	Duplicate forward() method	Superseded by decode-delegating forward

Files Changed

New Files

• src/diffusers/models/transformers/transformer_mage_flow.py
• src/diffusers/models/autoencoders/autoencoder_mage_vae.py
• src/diffusers/pipelines/mage_flow/pipeline_mage_flow.py
• src/diffusers/pipelines/mage_flow/pipeline_output.py
• src/diffusers/pipelines/mage_flow/__init__.py
• scripts/convert_mage_flow_to_diffusers.py
• tests/pipelines/mage_flow/test_pipeline_mage_flow.py
• tests/models/transformers/test_models_transformer_mage_flow.py
• tests/pipelines/mage_flow/__init__.py
• docs/source/en/api/pipelines/mage_flow.md
• docs/source/en/api/models/mage_flow_transformer2d.md

Modified Files

• src/diffusers/__init__.py — Register new components
• src/diffusers/models/__init__.py — Register models
• src/diffusers/models/transformers/__init__.py
• src/diffusers/models/autoencoders/__init__.py
• src/diffusers/pipelines/__init__.py — Register pipeline
• docs/source/en/_toctree.yml — Add doc entries

Fixes #14270

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@sayakpaul @yiyixuxu @DN6 @dg845 @asomoza

[Xinjie-Q]

We have reviewed the PR and confirmed that the implementation is correct. The Mage-Flow integration looks good, including the transformer, VAE, pipeline, weight conversion, tests. From our side, the PR is ready and can be merged into Diffusers. Thanks for the contribution!

[sergereview]

🤗 Serge says:

Thorough port, but there are several correctness gaps and repo-convention violations that need fixing before merge.

Correctness

• prompt_embeds_mask is computed but never used. encode_prompt returns a mask, check_inputs warns loudly about providing it, but the denoising loop never passes any attention_mask to the transformer — MageFlowAttnProcessor accepts attention_mask yet the pipeline calls self.transformer(...) without it. With batched prompts of differing lengths the padded text tokens will participate in joint attention. Either wire the mask through (joint_attention_kwargs/an explicit arg) or drop the mask plumbing and the warnings entirely (the repo rules forbid unused code paths).
• prompt_embeds_mask reshape is wrong. prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) is applied to a 2-D [B, S] tensor — repeat with 3 dims will unsqueeze to [1, B, S*num_images] and then the .view(batch*num_images, seq_len) only accidentally works for num_images_per_prompt == 1. QwenImage has the same latent bug but it is guarded differently; here it will raise/mis-shape for num_images_per_prompt > 1.
• negative_prompt_embeds_mask is discarded. In __call__ the negative branch assigns into local negative_prompt_embeds_mask but it is never used afterwards, same as above.
• VAE encode is dead code and off-convention. AutoencoderMageVAE.encode returns a bare tensor (all other autoencoders return EncoderOutput/AutoencoderKLOutput with return_dict, and are decorated with @apply_forward_hook), and no caller in the repo uses it. forward also just forwards to decode instead of the encode→decode round trip every other autoencoder implements — this will silently break generic tooling/tests that assume the standard VAE interface.
• freeze_adaln() is a manual, mutating post-load step. The docstring in the PR description requires users to call vae.freeze_adaln() after loading. This mutates module structure (replacing adaLN_modulation with a buffer holder), which breaks save_pretrained/reload round trips and offloading/device_map paths, and it is never called from the pipeline. Either fold this into the model at construction time based on config, or drop it.
• MageFlowPosEmbed.forward ignores img_ids except for .device, and the grid size is recovered via .item() on img_ids in the model forward — which is exactly why three TorchCompileTesterMixin tests had to be disabled. Passing latent height/width down from the pipeline (which already computes latent_h/latent_w) would remove both the host sync and the disabled tests.
• Hard-coded 4096 RoPE tables cap resolution silently: _compute_video_freqs slices freqs_pos[1][: height // 2], so heights/widths above 8192 latent units silently truncate instead of erroring. Also these non-persistent buffers inflate the tiny test model enough that test_model_parallelism had to be disabled — a symptom worth fixing rather than skipping.
• _apply_rotary_emb_complex forces fp32 complex math per block per step; fine functionally, but note x_pass is None-checked via a branch that the rules discourage — minor.

Repo conventions / release plumbing

• Missing dummy objects. MageFlowTransformer2DModel/AutoencoderMageVAE are not in utils/dummy_pt_objects.py and MageFlowPipeline/MageFlowPipelineOutput are not in utils/dummy_torch_and_transformers_objects.py. make fix-copies / python utils/check_dummies.py will fail CI.
• Alphabetical ordering broken in src/diffusers/__init__.py: MageFlowPipeline/MageFlowPipelineOutput are inserted before Lumina2Pipeline in both the _import_structure list and the TYPE_CHECKING block.
• AutoencoderMageVAE has no doc page but is exported from the public init — utils/check_repo.py::check_all_objects_are_documented will flag it (and there is no _toctree.yml entry for it).
• Missing AutoencoderMageVAE model tests. .ai/testing.md requires model-level tests for every new model class; only the transformer has one. No tests/models/autoencoders/test_models_autoencoder_mage_vae.py was added, contrary to the description's "tests (70 passed / 0 failed)" claim covering the new components.
• pretrained_model_name_or_path left as "" with a TODO in the transformer tester config, and the docstring example points at microsoft/Mage-Flow-4B which is not referenced anywhere else — the PR ships no usable checkpoint reference.
• Typing/import nits: pipeline_mage_flow.py imports numpy only for the np.linspace sigma default; transformer_mage_flow.py uses a mutable default axes_dim: list[int] = [16, 48, 48]. Docstring defaults (context_in_dim=3584, num_layers=32) disagree with the conversion script's real config (2560, 12).
• Import order in the pipeline test (Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen2Tokenizer) isn't isort-clean; make style hasn't been run.
• Conversion script style: key[len(VAE_ENCODER_PREFIX):] slices will be reformatted by ruff (E203-style spacing), and _diffusers_version: "0.37.0" is hard-coded rather than read from diffusers.__version__.

Minor

• docs/source/en/api/pipelines/mage_flow.md uses .md suffixes in doc links (../../using-diffusers/schedulers.md); the prevailing convention in this folder is extension-less links, and it uses <Tip> where newer pages use > blockquotes.
• MageFlowAttnProcessor.__init__ raising on missing F.scaled_dot_product_attention is defensive dead code — torch>=2 is a hard requirement.
• MageVAENerfEmbedder._pos_cache is an unbounded dict keyed by (patch_size, device, dtype); harmless in practice but it's mutable state on a module that won't move with .to().

serge v0.1.0 · model: claude-opus-5 · 24 LLM turns · 36 tool calls · 1069.6s · 2205133 in / 82286 out tokens

[sergereview]

prompt_embeds_mask is 2-D ([batch, seq_len]), so repeat(1, num_images_per_prompt, 1) prepends a broadcast dim and produces [1, batch, seq_len * num_images_per_prompt]. The following .view(batch_size * num_images_per_prompt, seq_len) then only happens to work when num_images_per_prompt == 1 and will mis-shape (or raise) otherwise.

Suggested change

	prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
	prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt)

[sergereview]

The masks returned by encode_prompt (prompt_embeds_mask, negative_prompt_embeds_mask) are never forwarded to the transformer, so padded text tokens fully participate in joint attention whenever a batch contains prompts of different lengths. MageFlowAttnProcessor already accepts attention_mask, and check_inputs warns users about supplying the mask — but nothing consumes it.

Either thread the mask through to the attention processor, or remove the mask plumbing and the associated check_inputs warnings entirely (the project rules disallow unused code paths).

[sergereview]

encode is not called anywhere in the repo (the pipeline is T2I-only and uses self.vae(...) for decode), and it deviates from the autoencoder contract used everywhere else: no return_dict, no EncoderOutput/AutoencoderKLOutput wrapper, and no @apply_forward_hook. Either bring it in line with e.g. AutoencoderRAE.encode (return EncoderOutput, support return_dict, add the forward hook) or drop it for the initial integration.

Relatedly, forward on line 709 just delegates to decode, whereas every other autoencoder's forward does encode→decode. Generic VAE tooling/tests that assume the standard interface will silently do the wrong thing here.

[sergereview]

freeze_adaln() mutates the module tree after load (swapping adaLN_modulation MLPs for a buffer holder). That breaks save_pretrained/from_pretrained round trips (the folded state dict no longer matches the class-constructed module), and interacts badly with device_map/offloading hooks. It's also never called from MageFlowPipeline, so users must remember the manual step documented in the PR description.

Please either perform the folding at construction time driven by config, or drop it — a post-hoc structural mutation isn't something the rest of the library supports.

[sergereview]

forward ignores img_ids entirely except for .device, taking the grid from the height/width args instead. Combined with _infer_grid_size calling .item() at line 601, this is what forces three torch.compile tests to be disabled and adds a host sync every step.

The pipeline already knows latent_h/latent_w — pass them through as explicit arguments (like Flux's img_shapes) and drop img_ids/_infer_grid_size rather than round-tripping the grid through a tensor and back.

[sergereview]

pretrained_model_name_or_path is still an unfilled TODO, which means any tester relying on it (e.g. the from_pretrained-based paths) can't actually load anything. Please point this at a tiny repo (a personal one is fine for review; a maintainer can move it under hf-internal-testing/ later) — .ai/testing.md calls this out explicitly.

Separately: .ai/testing.md requires model-level tests for every new model class, and AutoencoderMageVAE has none (tests/models/autoencoders/test_models_autoencoder_mage_vae.py is missing). The description's claim of full test coverage isn't supported by the diff.

[sergereview]

Disabling test_model_parallelism because the fixed 4096-entry RoPE buffers dominate the tiny model is treating a design issue as a test problem. Those buffers are non-persistent and sized independently of the config; computing frequencies on demand for the requested grid (or sizing the table from config) would both fix the parallelism split and remove the compile-time .item() graph break flagged below.

[sergereview]

Import order isn't isort-clean (Qwen2Tokenizer should precede Qwen3VLConfig), which suggests make style wasn't run before opening the PR.

Suggested change

	from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen2Tokenizer
	from transformers import Qwen2Tokenizer, Qwen3VLConfig, Qwen3VLForConditionalGeneration

[sergereview]

Doc links in this folder are written without the .md extension (../../using-diffusers/schedulers, ../../using-diffusers/loading#reuse-a-pipeline) — the .md suffix produces broken links in the built docs. Newer pages also use a > blockquote rather than <Tip>.

[sergereview]

_diffusers_version is hard-coded here and in VAE_DIFFUSERS_CONFIG/MODEL_INDEX; read it from diffusers.__version__ instead so converted repos don't claim a stale version.

Also key[len(VAE_ENCODER_PREFIX):] (lines 105/111) will be reformatted by ruff — please run make style over this script.

[alerikaisattera]

Is edit pipeline included?

[yiyixuxu]

thanks so much for the PR @chenyangzhu1 !

Could you convert the pipeline to Modular Diffusers?

We are currently trying to do modular-only releases whenever possible — as long as we're not under release time pressure, we prefer new integrations to be released in modular pipeline from the start (it should be a pretty straightforward task for your AI agent if you're using one! )

[alvarobartt]

Note that you should also include this in https://github.com/alvarobartt/diffusers/blob/a8c1c08f42f321ca853f8f4349ba3ede0ca4358b/src/diffusers/pipelines/auto_pipeline.py#L138-L191 so that the AutoPipelineForText2Image can be leveraged for this model; it's really just about adding this entry:

        ("mage-flow", MageFlowPipeline),

cc @asomoza for visibility