Views
No views yet
1import torch
2from diffusers import DiffusionPipeline
3
4repo = "bigshanedogg/Mage-Flow-Base"
5pipe = DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True,
6 torch_dtype=torch.bfloat16).to("cuda")
7image = pipe(prompt="a photograph of a red fox sitting in snow", num_inference_steps=30,
8 guidance_scale=5.0, height=1024, width=1024, seed=1234).images[0]1conditioning = pipe.encode_prompt(prompt="...", negative_prompt=" ")
2image = pipe(prompt_embeds=conditioning["prompt_embeds"],
3 prompt_embeds_mask=conditioning["prompt_embeds_mask"],
4 negative_prompt_embeds=conditioning["negative_prompt_embeds"],
5 negative_prompt_embeds_mask=conditioning["negative_prompt_embeds_mask"],
6 num_inference_steps=30, 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.| Path | What |
|---|---|
transformer/ | 4B NR-MMDiT (MageFlowTransformer2DModel), 12 blocks, hidden 3072, 24 heads |
vae/ | Mage-VAE (AutoencoderMageVAE), one-step diffusion tokenizer, 128 latent channels, /16 |
text_encoder/, tokenizer/ | Qwen3-VL-4B, read at its last hidden state (2560) |
scheduler/ | FlowMatchEulerDiscreteScheduler, shift 6.0 |
pipeline.py + component modules | Remote code: upstream's math with a diffusers surface |
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.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.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.| Part | License |
|---|---|
transformer/, vae/, scheduler/, pipeline.py and the component modules (Mage-Flow weights and vendored code) | MIT — LICENSE, (c) 2026 Microsoft |
text_encoder/, tokenizer/ (Qwen3-VL-4B-Instruct, redistributed unmodified) | Apache-2.0 — LICENSE.qwen3-vl-apache-2.0 |
generate_images and this pipeline's __call__ over the same component
instances gives bit-identical images (same prompt, seed, steps, guidance).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.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.MAGEFLOW_GS_KEY / default), kept so provenance detection still works.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).