MOSS-Image-heal: A Full Finetune That Makes fp16 Safe
KohakuBlueleaf's MOSS-Image-heal retrains MOSS-Image-Base so its residual stream stops spiking: fp16 no longer runs near the limit, and ComfyUI gets bf16, fp16 and int8 files.
Base and heal with the same prompt, seed, 16 steps, CFG 3 and shift 5, from the official model card.
The problem it trains away
MOSS-Image is a 5.08B single-stream diffusion transformer, and its residual stream carries a handful of very large channels. On the author's six-sample probe the largest stream value on the base model is 58,411 against an fp16 maximum of 65,504, so 1.12x of headroom. Anything that pushes activations a little harder, or any kernel that computes a slightly wider intermediate, overflows to inf, and the image comes back black or corrupt. That is why Comfy-Org's repack ships bf16 and int8 ConvRot files rather than an fp16 one.
Largest absolute stream value per block over six text-image samples at 1024² area. Chart from the model card.
| Measurement | MOSS-Image-Base | MOSS-Image-heal |
|---|---|---|
| Largest stream value | 58,411 | 3,083 |
| Headroom to the fp16 maximum | 1.12x | 21.2x |
| Stream RMS, mean over blocks | 27.75 | 3.70 |
| Stream RMS, last block | 65.6 | 12.0 |
| Energy share of the 32 largest channels | 0.837 | 0.708 |
The headroom comes from the same change: the model learned to spread what used to sit in 32 channels across more of them. Energy concentrating from 0.837 to 0.708 is the visible side of that, and it is what buys the fp16 space.
How it was trained
The run is a full finetune of every transformer weight except the frozen text refiner; the text encoder, tokenizer, VAE, scheduler and pipeline code are the official MOSS-Image-Base ones. Losses are flow matching plus a distillation term toward the frozen base, a stream write-concentration loss that penalises attention and FFN writes piling into few channels, and a stream RMS loss. Weights are separately decayed whenever they exceed a cap multiple of their starting RMS, with the cap tightening from 64 down to 16 over the first 10,000 steps.
| Steps | 50,000 x 64 samples, 3.2M samples |
| Resolution | 1024² area, each image at its own aspect ratio |
| Optimizer | AdamW, lr 1e-5 cosine to 1e-7 |
| Compute | About 200 GPU hours on B300, fp16 with a fp32 master copy |
| Data | 999,294 Danbooru posts, Graph-Based Captioning from Qwen3.8-27B and Qwen3.5-9B |
Weight outliers per 2-D tensor: 349 elements sat above 64x RMS in the base transformer and none do in heal. Chart from the model card.
Outputs stay where the base model was
Taming the stream is only useful if the model still draws like the base, and the card reports both numbers. Flow loss moves from 0.4280 to 0.4289, and the velocity cosine against the base on the probe rows is 0.9983. This is a numerical repair, not a style retrain.
The second grid from the card, portrait aspect, printed with both the caption prompt and the GBC layout prompt.
Files for ComfyUI
All three files are single-file checkpoints in Comfy-Org's key layout, so they load through the same Load Diffusion Model node as the official MOSS-Image repack. Nothing else changes: you keep the qwen3.5_9b text encoder and the MOSS VAE you already use.
| File | Weights | Size |
|---|---|---|
moss_image_heal_bf16.safetensors | bf16 | 10.2 GB |
moss_image_heal_fp16.safetensors | fp16 | 10.2 GB |
moss_image_heal_int8_convrot.safetensors | int8 ConvRot, group 256 | 5.4 GB |
The int8 file is written to match Comfy-Org's moss_image_base_int8_convrot layout exactly: same 757 keys, every non-quantized tensor identical, and 5e-8 of the quantization codes off by one. So it drops into the same slot as the official int8 checkpoint at half the disk.
On the same probe rows, fp16 heal reaches a velocity error of 2.1e-3 with a cosine of 0.999996, and peaks at 3,278 in the stream, 20x below the fp16 maximum. The base model on those rows peaks at 61,696, 6% below it.
Availability
- Weights and reports: KBlueLeaf/MOSS-Image-heal
- Trainer and write-ups: KohakuBlueleaf/KohakuMossTrainer
- Base model: OpenMOSS-Team/MOSS-Image-Base
Comments
Sign in with GitHub to join the discussion.