Mage-Flow

Mage-Flow (diffusers)

Mage-Flow's RL-aligned 20-step checkpoint packaged as a diffusers pipeline, so it loads from one path and splits into an encode half and a decode half.

An unofficial repackaging by an individual — not affiliated with, endorsed by, or released by Microsoft, Comfy-Org or the Qwen team. The weights are theirs, redistributed under their licenses; what is added here is the diffusers surface.

import torch
from diffusers import DiffusionPipeline

repo = "bigshanedogg/Mage-Flow"
pipe = DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True,
                                         torch_dtype=torch.bfloat16).to("cuda")
image = pipe(prompt="a photograph of a red fox sitting in snow", num_inference_steps=20,
             guidance_scale=5.0, height=1024, width=1024, seed=1234).images[0]

The two halves can also run apart — encode once, denoise elsewhere, or swap the text encoder:

conditioning = pipe.encode_prompt(prompt="...", negative_prompt=" ")
image = pipe(prompt_embeds=conditioning["prompt_embeds"],
             prompt_embeds_mask=conditioning["prompt_embeds_mask"],
             negative_prompt_embeds=conditioning["negative_prompt_embeds"],
             negative_prompt_embeds_mask=conditioning["negative_prompt_embeds_mask"],
             num_inference_steps=20, guidance_scale=5.0, seed=1234).images[0]

encode_prompt returns padded (batch, sequence, 2560) conditioning plus the mask marking what it padded; upstream keeps it packed end to end, and the mask is what lets it cross a process boundary.

Contents

Provenance

• Code: microsoft/Mage (mage_flow, MIT) at commit 76bec2bb3818, vendored with a diffusers wrapper. The packing, velocity/CFG combination, scheduler loop, noise construction and VAE are upstream's own functions rather than a reimplementation.
• Weights: Comfy-Org/Mage-Flow (MIT), which repackages microsoft's release as single files. The text encoder there is bit-identical to Qwen/Qwen3-VL-4B-Instruct (Mage-Flow keeps it frozen), and is shipped here so the checkpoint loads from one path.
• transformer/config.json is derived from the weights themselves — hidden_size, in_channels, out_channels and context_in_dim from the projection shapes, depth from the block count, num_heads from hidden_size / head_dim where head_dim is the q-norm width. The one value no weight carries is the RoPE axis split axes_dim; that value was read off ComfyUI's public Mage-Flow configuration — a hyperparameter, no code taken from it — and cross-checked to sum to head_dim.

Defaults

num_inference_steps follows the released recommendation for this variant — the model card and the technical report both state 30 steps for Base, 20 for the RL-aligned line and 4 for Turbo.

guidance_scale = 5.0 is upstream's implementation default (generate_images(..., cfg=5.0) and the CLI's --cfg), not a published per-variant recommendation: the report names 5.0 only for its own candidate-sampling procedure and 7.5 for the distillation objective, and gives no inference table. So treat it as a starting point, and for the distilled Turbo line in particular it is worth checking whether guidance is needed at all — above 1.0 it doubles the forward passes.

Licenses

Two licenses apply, and both texts ship with the repository:

Both permit commercial use. Keeping these two files alongside the weights is what satisfies their attribution requirements; the encoder weights are byte-identical to the upstream release, so there are no modifications to state.

Upstream adds a statement of intent that the license itself does not impose — microsoft/Mage's README says its models "are released for research purposes only and are not intended for product or service deployment", and asks that downstream users add their own safeguards. Worth reading before shipping anything built on this.

Equivalence with the reference implementation

Running upstream's own generate_images and this pipeline's __call__ over the same component instances gives bit-identical images (same prompt, seed, steps, guidance).

Loading the same weights twice is a different matter: two loads of the Qwen3-VL encoder — even both by upstream's own loader, from the same files with the same arguments — differ by about 1% of the pixel range (MAE ≈ 2.4/255, PSNR ≈ 32 dB) because the numerics depend on the load context. So expect image-level agreement rather than bitwise agreement when comparing across separate loads.

Notes

• attn_type defaults to sdpa, and the vendored attention shim falls back to it when flash-attn 2 is unavailable — required on pre-Ampere cards, where flash-attn 2 has no kernels.
• Upstream pins transformers>=5.3,<5.6; on newer versions its patched Qwen3-VL forward calls create_causal_mask(input_embeds=..., cache_position=...), which was renamed and trimmed. The translation is included here, so no caller-side shim is needed.
• The initial noise carries upstream's distribution-preserving watermark (Gaussian Shading, key MAGEFLOW_GS_KEY / default), kept so provenance detection still works.
• Upstream's mandatory content-policy gate is not part of this port. generate_images there runs an LLM-as-judge on the same encoder and returns a placeholder image for anything it flags; that path is absent here, so prompts are not screened. Apply moderation at your own layer if outputs are published.
• diffusers.AutoModel cannot load the components on its own; load through the pipeline (or import the classes from the component files directly).