MOSS-Image-SFT-4-steps

⚡ Demo: MOSS-Image 4-Step Turbo Space

v1.1 (2026-10-17): +400 continuation updates with a signage-lexicon prompt mix — fixes the v1.0 small-English-text failure class (e.g. "OPPEN 24H" → "OPEN 24H") with no regressions. v1.0 weights remain at checkpoint-update-1000/.

A 4-step distillation of OpenMOSS-Team/MOSS-Image-SFT (50 steps, CFG 3), warm-started from lodestones/MOSS-Image-SFT-8-steps-dmd and continued with data-space DMD (arXiv 2610.12102) for 1,000 generator updates. 11 s per 1024×1024 image on one A100 (2.7 s per fp32 transformer forward), no CFG, no negative prompt.

Both images above were generated by this model in 4 steps, no CFG — "TODAY'S SPECIAL / Flat White $4.50" and 深夜食堂 render correctly.

License: Apache-2.0 — commercial use OK.

Quickstart

Self-contained pipeline — this repo ships everything (merged fp32 transformer, text encoder, tokenizer, VAE, scheduler, pipeline code):

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "akhaliq/MOSS-Image-SFT-4-steps",
    custom_pipeline="pipeline_moss_image", trust_remote_code=True,
    torch_dtype={"transformer": torch.float32, "default": torch.bfloat16,
                 "vae": torch.float32},
).to("cuda")
# then run the data-space sampler loop below (NOT the standard scheduler)

Single-file CLI (downloads the base repo's bf16 transformer + this repo's fp32 adapters, ~31 GB VRAM):

pip install "torch>=2.9" "diffusers>=0.38" "transformers>=5.3" accelerate \
    safetensors huggingface_hub pillow
python inference.py --prompt "A neon sign reading '深夜食堂' over a rainy door" \
    --steps 4 --seed 1234 --out out.png

Lower VRAM: keep the text encoder off the GPU between prompts, or use pipe.enable_model_cpu_offload() — generation is latent-only (no VAE decode per step), so offloading the text encoder is the big win.

Files

• transformer/ — merged fp32 MOSSImageTransformer (load with the base repo's pipeline_moss_image.py); text encoder, tokenizer and VAE come from the base repo.
• lora.safetensors — fp32 LoRA, rank 128 / alpha 128 on every block's proj.{q,k,v,o,ffn_up_gate,ffn_down}, kohya keys diffusion_model.<module>.lora_down.weight / lora_up.weight / alpha; apply as y = W x + (alpha / rank) * up(down(x)).
• time_embedding.safetensors — fp32 replacements for time_embedding.* and shared_block_modulation.bias / video_head_modulation.bias. Keep these in fp32.
• dataspace.safetensors — latent_mix.Q [128, 128] and step_codes [4].
• checkpoint-update-{250,500,750,1000}/ — training checkpoints (u1000 = published). Note: training past ~1,000 updates washed out in our run; u1000 was selected from visual checkpoints as the final EMA weights.
• train_dmd.py — the training script.

Sampler (same loop as the 8-step checkpoint, 4 step codes)

import torch
from safetensors.torch import load_file

ds = load_file("dataspace.safetensors")            # latent_mix.Q [128,128], step_codes [4]
Q, codes = ds["latent_mix.Q"].cuda(), ds["step_codes"].cuda()
z = torch.randn(1, 128, H // 16, W // 16, device="cuda")   # fp32
x0 = torch.zeros_like(z)
for r in range(4):
    x = z + torch.einsum("oc,bchw->bohw", Q, x0)
    v = transformer(x, codes[r], prompt_embeds, prompt_mask).sample.float()
    x0 = z - v
# decode x0 with the release VAE (BN denormalize + 2x2 unpatchify)

Precision — read this

Like the 8-step checkpoint it builds on, the distillation delta is smaller than bf16's rounding step:

• Recommended: base repo's bf16 transformer + these fp32 adapters (lora.safetensors, time_embedding.safetensors, keep latent_mix.Q and the step codes fp32). This is the configuration the model was trained and evaluated in.
• Best quality: the merged fp32 transformer/ (~20 GB, 2× memory and ~2× latency of bf16).
• Not recommended: merging the LoRA into bf16 weights. It still produces coherent images, but fine text degrades first — in our test the main 深夜食堂 sign survived while small signage turned to gibberish.

ComfyUI

This repo uses the same file format and sampler as lodestones' 8-step checkpoint and darask0's rapid-moss: ComfyUI nodes built for those checkpoints load this model's lora.safetensors, time_embedding.safetensors and dataspace.safetensors unchanged (4 step codes). With the official MOSS-Image ComfyUI plugin (search MOSS-Image in ComfyUI Manager), load the merged fp32 transformer/ as the diffusion model and use the data-space sampler loop, 4 steps, CFG 1. Standard LoRA loaders won't work — the LoRA must stay unmerged in fp32 (see Precision).

Training

Method	data-space DMD (arXiv 2610.12102), latent space
Init	lodestones 8-step DMD checkpoint (LoRA + time-embedding + Q, first 4 step codes)
Teacher	frozen MOSS-Image-SFT (bf16), CFG 3.0, negative branch skip-layer 22
Fake score	base transformer (bf16) + fresh rank-128 LoRA, 1:1 update ratio, uniform τ ∈ [0.02, 0.98]
Generator τ	Beta(2.5, 1) importance sampling (high-noise skew)
Optimizer	Adam β₁=0, lr 1e-5 (gen) / 1e-5 (fake), grad clip 1.0, EMA 0.999
Updates	1,000 × batch 2 @ 1024px (v1.0; selected mid-run, >1,000 washed out) + 400 × batch 2 @ lr 5e-6 (v1.1)
Prompts	v1.0: 2,048 procedural scene-graph JSON + bilingual text prompts · v1.1: + signage lexicon (OPEN/CLOSED/SALE/prices…) + 3,000 real user prompts
Cost	5.6 h on 1× A100-80GB ($14) + 2.2 h ($6) for v1.1

Evaluation

Side-by-side strips (8 prompts, fixed seeds): truncated 8-step checkpoint at 4 steps vs this model at u750/u1000 vs the 8-step checkpoint at 8 steps: huggingface.co/datasets/akhaliq/moss-image-turbo-eval (eval_final/). At 4 steps this model matches the 8-step checkpoint on text accuracy and photographic detail, and clearly beats running the 8-step checkpoint truncated to 4 steps (sharper, better contrast, more accurate Chinese text).

Each strip below: 8-step checkpoint truncated to 4 steps | this model (u750) | this model (u1000, v1.0) | 8-step checkpoint at 8 steps — same prompt, same noise:

Clearest win: on the infographic prompt this model is the only one that gets both "62%" and "38%" right — the 8-step checkpoint at 8 steps invents "88% / 31%" on its second chart.

Head-to-head: the 4-step adapters

darask0/MOSS-Image_4-step_adapter_lora ("rapid-moss") is an independent 4-step adapter from the same 8-step warm-start, trained with DMAD — published within hours of this model. Same prompts, same seeds, both at 4 steps (this model left, rapid-moss right):

Scorecard (v1.1 vs rapid-moss): tied on headline text and menus; this model wins numeric fidelity (the infographic — theirs garbles the second chart); rapid-moss still produces livelier street-scene compositions. All 8 strips: vs_darask_v2/ (v1.0 strips: vs_darask/).

Citation

@misc{mossimagesft4steps2026,
  title = {MOSS-Image-SFT-4-steps: 4-step data-space DMD distillation of MOSS-Image-SFT},
  author = {akhaliq},
  year = {2026},
  howpublished = {\url{https://huggingface.co/akhaliq/MOSS-Image-SFT-4-steps}}
}

Builds on MOSS-Image (Apache 2.0), lodestones' 8-step DMD, and Few-Step Generation via Data-Space Iteration.