SearchingMan/Mage-Flow-W4A8-ARA

Mage-Flow W4A8 + ARA

A 4-bit build of microsoft/Mage-Flow that runs on ComfyUI's native quantized kernels, plus a small accuracy recovery adapter (ARA) that brings the 4-bit model back toward the BF16 original. Optionally, an INT8 text encoder cuts another ~3.7 GB of VRAM.

	Transformer	Text encoder	LPIPS ↓	PSNR ↑	Peak VRAM*	s / image
Mage-Flow BF16 (reference)	8.23 GB	BF16	–	–		16.1
Comfy-Org INT8 ConvRot	4.16 GB	BF16	0.037	26.5 dB		9.0
W4A8 without ARA	2.97 GB	BF16	0.194	18.8 dB		9.0
W4A8 + ARA	2.97 GB + 99 MB	BF16	0.117	21.5 dB	12.6 GB	10.0
W4A8 + ARA + INT8 encoder	2.97 GB + 99 MB	INT8 (this repo)	0.116	21.6 dB	8.9 GB	10.3

LPIPS (AlexNet) and PSNR against the BF16 model with the BF16 text encoder, same prompt and seed. ComfyUI, 1024 px class resolutions, 30 steps, 24 held-out prompts (12 short, 12 long) written after every choice about the adapter and the encoder was made. RTX 5090 Laptop GPU (24 GB). * ComfyUI's share of GPU memory while it keeps the text encoder and the transformer loaded together; the desktop's ~1.9 GB is excluded. Details below.

Pick this if you want Mage-Flow in the smallest transformer that runs on stock ComfyUI. Pick Comfy-Org INT8 if you have the extra 1.1 GB: it is much closer to BF16.

Examples

Left to right: BF16 · Comfy-Org INT8 · W4A8 + ARA with the BF16 text encoder · W4A8 + ARA with the INT8 text encoder. Same prompt, seed and sampler settings. The last two columns differ only in the text encoder.

Long, detailed prompts:

10 more

The 4-bit model on its own does not produce broken images. It renders a different image: other composition details, tone and contrast. The ARA pulls the result back toward what BF16 would draw.

Two ways to run it

Both use the same 4-bit model and the same ARA. They differ only in the text encoder file you put in the CLIPLoader node.

	Option 1: BF16 text encoder	Option 2: INT8 text encoder
Text encoder file	qwen3vl_4b_bf16.safetensors (Comfy-Org, 8.9 GB)	qwen3vl_4b_int8_convrot.safetensors (this repo, 4.9 GB)
Peak VRAM*	12.6 GB	8.9 GB
LPIPS, 24 held-out prompts	0.117	0.116
short prompts / long prompts	0.094 / 0.141	0.091 / 0.140
Choose it if	you want the encoder exactly as Mage-Flow ships it	you want the lower VRAM

The shipped workflow uses Option 2. To switch, pick the other file in CLIPLoader.

What the INT8 encoder costs, honestly. On its own it is not lossless: with the BF16 transformer it moves the image by LPIPS 0.043, about as much as Comfy-Org's INT8 transformer does (0.037). On top of the 4-bit model the difference mostly disappears: 0.116 vs 0.117 on the 24 held-out prompts. Single images still change in small ways, in either direction (per prompt from −0.038 to +0.024). On the 12 prompts used to choose the adapter, 8 of which are long, detailed descriptions, it cost more: 0.132 vs 0.119. Compare the last two columns of the examples above and decide for yourself.

Comfy-Org's FP8 encoder (qwen3vl_4b_fp8_scaled.safetensors, 5.2 GB) also works in the same node, but it loses more than INT8 (measured with an earlier version of the adapter: LPIPS 0.116 vs 0.100 with INT8 on 16 held-out prompts).

Files

File	Put it in
diffusion_models/mage_flow_w4a8.safetensors	ComfyUI/models/diffusion_models/	the 4-bit model (2.97 GB)
loras/mage_flow_w4a8_ara.safetensors	ComfyUI/models/loras/	the ARA (99 MB)
text_encoders/qwen3vl_4b_int8_convrot.safetensors	ComfyUI/models/text_encoders/	optional INT8 text encoder (4.9 GB, Apache-2.0, see NOTICE)
workflows/Mage_Flow_W4A8_ARA.json	drag into ComfyUI	text-to-image workflow
workflows/Mage_Flow_W4A8_ARA_API.json		same workflow, API format
ComfyUI-MageFlow-ARA/	ComfyUI/custom_nodes/	the ARA loader node (also on GitHub)

From Comfy-Org:

• mage_flow_vae_bf16.safetensors → ComfyUI/models/vae/
• for Option 1: qwen3vl_4b_bf16.safetensors → ComfyUI/models/text_encoders/

Quick start

1. Update ComfyUI. You need native Mage-Flow support and the asym_w4a8_int8 format (comfy-kitchen ≥ 0.2.37).
2. Install the node:

   cd ComfyUI/custom_nodes
   git clone https://github.com/kgonia/ComfyUI-MageFlow-ARA
   

3. Download the files above and restart ComfyUI.
4. Drag workflows/Mage_Flow_W4A8_ARA.json into ComfyUI. For Option 1, select qwen3vl_4b_bf16.safetensors in CLIPLoader.

The graph is the official Mage-Flow text-to-image workflow with one extra node:

Load Diffusion Model (mage_flow_w4a8) → Load Mage-Flow ARA (residual) → KSampler

Settings: 30 steps, CFG 5.0, euler / simple, 512–2048 px per side in multiples of 16. Keep the ARA strength at 1.0.

Why a custom node and not LoraLoader?

The ARA runs next to the 4-bit layer: y = Q(W)·x + up(down(x)). Its correction is smaller than one 4-bit step. ComfyUI's LoraLoader merges a LoRA into the weight and requantizes it, so most of the correction rounds away. The node adds the correction at runtime and leaves the 4-bit weights and ComfyUI's quantized kernels untouched. It costs about 10% generation time.

The ARA file is an ordinary LoRA file, so LoraLoader will load it, but most of the benefit is lost. Measured with an earlier version of the adapter on 16 held-out prompts:

Loaded with	LPIPS ↓	PSNR ↑	s / image
(no ARA)	0.211	18.3 dB	9.6
LoraLoaderModelOnly (merge + requantize)	0.170	19.1 dB	12.6
Load Mage-Flow ARA (residual)	0.106	22.0 dB	11.0

The node was better on all 16 prompts. LoraLoader keeps about 40% of the improvement and is also slower, because the patched 4-bit weights have to be requantized.

What is in the checkpoint

Mage-Flow's transformer has 12 blocks. The checkpoint uses ComfyUI's own formats, so the base model runs on stock ComfyUI:

Layers	Count	Format
Block linears (attention q/k/v/out, MLPs)	144	asym_w4a8_int8: 16-level codebook, groups of 16, ConvRot 256, INT8 activations
Modulation linears (img_mod.1, txt_mod.1)	24	int8_tensorwise with ConvRot 256
Everything else		BF16

The modulation layers hold a third of the parameters and are very sensitive. Putting them in 4-bit too cost far more quality than the 0.6 GB it saved (LPIPS 0.40 vs 0.24 at 512 px), so they stay in INT8.

How the ARA was made

The ARA covers all 144 4-bit block linears: rank 32 in blocks 0–9 and rank 64 in blocks 10–11 (99 MB in BF16). The node reads the rank of each layer from the file.

1. Closed form, no training. For each layer, the 4-bit quantization error E = W − Q(W) (exact ComfyUI codes) and the covariance C of that layer's real inputs give the best possible rank-32 correction directly: the SVD of E·C^½. Computed block by block on the quantized model, so each block sees inputs already corrected upstream. This alone removes about half of the 4-bit error.
2. Fine-tune. 1,800 steps starting from that solution, distilling from the BF16 model on its own generations: 250 images at 512–1024 px, from 95 short and medium prompts plus 60 long, caption-style prompts. A quarter of the states use the empty prompt (the CFG branch), and noise levels are drawn from ComfyUI's Mage-Flow schedule. No external images or captions. The training backend reproduces ComfyUI's quantized layers bit for bit.
3. Last two blocks averaged with a second adapter. In blocks 10 and 11, the correction is the average of the fine-tuned adapter and an earlier ARA trained from scratch: ΔW = ½·up₁down₁ + ½·up₂down₂. The two adapters come from different runs, so their low-rank factors cannot be averaged directly. The average is stored exactly by concatenating them, which is why these 24 layers have rank 64.

Step 3 is a judgment call, and it costs a little on the metric: the fine-tuned adapter alone scores LPIPS 0.106 on the 24 held-out prompts, the averaged one 0.117. In side-by-side comparisons, the averaged adapter brought back details that BF16 draws and the fine-tuned one dropped, on a part of the prompts (object layout, a pose, a building's windows). We preferred that.

Quality

ComfyUI, 30 steps, CFG 5, euler / simple, one seed per prompt, the same seed for every variant. Measured against the BF16 output with the BF16 text encoder.

Held-out set: 24 prompts written after the adapter and the encoder were chosen. 12 short prompts (9–16 words) and 12 long, detailed descriptions (65–81 words); photos, illustration styles, text in the image, aspect ratios from 2:3 to 3:2.

Variant	LPIPS ↓	short	long	PSNR ↑	PSNR at 64 px ↑
Comfy-Org INT8	0.037	0.024	0.049	26.46	29.20
BF16 transformer + INT8 text encoder	0.043	0.028	0.059	26.09	28.54
W4A8 without ARA	0.194	0.183	0.205	18.77	20.49
W4A8 + ARA	0.117	0.094	0.141	21.46	23.60
W4A8 + ARA + INT8 text encoder	0.116	0.091	0.140	21.55	23.71

PSNR at 64 px compares downscaled images. It ignores texture and shows whether the composition matches: where the subject is, which way it faces, the big shapes and colors.

Earlier sets, for comparison. Both were used while making choices, so they are less neutral:

Variant	16 prompts (chose the adapter)	12 prompts (chose between variants)
Comfy-Org INT8	0.025	0.032
W4A8 without ARA	0.211	0.219
W4A8 + ARA	0.110	0.119
W4A8 + ARA + INT8 text encoder	0.110	0.132

What did not help, for the record: rank 64 everywhere (more weight error captured, no better images), training the ARA from scratch on its own (LPIPS 0.147–0.156 on the 12-prompt set), a closed form fitted to the end-to-end layer gap (diverged), and a second LoRA trained to compensate the FP8 text encoder (recovered only about 30% of its loss).

Notes

• Tested on an RTX 5090 Laptop GPU (24 GB) only. On smaller cards ComfyUI offloads the text encoder to system RAM between prompts, so peak VRAM is lower than in the table.
• ComfyUI has no native 4-bit × 4-bit matmul on RTX 50 cards. Files labelled "W4A4" run with 8-bit activations there anyway. This checkpoint is W4A8 by design.
• Without the node, mage_flow_w4a8.safetensors works on its own at the "W4A8 without ARA" quality above.
• LPIPS against BF16 measures fidelity to the original model, not image quality on its own.

License and credits

The 4-bit model, the ARA, the node and the workflows: MIT, same as Mage-Flow.

The INT8 text encoder is a quantized version of Qwen/Qwen3-VL-4B-Instruct by the Qwen team and stays under the Apache License 2.0. See text_encoders/NOTICE for what was changed and text_encoders/LICENSE-Qwen3-VL for the license.

• Base model: microsoft/Mage-Flow
• Text encoder: Qwen/Qwen3-VL-4B-Instruct
• ComfyUI packaging, text encoder and VAE files: Comfy-Org/Mage-Flow
• Quantized formats and kernels: ComfyUI / comfy-kitchen