Duplicate from darask0/MOSS-Image_4-step_adapter_lora

df84732 16 days ago

16.9 kB

---
license: apache-2.0
base_model:
- OpenMOSS-Team/MOSS-Image-SFT
- lodestones/MOSS-Image-SFT-8-steps-dmd
library_name: diffusers
pipeline_tag: text-to-image
language:
- ja
- en
tags:
- moss-image
- lora
- distillation
- few-step
- dmad
- dmd
---
# 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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)
(50 steps × 2 forwards with CFG 3). It continues [lodestones/MOSS-Image-SFT-8-steps-dmd](https://huggingface.co/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`。)

## 使い方

```bash
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 を使い続ける):

```python
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](https://github.com/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](comfyui/rapid-moss/README.md) にあります。GPU での動作確認では、学習コードで作った画像とほぼ同じ絵になりました(下の図)。1024² 1 枚は 4.5〜6.9 秒(RTX PRO 6000、テキストエンコード・VAE 込み)です。

![ComfyUI check](images/comfyui-check.jpg)

## 注意点

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

- **元のモデルは 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`、学習には未使用)。

![compare-1](images/compare-1.jpg)
![compare-2](images/compare-2.jpg)
![compare-3](images/compare-3.jpg)

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

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

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

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

| init(lodestones) | 300 it | 500 it(公開版) | 800 it |
|---|---|---|---|
| ![](images/training/init.jpg) | ![](images/training/000300.jpg) | ![](images/training/000500.jpg) | ![](images/training/000800.jpg) |

- 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 ヘッドだけ)。

### 設定

| 項目 | 値 |
|---|---|
| バッチ | 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](https://huggingface.co/datasets/BootsofLagrangian/danbooru-multitier-captions-202606)(CC-BY-4.0、Giyeong Oh)、[CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq](https://huggingface.co/datasets/CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq)(CC-BY-SA-4.0) |
| タグ列 | 上記の Danbooru タグ、[FredZhang7/anime-prompts-180K](https://huggingface.co/datasets/FredZhang7/anime-prompts-180K)(CreativeML OpenRAIL-M)。絵師名は除去 |
| 一般 | [ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions](https://huggingface.co/datasets/ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions)(MIT)、[data-is-better-together/open-image-preferences-v1](https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1)(Apache-2.0)、[nateraw/parti-prompts](https://huggingface.co/datasets/nateraw/parti-prompts)(Apache-2.0) |
| 文字描画(日 60% / 英 25% / 中 15%)、JSON レイアウト、短文 | テンプレートから生成 |

## うまくいかなかったこと(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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)(Apache-2.0)、初期値: [lodestones/MOSS-Image-SFT-8-steps-dmd](https://huggingface.co/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、[コード](https://github.com/Yzmblog/DMAD))
  - Adversarial Distribution Matching / DMDX(arXiv 2507.18569)
- 前作: [Anima の DMDX・DMD2・PCM 蒸留の記録](https://note.com/darask/n/neef98f000087)

	---
	license: apache-2.0
	base_model:
	- OpenMOSS-Team/MOSS-Image-SFT
	- lodestones/MOSS-Image-SFT-8-steps-dmd
	library_name: diffusers
	pipeline_tag: text-to-image
	language:
	- ja
	- en
	tags:
	- moss-image
	- lora
	- distillation
	- few-step
	- dmad
	- dmd
	---

	# 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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)
	(50 steps × 2 forwards with CFG 3). It continues [lodestones/MOSS-Image-SFT-8-steps-dmd](https://huggingface.co/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`。)

	## 使い方

	```bash
	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 を使い続ける):

	```python
	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](https://github.com/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](comfyui/rapid-moss/README.md) にあります。GPU での動作確認では、学習コードで作った画像とほぼ同じ絵になりました(下の図)。1024² 1 枚は 4.5〜6.9 秒(RTX PRO 6000、テキストエンコード・VAE 込み)です。

	![ComfyUI check](images/comfyui-check.jpg)

	## 注意点

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

	- 元のモデルは 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`、学習には未使用)。

	![compare-1](images/compare-1.jpg)
	![compare-2](images/compare-2.jpg)
	![compare-3](images/compare-3.jpg)

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

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

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

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

	\| init(lodestones) \| 300 it \| 500 it(公開版) \| 800 it \|
	\|---\|---\|---\|---\|
	\| ![](images/training/init.jpg) \| ![](images/training/000300.jpg) \| ![](images/training/000500.jpg) \| ![](images/training/000800.jpg) \|

	- 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 ヘッドだけ)。

	### 設定

	\| 項目 \| 値 \|
	\|---\|---\|
	\| バッチ \| 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](https://huggingface.co/datasets/BootsofLagrangian/danbooru-multitier-captions-202606)(CC-BY-4.0、Giyeong Oh)、[CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq](https://huggingface.co/datasets/CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq)(CC-BY-SA-4.0) \|
	\| タグ列 \| 上記の Danbooru タグ、[FredZhang7/anime-prompts-180K](https://huggingface.co/datasets/FredZhang7/anime-prompts-180K)(CreativeML OpenRAIL-M)。絵師名は除去 \|
	\| 一般 \| [ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions](https://huggingface.co/datasets/ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions)(MIT)、[data-is-better-together/open-image-preferences-v1](https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1)(Apache-2.0)、[nateraw/parti-prompts](https://huggingface.co/datasets/nateraw/parti-prompts)(Apache-2.0) \|
	\| 文字描画(日 60% / 英 25% / 中 15%)、JSON レイアウト、短文 \| テンプレートから生成 \|

	## うまくいかなかったこと(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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)(Apache-2.0)、初期値: [lodestones/MOSS-Image-SFT-8-steps-dmd](https://huggingface.co/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、[コード](https://github.com/Yzmblog/DMAD))
	- Adversarial Distribution Matching / DMDX(arXiv 2507.18569)
	- 前作: [Anima の DMDX・DMD2・PCM 蒸留の記録](https://note.com/darask/n/neef98f000087)