rapid-moss: MOSS-Image-SFT 4-step adapter (data-space DMD → DMAD)

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 のみの計測。テキストエンコーダと 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 と同じ形式です。

サンプラー(ノイズの再注入もタイムステップのスケジュールもなく、同じ 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

カスタムノード comfyui/rapid-moss/ を同梱しています。ComfyUI 標準の LoraLoaderModelOnly は LoRA を重みにマージするので、このアダプタには使えません。このノードを使ってください。

1. MOSS-Image 公式ノード(OpenMOSS/MOSS-Image の comfyui/moss-image/)を入れます。
2. このリポジトリの comfyui/rapid-moss/ を ComfyUI/custom_nodes/rapid-moss/ にコピーします。
3. 上の 3 ファイルを ComfyUI/models/rapid_moss/rapid-moss-4step/ に置きます。
4. 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 を返します。

• LoRA をマージせず fp32 で足す
• time path を fp32 に置き換える

これを BasicGuider(CFG なし)と SamplerCustomAdvanced につなぎます。lodestones の 8 ステップ版も、同じ 3 ファイルを別フォルダに置けば同じノードで使えます。

詳しい手順と注意点は comfyui/rapid-moss/README.md にあります。GPU での動作確認では、学習コードで作った画像とほぼ同じ絵になりました(下の図)。1024² 1 枚は 4.5〜6.9 秒(RTX PRO 6000、テキストエンコード・VAE 込み)です。

注意点

必ず守ること(守らないと効かない、または絵が崩れる)

• 元のモデルは MOSS-Image-SFT です。
  • 学習の teacher も初期値も SFT です。MOSS-Image-Base と組み合わせたときの動作は確かめていません。
• LoRA はマージせず、fp32 のまま足してください。
  • 差分が bf16 の丸め幅より小さいため、bf16 の重みにマージすると効果がほぼ消えます。
  • そのため、マージ方式のローダーでは効きません。ComfyUI の LoraLoaderModelOnly、diffusers の fuse_lora、マージ済みの bf16 重みを書き出すツールなどがこれに当たります。
  • inference.py か、ComfyUI のカスタムノード(comfyui/rapid-moss/)を使ってください。
• time_embedding.safetensors の 3 モジュールも置き換え、fp32 で計算します。
  • LoRA だけを当てても正しく動きません。
  • t=1 の時間ベクトルがバイアスに畳み込まれているので、3 つは必ずセットで使います。
• 専用のサンプラーを使います。
  • x = z + Q·x0、x0 = z − v を 4 回繰り返し、timestep には step_codes を渡します。
  • Euler などの標準サンプラーや、通常のスケジューラ(shift 5 など)では正しく動きません。
• CFG とネガティブプロンプトは使いません。
  • CFG の効果は蒸留で取り込み済みです。CFG を 1 より大きくすると計算が倍になり、効果が二重にかかって色や形が崩れる可能性が高いです(試していません)。
  • 負側 block 22 skip や SLG も入れないでください。
• ステップ数は 1〜4 です。
  • 4 を推奨します。1〜3 回で止めても画像になります(途中の各段階も学習しています)が、その品質は確かめていません。5 回以上にはできません(step code が 4 つしかないため)。
• base の重みは bf16 か fp32 で使ってください。
  • fp8 などに量子化した重みでは、LoRA の差分が失われるか、time path の置き換えが正しくできません。ComfyUI のノードは fp8 だとエラーにしています。

おすすめの設定

• 解像度: 16 の倍数で、約 100 万画素が目安です。
  • 学習した解像度は 1024×1024、896×1152、832×1216、768×1344 と、それぞれの横向きです。これより大きい・小さい解像度での品質は確かめていません。
• プロンプト: MOSS-Image と同じです(自然文、タグ列、JSON レイアウト)。
  • 描かせたい文字は "..." で囲みます。
  • 日本語の文字描画は、付箋・看板・ポスター・吹き出しで確認しました。
• seed: 初期ノイズは CPU の torch.randn(fp32)です。inference.py と ComfyUI のノードは、同じ seed なら同じノイズになります。
  • ただしテキストエンコーダの実装(transformers と ComfyUI)の違いで、細部は少し変わります。
• VRAM: base に加えて約 1.5 GB 使います。fp32 の LoRA が 1.19 GB、fp32 の time path と modulation が約 0.34 GB です。

分かっている弱点

• 厚塗り・絵画調の絵に、頼んでいない題字や漢字が入ることがあります。lodestones の 8 ステップ版から引き継いだ癖で、teacher は入れません。
• 暖簾・看板の小さな文字や、白黒漫画の吹き出しは、意味の通らない文字になりやすいです。
• 頼んでいない人物が増えることがあります。JSON レイアウトの 2 コマ漫画では、下のコマに 2 人目が出ました。
• lodestones の 4 ステップより描き込みは増えますが、teacher(50 ステップ)の絵画的な質感や多様性には届きません。
• 学習を続けると、プロンプトへの従い方と文字が崩れていきました(800 イテレーション目)。公開しているのは崩れる前の 500 イテレーション目です。

検証していないこと

• 他の LoRA との組み合わせ: キャラクター LoRA や画風 LoRA を重ねたときの挙動は試していません。4 ステップ化で効きが弱まる可能性があります。
• 評価の範囲: 評価は 17 プロンプトの目視だけです。FID や好みの評価などの数値指標は取っていません。
• 学習データの規模: 学習プロンプトは 3,199 件で、teacher サンプルは 1 プロンプトにつき 1 枚です。
• 学習データの内容: 性的に露骨な内容は学習プロンプトから除いています。そうしたプロンプトでの挙動は確かめていません。

結果

左から teacher(50 step)/ lodestones 8 step 版を 4 step で使ったもの / このアダプタ(4 step)。同じ seed、評価用 17 プロンプト(eval/eval_prompts.jsonl、学習には未使用)。

lodestones の 4 step と比べた所見(目視のみ。数値指標では評価していません):

• 良くなった点
  • 背景や小物の描き込みが増え、構図に動きが出た(teacher の雰囲気に近づいた)
  • 縦書きポスターで「夏」が二重になる崩れが直った
  • 付箋の 3 行、看板、吹き出しの日本語は正確なまま
• 悪くなった点
  • ライトノベル表紙に頼んでいない英字が増えた
  • 2 コマ漫画(JSON レイアウト)の下のコマに 2 人目が出た
  • 暖簾などの小さな文字の崩れが増えた
• 変わらない点
  • 厚塗りの侍の絵に頼んでいない題字が入る(lodestones から引き継いだ癖。teacher は入れない)
  • 白黒漫画の吹き出しが意味の通らない文字になる

結論として、細部は良くなったが一長一短です。

学習の経過(EMA、4 step、同じ seed)

• 300 it では、桜の女子高生の絵が顔の寄せ集めになる崩れが一度出ました。500 it では元に戻っています。
• 500 → 800 it ではプロンプトへの従い方(1girl, solo なのに 2 人)と文字が崩れていきました。841〜850 it で NaN になり、学習を止めました。teacher サンプルに非有限値はなく、原因は特定できていません。現在のコードは非有限の勾配で更新を飛ばします。
• 公開版は 500 it の EMA です。

手法

生成器(data-space iteration、arXiv 2610.12102)

lodestones の形式そのままです。x = z + Q·x0, v = net(x, code_r), x0 ← z − v。

• 学習対象は LoRA(rank 128)、time path(fp32)、Q の 3 つです。
• lodestones update 2600 から初期化し、step codes は先頭 4 つを使います。
• 反復番号 r は 0〜3 を順番に回します。r 回目までは勾配なしで進め、最後の 1 回だけ勾配を流します。

識別器と損失(DMAD、arXiv 2610.02188)

• critic の構成: 凍結した teacher(MOSS-Image-SFT の bf16)に、block 21〜41 の LoRA(rank 128)を付けます。block 41 の画像トークンにヘッドを付け、1 枚につき 1 つのロジットを出します。ヘッドは 4×4 stride-2 conv(3072→512)、GroupNorm、SiLU を 2 回重ね、空間平均を取ってから線形層に通します。
• critic の損失: softplus(−d(teacher サンプル)) + softplus(d(生徒の出力))。ノイズ量は σ = 5u / (1 + 4u)、u ~ U[0.02, 0.98] で、teacher と生徒の組で同じ σ を使います。
• 生成器の損失: −d(ノイズを足した生徒の x0)。teacher の CFG 推論も fake score も、学習中は使いません。
• teacher サンプル: 3,199 枚。1 プロンプトにつき 1 枚で、50 step・CFG 3・負側 block 22 skip で作りました。解像度は約 1MP で、正方形・縦長・横長を混ぜています。学習時のプロンプトはこの 3,199 件から選びます。
• 本物の画像を使うヘッドは使っていません(DMAD は 2 ヘッドですが、こちらは teacher ヘッドだけ)。

設定

学習プロンプト

24,000 件(再配布はしていません)から、teacher サンプルを作った 3,199 件を学習に使いました。

うまくいかなかったこと(DMDX)

最初は DMDX(arXiv 2507.18569)の ADM を data-space DMD に足す計画でした。ADM は、生成器の出力にノイズをかけた同じ x_τ から、teacher(CFG あり)と fake score の ODE を Δt = 1/64 だけ進め、その 2 つを識別器に見分けさせる方式です。

• 結果: 識別器は 130 回更新しても正解率 0.500〜0.506 のままで、学習しませんでした(logs/probe_adm_discriminator.jsonl)。2 つの入力の差は Δt·(v_real − v_fake) だけで、1 枚ずつの判定ではほぼ見分けられません。
• 比較: 同じヘッドで本物側を teacher サンプルにすると、60 回で 0.56 → 0.68 まで上がりました(logs/probe_teacher_discriminator.jsonl)。
• この結果と、その後に出た DMAD を受けて、fake score も teacher の推論もいらない DMAD に切り替えました。

ライセンスとクレジット

• 重みとコード: Apache-2.0
• 元のモデル: OpenMOSS-Team/MOSS-Image-SFT(Apache-2.0)、初期値: lodestones/MOSS-Image-SFT-8-steps-dmd(Apache-2.0)
• 参考にした論文:
  • Few-Step Generation via Data-Space Iteration(arXiv 2610.12102)
  • DMAD: Distribution Matching as Adversarial Distillation(arXiv 2610.02188、コード)
  • Adversarial Distribution Matching / DMDX(arXiv 2507.18569)
• 前作: Anima の DMDX・DMD2・PCM 蒸留の記録