MOSS-Image-heal

A full finetune of MOSS-Image-Base with a healed residual stream: the largest stream value is 3,083 (base: 58,411; fp16 max: 65,504), the stream RMS is 7.5x lower, and every weight element is within 64x its tensor's RMS (base: 349 elements above). The outputs stay close to the base model: flow loss 0.4289 vs 0.4280, velocity cosine 0.9983 on the probe set below.

The transformer is the trained part. Text encoder, tokenizer, VAE, scheduler and pipeline code are the official MOSS-Image-Base ones.

The concepts behind this model are written up in KohakuMoss:

Massive residual channels in MOSS-Image: the few huge channels of the base model's residual stream, what they do, and how a heal shrinks them.
Healing the stream and the weights during a finetune: the stream losses (write concentration, RMS) and the outlier decay of the weights used in this run, with a toy model and their combination.
Numerics: the fp16 residual-stream fold the training ran with.

file	for
`transformer/`	diffusers `MOSSImageTransformer` (official config, bf16)
`comfyui/moss_image_heal_bf16.safetensors`	ComfyUI single file, bf16 (fused qkv, the Comfy-Org layout)
`comfyui/moss_image_heal_fp16.safetensors`	ComfyUI single file, fp16
`comfyui/moss_image_heal_int8_convrot.safetensors`	ComfyUI single file, int8 ConvRot (group 256), the layout of the Comfy-Org `int8_convrot` files
`reports/`	every number and chart value on this page

Samples

Base | heal, with the caption prompt (tags + natural-language caption) and the GBC prompt (gbc.soft). 16 steps, CFG 3, shift 5, seed 0, same negative prompt for all.

5057091

7637375

The caption and GBC prompts of each grid are in reports/prompts/ (caption.txt, gbc.soft.txt). Negative prompt:

worst quality, low quality, bad quality, old, lowres, blurry, cropped, jpeg artifacts, watermark, signature, bad anatomy, bad hands, extra digits, fewer digits, nsfw, explicit

Residual stream

The stream statistics below follow massive-channels.md. Measured on 6 text-image samples at 1024^2 area, sigma 0.3 / 0.5 / 0.7 / 0.9 / 0.5 / 0.7 with fixed noise, bf16 forward, fp32 statistics; each value is the mean over the 6 rows of the per-block statistic.

residual stream per block

	base	heal
largest \|x\| in the stream (base: block 26, heal: block 27)	58,411	3,083
headroom to the fp16 max	1.12x	21.2x
mean of the top 1e-4 of \|x\|, worst block	1,603	258
stream RMS, mean over blocks	27.75	3.70
stream RMS, last block	65.6	12.0
energy share of the 32 largest channels, mean over blocks	0.837	0.708
flow loss	0.4280	0.4289
velocity cosine to base	1	0.9983

What each block writes into the stream (the attention + FFN sublayer outputs after their gates); the write-concentration loss (stream-healing.md section 4) acts on these writes:

writes per block

	base	heal
largest write value (block 0)	17,250	1,113
top-32 channel share of the write, mean over blocks	0.553	0.316

Weights

Elements above a multiple of their tensor's RMS, over every 2-D tensor of the transformer (185 tensors, about 5.08 billion elements). The outlier decay (stream-healing.md section 5) pulls each weight's excess over a capped multiple of its start RMS:

above	base	heal
8x RMS	171,936	184,750
16x RMS	15,310	16,597
32x RMS	2,621	345
64x RMS	349	0
largest max / RMS	271.7 (`time_embedding.2`)	63.5 (`time_embedding.2`)

The text refiner (text_embedder.*) is frozen and keeps its base values; its blocks.0.ffn_down (54.7x) and proj (37.4x) are the third- and fourth-largest tensors of heal.

largest weight / RMS per tensor

Precision of the files

Each file against the same weights in fp32, on the 6 probe rows above (RTX 3090, 16-bit matmuls with fp32 accumulation): relative L2 error and cosine of the predicted velocity.

file	compute	velocity rel. error (mean / max)	cosine (min)	largest stream value
`moss_image_heal_fp16`	fp16	2.1e-3 / 2.7e-3	0.999996	3,278
`moss_image_heal_bf16`	bf16	1.7e-2 / 2.1e-2	0.999787	3,248
`transformer/`	bf16	1.7e-2 / 2.1e-2	0.999787	3,248
`moss_image_heal_int8_convrot`	int8 weights, fp16 compute	1.2e-2 / 2.0e-2	0.999804	3,278
`moss_image_heal_int8_convrot`	int8 weights, bf16 compute	2.1e-2 / 2.9e-2	0.999583	3,248
MOSS-Image-Base, for reference	fp16	2.0e-3 / 2.7e-3	0.999996	61,696
MOSS-Image-Base, for reference	bf16	1.7e-2 / 2.2e-2	0.999767	60,672

In fp16 the heal stream peaks at 3,278, 20x below the fp16 max; the base stream peaks at 61,696 on the same rows, 6% below it.
The int8 file's dequantized weights are within 0.85-0.98% (relative L2) of the bf16 weights per layer; 168 block linears are int8, every other tensor bf16.
The int8 writer reproduces Comfy-Org's moss_image_base_int8_convrot.safetensors from the base bf16 weights: same 757 keys, every record and non-quantized tensor equal, 5e-8 of the codes differ by 1 (report).

Training


start	MOSS-Image-Base
method	full finetune (every transformer weight except the text refiner)
steps	50,000 x 64 samples (4 GPUs x batch 4 x 4 accumulation; 3.2M samples), 1024^2 area at each image's own aspect ratio
optimizer	AdamW, lr 1e-5, cosine to 1e-7
compute	fp16 with the residual-stream fold of MOSS-Image-FP16-fixed, fp32 master weights
losses	flow matching + distillation toward the frozen base (weight 1.0) + stream write-concentration (5e-3, target share 0.5) + stream RMS (5e-5)
weight outliers	decoupled decay of each weight's excess over a cap x its start RMS (rate 12, cap 64 -> 16 over the first 10k steps)
cost	about 200 GPU hours on B300

Data

Images: about 1M danbooru posts (999,294), chosen by a normalized quality score q_soft: a post's favourites and score, ranked against posts of the same rating uploaded within ±7 days, then corrected for its aspect ratio and its franchise. The best-scoring posts are kept per era, rating group and aspect ratio. The run reads 3.2M samples, about 3.2 passes over the set.
Captions: each image has a Graph-Based Captioning (GBC) graph made with KohakuGBC: an image caption, the entities and compositions in it with their boxes and captions, the relations between them, and tags placed on the vertices. The VLM captions come from Qwen3.8 27B and Qwen3.5 9B.
Prompt formats: every sample is drawn into one of several prompt formats: artist + quality, tags, tags + natural-language caption, natural-language caption, and the GBC graph (gbc.soft in the samples above), each with artist, copyright, character, quality, era and rating sections and random section and tag dropout.

The losses and the decay are KohakuMoss's STREAM_LOSS and OUTLIER_DECAY: the write term in stream-healing.md section 4, the outlier decay in section 5, the combination used here in section 6, and the fp16 fold in numerics.md.

Usage

diffusers

import importlib

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "OpenMOSS-Team/MOSS-Image-Base",
    custom_pipeline="pipeline_moss_image",
    trust_remote_code=True,
    transformer=None,
    torch_dtype=torch.bfloat16,
)
pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)
MOSSImageTransformer = importlib.import_module(type(pipe).__module__).MOSSImageTransformer
pipe.transformer = MOSSImageTransformer.from_pretrained(
    "KBlueLeaf/MOSS-Image-heal", subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe.to("cuda")

ComfyUI

One of the comfyui/ files into models/diffusion_models/; the Comfy-Org qwen3.5_9b text encoder and flux2-vae_small.safetensors VAE as for the official model.
Load it with Load Diffusion Model in the official MOSS-Image workflow. The fp16 file runs in fp16 with ComfyUI's --fp16-unet.

KohakuMoss

TRANSFORMER (sampling) / WEIGHTS (training) set to any of the files, PRECISION bf16 or fp16. See KohakuMoss.

License

Apache-2.0, as the official MOSS-Image weights.