MageTrail-Shared-AdaLN-SVD

MageTrail — Shared AdaLN SVD, Rank 256

An experimental rank-256 compression of RicemanT/MageTrail, a full finetune of Mage-Flow. This release reduces its AdaLN parameter count to make full finetuning more accessible. Built with mage-flow-trainer.

Shared low-rank conditioning, separate per-block AdaLN heads: a single 3072→256 projection feeds independent image/text modulation heads in all 12 blocks. The blocks retain their own modulation; they do not all share identical AdaLN outputs.

The conversion uses calibrated timestep features and FP64 QR-based SVD (gesvd). This checkpoint is the SVD-converted model, without subsequent image-training or gradient-based distillation. Saved modulation factors are FP32; the remaining transformer weights are BF16.

Visual comparison

Click either image to view it at full resolution.

Using it in ComfyUI

1. Install ComfyUI-MageFlow-Compressed and restart ComfyUI.
2. Put rank256-fp64-qr.safetensors in ComfyUI/models/diffusion_models/.
3. Select it in Load Compressed Mage-Flow and connect its MODEL output to your sampler.
4. Use the usual Mage-Flow text encoder and VAE: Qwen3-VL 4B with Load CLIP → type: mage, and the Mage-Flow VAE.

For a starting workflow, use Empty Flux 2 Latent, Euler, 32 steps, CFG 4, and 1024×1024. The loader supplies Mage-Flow's shift 6.

This is a complete compressed transformer checkpoint, not a LoRA or adapter. It includes the trunk weights and replaces the original block modulation projections. This repository contains the transformer only. It requires the dedicated loader; it is not a drop-in checkpoint for the standard diffusion-model loader. The custom node does not modify ComfyUI's built-in loaders or model detection.

Finetuning

Use the compressed-checkpoint support in mage-flow-trainer. The shared projection and separate heads remain trainable. Their weights and gradients default to FP32; train.compressed_adaln_dtype = "bfloat16" opts into BF16. SDNQ can quantize the supported trunk weights while leaving these factors unquantized.

A 10-step rank-256 full-finetuning smoke test with INT8 SDNQ and quantized optimizer state peaked at approximately 16.08 GiB allocated VRAM, without CPU optimizer offload. This is configuration-specific, not a general memory requirement or a demonstration of long-run stability.

Status and credits

Experimental. Held-out prediction error against the original transformer was approximately 1.18% in the local calibration evaluation; this does not establish identical image quality. ComfyUI loading, offloading, and a complete 1024px sampling run were tested. Evaluate your own prompts and finetuning results.

Credits to RicemanT for the MageTrail finetune and Microsoft for the original Mage project. MageTrail is distributed under MIT; this derivative retains that license. The separate custom node is GPL-3.0-or-later.