arxiv:2607.19064

Mage-Flow: An Efficient Native-Resolution Foundation Model for Image Generation and Editing

Published on Jul 21

· Submitted by

Xinjie on Jul 22

Microsoft

Upvote

Authors:

Yifei Shen ,

Yuxuan Luo ,

Fanyi Pu ,

Tianci Bi ,

Abstract

Large-scale visual generators are increasingly capable but costly to train, fine-tune, and deploy. We introduce Mage-Flow, a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing. The stack is built from two co-designed components: Mage-VAE, a lightweight high-fidelity latent tokenizer, and a Native-Resolution Multimodal Diffusion Transformer trained with rectified flow matching. Mage-VAE uses one-step diffusion-style encoding and decoding with anchor-latent regularization, preserving the reconstruction quality of strong public VAEs while reducing tokenization cost by more than an order of magnitude. Together with native-resolution packing and stack-level CUDA kernel fusion, the stack supports flexible-resolution training and improves end-to-end training throughput by about 2.5times. Built on this foundation, we develop a complete model family with Base, RL-aligned, and Turbo variants for both generation and editing. Diffusion-NFT improves prompt following, text rendering, aesthetic quality, and editing fidelity, while few-step distillation with adversarial perceptual guidance produces 4-step Turbo models for low-latency inference. Despite its compact scale, Mage-Flow and Mage-Flow-Edit achieves competitive performance across standard generation and editing benchmarks. More importantly, the Turbo variants make high-resolution generation and editing practical for interactive use: at 1024^2 resolution on a single NVIDIA A100 GPU, Mage-Flow-Turbo generates an image in 0.59s, and Mage-Flow-Edit-Turbo edits an image in 1.02s, while maintaining a small memory footprint. These results show that careful tokenizer--backbone--system co-design can deliver strong high-resolution generation and editing within an efficient 4B model family.

View arXiv page View PDF Project page GitHub 189 Add to collection

Community

Xinjie-Q

Paper submitter 1 day ago

•

edited about 22 hours ago

Mage-Flow is a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing. Instead of scaling to tens of billions of parameters, Mage-Flow reaches state-of-the-art-competitive quality through careful tokenizer–backbone–system co-design, so it stays fast, memory-light, and easy to fine-tune under realistic compute budgets.

The stack is built from two shared, co-designed components:

Mage-VAE — a lightweight, high-fidelity latent tokenizer (one-step diffusion encode/decode with anchor-latent KL regularization).
NR-MMDiT — a shared 4B Native-Resolution Multimodal Diffusion Transformer, trained with rectified flow matching in the Mage-VAE latent space.

Together with native-resolution packing and a fused-kernel training infrastructure, this shared stack powers two model instantiations: Mage-Flow for text-to-image generation and Mage-Flow-Edit for instruction-based image editing. Each ships in Base, RL-aligned, and 4-step Turbo variants.

✨ Highlights

Compact & competitive. A single 4B family for generation and editing that matches or beats much larger open systems (Qwen-Image 20B, Z-Image 6B, FLUX.2 32B, FireRed-Image-Edit 20B).
Efficient tokenizer. Mage-VAE matches FLUX.2-VAE reconstruction fidelity while using ~12× / ~22× fewer encode / decode MACs per pixel, removing the VAE as the high-resolution bottleneck.
Native resolution. One checkpoint generates from 512 to 2048 on any aspect ratio, including extreme 4:1 (e.g. 512×2048, 2048×512).
System-level speed. Native-resolution packing (FlashAttention var-len + per-sample 2D RoPE) + fused CUDA kernels achieve ~2.5× faster training; CFG's conditional/unconditional branches run in one packed forward.
Full family. Base, RL-aligned, and 4-step Turbo variants for both generation and editing.
Versatile editing. Mage-Flow-Edit supports semantic content editing, appearance transformation, image restoration, and structure-aware outputs within a unified image-and-text-conditioned model. See the report's editing galleries.
Interactive latency. At 1024² on a single A100: Mage-Flow-Turbo 0.59 s/image, Mage-Flow-Edit-Turbo 1.02 s/edit, peak memory ~18–20 GB (lowest among compared systems).