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.

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.

Prompts

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.

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:

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:

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.

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.

• 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

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

1. 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.
2. 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.