Model
English summary. A 4-step, CFG-free adapter for OpenMOSS-Team/MOSS-Image-SFT
(50 steps × 2 forwards with CFG 3). It continues lodestones/MOSS-Image-SFT-8-steps-dmd
(data-space DMD, arXiv 2610.12102) restricted to 4 iterations, trained with DMAD (arXiv 2610.02188): a LoRA critic on the
frozen teacher separates teacher samples from student samples, and the student minimises the negative critic logit.
DiT time per 1024² image on one H100: 16.2 s (teacher) → 0.81 s. Compared with the lodestones checkpoint run for 4 steps,
images have more detail and livelier composition, with some regressions (see Results). Same file format and sampler as the
lodestones checkpoint; inference.py runs it with the release pipeline. Training code is in code/.
MOSS-Image-SFT(50 ステップ・CFG 3 で 1 枚 100 回の DiT 計算)を、4 回の計算・CFG なしで画像を出せるようにする高速化アダプタです。 lodestones の 8 ステップ版(data-space DMD、update 2600)を出発点に、4 ステップに絞って DMAD で追加学習しました。
| DiT 時間 / 1024² 1 枚(H100) | DiT の計算回数 | |
|---|---|---|
| MOSS-Image-SFT(50 step、CFG 3、負側 block 22 skip) | 16.2 s | 100 |
| lodestones 8 step | 1.54 s | 8 |
| このアダプタ(4 step) | 0.81 s | 4 |
(DiT のみの計測。テキストエンコーダと VAE は含まない。実装は code/rapid_moss/model.py。)
pip install "diffusers>=0.38" "transformers>=5.3" accelerate safetensors huggingface_hub pillow
python inference.py --prompt "an anime girl in a yellow raincoat at a rainy bus stop at night" --width 1024 --height 1024 --out out.png
ファイルは lodestones と同じ形式です。
| ファイル | 中身 |
|---|---|
lora.safetensors | fp32 LoRA、rank 128 / alpha 128、全 42 ブロックの proj.{q,k,v,o,ffn_up_gate,ffn_down}。kohya 形式のキー(diffusion_model.<module>.lora_down.weight / lora_up.weight / alpha) |
time_embedding.safetensors | fp32 の time_embedding.* と shared_block_modulation.bias / video_head_modulation.bias の置き換え |
dataspace.safetensors | latent_mix.Q [128, 128] と step_codes [1000, 968, 926, 871] |
サンプラー(ノイズの再注入もタイムステップのスケジュールもなく、同じ z を使い続ける):
x0 = torch.zeros_like(z) # z: fp32 noise [1, 128, H/16, W/16]
for r in range(4): # 1〜3 回で止めても画像になる(各段階を学習している)
x = z + torch.einsum("oc,bchw->bohw", Q, x0)
x0 = z - transformer(x, codes[r], prompt_embeds, prompt_mask).sample.float()
# x0 を release の VAE でデコード(BN の逆正規化 + 2x2 unpatchify)
LoRA の適用方法、CFG、解像度などの決まりは、下の「注意点」を読んでください。
カスタムノード comfyui/rapid-moss/ を同梱しています。ComfyUI 本体も MOSS-Image に対応していますが、専用サンプラーと time path の置き換えはなく、標準の LoraLoaderModelOnly は LoRA を bf16 の重みに足し込みます(当たる差分の相対誤差 81%)。標準ノードだけでも絵は出ますが、学習した生成器とは別物になります(下の比較)。このノードを使ってください。
ComfyUI 標準ローダーにも対応しています。 comfyui/rapid-moss/example_workflows/rapid-moss-4step-native.json を使い、Load Diffusion Model で moss_image_sft_bf16.safetensors、Load CLIP で Qwen3.5-9B (type: moss_image)、Load VAE で FLUX.2 VAE を選んでください。この構成では OpenMOSS の HF ローダーノードは不要です。新しい ComfyUI で出ていた aimdo memory compile error(2026-10-19 修正)と、QKV が結合されたモデルで出ていた ModuleDict ... has no attribute q は、ノードを更新すると直ります。
従来の HF ローダーを使う場合:
comfyui/moss-image/)を入れます。comfyui/rapid-moss/ を ComfyUI/custom_nodes/rapid-moss/ にコピーします。ComfyUI/models/rapid_moss/rapid-moss-4step/ に置きます。comfyui/rapid-moss/example_workflows/rapid-moss-4step.json を画面にドラッグ&ドロップし、HF Loader の model_path を MOSS-Image-SFT のフォルダにして実行します。ノード MOSS-Image Few-Step Adapter (rapid-moss) は、次の処理をして MODEL と、専用の SAMPLER・SIGMAS を返します。
これを BasicGuider(CFG なし)と SamplerCustomAdvanced につなぎます。lodestones の 8 ステップ版も、同じ 3 ファイルを別フォルダに置けば同じノードで使えます。
詳しい手順と注意点は comfyui/rapid-moss/README.md にあります。HF ローダーでの動作確認では、学習コードで作った画像とほぼ同じ絵になりました(下の図)。標準ローダーでは、2026-10-19 の ComfyUI で 1024² 1 枚約 1.1 秒(2 枚目以降、RTX PRO 6000、VAE 込み)です。
inference.py は fp32)。LoraLoaderModelOnly、diffusers の fuse_lora、マージ済みの bf16 重みを書き出すツールなどがこれに当たります。inference.py か、ComfyUI のカスタムノード(comfyui/rapid-moss/)を使ってください。time_embedding.safetensors の 3 モジュールも置き換え、fp32 で計算します。
x = z + Q·x0、x0 = z − v を 4 回繰り返し、timestep には step_codes を渡します。"..." で囲みます。torch.randn(fp32)です。inference.py と ComfyUI のノードは、同じ seed なら同じノイズになります。
inference.py と ComfyUI の HF ローダーでは、base に加えて約 1.5 GB 使います。fp32 の LoRA が 1.19 GB、fp32 の time path と modulation が約 0.34 GB です。ComfyUI の標準ローダーでは約 0.65 GB(bf16 の LoRA と fp32 の time path)です。左から teacher(50 step)/ lodestones 8 step 版を 4 step で使ったもの / このアダプタ(4 step)。同じ seed、評価用 17 プロンプト(eval/eval_prompts.jsonl、学習には未使用)。
lodestones の 4 step と比べた所見(目視のみ。数値指標では評価していません):
結論として、細部は良くなったが一長一短です。
| init(lodestones) | 300 it | 500 it(公開版) | 800 it |
|---|---|---|---|
 |  |  |  |
1girl, solo なのに 2 人)と文字が崩れていきました。841〜850 it で NaN になり、学習を止めました。teacher サンプルに非有限値はなく、原因は特定できていません。現在のコードは非有限の勾配で更新を飛ばします。lodestones の形式そのままです。x = z + Q·x0, v = net(x, code_r), x0 ← z − v。
softplus(−d(teacher サンプル)) + softplus(d(生徒の出力))。ノイズ量は σ = 5u / (1 + 4u)、u ~ U[0.02, 0.98] で、teacher と生徒の組で同じ σ を使います。−d(ノイズを足した生徒の x0)。teacher の CFG 推論も fake score も、学習中は使いません。| 項目 | 値 |
|---|---|
| バッチ | 16(8 GPU × 2、1 サンプルずつ順に処理) |
| 生成器 | AdamW lr 2e-5、betas (0, 0.99)、weight decay 0、EMA 0.998 |
| critic | AdamW lr 4e-5、betas (0, 0.99)、weight decay 0.01、生成器を更新しない 50 it のウォームアップ |
| 反復 | 500 it(生成器の更新 450 回)、1 it 約 5.7 秒 |
| 計算機 | RunPod 8×H100 SXM。teacher サンプル生成約 1.9 時間、学習 500 it 約 47 分 |
| 費用 | 試運転と比較実験を含めて計約 $126 |
24,000 件(再配布はしていません)から、teacher サンプルを作った 3,199 件を学習に使いました。
| 種類 | 出典(ライセンス) |
|---|---|
| アニメの自然文キャプション | BootsofLagrangian/danbooru-multitier-captions-202606(CC-BY-4.0、Giyeong Oh)、CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq(CC-BY-SA-4.0) |
| タグ列 | 上記の Danbooru タグ、FredZhang7/anime-prompts-180K(CreativeML OpenRAIL-M)。絵師名は除去 |
| 一般 | ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions(MIT)、data-is-better-together/open-image-preferences-v1(Apache-2.0)、nateraw/parti-prompts(Apache-2.0) |
| 文字描画(日 60% / 英 25% / 中 15%)、JSON レイアウト、短文 | テンプレートから生成 |
最初は DMDX(arXiv 2507.18569)の ADM を data-space DMD に足す計画でした。ADM は、生成器の出力にノイズをかけた同じ x_τ から、teacher(CFG あり)と fake score の ODE を Δt = 1/64 だけ進め、その 2 つを識別器に見分けさせる方式です。
logs/probe_adm_discriminator.jsonl)。2 つの入力の差は Δt·(v_real − v_fake) だけで、1 枚ずつの判定ではほぼ見分けられません。logs/probe_teacher_discriminator.jsonl)。