rapid-moss: MOSS-Image-SFT 4-step adapter (data-space DMD + DMAD, it 500 EMA), code, logs, eval

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/compare-1.jpg filter=lfs diff=lfs merge=lfs -text
+images/compare-2.jpg filter=lfs diff=lfs merge=lfs -text
+images/compare-3.jpg filter=lfs diff=lfs merge=lfs -text
+images/training/000300.jpg filter=lfs diff=lfs merge=lfs -text
+images/training/000500.jpg filter=lfs diff=lfs merge=lfs -text
+images/training/000800.jpg filter=lfs diff=lfs merge=lfs -text
+images/training/init.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,168 @@
+---
+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)
+```
+
+- base の transformer は bf16 のまま、LoRA は**マージせずに** fp32 で足してください。差分が bf16 の丸め幅より小さく、bf16 の重みにマージすると効果がほぼ消えます。
+- time path の 3 モジュールは fp32 のまま使います。
+- CFG とネガティブプロンプトは使いません。プロンプトの形式(自然文・タグ列・JSON レイアウト、日本語の文字指定)は MOSS-Image と同じです。
+
+## 結果
+
+左から 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 に切り替えました。
+
+## 制限
+
+- 評価は目視だけで、プロンプトは 17 件です。FID などの数値指標は取っていません。
+- キャラクター LoRA など、他の LoRA と重ねたときの挙動は試していません。
+- 学習プロンプトは 3,199 件と少なく、teacher サンプルは 1 プロンプトにつき 1 枚です。
+
+## ライセンスとクレジット
+
+- 重みとコード: 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)

code/README.md

@@ -0,0 +1,30 @@
+# rapid-moss training code
+
+Distillation of MOSS-Image-SFT into an R-step data-space generator. Python 3.10+, torch >= 2.9, diffusers >= 0.38,
+transformers >= 5.3. The release `pipeline_moss_image.py` (shipped in the MOSS-Image-SFT repo) is imported for the
+transformer class only; the forward pass is reimplemented in `rapid_moss/model.py` (fp32 residual + modulation,
+bf16 matmuls, unmerged LoRA sets, per-block gradient checkpointing, feature taps).
+
+| file | role |
+|---|---|
+| `rapid_moss/model.py` | MOSS DiT forward shared by teacher / generator / fake score / critic (one frozen bf16 base, per-path LoRA and time path); release sampler (`sample_teacher`) |
+| `rapid_moss/nets.py` | `Generator` (data-space iteration, lodestones format I/O), `FakeScore`, DMDX `Discriminator`, DMAD `Critic` |
+| `rapid_moss/textenc.py` | prompt embedding cache (release Qwen3.5 text encoder, same chat template as the pipeline) |
+| `rapid_moss/teachergen.py` | teacher samples (50 steps, CFG 3, negative-branch block-22 skip) for the critic's real side |
+| `rapid_moss/train.py` | training loop: `method=dmad` (used for the release) or `method=dmd` (data-space DMD + DMDX discriminator) |
+| `rapid_moss/sample.py` | side-by-side evaluation (teacher vs lodestones-format adapters) |
+| `tests/test_equiv.py` | CPU fp32 check: forward, skip-22 and the lodestones generator path match the release transformer exactly |
+| `tests/test_dmad.py` | CPU check of the critic wiring (generator phase / critic phase gradients) |
+| `scripts/` | RunPod GraphQL client, pod setup and run scripts, balance watchdog |
+| `configs/dmad.json` | the run that produced the released adapter (best checkpoint = EMA export at iteration 500) |
+
+```
+torchrun --nproc-per-node 8 -m rapid_moss.textenc --model MOSS-Image-SFT --prompts train.jsonl eval.jsonl --out cache
+torchrun --nproc-per-node 8 -m rapid_moss.teachergen --model MOSS-Image-SFT --cache cache --prompts train.jsonl --out teacher --count 2000
+torchrun --nproc-per-node 8 -m rapid_moss.train --config configs/dmad.json --set teacher_dir=teacher
+python -m rapid_moss.sample --model MOSS-Image-SFT --cache cache --prompts eval.jsonl --out eval \
+    --method teacher:50 --method dmd:LODESTONES_DIR:4 --method dmd:runs/dmad/export/000500:4
+```
+
+The prompt set (24,000 prompts: anime captions and tags, general captions, Japanese / English / Chinese text rendering,
+MOSS JSON layouts, short prompts) is not redistributed; see the model card for its sources.

code/configs/dmad.json

@@ -0,0 +1,16 @@
+{
+ "method": "dmad",
+ "out_dir": "/workspace/runs/dmad",
+ "teacher_dir": ["/workspace/teacher", "/workspace/teacher2"],
+ "iterations": 2500,
+ "micro": 2,
+ "critic_warmup": 50,
+ "lr_gen": 2e-5, "lr_time": 2e-5, "lr_q": 2e-5,
+ "lr_critic": 4e-5,
+ "ema": 0.998,
+ "log_every": 10,
+ "sample_every": 100,
+ "save_every": 200,
+ "export_every": 500,
+ "eval_count": 8
+}

code/configs/main.json

@@ -0,0 +1,19 @@
+{
+ "out_dir": "/workspace/runs/main",
+ "iterations": 2000,
+ "disc_real": "teacher",
+ "teacher_dir": "/workspace/teacher",
+ "p_gan": 0.5,
+ "micro": 2,
+ "fake_per_gen": 2,
+ "fake_warmup": 150,
+ "adv_start": 50,
+ "adv_warmup": 200,
+ "lr_gen": 1.5e-5, "lr_time": 1.5e-5, "lr_q": 1.5e-5,
+ "ema": 0.998,
+ "log_every": 10,
+ "sample_every": 100,
+ "save_every": 200,
+ "export_every": 400,
+ "eval_count": 8
+}

code/configs/smoke.json

@@ -0,0 +1,14 @@
+{
+ "out_dir": "/workspace/runs/smoke",
+ "iterations": 24,
+ "micro": 1,
+ "fake_per_gen": 2,
+ "fake_warmup": 4,
+ "adv_start": 2,
+ "adv_warmup": 4,
+ "log_every": 2,
+ "sample_every": 12,
+ "save_every": 24,
+ "export_every": 24,
+ "eval_count": 4
+}

code/rapid_moss/model.py

@@ -0,0 +1,234 @@
+"""MOSS-Image DiT for distillation.
+
+One frozen bf16 base (the release MOSSImageTransformer) is shared by the teacher, the
+data-space generator, the fake score and the discriminator backbone; each path only
+differs by its LoRA, its time path and optional extras. The residual stream and the
+modulation run in fp32 (the release trains fp32-master + autocast), matmuls in bf16.
+"""
+from __future__ import annotations
+
+import importlib.util
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+BF16 = torch.bfloat16
+PROJ = ('q', 'k', 'v', 'o', 'ffn_up_gate', 'ffn_down')
+
+
+def load_release_module(model_dir):
+    """Import the release pipeline_moss_image.py (read and reviewed; no side effects on import)."""
+    path = Path(model_dir) / 'pipeline_moss_image.py'
+    spec = importlib.util.spec_from_file_location('pipeline_moss_image', path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def load_base(model_dir, device):
+    release = load_release_module(model_dir)
+    tr = release.MOSSImageTransformer.from_pretrained(Path(model_dir) / 'transformer', torch_dtype=BF16,
+                                                      local_files_only=True)
+    return Moss(tr.to(device), release)
+
+
+def rms(x, weight, eps=1e-6):
+    y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps)
+    return y.to(x.dtype) * weight
+
+
+def situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def proj_dims(dim, ffn_dim):
+    return {'q': (dim, dim), 'k': (dim, dim), 'v': (dim, dim), 'o': (dim, dim),
+            'ffn_up_gate': (dim, 2 * ffn_dim), 'ffn_down': (ffn_dim, dim)}
+
+
+class LoRA(nn.Module):
+    """Stacked LoRA for the projections of blocks [start, layers): A [n, r, in], B [n, out, r]; y += (alpha/r) B A x."""
+
+    def __init__(self, dims, layers, rank, alpha, start=0):
+        super().__init__()
+        self.rank, self.alpha, self.scale, self.start = rank, alpha, alpha / rank, start
+        layers -= start
+        self.A = nn.ParameterDict()
+        self.B = nn.ParameterDict()
+        for name, (fan_in, fan_out) in dims.items():
+            bound = 1 / math.sqrt(fan_in)
+            self.A[name] = nn.Parameter(torch.empty(layers, rank, fan_in).uniform_(-bound, bound))
+            self.B[name] = nn.Parameter(torch.zeros(layers, fan_out, rank))
+        self._cache = None
+
+    def delta(self, x, layer, name):
+        layer -= self.start
+        if self._cache is not None and not torch.is_grad_enabled():
+            a, b = self._cache[name]
+            a, b = a[layer], b[layer]
+        else:
+            a, b = self.A[name][layer].to(x.dtype), self.B[name][layer].to(x.dtype)
+        return F.linear(F.linear(x, a), b) * self.scale
+
+    def freeze_cache(self):
+        """bf16 copies used by no-grad forwards only; drop_cache() before the next optimizer step."""
+        with torch.no_grad():
+            self._cache = {n: (self.A[n].to(BF16), self.B[n].to(BF16)) for n in self.A}
+
+    def drop_cache(self):
+        self._cache = None
+
+
+class TimePath(nn.Module):
+    """time_embedding MLP + both modulation biases (fp32); modulation weights stay the base's."""
+
+    NAMES = ('time_embedding.0.weight', 'time_embedding.0.bias', 'time_embedding.2.weight',
+             'time_embedding.2.bias', 'shared_block_modulation.bias', 'video_head_modulation.bias')
+
+    def __init__(self, tensors, trainable):
+        super().__init__()
+        self.p = nn.ParameterDict({k.replace('.', '_'): nn.Parameter(tensors[k].detach().float().clone(),
+                                                                     requires_grad=trainable)
+                                   for k in self.NAMES})
+
+    def __getitem__(self, name):
+        return self.p[name.replace('.', '_')]
+
+    def state(self):
+        return {k: self[k].detach() for k in self.NAMES}
+
+
+@dataclass
+class Context:
+    tokens: torch.Tensor  # [1, L, D] text_embedder output (bf16)
+    length: int
+
+
+class Moss(nn.Module):
+    def __init__(self, tr, release):
+        super().__init__()
+        self.tr = tr.requires_grad_(False).eval()
+        self.release = release
+        cfg = tr.config
+        self.dim, self.heads, self.depth = cfg.dim, cfg.num_heads, cfg.num_layers
+        self.dims = proj_dims(cfg.dim, cfg.ffn_dim)
+        sd = dict(tr.named_parameters())
+        self.base_time = TimePath(sd, trainable=False)
+        self.register_buffer('mod_w', tr.shared_block_modulation.weight.detach().float(), persistent=False)
+        self.register_buffer('head_w', tr.video_head_modulation.weight.detach().float(), persistent=False)
+        self._rope = {}
+
+    # ---- conditioning ----
+    def rope(self, h, w, length):
+        key = (h, w, length)
+        if key not in self._rope:
+            if len(self._rope) > 64:
+                self._rope.clear()
+            with torch.no_grad():
+                self._rope[key] = self.tr.rope(h, w, length, self.mod_w.device).to(BF16)
+        return self._rope[key]
+
+    @torch.no_grad()
+    def context(self, emb):
+        """emb [L, text_dim] text-encoder hidden states -> Context (text refinement is prompt-only)."""
+        length = emb.shape[0]
+        rope = self.tr.rope(1, 1, length, self.mod_w.device).to(BF16)[:, 1:]
+        tokens = self.tr.text_embedder(emb[None].to(device=self.mod_w.device, dtype=BF16), rope)
+        return Context(tokens, length)
+
+    def time_mod(self, t, tp, extra=None):
+        half = self.tr.config.freq_dim // 2
+        t = torch.as_tensor(t, device=self.mod_w.device, dtype=torch.float32).reshape(-1)
+        freq = 10000 ** (-torch.arange(half, device=t.device, dtype=torch.float64) / half)
+        phase = t.double()[:, None] * freq
+        emb = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = F.linear(emb, tp['time_embedding.0.weight'], tp['time_embedding.0.bias'])
+        time = F.linear(F.silu(time), tp['time_embedding.2.weight'], tp['time_embedding.2.bias'])
+        if extra is not None:
+            time = time + extra
+        mod = F.linear(time, self.mod_w, tp['shared_block_modulation.bias']).view(-1, 6, self.dim)
+        head = F.linear(time, self.head_w, tp['video_head_modulation.bias']).view(-1, 2, self.dim)
+        return mod, head
+
+    # ---- blocks ----
+    def lin(self, y, i, name, lora):
+        out = F.linear(y, self.tr.layers[i].proj[name].weight)
+        if lora is not None and i >= lora.start:
+            out = out + lora.delta(y, i, name)
+        return out
+
+    def block(self, i, x, rope, mod, lora):
+        layer = self.tr.layers[i]
+        heads = self.heads
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.unbind(1)
+        y = (rms(x, layer.norm1.weight) * (1 + scale_a[:, None]) + shift_a[:, None]).to(BF16)
+        q = rms(self.lin(y, i, 'q', lora).unflatten(-1, (heads, -1)), layer.q_norm.weight).flatten(2)
+        k = rms(self.lin(y, i, 'k', lora).unflatten(-1, (heads, -1)), layer.k_norm.weight).flatten(2)
+        a = self.release._attention(q, k, self.lin(y, i, 'v', lora), rope, heads)
+        x = x + self.lin(a, i, 'o', lora).float() * gate_a[:, None]
+        y = (rms(x, layer.norm2.weight) * (1 + scale_f[:, None]) + shift_f[:, None]).to(BF16)
+        hidden, gate = self.lin(y, i, 'ffn_up_gate', lora).chunk(2, -1)
+        return x + self.lin(situ(hidden, gate), i, 'ffn_down', lora).float() * gate_f[:, None]
+
+    def forward(self, latent, t, ctx, *, lora=None, tp=None, time_extra=None, skip=(), taps=None, ckpt=False):
+        """latent [1, C, h, w] fp32, t on the 0..1000 scale. Returns velocity [1, C, h, w] fp32,
+        or {block: image-token features [1, h*w, D]} when taps is given (stops after the last tap)."""
+        _, _, h, w = latent.shape
+        hw = h * w
+        mod, head = self.time_mod(t, tp if tp is not None else self.base_time, time_extra)
+        rope = self.rope(h, w, ctx.length)
+        video = self.tr.video_embedder(latent.flatten(2).transpose(1, 2).to(BF16))
+        x = torch.cat((video, ctx.tokens), 1).float()
+        last = max(taps) if taps else self.depth - 1
+        use_ckpt = ckpt and torch.is_grad_enabled()
+        feats = {}
+        for i in range(last + 1):
+            if i not in skip:
+                if use_ckpt:
+                    x = checkpoint(self.block, i, x, rope, mod, lora, use_reentrant=False)
+                else:
+                    x = self.block(i, x, rope, mod, lora)
+            if taps and i in taps:
+                feats[i] = x[:, :hw]
+        if taps:
+            return feats
+        x = rms(x[:, :hw], self.tr.final_norm_video.weight)
+        x = x * (1 + head[:, 1, None]) + head[:, 0, None]
+        x = F.linear(x.to(BF16), self.tr.final_linear_video.weight).float()
+        return x.transpose(1, 2).reshape(1, -1, h, w)
+
+    @torch.no_grad()
+    def teacher(self, latent, t, pos, neg, cfg, skip_neg=(22,)):
+        """Release guidance: CFG with block 22 skipped in the negative branch only."""
+        vp = self(latent, t, pos)
+        vn = self(latent, t, neg, skip=skip_neg)
+        return vn + cfg * (vp - vn)
+
+
+def release_sigmas(steps, shift=5.0, total=1000, training_shift=1.0):
+    """MOSSImageScheduler.set_timesteps sigmas (shifted grid snapped to the training grid)."""
+    def grid(count, s):
+        lower = s / (total + s - 1)
+        values = torch.linspace(1, lower, count)
+        return s * values / (1 + (s - 1) * values)
+    sigmas = grid(steps, shift)
+    training = grid(total, training_shift).flip(0)
+    upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+    lower = (upper - 1).clamp(min=0)
+    nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+    return torch.unique_consecutive(training[nearest]).tolist() + [0.0]
+
+
+@torch.no_grad()
+def sample_teacher(model, noise, pos, neg, steps=50, cfg=3.0, shift=5.0):
+    x = noise.float()
+    sigmas = release_sigmas(steps, shift)
+    for s, s_next in zip(sigmas[:-1], sigmas[1:]):
+        v = model.teacher(x, s * 1000, pos, neg, cfg)
+        x = x + (s_next - s) * v
+    return x

code/rapid_moss/nets.py

@@ -0,0 +1,139 @@
+"""Trainable parts: data-space generator, fake score, DMDX discriminator heads; lodestones-format I/O."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+from torch.nn import functional as F
+
+from .model import PROJ, LoRA, TimePath
+
+
+class Generator(nn.Module):
+    """Data-space iteration (arXiv 2610.12102) in the lodestones MOSS form:
+    x = z + Q x0 (channel mix), v = net(x, code_r), x0' = z - v, x0 starts at zero."""
+
+    def __init__(self, model, rank, alpha, steps):
+        super().__init__()
+        self.lora = LoRA(model.dims, model.depth, rank, alpha)
+        self.time = TimePath(model.base_time.state(), trainable=True)
+        channels = model.tr.config.latent_channels
+        self.Q = nn.Parameter(torch.zeros(channels, channels))
+        self.register_buffer('codes', torch.full((steps,), 1000.0))
+        self.steps = steps
+
+    def step(self, model, z, x0, r, ctx, ckpt=False):
+        x = z + torch.einsum('oc,bchw->bohw', self.Q, x0)
+        return z - model(x, self.codes[r], ctx, lora=self.lora, tp=self.time, ckpt=ckpt)
+
+    @torch.no_grad()
+    def rollout(self, model, z, ctx, n):
+        """n generator calls from x0 = 0 (no grad); returns x0 after n steps."""
+        x0 = torch.zeros_like(z)
+        for r in range(n):
+            x0 = self.step(model, z, x0, r, ctx)
+        return x0
+
+    @torch.no_grad()
+    def load_lodestones(self, path):
+        path = Path(path)
+        ds = load_file(path / 'dataspace.safetensors')
+        self.Q.copy_(ds['latent_mix.Q'])
+        self.codes.copy_(ds['step_codes'][:self.steps])
+        te = load_file(path / 'time_embedding.safetensors')
+        for k in TimePath.NAMES:
+            self.time[k].copy_(te[k])
+        lora = load_file(path / 'lora.safetensors')
+        for i in range(self.lora.A[PROJ[0]].shape[0]):
+            for name in PROJ:
+                p = f'diffusion_model.layers.{i}.proj.{name}'
+                down, up = lora[p + '.lora_down.weight'], lora[p + '.lora_up.weight']
+                alpha = float(lora[p + '.alpha'])
+                if down.shape[0] != self.lora.rank:
+                    raise ValueError(f'rank mismatch: file {down.shape[0]} vs {self.lora.rank}')
+                # fold the file's alpha/rank into B so our own scale applies unchanged
+                self.lora.A[name][i].copy_(down)
+                self.lora.B[name][i].copy_(up * (alpha / down.shape[0]) / self.lora.scale)
+
+    @torch.no_grad()
+    def export(self, out_dir, meta=None):
+        """Write lodestones-compatible files (kohya LoRA, fp32 time path, Q + step codes)."""
+        out = Path(out_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        rank = self.lora.rank
+        lora = {}
+        for i in range(self.lora.A[PROJ[0]].shape[0]):
+            for name in PROJ:
+                p = f'diffusion_model.layers.{i}.proj.{name}'
+                lora[p + '.lora_down.weight'] = self.lora.A[name][i].float().contiguous()
+                lora[p + '.lora_up.weight'] = (self.lora.B[name][i].float() * self.lora.scale).contiguous()
+                lora[p + '.alpha'] = torch.tensor(float(rank))
+        info = {'rank': str(rank), 'alpha': f'{float(rank)}', 'steps': str(self.steps), **(meta or {})}
+        save_file(lora, out / 'lora.safetensors', metadata=info)
+        save_file({k: v.float().contiguous() for k, v in self.time.state().items()}, out / 'time_embedding.safetensors')
+        save_file({'latent_mix.Q': self.Q.float().contiguous(), 'step_codes': self.codes.float().contiguous()},
+                  out / 'dataspace.safetensors', metadata=info)
+
+
+class FakeScore(nn.Module):
+    """Generator-distribution velocity: base + LoRA, conditioned on the refinement index r."""
+
+    def __init__(self, model, rank, alpha, steps):
+        super().__init__()
+        self.lora = LoRA(model.dims, model.depth, rank, alpha)
+        self.r_emb = nn.Parameter(torch.zeros(steps, model.dim))
+
+    def __call__(self, model, x, t, ctx, r, ckpt=False):
+        return model(x, t, ctx, lora=self.lora, time_extra=self.r_emb[r][None], ckpt=ckpt)
+
+
+class Head(nn.Module):
+    """DMDX 2D head: 4x4 stride-2 conv, GroupNorm(32), SiLU, then 1x1 to patch logits."""
+
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.conv = nn.Conv2d(dim, hidden, 4, 2, 1)
+        self.norm = nn.GroupNorm(32, hidden)
+        self.out = nn.Conv2d(hidden, 1, 1)
+
+    def forward(self, f):
+        return self.out(F.silu(self.norm(self.conv(f))))
+
+
+class Discriminator(nn.Module):
+    """Trainable heads on frozen-teacher block features (backbone = base, no LoRA)."""
+
+    def __init__(self, dim, taps, hidden):
+        super().__init__()
+        self.taps = tuple(sorted(taps))
+        self.heads = nn.ModuleList([Head(dim, hidden) for _ in self.taps])
+
+    def forward(self, feats, h, w):
+        logits = []
+        for i, head in zip(self.taps, self.heads):
+            f = feats[i].float()
+            f = f * torch.rsqrt(f.square().mean(-1, keepdim=True) + 1e-6)
+            logits.append(head(f.transpose(1, 2).reshape(1, -1, h, w)).flatten(1))
+        return torch.cat(logits, 1)
+
+
+class Critic(nn.Module):
+    """DMAD critic (arXiv 2610.02188), teacher head only: base + LoRA on blocks [start:], a conv head on the
+    last block's image tokens giving one logit per sample (teacher sample vs generator sample)."""
+
+    def __init__(self, model, rank, alpha, start, hidden):
+        super().__init__()
+        self.lora = LoRA(model.dims, model.depth, rank, alpha, start=start)
+        self.last = model.depth - 1
+        self.head = nn.Sequential(nn.Conv2d(model.dim, hidden, 4, 2, 1), nn.GroupNorm(32, hidden), nn.SiLU(),
+                                  nn.Conv2d(hidden, hidden, 4, 2, 1), nn.GroupNorm(32, hidden), nn.SiLU())
+        self.out = nn.Linear(hidden, 1)
+
+    def forward(self, model, x, t, ctx, ckpt=False):
+        _, _, h, w = x.shape
+        f = model(x, t, ctx, lora=self.lora, taps=(self.last,), ckpt=ckpt)[self.last].float()
+        f = f * torch.rsqrt(f.square().mean(-1, keepdim=True) + 1e-6)
+        f = self.head(f.transpose(1, 2).reshape(1, -1, h, w))
+        return self.out(f.mean((2, 3)))[:, 0]

code/rapid_moss/sample.py

@@ -0,0 +1,110 @@
+"""Side-by-side evaluation: release teacher vs data-space generators (lodestones-format dirs).
+
+    python -m rapid_moss.sample --model DIR --cache CACHE --prompts eval.jsonl --out OUT \
+        --method teacher:50 --method dmd:/path/lodestones:8 --method dmd:/path/lodestones:4
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+from .model import load_base, sample_teacher
+from .nets import Generator
+from .textenc import Cache, read_prompts
+from .train import decode
+
+
+def load_generator(model, path, steps, device):
+    with safe_open(Path(path) / 'lora.safetensors', 'pt') as f:
+        key = next(k for k in f.keys() if k.endswith('.lora_down.weight'))
+        rank = f.get_slice(key).get_shape()[0]
+    gen = Generator(model, rank, rank, steps).to(device).requires_grad_(False)
+    gen.load_lodestones(path)
+    gen.lora.freeze_cache()
+    return gen
+
+
+def label_grid(columns, labels, ids, path, cell=512):
+    from PIL import Image, ImageDraw
+    pad = 28
+    sheet = Image.new('RGB', (len(columns) * cell, len(ids) * (cell + pad)), 'white')
+    draw = ImageDraw.Draw(sheet)
+    for c, (images, label) in enumerate(zip(columns, labels)):
+        for r, im in enumerate(images):
+            im = Image.fromarray(im)
+            im.thumbnail((cell, cell))
+            y = r * (cell + pad)
+            draw.text((c * cell + 6, y + 6), f'{label} | {ids[r]}', fill='black')
+            sheet.paste(im, (c * cell + (cell - im.width) // 2, y + pad + (cell - im.height) // 2))
+    sheet.save(path, quality=88)
+
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--model', required=True)
+    ap.add_argument('--cache', required=True)
+    ap.add_argument('--prompts', required=True)
+    ap.add_argument('--out', required=True)
+    ap.add_argument('--method', action='append', required=True)
+    ap.add_argument('--seed', type=int, default=1234)
+    ap.add_argument('--limit', type=int, default=0)
+    args = ap.parse_args()
+    dev = torch.device('cuda')
+    torch.backends.cuda.matmul.allow_tf32 = True
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    model = load_base(args.model, dev)
+    from diffusers import AutoencoderKLFlux2
+    vae = AutoencoderKLFlux2.from_pretrained(Path(args.model) / 'vae', torch_dtype=torch.float32,
+                                             local_files_only=True).to(dev).eval()
+    cache = Cache(args.cache)
+    rows = read_prompts([args.prompts])
+    rows = rows[:args.limit] if args.limit else rows
+    neg = model.context(cache.get('').to(dev))
+    columns, labels, timings = [], [], {}
+    for spec in args.method:
+        kind, *rest = spec.split(':')
+        if kind == 'teacher':
+            steps = int(rest[0]) if rest else 50
+            cfg = float(rest[1]) if len(rest) > 1 else 3.0
+            tag, gen = f'teacher{steps}-cfg{cfg:g}', None
+        elif kind == 'dmd':
+            path, steps = rest[0], int(rest[1])
+            tag = f'{Path(path).name}-{steps}step'
+            gen = load_generator(model, path, steps, dev)
+        else:
+            raise SystemExit(f'unknown method {spec}')
+        (out / tag).mkdir(exist_ok=True)
+        images, secs = [], []
+        for row in rows:
+            ctx = model.context(cache.get(row['id']).to(dev))
+            g = torch.Generator('cpu').manual_seed(args.seed)
+            z = torch.randn(1, model.tr.config.latent_channels, row['h'] // 16, row['w'] // 16, generator=g).to(dev)
+            torch.cuda.synchronize()
+            t0 = time.time()
+            x0 = sample_teacher(model, z, ctx, neg, steps, cfg) if gen is None else gen.rollout(model, z, ctx, steps)
+            torch.cuda.synchronize()
+            secs.append(time.time() - t0)
+            images.append(decode(vae, x0))
+            from PIL import Image
+            Image.fromarray(images[-1]).save(out / tag / f"{row['id']}.png")
+        del gen
+        columns.append(images)
+        labels.append(tag)
+        timings[tag] = round(sum(secs) / len(secs), 3)
+        print(json.dumps({'method': tag, 'dit_s_per_image': timings[tag]}), flush=True)
+    (out / 'timings.json').write_text(json.dumps(timings, indent=1))
+    ids = [r['id'] for r in rows]
+    for start in range(0, len(ids), 6):
+        label_grid([c[start:start + 6] for c in columns], labels, ids[start:start + 6],
+                   out / f'grid-{start // 6:02d}.jpg')
+
+
+if __name__ == '__main__':
+    main()

code/rapid_moss/teachergen.py

@@ -0,0 +1,113 @@
+"""Teacher samples for the discriminator's real side (DMDX ADP-style: teacher ODE endpoints).
+
+    torchrun --nproc-per-node N -m rapid_moss.teachergen --model DIR --cache CACHE --prompts train.jsonl \
+        --out TEACHER --count 2400 [--steps 50 --cfg 3]
+Picks `count` prompts stratified by src (fixed seed), a resolution per prompt (its own w/h for
+layouts, else a bucket), runs the release sampler (CFG 3, negative skips block 22) and stores
+the final latent (fp16) per prompt id with index.json {id: {file, w, h}}.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import time
+import zlib
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from .model import load_base, sample_teacher
+from .textenc import Cache, read_prompts
+from .train import DEFAULTS
+
+
+def pick(rows, count, seed):
+    rng = random.Random(seed)
+    by_src = defaultdict(list)
+    for r in rows:
+        by_src[r.get('src', '')].append(r)
+    total = len(rows)
+    chosen = []
+    for src, items in sorted(by_src.items()):
+        rng.shuffle(items)
+        chosen += items[:round(count * len(items) / total)]
+    rng.shuffle(chosen)
+    buckets = DEFAULTS['buckets']
+    for r in chosen:
+        if 'w' not in r or 'h' not in r:
+            r['w'], r['h'], _ = rng.choices(buckets, weights=[b[2] for b in buckets])[0]
+    return chosen[:count]
+
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--model', required=True)
+    ap.add_argument('--cache', required=True)
+    ap.add_argument('--prompts', required=True)
+    ap.add_argument('--out', required=True)
+    ap.add_argument('--count', type=int, default=2400)
+    ap.add_argument('--steps', type=int, default=50)
+    ap.add_argument('--cfg', type=float, default=3.0)
+    ap.add_argument('--seed', type=int, default=7)
+    ap.add_argument('--shard-size', type=int, default=200)
+    ap.add_argument('--exclude', nargs='*', default=[], help='index.json files of earlier runs; their prompts are skipped')
+    args = ap.parse_args()
+    distributed = 'RANK' in os.environ
+    if distributed:
+        dist.init_process_group('nccl')
+    rank, world = (dist.get_rank(), dist.get_world_size()) if distributed else (0, 1)
+    dev = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0)))
+    torch.cuda.set_device(dev)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    cache = Cache(args.cache)
+    skip = set()
+    for path in args.exclude:
+        skip.update(json.loads(Path(path).read_text()))
+    rows = pick([r for r in read_prompts([args.prompts]) if r['id'] in cache and r['id'] not in skip],
+                args.count, args.seed)
+    meta = {r['id']: {'w': r['w'], 'h': r['h']} for r in rows}
+    done = set()
+    for p in out.glob(f'shard-{rank}-*.safetensors'):
+        with safe_open(p, 'pt') as f:
+            done.update(f.keys())
+    mine = [r for i, r in enumerate(rows) if i % world == rank and r['id'] not in done]
+    model = load_base(args.model, dev)
+    neg = model.context(cache.get('').to(dev))
+    shard = len(list(out.glob(f'shard-{rank}-*.safetensors')))
+    buf, t0 = {}, time.time()
+    for n, row in enumerate(mine):
+        g = torch.Generator('cpu').manual_seed(zlib.crc32(f"{args.seed}:{row['id']}".encode()))
+        z = torch.randn(1, model.tr.config.latent_channels, row['h'] // 16, row['w'] // 16, generator=g).to(dev)
+        x0 = sample_teacher(model, z, model.context(cache.get(row['id']).to(dev)), neg, args.steps, args.cfg)
+        buf[row['id']] = x0[0].to('cpu', torch.float16).contiguous()
+        if len(buf) >= args.shard_size or n == len(mine) - 1:
+            save_file(buf, out / f'shard-{rank}-{shard:05d}.safetensors')
+            shard += 1
+            buf = {}
+        if rank == 0 and n % 20 == 0:
+            print(json.dumps({'done': n, 'of': len(mine), 's': round(time.time() - t0, 1)}), flush=True)
+    if distributed:
+        dist.barrier()
+    if rank == 0:
+        index = {}
+        for p in sorted(out.glob('shard-*.safetensors')):
+            with safe_open(p, 'pt') as f:
+                for k in f.keys():
+                    index[k] = {'file': p.name, **meta[k]}
+        (out / 'index.json').write_text(json.dumps(index))
+        print(json.dumps({'index': len(index), 'wanted': len(rows)}), flush=True)
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()

code/rapid_moss/textenc.py

@@ -0,0 +1,122 @@
+"""Prompt embedding cache: release text encoder (Qwen3.5) hidden states, one tensor per prompt.
+
+    torchrun --nproc-per-node N -m rapid_moss.textenc --model DIR --prompts train.jsonl eval.jsonl --out CACHE
+Each rank writes cache/shard-<rank>-<k>.safetensors; rank 0 merges cache/index.json at the end.
+Key "" is the empty negative prompt.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+def read_prompts(paths):
+    rows = []
+    for p in paths:
+        with open(p, encoding='utf-8') as f:
+            rows += [json.loads(line) for line in f if line.strip()]
+    return rows
+
+
+class Encoder:
+    def __init__(self, model_dir, device):
+        from transformers import AutoModel, AutoTokenizer
+        self.tok = AutoTokenizer.from_pretrained(Path(model_dir) / 'tokenizer', local_files_only=True)
+        self.enc = AutoModel.from_pretrained(Path(model_dir) / 'text_encoder', dtype=torch.bfloat16,
+                                             local_files_only=True).to(device).eval()
+        self.device = device
+
+    @torch.no_grad()
+    def __call__(self, prompt, max_len=4096):
+        # Same as MOSSImagePipeline.encode_prompt for one prompt.
+        text = self.tok.apply_chat_template([{'role': 'user', 'content': prompt}], tokenize=False,
+                                            add_generation_prompt=True, enable_thinking=True)
+        inputs = self.tok(text, return_tensors='pt', truncation=True, max_length=max_len).to(self.device)
+        return self.enc(**inputs, use_cache=False, return_dict=True).last_hidden_state[0]
+
+
+class Cache:
+    """Read side: id -> [L, 4096] bf16 (CPU, memory-mapped shards)."""
+
+    def __init__(self, root):
+        self.root = Path(root)
+        self.index = json.loads((self.root / 'index.json').read_text())
+        self._files = {}
+
+    def __contains__(self, key):
+        return key in self.index
+
+    def get(self, key):
+        name = self.index[key]
+        if name not in self._files:
+            self._files[name] = safe_open(self.root / name, 'pt')
+        return self._files[name].get_tensor(key)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--model', required=True)
+    ap.add_argument('--prompts', nargs='+', required=True)
+    ap.add_argument('--out', required=True)
+    ap.add_argument('--shard-size', type=int, default=500)
+    args = ap.parse_args()
+    distributed = 'RANK' in os.environ
+    if distributed:
+        dist.init_process_group('nccl')
+    rank, world = (dist.get_rank(), dist.get_world_size()) if distributed else (0, 1)
+    device = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0)))
+    torch.cuda.set_device(device)
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    rows = read_prompts(args.prompts)
+    keys = [''] + [r['id'] for r in rows]
+    texts = dict([('', '')] + [(r['id'], r['prompt']) for r in rows])
+    if len(set(keys)) != len(keys):
+        raise SystemExit('duplicate prompt ids')
+    done = set()
+    for p in out.glob(f'shard-{rank}-*.safetensors'):
+        with safe_open(p, 'pt') as f:
+            done.update(f.keys())
+    mine = [k for i, k in enumerate(keys) if i % world == rank and k not in done]
+    enc = Encoder(args.model, device)
+    buf, shard, t0 = {}, len(list(out.glob(f'shard-{rank}-*.safetensors'))), time.time()
+
+    def flush():
+        nonlocal buf, shard
+        if buf:
+            save_file(buf, out / f'shard-{rank}-{shard:05d}.safetensors')
+            shard += 1
+            buf = {}
+
+    for n, key in enumerate(mine):
+        buf[key] = enc(texts[key]).to('cpu', torch.bfloat16).contiguous()
+        if len(buf) >= args.shard_size:
+            flush()
+        if rank == 0 and n % 200 == 0:
+            print(json.dumps({'encoded': n, 'of': len(mine), 's': round(time.time() - t0, 1)}), flush=True)
+    flush()
+    if distributed:
+        dist.barrier()
+    if rank == 0:
+        index = {}
+        for p in sorted(out.glob('shard-*.safetensors')):
+            with safe_open(p, 'pt') as f:
+                for k in f.keys():
+                    index[k] = p.name
+        missing = [k for k in keys if k not in index]
+        (out / 'index.json').write_text(json.dumps(index))
+        print(json.dumps({'index': len(index), 'missing': len(missing)}), flush=True)
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()

code/rapid_moss/train.py

@@ -0,0 +1,499 @@
+"""Distillation of MOSS-Image-SFT to an R-step data-space generator (arXiv 2610.12102).
+
+    torchrun --nproc-per-node N -m rapid_moss.train --config configs/main.json [--set key=value ...]
+
+method=dmad (arXiv 2610.02188): no fake score and no teacher evaluation during training. A critic
+(base + LoRA on the later blocks + a head on the last block) separates noised teacher samples
+(rapid_moss.teachergen) from noised generator samples of the same prompts with a logistic loss;
+the generator minimises -critic(noised x0). Prompts are drawn from the teacher-sample pool.
+
+method=dmd, per iteration: fake_per_gen fake-score updates, then (after fake_warmup) one generator update
+with the DMD direction sg[v_real - v_fake] on x0 (arXiv 2610.12102) plus the ADM generator loss
+-D(x_fake at tau - dt) (arXiv 2507.18569), and one discriminator hinge update on the same pair
+(x_real = x_tau - dt * v_teacher_cfg, x_fake = x_tau - dt * v_fake). Each GPU runs `micro`
+single-sample micro-steps per update (samples have different resolutions and text lengths).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import time
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch.nn import functional as F
+
+from .model import load_base
+from .nets import Critic, Discriminator, FakeScore, Generator
+from .textenc import Cache, read_prompts
+
+DEFAULTS = {
+    'model_dir': '/workspace/models/MOSS-Image-SFT',
+    'lodestones_dir': '/workspace/models/MOSS-Image-SFT-8-steps-dmd',
+    'cache_dir': '/workspace/cache',
+    'prompts': '/workspace/prompts/train.jsonl',
+    'eval_prompts': '/workspace/prompts/eval.jsonl',
+    'out_dir': '/workspace/runs/main',
+    'steps': 4,
+    'gen_rank': 128, 'gen_alpha': 128,
+    'fake_rank': 64, 'fake_alpha': 64,
+    'disc_taps': [5, 11, 17, 23, 29, 35], 'disc_hidden': 256,
+    'iterations': 3000,
+    'micro': 2,
+    'fake_per_gen': 2,
+    'fake_warmup': 200,
+    'lr_gen': 1e-5, 'lr_time': 1e-5, 'lr_q': 1e-5, 'lr_fake': 1e-4, 'lr_disc': 5e-5,
+    'betas': [0.0, 0.99],
+    'ema': 0.999,
+    'cfg': [3.0, 3.0],
+    'skip_neg': [22],
+    'tau_beta': [2.5, 1.0], 'tau_clamp': [0.03, 0.985],
+    'dt': 1 / 64,
+    'adv_weight': 0.5, 'adv_start': 50, 'adv_warmup': 200,
+    # discriminator real side: 'adm' = teacher ODE step from the same noised generator sample
+    # (ADM, x_real/x_fake at tau - dt); 'teacher' = noised teacher samples of the same prompt
+    # (ADP-style endpoints from rapid_moss.teachergen), drawn on a p_gan share of generator steps
+    'disc_real': 'adm', 'teacher_dir': '/workspace/teacher', 'p_gan': 0.5, 'gan_t': [0.02, 0.98],
+    'clip_gen': 0.0, 'clip_critic': 0.0,  # 0 = no clipping; a non-finite gradient norm always skips the step
+    'buckets': [[1024, 1024, 0.35], [896, 1152, 0.15], [832, 1216, 0.15], [1152, 896, 0.1],
+                [1216, 832, 0.1], [768, 1344, 0.075], [1344, 768, 0.075]],
+    'method': 'dmd',
+    'critic_rank': 128, 'critic_alpha': 128, 'critic_start': 21, 'critic_hidden': 512,
+    'critic_shift': 5.0, 'critic_sigma': [0.02, 0.98], 'lr_critic': 4e-5, 'wd_critic': 0.01,
+    'critic_warmup': 50,
+    'ckpt': True,
+    'seed': 0,
+    'log_every': 10, 'sample_every': 100, 'save_every': 250, 'export_every': 500,
+    'eval_count': 8, 'eval_seed': 1234,
+}
+
+
+def parse_config():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--config')
+    ap.add_argument('--set', nargs='*', default=[])
+    args = ap.parse_args()
+    cfg = dict(DEFAULTS)
+    if args.config:
+        cfg.update(json.loads(Path(args.config).read_text()))
+    for item in args.set:
+        key, value = item.split('=', 1)
+        if key not in DEFAULTS:
+            raise SystemExit(f'unknown key {key}')
+        try:
+            cfg[key] = json.loads(value)
+        except json.JSONDecodeError:
+            cfg[key] = value
+    return cfg
+
+
+class Dist:
+    def __init__(self):
+        self.on = 'RANK' in os.environ
+        if self.on:
+            dist.init_process_group('nccl')
+        self.rank = dist.get_rank() if self.on else 0
+        self.world = dist.get_world_size() if self.on else 1
+        self.device = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0)))
+        torch.cuda.set_device(self.device)
+
+    def mean_grads(self, params):
+        grads = [p.grad for p in params]
+        if self.world == 1:
+            return
+        flat = torch._utils._flatten_dense_tensors(grads)
+        dist.all_reduce(flat)
+        flat /= self.world
+        for g, f in zip(grads, torch._utils._unflatten_dense_tensors(flat, grads)):
+            g.copy_(f)
+
+    def mean(self, values):
+        t = torch.tensor(values, device=self.device, dtype=torch.float64)
+        if self.on:
+            dist.all_reduce(t)
+            t /= self.world
+        return t.tolist()
+
+    def barrier(self):
+        if self.on:
+            dist.barrier()
+
+
+def zero_grads(params):
+    for p in params:
+        if p.grad is None:
+            p.grad = torch.zeros_like(p)
+        else:
+            p.grad.zero_()
+
+
+def load_teacher(dirs, rows, dev):
+    """Scan teachergen shards (any number of dirs); returns the prompt rows that have a teacher latent
+    (with w/h from the latent shape) and a getter id -> [1, C, h, w] fp32 on dev."""
+    from safetensors import safe_open
+    files, where = [], {}
+    for d in [dirs] if isinstance(dirs, str) else dirs:
+        for path in sorted(Path(d).glob('shard-*.safetensors')):
+            try:
+                f = safe_open(path, 'pt')
+            except Exception as e:  # a shard cut off when its teachergen run was stopped
+                print(json.dumps({'skip_shard': str(path), 'error': str(e)[:120]}), flush=True)
+                continue
+            files.append(f)
+            for k in f.keys():
+                _, h, w = f.get_slice(k).get_shape()
+                where[k] = (len(files) - 1, w * 16, h * 16)
+    pool = [dict(r, w=where[r['id']][1], h=where[r['id']][2]) for r in rows if r['id'] in where]
+
+    def get(key):
+        return files[where[key][0]].get_tensor(key)[None].to(dev, torch.float32)
+    return pool, get
+
+
+def decode(vae, x0):
+    z = x0.to(vae.dtype)
+    bn, eps = vae.bn, vae.config.batch_norm_eps
+    z = z * (bn.running_var + eps).sqrt().view(1, -1, 1, 1) + bn.running_mean.view(1, -1, 1, 1)
+    b, c, h, w = z.shape
+    z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+    img = vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+    return ((img[0].float() + 1) * 127.5).round().byte().permute(1, 2, 0).cpu().numpy()
+
+
+def contact_sheet(images, path, cell=384, cols=4):
+    from PIL import Image
+    rows = math.ceil(len(images) / cols)
+    sheet = Image.new('RGB', (cols * cell, rows * cell), 'white')
+    for i, im in enumerate(images):
+        im = Image.fromarray(im)
+        im.thumbnail((cell, cell))
+        sheet.paste(im, ((i % cols) * cell + (cell - im.width) // 2, (i // cols) * cell + (cell - im.height) // 2))
+    sheet.save(path, quality=90)
+
+
+def main():
+    C = parse_config()
+    D = Dist()
+    dev = D.device
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    out = Path(C['out_dir'])
+    if D.rank == 0:
+        (out / 'samples').mkdir(parents=True, exist_ok=True)
+        (out / 'ckpt').mkdir(exist_ok=True)
+        (out / 'config.json').write_text(json.dumps(C, indent=1))
+    R = C['steps']
+
+    model = load_base(C['model_dir'], dev)
+    gen = Generator(model, C['gen_rank'], C['gen_alpha'], R).to(dev)
+    gen.load_lodestones(C['lodestones_dir'])
+    ema = Generator(model, C['gen_rank'], C['gen_alpha'], R).to(dev).requires_grad_(False)
+    ema.load_state_dict(gen.state_dict())
+    betas = tuple(C['betas'])
+    gen_params = list(gen.parameters())
+    opt_gen = torch.optim.AdamW([{'params': list(gen.lora.parameters()), 'lr': C['lr_gen']},
+                                 {'params': list(gen.time.parameters()), 'lr': C['lr_time']},
+                                 {'params': [gen.Q], 'lr': C['lr_q']}], betas=betas, weight_decay=0.0)
+    parts = {'gen': gen, 'ema': ema, 'opt_gen': opt_gen}
+    dmad = C['method'] == 'dmad'
+    if dmad:
+        critic = Critic(model, C['critic_rank'], C['critic_alpha'], C['critic_start'], C['critic_hidden']).to(dev)
+        opt_critic = torch.optim.AdamW(critic.parameters(), lr=C['lr_critic'], betas=betas,
+                                       weight_decay=C['wd_critic'])
+        parts.update(critic=critic, opt_critic=opt_critic)
+        trainables = (gen_params, list(critic.parameters()))
+    elif C['method'] == 'dmd':
+        fake = FakeScore(model, C['fake_rank'], C['fake_alpha'], R).to(dev)
+        disc = Discriminator(model.dim, C['disc_taps'], C['disc_hidden']).to(dev)
+        opt_fake = torch.optim.AdamW(fake.parameters(), lr=C['lr_fake'], betas=betas, weight_decay=0.0)
+        opt_disc = torch.optim.AdamW(disc.parameters(), lr=C['lr_disc'], betas=betas, weight_decay=0.0)
+        parts.update(fake=fake, disc=disc, opt_fake=opt_fake, opt_disc=opt_disc)
+        trainables = (gen_params, list(fake.parameters()), list(disc.parameters()))
+    else:
+        raise SystemExit(f"method must be dmd or dmad, got {C['method']}")
+    for params in trainables:
+        zero_grads(params)
+
+    cache = Cache(C['cache_dir'])
+    rows = [r for r in read_prompts([C['prompts']]) if r['id'] in cache]
+    evals = [r for r in read_prompts([C['eval_prompts']]) if r['id'] in cache][:C['eval_count']]
+    neg = model.context(cache.get('').to(dev))
+    buckets = C['buckets']
+    teacher = None
+    if dmad or C['disc_real'] == 'teacher':
+        trows, teacher = load_teacher(C['teacher_dir'], rows, dev)
+        if D.rank == 0:
+            print(json.dumps({'teacher_samples': len(trows)}), flush=True)
+    elif C['disc_real'] != 'adm':
+        raise SystemExit(f"disc_real must be adm or teacher, got {C['disc_real']}")
+
+    start, gen_updates = 0, 0
+    latest = out / 'ckpt' / 'latest.pt'
+    if latest.exists():
+        state = torch.load(latest, map_location=dev, weights_only=False)
+        for key, mod in parts.items():
+            mod.load_state_dict(state[key])
+        start, gen_updates = state['it'], state['gen_updates']
+        del state
+        if D.rank == 0:
+            print(json.dumps({'resume': start, 'gen_updates': gen_updates}), flush=True)
+    rng = random.Random(C['seed'] * 1_000_003 + D.rank * 7919 + start)
+    torch.manual_seed(C['seed'] * 1_000_003 + D.rank * 7919 + start)
+    beta = torch.distributions.Beta(*[torch.tensor(float(v)) for v in C['tau_beta']])
+    lo, hi = C['tau_clamp']
+    dt = C['dt']
+    vae = None
+
+    def draw(pool=None):
+        pool = pool or rows
+        row = pool[rng.randrange(len(pool))]
+        if 'w' in row and 'h' in row:
+            w, h = row['w'], row['h']
+        else:
+            w, h, _ = rng.choices(buckets, weights=[b[2] for b in buckets])[0]
+        ctx = model.context(cache.get(row['id']).to(dev))
+        z = torch.randn(1, model.tr.config.latent_channels, h // 16, w // 16, device=dev)
+        return ctx, z, row
+
+    def tau():
+        return float(beta.sample().clamp(lo, hi))
+
+    def r_index(phase, m):
+        return ((phase * D.world + D.rank) * C['micro'] + m) % R
+
+    def fake_step(phase):
+        losses = []
+        gen.lora.freeze_cache()
+        for m in range(C['micro']):
+            ctx, z, _ = draw()
+            r = r_index(phase, m)
+            x0 = gen.rollout(model, z, ctx, r + 1)
+            t = tau()
+            eps = torch.randn_like(z)
+            v = fake(model, (1 - t) * x0 + t * eps, t * 1000, ctx, r, ckpt=C['ckpt'])
+            loss = F.mse_loss(v, eps - x0)
+            (loss / C['micro']).backward()
+            losses.append(loss.item())
+        gen.lora.drop_cache()
+        D.mean_grads(list(fake.parameters()))
+        opt_fake.step()
+        opt_fake.zero_grad(set_to_none=False)
+        return sum(losses) / len(losses)
+
+    def gen_step(phase):
+        ramp = 0.0 if gen_updates < C['adv_start'] else min(1.0, (gen_updates - C['adv_start'] + 1) / max(1, C['adv_warmup']))
+        stats = {k: 0.0 for k in ('dmd_rms', 'adv_ratio', 'lam', 'd_loss', 'd_real', 'd_fake', 'd_acc')}
+        n_gan = 0
+        fake.lora.freeze_cache()
+        gen.lora.freeze_cache()
+        for m in range(C['micro']):
+            use_teacher = teacher is not None and rng.random() < C['p_gan']
+            ctx, z, row = draw(trows if use_teacher else None)
+            r = r_index(phase, m)
+            x_prev = gen.rollout(model, z, ctx, r)
+            x0_hat = gen.step(model, z, x_prev, r, ctx, ckpt=C['ckpt'])
+            x0d = x0_hat.detach()
+            t = tau()
+            eps = torch.randn_like(z)
+            xt = (1 - t) * x0d + t * eps
+            w = rng.uniform(*C['cfg'])
+            with torch.no_grad():
+                v_real = model.teacher(xt, t * 1000, ctx, neg, w, tuple(C['skip_neg']))
+                v_fake = fake(model, xt, t * 1000, ctx, r)
+            g_dmd = (v_real - v_fake) / x0d.numel()
+            stats['dmd_rms'] += (v_real - v_fake).square().mean().sqrt().item()
+            _, _, h, wd = z.shape
+            leaf = x0d.clone().requires_grad_(True)
+            if teacher is None:
+                # ADM: generator loss -D(x_tau - dt * v_fake) at tau - dt; gradient taken w.r.t. x0 only
+                s = t - dt
+                x_fake = (1 - t) * leaf + t * eps - dt * v_fake
+                x_real = xt - dt * v_real
+            elif use_teacher:
+                s = rng.uniform(*C['gan_t'])
+                x_fake = (1 - s) * leaf + s * torch.randn_like(z)
+                x_real = (1 - s) * teacher(row['id']) + s * torch.randn_like(z)
+            else:
+                x0_hat.backward(g_dmd / C['micro'])
+                del x0_hat
+                continue
+            n_gan += 1
+            feats = model(x_fake, s * 1000, ctx, taps=disc.taps, ckpt=C['ckpt'])
+            logits_f = disc(feats, h, wd)
+            g_adv, = torch.autograd.grad(-logits_f.mean(), leaf)
+            ratio = (g_adv.norm() / (g_dmd.norm() + 1e-30)).item()
+            lam = C['adv_weight'] * ramp / max(ratio, 1e-12)
+            stats['adv_ratio'] += ratio
+            stats['lam'] += lam
+            x0_hat.backward((g_dmd + lam * g_adv) / C['micro'])
+            # discriminator hinge on the same pair (fake features reused from the pass above)
+            f_fake = {k: v.detach() for k, v in feats.items()}
+            del x0_hat, feats, logits_f
+            with torch.no_grad():
+                f_real = model(x_real, s * 1000, ctx, taps=disc.taps)
+            d_real, d_fake = disc(f_real, h, wd), disc(f_fake, h, wd)
+            d_loss = F.relu(1 - d_real).mean() + F.relu(1 + d_fake).mean()
+            (d_loss / C['micro']).backward()
+            stats['d_loss'] += d_loss.item()
+            stats['d_real'] += d_real.mean().item()
+            stats['d_fake'] += d_fake.mean().item()
+            stats['d_acc'] += 0.5 * ((d_real > 0).float().mean() + (d_fake < 0).float().mean()).item()
+        gen.lora.drop_cache()
+        fake.lora.drop_cache()
+        D.mean_grads(gen_params)
+        D.mean_grads(list(disc.parameters()))
+        gnorm = torch.nn.utils.clip_grad_norm_(gen_params, C['clip_gen'] or 1e12)
+        if torch.isfinite(gnorm):
+            opt_gen.step()
+            opt_disc.step()
+        opt_gen.zero_grad(set_to_none=False)
+        opt_disc.zero_grad(set_to_none=False)
+        with torch.no_grad():
+            for pe, p in zip(ema.parameters(), gen.parameters()):
+                pe.lerp_(p, 1 - C['ema'])
+        stats = {k: v / (C['micro'] if k == 'dmd_rms' else max(n_gan, 1)) for k, v in stats.items()}
+        stats['gan_frac'] = n_gan / C['micro']
+        stats['gnorm'] = gnorm.item()
+        stats['ramp'] = ramp
+        return stats
+
+    def critic_sigma():
+        u, k = rng.uniform(*C['critic_sigma']), C['critic_shift']
+        return k * u / (1 + (k - 1) * u)
+
+    def dmad_step(phase, gen_on):
+        stats = {k: 0.0 for k in ('g_logit', 'c_loss', 'd_teacher', 'd_gen', 'c_acc')}
+        stash = []
+        gen.lora.freeze_cache()
+        critic.requires_grad_(False)
+        for m in range(C['micro']):
+            ctx, z, row = draw(trows)
+            r = r_index(phase, m)
+            if gen_on:
+                x0_hat = gen.step(model, z, gen.rollout(model, z, ctx, r), r, ctx, ckpt=C['ckpt'])
+                s = critic_sigma()
+                logit = critic(model, (1 - s) * x0_hat + s * torch.randn_like(z), s * 1000, ctx, ckpt=C['ckpt'])
+                (-logit.mean() / C['micro']).backward()
+                stats['g_logit'] += logit.item()
+                x0 = x0_hat.detach()
+                del x0_hat, logit
+            else:
+                x0 = gen.rollout(model, z, ctx, r + 1)
+            stash.append((ctx, row, x0))
+        gen.lora.drop_cache()
+        if gen_on:
+            D.mean_grads(gen_params)
+            gnorm = torch.nn.utils.clip_grad_norm_(gen_params, C['clip_gen'] or 1e12)
+            if torch.isfinite(gnorm):
+                opt_gen.step()
+                with torch.no_grad():
+                    for pe, p in zip(ema.parameters(), gen.parameters()):
+                        pe.lerp_(p, 1 - C['ema'])
+                stats['gnorm'] = gnorm.item() * C['micro']
+            else:
+                stats['skip_gen'] = C['micro']
+            opt_gen.zero_grad(set_to_none=False)
+        critic.requires_grad_(True)
+        for ctx, row, x0 in stash:
+            s = critic_sigma()
+            d_q = critic(model, (1 - s) * teacher(row['id']) + s * torch.randn_like(x0), s * 1000, ctx, ckpt=C['ckpt'])
+            d_g = critic(model, (1 - s) * x0 + s * torch.randn_like(x0), s * 1000, ctx, ckpt=C['ckpt'])
+            loss = F.softplus(-d_q).mean() + F.softplus(d_g).mean()
+            (loss / C['micro']).backward()
+            stats['c_loss'] += loss.item()
+            stats['d_teacher'] += d_q.item()
+            stats['d_gen'] += d_g.item()
+            stats['c_acc'] += 0.5 * ((d_q > 0).float() + (d_g < 0).float()).item()
+        critic_params = list(critic.parameters())
+        D.mean_grads(critic_params)
+        cnorm = torch.nn.utils.clip_grad_norm_(critic_params, C['clip_critic'] or 1e12)
+        if torch.isfinite(cnorm):
+            opt_critic.step()
+            stats['cnorm'] = cnorm.item() * C['micro']
+        else:
+            stats['skip_critic'] = C['micro']
+        opt_critic.zero_grad(set_to_none=False)
+        return {k: v / C['micro'] for k, v in stats.items()}
+
+    @torch.no_grad()
+    def sample(tag):
+        nonlocal vae
+        if D.rank == 0:
+            if vae is None:
+                from diffusers import AutoencoderKLFlux2
+                vae = AutoencoderKLFlux2.from_pretrained(Path(C['model_dir']) / 'vae', torch_dtype=torch.float32,
+                                                         local_files_only=True).to(dev).eval()
+            d = out / 'samples' / tag
+            d.mkdir(parents=True, exist_ok=True)
+            images = []
+            for row in evals:
+                g = torch.Generator('cpu').manual_seed(C['eval_seed'])
+                z = torch.randn(1, model.tr.config.latent_channels, row['h'] // 16, row['w'] // 16, generator=g).to(dev)
+                x0 = ema.rollout(model, z, model.context(cache.get(row['id']).to(dev)), R)
+                images.append(decode(vae, x0))
+                from PIL import Image
+                Image.fromarray(images[-1]).save(d / f"{row['id']}.png")
+            contact_sheet(images, out / 'samples' / f'{tag}.jpg')
+        D.barrier()
+
+    def save(it):
+        if D.rank == 0:
+            state = {key: mod.state_dict() for key, mod in parts.items()}
+            state.update(it=it, gen_updates=gen_updates, config=C)
+            tmp = out / 'ckpt' / 'latest.tmp'
+            torch.save(state, tmp)
+            if latest.exists():
+                latest.replace(out / 'ckpt' / 'prev.pt')  # one older state survives a bad save
+            tmp.replace(latest)
+        D.barrier()
+
+    if start == 0 and C['sample_every']:
+        sample('init')
+    t_last, log = time.time(), []
+    for it in range(start, C['iterations']):
+        entry = {'it': it}
+        if dmad:
+            gen_on = it >= C['critic_warmup']
+            entry.update(dmad_step(it, gen_on))
+            gen_updates += gen_on
+        else:
+            entry['fake'] = sum(fake_step(it * C['fake_per_gen'] + k) for k in range(C['fake_per_gen'])) / C['fake_per_gen']
+            if it >= C['fake_warmup']:
+                entry.update(gen_step(it))
+                gen_updates += 1
+        log.append(entry)
+        done = it + 1
+        if done % C['log_every'] == 0:
+            keys = sorted({k for e in log for k in e if k != 'it'})
+            means = D.mean([sum(e.get(k, 0.0) for e in log) / len(log) for k in keys])
+            now = time.time()
+            if D.rank == 0:
+                rec = {'it': done, 'gen_updates': gen_updates, 's_per_it': round((now - t_last) / len(log), 2),
+                       'mem_gib': round(torch.cuda.max_memory_allocated() / 2**30, 1),
+                       **{k: round(v, 6) for k, v in zip(keys, means)}}
+                print(json.dumps(rec), flush=True)
+                with open(out / 'log.jsonl', 'a') as f:
+                    f.write(json.dumps(rec) + '\n')
+            t_last, log = now, []
+        if C['sample_every'] and done % C['sample_every'] == 0:
+            sample(f'{done:06d}')
+        if C['save_every'] and done % C['save_every'] == 0:
+            save(done)
+        if C['export_every'] and done % C['export_every'] == 0 and D.rank == 0:
+            ema.export(out / 'export' / f'{done:06d}', {'iteration': str(done)})
+    save(C['iterations'])
+    if D.rank == 0:
+        ema.export(out / 'export' / 'final', {'iteration': str(C['iterations'])})
+        gen.export(out / 'export' / 'final-raw', {'iteration': str(C['iterations'])})
+    if C['sample_every']:
+        sample('final')
+    if D.on:
+        dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()

code/scripts/main_run.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+# Runs on the main pod (all GPUs): setup -> text cache -> [teacher samples] -> training.
+#   DISC_REAL=adm|teacher TEACHER_COUNT=2400 bash scripts/main_run.sh [--set key=value ...]
+# Expects code at /workspace/code and prompts at /workspace/prompts/{train,eval}.jsonl.
+set -euo pipefail
+cd /workspace/code
+N=$(nvidia-smi -L | wc -l)
+M=/workspace/models/MOSS-Image-SFT
+echo "== setup $(date -u +%T)"
+bash scripts/pod_setup.sh
+echo "== text cache $(date -u +%T)"
+torchrun --nproc-per-node "$N" -m rapid_moss.textenc --model $M \
+    --prompts /workspace/prompts/train.jsonl /workspace/prompts/eval.jsonl --out /workspace/cache
+EXTRA=(disc_real=adm)
+if [[ "${DISC_REAL:-teacher}" == teacher ]]; then
+    echo "== teacher samples $(date -u +%T)"
+    torchrun --nproc-per-node "$N" -m rapid_moss.teachergen --model $M --cache /workspace/cache \
+        --prompts /workspace/prompts/train.jsonl --out /workspace/teacher --count "${TEACHER_COUNT:-2000}"
+    EXTRA=()
+fi
+echo "== train $(date -u +%T)"
+torchrun --nproc-per-node "$N" -m rapid_moss.train --config configs/main.json --set "${EXTRA[@]}" "$@"
+echo "== done $(date -u +%T)"

code/scripts/pod_dmad.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+# On the pod: let the second teacher run go until DEADLINE (UTC, e.g. 21:50), stop it, then train DMAD on all GPUs.
+DEADLINE=${1:?usage: pod_dmad.sh HH:MM}
+until [[ "$(date -u +%H:%M)" > "$DEADLINE" || "$(date -u +%H:%M)" == "$DEADLINE" ]]; do sleep 20; done
+pkill -f "rapid_moss.teachergen"
+while pgrep -f "rapid_moss.teachergen" > /dev/null; do sleep 5; done
+echo "== teacher samples: $(ls /workspace/teacher /workspace/teacher2 | grep -c safetensors) shards $(date -u +%T)"
+cd /workspace/code
+N=$(nvidia-smi -L | wc -l)
+torchrun --nproc-per-node "$N" -m rapid_moss.train --config configs/dmad.json
+echo "== done $(date -u +%T)"

code/scripts/pod_setup.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# Runs on the pod: Python deps + model downloads into /workspace (idempotent).
+set -euo pipefail
+cd /workspace
+export HF_HUB_ENABLE_HF_TRANSFER=0 HF_XET_HIGH_PERFORMANCE=1 PIP_BREAK_SYSTEM_PACKAGES=1
+pip install -q --root-user-action=ignore "diffusers==0.38.0" "transformers==5.17.0" "accelerate>=1.10" \
+    "safetensors>=0.6" "huggingface_hub[hf_xet]>=1.0" pillow flash-linear-attention 2>&1 | tail -3
+python - <<'EOF'
+import torch, diffusers, transformers
+print('torch', torch.__version__, 'cuda', torch.version.cuda, 'diffusers', diffusers.__version__,
+      'transformers', transformers.__version__, [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])
+EOF
+mkdir -p models
+hf download OpenMOSS-Team/MOSS-Image-SFT --local-dir models/MOSS-Image-SFT --exclude "assets/*" > /dev/null
+hf download lodestones/MOSS-Image-SFT-8-steps-dmd lora.safetensors time_embedding.safetensors dataspace.safetensors README.md \
+    --local-dir models/MOSS-Image-SFT-8-steps-dmd > /dev/null
+du -sh models/*
+md5sum models/MOSS-Image-SFT/pipeline_moss_image.py

code/scripts/pod_stop_train.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# On the pod: stop the training run (torchrun + workers) and its pod_dmad.sh wrapper.
+pkill -f "scripts/pod_dmad.sh"
+pkill -f "rapid_moss.train"
+sleep 5
+pgrep -fa "rapid_moss.train" | head -3 || true

code/scripts/pod_teacher2.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+# On the pod: after the first teacher run finishes, generate more teacher samples for other prompts
+# (small shards so a stop loses little; train.py scans shards directly).
+while pgrep -f "rapid_moss.teachergen --model" > /dev/null; do sleep 15; done
+cd /workspace/code
+torchrun --nproc-per-node 8 -m rapid_moss.teachergen --model /workspace/models/MOSS-Image-SFT --cache /workspace/cache \
+    --prompts /workspace/prompts/train.jsonl --out /workspace/teacher2 --count 4000 --seed 8 --shard-size 25 \
+    --exclude /workspace/teacher/index.json

code/scripts/pssh.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+# ssh/rsync helpers for the pod in scripts/pod.env:  pssh.sh CMD...  |  pssh.sh --push SRC DST  |  pssh.sh --pull SRC DST
+set -euo pipefail
+source "$(dirname "$0")/pod.env"
+KEY=$HOME/.runpod/ssh/runpodctl-ssh-key
+OPTS=(-i "$KEY" -p "$POD_PORT" -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=$HOME/.runpod/ssh/known_hosts -o ConnectTimeout=20 -o ServerAliveInterval=30)
+if [[ "${1:-}" == --push ]]; then rsync -rlptz --no-owner --no-group --exclude .git -e "ssh ${OPTS[*]}" "$2" "root@$POD_HOST:$3"
+elif [[ "${1:-}" == --pull ]]; then rsync -rlptz --no-owner --no-group --exclude .git -e "ssh ${OPTS[*]}" "root@$POD_HOST:$2" "$3"
+else exec ssh "${OPTS[@]}" "root@$POD_HOST" "$@"; fi

code/scripts/runpod.py

@@ -0,0 +1,103 @@
+"""Minimal RunPod GraphQL client (API key from ~/.runpod/config.toml).
+
+    python scripts/runpod.py balance
+    python scripts/runpod.py gpus
+    python scripts/runpod.py create NAME --gpu "NVIDIA H100 80GB HBM3" --count 1 [--cloud COMMUNITY] [--volume 300]
+    python scripts/runpod.py status POD_ID        # prints ssh target when running
+    python scripts/runpod.py stop POD_ID | terminate POD_ID
+"""
+import argparse
+import json
+import re
+import sys
+import time
+import urllib.request
+from pathlib import Path
+
+IMAGE = 'runpod/pytorch:1.4.1-cu1281-torch2130-ubuntu2404'
+KEY_FILE = Path.home() / '.runpod' / 'ssh' / 'runpodctl-ssh-key'
+
+
+def api_key():
+    text = (Path.home() / '.runpod' / 'config.toml').read_text()
+    return re.search(r"apikey\s*=\s*['\"]([^'\"]+)", text).group(1)
+
+
+def gql(query, variables=None):
+    req = urllib.request.Request('https://api.runpod.io/graphql',
+                                 data=json.dumps({'query': query, 'variables': variables or {}}).encode(),
+                                 headers={'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key()}',
+                                          'User-Agent': 'rapid-moss'})
+    with urllib.request.urlopen(req, timeout=60) as r:
+        out = json.load(r)
+    if out.get('errors'):
+        raise SystemExit(json.dumps(out['errors'], indent=1))
+    return out['data']
+
+
+def pod(pod_id):
+    return gql('query($id: String!) { pod(input: {podId: $id}) { id name desiredStatus costPerHr gpuCount '
+               'machine { gpuDisplayName location } runtime { uptimeInSeconds ports { ip isIpPublic privatePort '
+               'publicPort type } } } }', {'id': pod_id})['pod']
+
+
+def ssh_target(p):
+    for port in (p.get('runtime') or {}).get('ports') or []:
+        if port['privatePort'] == 22 and port['isIpPublic']:
+            return port['ip'], port['publicPort']
+    return None
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    sub = ap.add_subparsers(dest='cmd', required=True)
+    sub.add_parser('balance')
+    sub.add_parser('gpus')
+    c = sub.add_parser('create')
+    c.add_argument('name')
+    c.add_argument('--gpu', default='NVIDIA H100 80GB HBM3')
+    c.add_argument('--count', type=int, default=1)
+    c.add_argument('--cloud', default='COMMUNITY')
+    c.add_argument('--volume', type=int, default=300)
+    c.add_argument('--disk', type=int, default=60)
+    for name in ('status', 'stop', 'terminate', 'wait'):
+        sub.add_parser(name).add_argument('pod_id')
+    a = ap.parse_args()
+    if a.cmd == 'balance':
+        print(json.dumps(gql('query { myself { clientBalance currentSpendPerHr pods { id name desiredStatus costPerHr } } }')['myself'], indent=1))
+    elif a.cmd == 'gpus':
+        data = gql('query { gpuTypes { id memoryInGb securePrice communityPrice lowestPrice(input:{gpuCount:1}) { stockStatus } } }')
+        for g in sorted(data['gpuTypes'], key=lambda g: -(g['memoryInGb'] or 0)):
+            if (g['memoryInGb'] or 0) >= 80:
+                print(f"{g['id']:<60} {g['memoryInGb']:>4} GB secure {g['securePrice']} community {g['communityPrice']} stock {(g['lowestPrice'] or {}).get('stockStatus')}")
+    elif a.cmd == 'create':
+        pub = KEY_FILE.with_suffix('.pub').read_text().strip()
+        q = ('mutation($in: PodFindAndDeployOnDemandInput) { podFindAndDeployOnDemand(input: $in) '
+             '{ id costPerHr machine { gpuDisplayName location } } }')
+        inp = {'cloudType': a.cloud, 'gpuCount': a.count, 'volumeInGb': a.volume, 'containerDiskInGb': a.disk,
+               'gpuTypeId': a.gpu, 'name': a.name, 'imageName': IMAGE, 'ports': '22/tcp',
+               'volumeMountPath': '/workspace', 'startSsh': True, 'supportPublicIp': True,
+               'minCudaVersion': '12.8', 'env': [{'key': 'PUBLIC_KEY', 'value': pub}]}
+        print(json.dumps(gql(q, {'in': inp})['podFindAndDeployOnDemand'], indent=1))
+    elif a.cmd == 'status':
+        p = pod(a.pod_id)
+        print(json.dumps(p, indent=1))
+        t = ssh_target(p)
+        if t:
+            print(f'ssh -i {KEY_FILE} -p {t[1]} root@{t[0]}')
+    elif a.cmd == 'wait':
+        for _ in range(120):
+            t = ssh_target(pod(a.pod_id))
+            if t:
+                print(f'{t[0]} {t[1]}')
+                return
+            time.sleep(10)
+        sys.exit('pod did not expose ssh in 20 min')
+    elif a.cmd == 'stop':
+        print(gql('mutation($id: String!) { podStop(input: {podId: $id}) { id desiredStatus } }', {'id': a.pod_id}))
+    elif a.cmd == 'terminate':
+        print(gql('mutation($id: String!) { podTerminate(input: {podId: $id}) }', {'id': a.pod_id}))
+
+
+if __name__ == '__main__':
+    main()

code/scripts/watchdog.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+# Local: stop the pod (keeps its volume) when the RunPod balance drops below $MIN (default 6).
+cd "$(dirname "$0")/.."
+source scripts/pod.env
+MIN=${MIN:-6}
+while true; do
+  OUT=$(nix shell nixpkgs#python3 -c python3 scripts/runpod.py balance 2>/dev/null) || { sleep 120; continue; }
+  BAL=$(echo "$OUT" | grep clientBalance | grep -o '[0-9.]*')
+  RUN=$(echo "$OUT" | grep -c '"RUNNING"')
+  echo "$(date +%T) balance=$BAL running=$RUN"
+  [[ "$RUN" == 0 ]] && { echo "no running pod, exit"; exit 0; }
+  if awk "BEGIN{exit !($BAL < $MIN)}"; then
+    echo "balance below $MIN, stopping $POD_ID"; nix shell nixpkgs#python3 -c python3 scripts/runpod.py stop "$POD_ID"; exit 0
+  fi
+  sleep 300
+done

code/tests/test_dmad.py

@@ -0,0 +1,52 @@
+"""CPU check of the DMAD critic wiring (fp32, tiny latent).
+
+    python tests/test_dmad.py MODEL_DIR LODESTONES_DIR
+Generator phase (critic frozen): gradients reach the generator and nothing reaches the critic.
+Critic phase: gradients reach the critic head and its LoRA B on blocks >= start only.
+"""
+import sys
+
+import torch
+from torch.nn import functional as F
+
+import rapid_moss.model as M
+from rapid_moss.nets import Critic, Generator
+
+torch.manual_seed(0)
+model_dir, lode_dir = sys.argv[1], sys.argv[2]
+M.BF16 = torch.float32
+release = M.load_release_module(model_dir)
+tr = release.MOSSImageTransformer.from_pretrained(f'{model_dir}/transformer', torch_dtype=torch.float32,
+                                                  local_files_only=True)
+model = M.Moss(tr, release)
+ctx = model.context(torch.randn(10, 4096) * 2)
+z = torch.randn(1, 128, 4, 6)
+
+gen = Generator(model, 128, 128, 4)
+gen.load_lodestones(lode_dir)
+critic = Critic(model, 16, 16, start=30, hidden=64)
+assert critic.lora.A['q'].shape[0] == 12
+
+# generator phase
+critic.requires_grad_(False)
+x0 = gen.step(model, z, gen.rollout(model, z, ctx, 1), 1, ctx, ckpt=True)
+logit = critic(model, 0.6 * x0 + 0.4 * torch.randn_like(z), 400.0, ctx, ckpt=True)
+assert logit.shape == (1,)
+(-logit.mean()).backward()
+assert all(p.grad is not None and p.grad.abs().sum() > 0 for p in gen.parameters()), 'generator grads missing'
+assert all(p.grad is None for p in critic.parameters()), 'critic got generator-phase grads'
+print('generator phase OK, logit', logit.item())
+
+# critic phase
+critic.requires_grad_(True)
+for p in gen.parameters():
+    p.grad = None
+d_q = critic(model, 0.5 * torch.randn_like(z) + 0.5 * z, 300.0, ctx, ckpt=True)
+d_g = critic(model, 0.5 * x0.detach() + 0.5 * torch.randn_like(z), 300.0, ctx, ckpt=True)
+(F.softplus(-d_q) + F.softplus(d_g)).mean().backward()
+assert all(p.grad is None for p in gen.parameters()), 'generator got critic-phase grads'
+assert all(p.grad.abs().sum() > 0 for p in critic.head.parameters()) and critic.out.weight.grad.abs().sum() > 0
+b_rows = (critic.lora.B['q'].grad.flatten(1).abs().sum(1) > 0).tolist()
+print('critic LoRA B grad per block (30..41):', b_rows)
+assert all(b_rows)
+print('OK')

code/tests/test_equiv.py

@@ -0,0 +1,93 @@
+"""CPU check of rapid_moss.model against the release MOSSImageTransformer.
+
+    python tests/test_equiv.py MODEL_DIR LODESTONES_DIR
+Runs everything in fp32 (rapid_moss.model.BF16 patched) so the logic must match to ~1e-5:
+plain forward, negative-branch skip 22, the lodestones generator path (fp32 time path +
+unmerged LoRA hooks, as its model card describes), discriminator taps and a gradient check.
+"""
+import sys
+import time
+
+import torch
+from safetensors.torch import load_file
+
+import rapid_moss.model as M
+from rapid_moss.nets import Discriminator, FakeScore, Generator
+
+torch.manual_seed(0)
+model_dir, lode_dir = sys.argv[1], sys.argv[2]
+M.BF16 = torch.float32
+release = M.load_release_module(model_dir)
+t0 = time.time()
+tr = release.MOSSImageTransformer.from_pretrained(f'{model_dir}/transformer', torch_dtype=torch.float32,
+                                                  local_files_only=True)
+model = M.Moss(tr, release)
+print(f'load {time.time() - t0:.1f}s')
+
+h, w, L = 6, 10, 12
+latent = torch.randn(1, 128, h, w)
+emb = torch.randn(L, 4096) * 2
+ctx = model.context(emb)
+
+
+def rel(a, b):
+    return ((a - b).norm() / b.norm()).item()
+
+
+with torch.no_grad():
+    for t in (999.0, 500.0, 37.0):
+        ref = tr(latent, t, emb[None]).sample.float()
+        got = model(latent, t, ctx)
+        print(f'forward t={t}: rel {rel(got, ref):.2e}')
+        assert rel(got, ref) < 1e-4
+    ref = tr(latent, 600.0, emb[None], skip_layers=[22]).sample.float()
+    got = model(latent, 600.0, ctx, skip=(22,))
+    print(f'skip22: rel {rel(got, ref):.2e} (vs no-skip {rel(got, model(latent, 600.0, ctx)):.2e})')
+    assert rel(got, ref) < 1e-4
+
+    # lodestones generator: release model + replaced time path + LoRA hooks (model card recipe)
+    gen = Generator(model, 128, 128, 4)
+    gen.load_lodestones(lode_dir)
+    te = load_file(f'{lode_dir}/time_embedding.safetensors')
+    lora = load_file(f'{lode_dir}/lora.safetensors')
+    import copy
+    tr2 = copy.deepcopy(tr)
+    tr2.load_state_dict(te, strict=False)
+    for i, layer in enumerate(tr2.layers):
+        for name, lin in layer.proj.items():
+            p = f'diffusion_model.layers.{i}.proj.{name}'
+            down, up, alpha = lora[p + '.lora_down.weight'], lora[p + '.lora_up.weight'], lora[p + '.alpha']
+            scale = float(alpha) / down.shape[0]
+            lin.register_forward_hook(lambda m, inp, out, d=down, u=up, s=scale: out + s * (inp[0] @ d.t()) @ u.t())
+    Q = load_file(f'{lode_dir}/dataspace.safetensors')['latent_mix.Q']
+    codes = load_file(f'{lode_dir}/dataspace.safetensors')['step_codes']
+    z = torch.randn(1, 128, h, w)
+    x0_ref = torch.zeros_like(z)
+    x0_got = torch.zeros_like(z)
+    for r in range(3):
+        x = z + torch.einsum('oc,bchw->bohw', Q, x0_ref)
+        x0_ref = z - tr2(x, codes[r], emb[None]).sample.float()
+        x0_got = gen.step(model, z, x0_got, r, ctx)
+        print(f'generator r={r}: rel {rel(x0_got, x0_ref):.2e}')
+        assert rel(x0_got, x0_ref) < 1e-4
+    del tr2
+
+    disc = Discriminator(model.dim, (5, 11), 256)
+    feats = model(latent, 300.0, ctx, taps=disc.taps)
+    print('taps', {k: tuple(v.shape) for k, v in feats.items()}, 'logits', tuple(disc(feats, h, w).shape))
+
+# gradients: generator step through checkpointed blocks, fake score, D input-gradient
+fake = FakeScore(model, 16, 16, 4)
+x0 = gen.step(model, z, torch.randn_like(z), 1, ctx, ckpt=True)
+x0.backward(torch.randn_like(x0))
+nz = {n: p.grad.abs().sum().item() > 0 for n, p in gen.named_parameters()}
+print('gen grads nonzero:', all(nz.values()), [n for n, v in nz.items() if not v][:5])
+assert all(nz.values())
+v = fake(model, latent, 400.0, ctx, 2, ckpt=True)
+v.square().mean().backward()
+print('fake r_emb grad rows nonzero:', (fake.r_emb.grad.abs().sum(1) > 0).tolist())
+leaf = latent.clone().requires_grad_(True)
+g, = torch.autograd.grad(-disc(model(leaf, 300.0, ctx, taps=disc.taps, ckpt=True), h, w).mean(), leaf)
+print('D input grad norm', g.norm().item())
+assert g.norm() > 0
+print('OK')

dataspace.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d19f0e695bb26886cfce57441fd412f358e159d6eec8555ec1a71daaf30a7968
+size 65784

eval/eval_prompts.jsonl

@@ -0,0 +1,17 @@
+{"id": "anime-girl", "src": "eval", "prompt": "A high-quality anime illustration of a girl with long silver hair and teal eyes, wearing a yellow raincoat, standing at a rainy bus stop at night, neon reflections on wet asphalt, detailed cel shading, Japanese anime style", "w": 1024, "h": 1024}
+{"id": "anime-tags", "src": "eval", "prompt": "1girl, solo, school uniform, cherry blossoms, wind, smile, looking at viewer, anime style, masterpiece, best quality, detailed eyes, spring, blue sky", "w": 1024, "h": 1024}
+{"id": "chibi", "src": "eval", "prompt": "A cute chibi anime character, a small fox girl with big round eyes holding a strawberry daifuku, pastel colors, white background, sticker style illustration", "w": 1024, "h": 1024}
+{"id": "thick-paint", "src": "eval", "prompt": "Painterly semi-realistic anime illustration of a swordsman in a bamboo forest at dawn, dramatic lighting, thick brush strokes, Japanese fantasy key visual", "w": 1024, "h": 1024}
+{"id": "manga-mono", "src": "eval", "prompt": "Black and white Japanese manga panel, screentone shading, a high school boy running on a rooftop with speed lines, dynamic perspective, ink lineart", "w": 1024, "h": 1024}
+{"id": "jp-sign", "src": "eval", "prompt": "A traditional Japanese ramen shop at night, a wooden sign above the door clearly reads \"麺処 はやて\" in bold brush calligraphy, a red paper lantern with the text \"営業中\"", "w": 1024, "h": 1024}
+{"id": "jp-bubble", "src": "eval", "prompt": "An anime illustration of a girl surprised in a classroom, with a white manga speech bubble that says \"えっ、本当に！？\" in clear Japanese text", "w": 1024, "h": 1024}
+{"id": "jp-note", "src": "eval", "prompt": "A close-up of a handwritten sticky note on a refrigerator, written in Japanese: \"牛乳を買う\n卵は二パック\n夜七時に帰ります\"", "w": 1024, "h": 1024}
+{"id": "jp-cover", "src": "eval", "prompt": "A light novel cover illustration, anime girl mage casting a spell, the title text \"疾風の魔導書\" in large stylized Japanese letters at the top, and \"第一巻\" in the corner", "w": 1024, "h": 1024}
+{"id": "jp-poster", "src": "eval", "prompt": "A Japanese summer festival poster, fireworks over a river, vertical text on the right side reading \"夏祭り 花火大会\" and the date \"八月十五日\" at the bottom", "w": 1024, "h": 1024}
+{"id": "layout-manga", "src": "eval", "prompt": "{\"summary\": {\"meta\": [\"manga\"], \"canvas\": \"single\", \"caption\": \"A two-panel anime manga page in color. Top panel: a girl with short black hair waves at a friend in front of a school gate in the morning. Bottom panel: her friend, a girl with long brown hair, smiles back. Each panel has a white speech bubble with Japanese text.\"}, \"items\": [{\"label\": \"object\", \"caption\": \"Top panel: a girl with short black hair in a sailor school uniform waving, school gate behind her.\", \"bbox\": [-512, -512, -16, 512]}, {\"label\": \"object\", \"caption\": \"Bottom panel: a girl with long brown hair in a sailor school uniform smiling.\", \"bbox\": [16, -512, 512, 512]}, {\"label\": \"text\", \"caption\": \"Black Japanese text inside a white round speech bubble at the top right of the top panel.\", \"text\": \"おはよう！\", \"bbox\": [-460, 180, -330, 460]}, {\"label\": \"text\", \"caption\": \"Black Japanese text inside a white round speech bubble at the bottom left of the bottom panel.\", \"text\": \"今日もいい天気だね\", \"bbox\": [300, -470, 440, -80]}], \"bbox_coord_space\": \"abs_centralized\", \"bbox_format\": \"yxyx\"}", "w": 1024, "h": 1024}
+{"id": "rain-neon", "src": "eval", "prompt": "masterpiece, best quality, highres, 1girl, solo, long silver hair, glowing blue eyes, oversized black hoodie, holding a transparent umbrella, neon-lit Tokyo alley at night, heavy rain, reflections on wet ground, cinematic lighting, depth of field, light particles, detailed background, vibrant colors", "w": 832, "h": 1216}
+{"id": "sky-flight", "src": "eval", "prompt": "A breathtaking anime key visual of a girl in a white sundress jumping off a cliff into a vast summer sky, dynamic low-angle perspective, towering cumulonimbus clouds, sunbeams, flying birds, hair and dress fluttering in the wind, ultra detailed, vivid saturated colors, sense of freedom", "w": 832, "h": 1216}
+{"id": "duo-fantasy", "src": "eval", "prompt": "masterpiece, best quality, 2girls, fantasy, a knight girl with golden armor and a witch girl with a large hat standing back to back, magic circles glowing around them, battlefield with ruins, dramatic rim lighting, sparks, highly detailed armor and fabric, epic composition", "w": 832, "h": 1216}
+{"id": "wafu", "src": "eval", "prompt": "Japanese-style fantasy illustration of a fox spirit shrine maiden with white hair and red eyes, nine fluffy tails, holding a paper lantern, red torii gates at night, falling maple leaves, blue will-o'-the-wisps, intricate kimono patterns, highly detailed, soft glowing light", "w": 832, "h": 1216}
+{"id": "transparent", "src": "eval", "prompt": "A delicate anime illustration with a transparent watercolor feeling, a girl with light blue hair lying in shallow water surrounded by floating flowers and small fish, sunlight caustics, pastel colors, soft focus, dreamy atmosphere, extremely detailed eyes", "w": 832, "h": 1216}
+{"id": "scenery", "src": "eval", "prompt": "best quality, highres, scenery, 1girl, small figure, standing on a railway platform in the countryside at sunset, an old train arriving, power lines, summer clouds, orange and purple sky, lens flare, nostalgic atmosphere, highly detailed background art", "w": 832, "h": 1216}

inference.py

@@ -0,0 +1,88 @@
+"""4-step MOSS-Image-SFT with the rapid-moss adapter (data-space sampler, no CFG).
+
+    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" --out out.png
+
+The base stays bf16; the LoRA is applied unmerged in fp32 (its delta is smaller than bf16 rounding),
+and the step embedder + modulation biases are replaced by fp32 tensors. Sampler:
+    x0 = 0; for r in range(steps): x = z + Q x0; x0 = z - transformer(x, codes[r], text)
+with one fixed noise z, no noise re-injection and no timestep schedule.
+"""
+import argparse
+import os
+
+import torch
+from diffusers import DiffusionPipeline
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+
+def load(base, adapter, device):
+    custom = os.path.join(base, 'pipeline_moss_image.py') if os.path.isdir(base) else 'pipeline_moss_image'
+    pipe = DiffusionPipeline.from_pretrained(base, custom_pipeline=custom, trust_remote_code=True,
+                                             torch_dtype={'default': torch.bfloat16, 'vae': torch.float32})
+    pipe.to(device)
+    tr = pipe.transformer
+    if not os.path.isdir(adapter):
+        adapter = snapshot_download(adapter, allow_patterns=['*.safetensors'])
+    for name in ('time_embedding', 'shared_block_modulation', 'video_head_modulation'):
+        getattr(tr, name).float()
+    time_path = load_file(os.path.join(adapter, 'time_embedding.safetensors'), device=str(device))
+    result = tr.load_state_dict(time_path, strict=False)
+    assert not result.unexpected_keys, result.unexpected_keys
+    lora = load_file(os.path.join(adapter, 'lora.safetensors'), device=str(device))
+    for i, layer in enumerate(tr.layers):
+        for name, lin in layer.proj.items():
+            key = f'diffusion_model.layers.{i}.proj.{name}'
+            down, up = lora[key + '.lora_down.weight'], lora[key + '.lora_up.weight']
+            scale = float(lora[key + '.alpha']) / down.shape[0]
+
+            def hook(module, inputs, output, down=down, up=up, scale=scale):
+                return output + (scale * ((inputs[0].float() @ down.t()) @ up.t())).to(output.dtype)
+            lin.register_forward_hook(hook)
+    ds = load_file(os.path.join(adapter, 'dataspace.safetensors'), device=str(device))
+    return pipe, ds['latent_mix.Q'], ds['step_codes']
+
+
+@torch.no_grad()
+def generate(pipe, Q, codes, prompt, width, height, steps, seed):
+    device = pipe._execution_device
+    emb, mask = pipe.encode_prompt(prompt)
+    z = torch.randn(1, 128, height // 16, width // 16, generator=torch.Generator('cpu').manual_seed(seed)).to(device)
+    x0 = torch.zeros_like(z)
+    for r in range(steps):
+        x = z + torch.einsum('oc,bchw->bohw', Q, x0)
+        x0 = z - pipe.transformer(x, codes[r], emb, mask).sample.float()
+    vae = pipe.vae
+    lat = x0.to(vae.dtype)
+    bn, eps = vae.bn, vae.config.batch_norm_eps
+    lat = lat * (bn.running_var + eps).sqrt().view(1, -1, 1, 1) + bn.running_mean.view(1, -1, 1, 1)
+    b, c, h, w = lat.shape
+    lat = lat.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+    image = vae.decode(lat, return_dict=False)[0].clamp(-1, 1)
+    return pipe.image_processor.postprocess(image, output_type='pil')[0]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--prompt', required=True, help='natural language, tags or a MOSS JSON layout string')
+    ap.add_argument('--width', type=int, default=1024)
+    ap.add_argument('--height', type=int, default=1024)
+    ap.add_argument('--steps', type=int, default=4, help='1-4; each prefix of the 4 steps is trained')
+    ap.add_argument('--seed', type=int, default=42)
+    ap.add_argument('--out', default='out.png')
+    ap.add_argument('--base', default='OpenMOSS-Team/MOSS-Image-SFT')
+    ap.add_argument('--adapter', default='darask0/rapid-moss')
+    ap.add_argument('--device', default='cuda')
+    args = ap.parse_args()
+    if args.width % 16 or args.height % 16:
+        raise SystemExit('width and height must be multiples of 16')
+    pipe, Q, codes = load(args.base, args.adapter, torch.device(args.device))
+    if not 1 <= args.steps <= len(codes):
+        raise SystemExit(f'steps must be 1..{len(codes)}')
+    generate(pipe, Q, codes, args.prompt, args.width, args.height, args.steps, args.seed).save(args.out)
+    print(args.out)
+
+
+if __name__ == '__main__':
+    main()

logs/baseline_timings_h100.json

@@ -0,0 +1,5 @@
+{
+ "MOSS-Image-SFT-8-steps-dmd-4step": 0.805,
+ "MOSS-Image-SFT-8-steps-dmd-8step": 1.544,
+ "teacher50-cfg3": 16.188
+}

logs/dmad_config.json

@@ -0,0 +1,124 @@
+{
+ "model_dir": "/workspace/models/MOSS-Image-SFT",
+ "lodestones_dir": "/workspace/models/MOSS-Image-SFT-8-steps-dmd",
+ "cache_dir": "/workspace/cache",
+ "prompts": "/workspace/prompts/train.jsonl",
+ "eval_prompts": "/workspace/prompts/eval.jsonl",
+ "out_dir": "/workspace/runs/dmad",
+ "steps": 4,
+ "gen_rank": 128,
+ "gen_alpha": 128,
+ "fake_rank": 64,
+ "fake_alpha": 64,
+ "disc_taps": [
+  5,
+  11,
+  17,
+  23,
+  29,
+  35
+ ],
+ "disc_hidden": 256,
+ "iterations": 2500,
+ "micro": 2,
+ "fake_per_gen": 2,
+ "fake_warmup": 200,
+ "lr_gen": 2e-05,
+ "lr_time": 2e-05,
+ "lr_q": 2e-05,
+ "lr_fake": 0.0001,
+ "lr_disc": 5e-05,
+ "betas": [
+  0.0,
+  0.99
+ ],
+ "ema": 0.998,
+ "cfg": [
+  3.0,
+  3.0
+ ],
+ "skip_neg": [
+  22
+ ],
+ "tau_beta": [
+  2.5,
+  1.0
+ ],
+ "tau_clamp": [
+  0.03,
+  0.985
+ ],
+ "dt": 0.015625,
+ "adv_weight": 0.5,
+ "adv_start": 50,
+ "adv_warmup": 200,
+ "disc_real": "adm",
+ "teacher_dir": [
+  "/workspace/teacher",
+  "/workspace/teacher2"
+ ],
+ "p_gan": 0.5,
+ "gan_t": [
+  0.02,
+  0.98
+ ],
+ "clip_gen": 0.0,
+ "buckets": [
+  [
+   1024,
+   1024,
+   0.35
+  ],
+  [
+   896,
+   1152,
+   0.15
+  ],
+  [
+   832,
+   1216,
+   0.15
+  ],
+  [
+   1152,
+   896,
+   0.1
+  ],
+  [
+   1216,
+   832,
+   0.1
+  ],
+  [
+   768,
+   1344,
+   0.075
+  ],
+  [
+   1344,
+   768,
+   0.075
+  ]
+ ],
+ "method": "dmad",
+ "critic_rank": 128,
+ "critic_alpha": 128,
+ "critic_start": 21,
+ "critic_hidden": 512,
+ "critic_shift": 5.0,
+ "critic_sigma": [
+  0.02,
+  0.98
+ ],
+ "lr_critic": 4e-05,
+ "wd_critic": 0.01,
+ "critic_warmup": 50,
+ "ckpt": true,
+ "seed": 0,
+ "log_every": 10,
+ "sample_every": 100,
+ "save_every": 200,
+ "export_every": 500,
+ "eval_count": 8,
+ "eval_seed": 1234
+}

logs/dmad_train_log.jsonl

@@ -0,0 +1,102 @@
+{"it": 10, "gen_updates": 0, "s_per_it": 3.51, "mem_gib": 20.4, "c_acc": 0.503125, "c_loss": 1.472324, "d_gen": -0.026185, "d_teacher": 0.02548, "g_logit": 0.0}
+{"it": 20, "gen_updates": 0, "s_per_it": 3.4, "mem_gib": 20.4, "c_acc": 0.55, "c_loss": 1.393549, "d_gen": -0.070432, "d_teacher": 0.090518, "g_logit": 0.0}
+{"it": 30, "gen_updates": 0, "s_per_it": 3.41, "mem_gib": 20.4, "c_acc": 0.5875, "c_loss": 1.374437, "d_gen": -0.295139, "d_teacher": 0.018142, "g_logit": 0.0}
+{"it": 40, "gen_updates": 0, "s_per_it": 3.53, "mem_gib": 20.4, "c_acc": 0.640625, "c_loss": 1.303367, "d_gen": -0.228213, "d_teacher": 0.229624, "g_logit": 0.0}
+{"it": 50, "gen_updates": 0, "s_per_it": 3.44, "mem_gib": 20.5, "c_acc": 0.65625, "c_loss": 1.230491, "d_gen": -0.380138, "d_teacher": 0.359761, "g_logit": 0.0}
+{"it": 60, "gen_updates": 10, "s_per_it": 5.85, "mem_gib": 24.9, "c_acc": 0.60625, "c_loss": 1.320506, "d_gen": -0.204005, "d_teacher": 0.280464, "g_logit": -0.257044, "gnorm": 344.827544}
+{"it": 70, "gen_updates": 20, "s_per_it": 5.89, "mem_gib": 24.9, "c_acc": 0.6125, "c_loss": 1.321517, "d_gen": -0.27649, "d_teacher": 0.113134, "g_logit": -0.267412, "gnorm": 169.696696}
+{"it": 80, "gen_updates": 30, "s_per_it": 5.89, "mem_gib": 24.9, "c_acc": 0.61875, "c_loss": 1.258887, "d_gen": -0.346907, "d_teacher": 0.261502, "g_logit": -0.35437, "gnorm": 634.010178}
+{"it": 90, "gen_updates": 40, "s_per_it": 5.76, "mem_gib": 24.9, "c_acc": 0.678125, "c_loss": 1.195195, "d_gen": -0.554969, "d_teacher": 0.365944, "g_logit": -0.459122, "gnorm": 157.662113}
+{"it": 100, "gen_updates": 50, "s_per_it": 5.86, "mem_gib": 24.9, "c_acc": 0.65, "c_loss": 1.271955, "d_gen": -0.339471, "d_teacher": 0.444618, "g_logit": -0.380601, "gnorm": 398.204028}
+{"it": 110, "gen_updates": 60, "s_per_it": 6.82, "mem_gib": 24.9, "c_acc": 0.703125, "c_loss": 1.135317, "d_gen": -0.523898, "d_teacher": 0.40571, "g_logit": -0.523776, "gnorm": 515.074624}
+{"it": 120, "gen_updates": 70, "s_per_it": 5.75, "mem_gib": 24.9, "c_acc": 0.7, "c_loss": 1.107486, "d_gen": -0.584106, "d_teacher": 0.698363, "g_logit": -0.60042, "gnorm": 555.140576}
+{"it": 130, "gen_updates": 80, "s_per_it": 5.81, "mem_gib": 24.9, "c_acc": 0.70625, "c_loss": 1.103744, "d_gen": -0.60869, "d_teacher": 0.821214, "g_logit": -0.620043, "gnorm": 425.754166}
+{"it": 140, "gen_updates": 90, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.784375, "c_loss": 0.8933, "d_gen": -1.041128, "d_teacher": 0.990952, "g_logit": -0.890582, "gnorm": 666.534967}
+{"it": 150, "gen_updates": 100, "s_per_it": 5.76, "mem_gib": 24.9, "c_acc": 0.78125, "c_loss": 0.887413, "d_gen": -1.197625, "d_teacher": 0.985895, "g_logit": -1.082821, "gnorm": 379.410664}
+{"it": 160, "gen_updates": 110, "s_per_it": 5.72, "mem_gib": 24.9, "c_acc": 0.790625, "c_loss": 0.856912, "d_gen": -1.541816, "d_teacher": 1.098191, "g_logit": -1.66994, "gnorm": 449.519646}
+{"it": 170, "gen_updates": 120, "s_per_it": 5.7, "mem_gib": 24.9, "c_acc": 0.790625, "c_loss": 0.91846, "d_gen": -1.285428, "d_teacher": 1.172357, "g_logit": -1.261887, "gnorm": 269.401883}
+{"it": 180, "gen_updates": 130, "s_per_it": 5.77, "mem_gib": 24.9, "c_acc": 0.84375, "c_loss": 0.797819, "d_gen": -1.637835, "d_teacher": 1.326887, "g_logit": -1.754513, "gnorm": 2882.158878}
+{"it": 190, "gen_updates": 140, "s_per_it": 5.77, "mem_gib": 24.9, "c_acc": 0.825, "c_loss": 0.761074, "d_gen": -1.504383, "d_teacher": 1.440473, "g_logit": -1.562236, "gnorm": 861.8714}
+{"it": 200, "gen_updates": 150, "s_per_it": 5.74, "mem_gib": 24.9, "c_acc": 0.91875, "c_loss": 0.463252, "d_gen": -2.063986, "d_teacher": 1.998507, "g_logit": -1.98676, "gnorm": 318.910345}
+{"it": 210, "gen_updates": 160, "s_per_it": 7.69, "mem_gib": 24.9, "c_acc": 0.8375, "c_loss": 0.818473, "d_gen": -1.431804, "d_teacher": 1.621684, "g_logit": -1.605268, "gnorm": 399.719297}
+{"it": 220, "gen_updates": 170, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.85, "c_loss": 0.717284, "d_gen": -1.603106, "d_teacher": 1.577599, "g_logit": -1.603056, "gnorm": 452.001189}
+{"it": 230, "gen_updates": 180, "s_per_it": 5.69, "mem_gib": 24.9, "c_acc": 0.875, "c_loss": 0.64667, "d_gen": -1.80143, "d_teacher": 1.676809, "g_logit": -1.889914, "gnorm": 366.318943}
+{"it": 240, "gen_updates": 190, "s_per_it": 5.66, "mem_gib": 24.9, "c_acc": 0.8625, "c_loss": 0.639551, "d_gen": -1.8087, "d_teacher": 1.857835, "g_logit": -1.802448, "gnorm": 605.091997}
+{"it": 250, "gen_updates": 200, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.90625, "c_loss": 0.545975, "d_gen": -2.125394, "d_teacher": 1.798105, "g_logit": -2.117861, "gnorm": 649.524585}
+{"it": 260, "gen_updates": 210, "s_per_it": 5.74, "mem_gib": 24.9, "c_acc": 0.9125, "c_loss": 0.484202, "d_gen": -2.16631, "d_teacher": 2.169687, "g_logit": -2.021578, "gnorm": 499.114932}
+{"it": 270, "gen_updates": 220, "s_per_it": 5.7, "mem_gib": 24.9, "c_acc": 0.959375, "c_loss": 0.270355, "d_gen": -2.670422, "d_teacher": 2.456078, "g_logit": -2.582783, "gnorm": 495.963659}
+{"it": 280, "gen_updates": 230, "s_per_it": 5.67, "mem_gib": 24.9, "c_acc": 0.925, "c_loss": 0.35589, "d_gen": -2.598104, "d_teacher": 2.468642, "g_logit": -2.570019, "gnorm": 357.345479}
+{"it": 290, "gen_updates": 240, "s_per_it": 5.72, "mem_gib": 24.9, "c_acc": 0.978125, "c_loss": 0.224288, "d_gen": -3.003006, "d_teacher": 2.784076, "g_logit": -2.92619, "gnorm": 74.857542}
+{"it": 300, "gen_updates": 250, "s_per_it": 5.69, "mem_gib": 24.9, "c_acc": 0.9875, "c_loss": 0.142997, "d_gen": -3.190769, "d_teacher": 3.069393, "g_logit": -3.20608, "gnorm": 69.31735}
+{"it": 310, "gen_updates": 260, "s_per_it": 6.64, "mem_gib": 24.9, "c_acc": 0.9, "c_loss": 0.551212, "d_gen": -2.347602, "d_teacher": 2.404747, "g_logit": -2.188123, "gnorm": 342.295467}
+{"it": 320, "gen_updates": 270, "s_per_it": 5.68, "mem_gib": 24.9, "c_acc": 0.91875, "c_loss": 0.433934, "d_gen": -2.529411, "d_teacher": 2.414225, "g_logit": -2.553354, "gnorm": 206.773278}
+{"it": 330, "gen_updates": 280, "s_per_it": 5.7, "mem_gib": 24.9, "c_acc": 0.84375, "c_loss": 0.779442, "d_gen": -1.927769, "d_teacher": 1.969032, "g_logit": -2.008004, "gnorm": 551.766666}
+{"it": 340, "gen_updates": 290, "s_per_it": 5.64, "mem_gib": 24.9, "c_acc": 0.85625, "c_loss": 0.71635, "d_gen": -1.580061, "d_teacher": 1.732232, "g_logit": -1.535344, "gnorm": 572.014401}
+{"it": 350, "gen_updates": 300, "s_per_it": 5.7, "mem_gib": 24.9, "c_acc": 0.784375, "c_loss": 0.976114, "d_gen": -1.113697, "d_teacher": 0.95776, "g_logit": -1.236905, "gnorm": 1454.895154}
+{"it": 360, "gen_updates": 310, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.765625, "c_loss": 0.824124, "d_gen": -1.330074, "d_teacher": 1.316943, "g_logit": -1.403452, "gnorm": 1251.758484}
+{"it": 370, "gen_updates": 320, "s_per_it": 5.64, "mem_gib": 24.9, "c_acc": 0.696875, "c_loss": 1.100472, "d_gen": -0.896046, "d_teacher": 0.928468, "g_logit": -0.96601, "gnorm": 605.353458}
+{"it": 380, "gen_updates": 330, "s_per_it": 5.68, "mem_gib": 24.9, "c_acc": 0.778125, "c_loss": 0.887347, "d_gen": -1.422161, "d_teacher": 1.172825, "g_logit": -1.269689, "gnorm": 1429.191895}
+{"it": 390, "gen_updates": 340, "s_per_it": 5.67, "mem_gib": 24.9, "c_acc": 0.8625, "c_loss": 0.707317, "d_gen": -1.915341, "d_teacher": 1.878058, "g_logit": -1.860343, "gnorm": 797.769325}
+{"it": 400, "gen_updates": 350, "s_per_it": 5.67, "mem_gib": 24.9, "c_acc": 0.734375, "c_loss": 0.974438, "d_gen": -1.03391, "d_teacher": 1.058606, "g_logit": -0.970323, "gnorm": 591.646497}
+{"it": 410, "gen_updates": 360, "s_per_it": 7.69, "mem_gib": 24.9, "c_acc": 0.79375, "c_loss": 0.89193, "d_gen": -1.427567, "d_teacher": 1.424491, "g_logit": -1.53325, "gnorm": 596.328967}
+{"it": 420, "gen_updates": 370, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.778125, "c_loss": 0.843672, "d_gen": -1.146293, "d_teacher": 1.200306, "g_logit": -1.091098, "gnorm": 491.311426}
+{"it": 430, "gen_updates": 380, "s_per_it": 5.63, "mem_gib": 24.9, "c_acc": 0.80625, "c_loss": 0.798015, "d_gen": -1.562343, "d_teacher": 1.305641, "g_logit": -1.485226, "gnorm": 969.241519}
+{"it": 440, "gen_updates": 390, "s_per_it": 5.67, "mem_gib": 24.9, "c_acc": 0.80625, "c_loss": 0.867097, "d_gen": -1.602022, "d_teacher": 1.307751, "g_logit": -1.710642, "gnorm": 1851.946349}
+{"it": 450, "gen_updates": 400, "s_per_it": 5.66, "mem_gib": 24.9, "c_acc": 0.84375, "c_loss": 0.640278, "d_gen": -1.727461, "d_teacher": 1.752258, "g_logit": -1.790782, "gnorm": 1126.413405}
+{"it": 460, "gen_updates": 410, "s_per_it": 5.63, "mem_gib": 24.9, "c_acc": 0.859375, "c_loss": 0.608079, "d_gen": -1.933124, "d_teacher": 1.953212, "g_logit": -2.132834, "gnorm": 316.285908}
+{"it": 470, "gen_updates": 420, "s_per_it": 5.66, "mem_gib": 24.9, "c_acc": 0.915625, "c_loss": 0.396266, "d_gen": -2.375104, "d_teacher": 2.246049, "g_logit": -2.372623, "gnorm": 426.73709}
+{"it": 480, "gen_updates": 430, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.909375, "c_loss": 0.394121, "d_gen": -2.525033, "d_teacher": 2.411426, "g_logit": -2.550583, "gnorm": 580.174043}
+{"it": 490, "gen_updates": 440, "s_per_it": 5.6, "mem_gib": 24.9, "c_acc": 0.865625, "c_loss": 0.593566, "d_gen": -2.148674, "d_teacher": 2.374551, "g_logit": -2.169901, "gnorm": 716.07439}
+{"it": 500, "gen_updates": 450, "s_per_it": 5.6, "mem_gib": 24.9, "c_acc": 0.846875, "c_loss": 0.729606, "d_gen": -2.290267, "d_teacher": 1.572728, "g_logit": -2.285638, "gnorm": 392.940741}
+{"it": 510, "gen_updates": 460, "s_per_it": 6.78, "mem_gib": 24.9, "c_acc": 0.921875, "c_loss": 0.370543, "d_gen": -2.736271, "d_teacher": 2.455694, "g_logit": -2.707942, "gnorm": 537.985992}
+{"it": 520, "gen_updates": 470, "s_per_it": 5.67, "mem_gib": 24.9, "c_acc": 0.965625, "c_loss": 0.252394, "d_gen": -2.877263, "d_teacher": 2.850073, "g_logit": -2.739951, "gnorm": 515.508899}
+{"it": 530, "gen_updates": 480, "s_per_it": 5.64, "mem_gib": 24.9, "c_acc": 0.890625, "c_loss": 0.563771, "d_gen": -2.414951, "d_teacher": 2.277598, "g_logit": -2.654116, "gnorm": 709.700902}
+{"it": 540, "gen_updates": 490, "s_per_it": 5.7, "mem_gib": 24.9, "c_acc": 0.94375, "c_loss": 0.285931, "d_gen": -2.918549, "d_teacher": 2.788963, "g_logit": -2.851511, "gnorm": 271.158658}
+{"it": 550, "gen_updates": 500, "s_per_it": 5.64, "mem_gib": 24.9, "c_acc": 0.93125, "c_loss": 0.386427, "d_gen": -3.007278, "d_teacher": 2.665188, "g_logit": -2.941659, "gnorm": 374.554742}
+{"it": 560, "gen_updates": 510, "s_per_it": 5.59, "mem_gib": 24.9, "c_acc": 0.953125, "c_loss": 0.253991, "d_gen": -2.992537, "d_teacher": 2.871128, "g_logit": -2.8083, "gnorm": 421.902652}
+{"it": 570, "gen_updates": 520, "s_per_it": 5.74, "mem_gib": 24.9, "c_acc": 0.94375, "c_loss": 0.296528, "d_gen": -3.032215, "d_teacher": 2.960297, "g_logit": -2.982546, "gnorm": 367.955676}
+{"it": 580, "gen_updates": 530, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.9, "c_loss": 0.493329, "d_gen": -2.62191, "d_teacher": 2.719897, "g_logit": -2.698434, "gnorm": 780.508957}
+{"it": 590, "gen_updates": 540, "s_per_it": 5.69, "mem_gib": 24.9, "c_acc": 0.946875, "c_loss": 0.312489, "d_gen": -3.124627, "d_teacher": 2.567042, "g_logit": -3.063663, "gnorm": 672.813906}
+{"it": 600, "gen_updates": 550, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.9625, "c_loss": 0.203462, "d_gen": -3.358016, "d_teacher": 3.01446, "g_logit": -3.541651, "gnorm": 244.839248}
+{"it": 610, "gen_updates": 560, "s_per_it": 7.86, "mem_gib": 24.9, "c_acc": 0.959375, "c_loss": 0.246468, "d_gen": -3.294766, "d_teacher": 2.899352, "g_logit": -3.391545, "gnorm": 274.389373}
+{"it": 620, "gen_updates": 570, "s_per_it": 5.6, "mem_gib": 24.9, "c_acc": 0.9375, "c_loss": 0.389626, "d_gen": -3.036653, "d_teacher": 2.82729, "g_logit": -2.993505, "gnorm": 602.330268}
+{"it": 630, "gen_updates": 580, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.953125, "c_loss": 0.236319, "d_gen": -3.326519, "d_teacher": 2.977125, "g_logit": -3.006266, "gnorm": 275.102869}
+{"it": 640, "gen_updates": 590, "s_per_it": 5.63, "mem_gib": 24.9, "c_acc": 0.9125, "c_loss": 0.468886, "d_gen": -2.839741, "d_teacher": 2.540864, "g_logit": -2.657633, "gnorm": 181.273212}
+{"it": 650, "gen_updates": 600, "s_per_it": 5.57, "mem_gib": 24.9, "c_acc": 0.953125, "c_loss": 0.311223, "d_gen": -3.182146, "d_teacher": 2.863069, "g_logit": -3.080172, "gnorm": 173.449964}
+{"it": 660, "gen_updates": 610, "s_per_it": 5.72, "mem_gib": 24.9, "c_acc": 0.94375, "c_loss": 0.354878, "d_gen": -3.160589, "d_teacher": 2.84673, "g_logit": -2.883895, "gnorm": 240.966786}
+{"it": 670, "gen_updates": 620, "s_per_it": 5.82, "mem_gib": 24.9, "c_acc": 0.90625, "c_loss": 0.406371, "d_gen": -2.754415, "d_teacher": 2.638921, "g_logit": -2.808692, "gnorm": 298.15805}
+{"it": 680, "gen_updates": 630, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.878125, "c_loss": 0.531884, "d_gen": -2.432082, "d_teacher": 2.578791, "g_logit": -2.482742, "gnorm": 271.599031}
+{"it": 690, "gen_updates": 640, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.9125, "c_loss": 0.4271, "d_gen": -2.860206, "d_teacher": 2.463032, "g_logit": -2.752373, "gnorm": 355.511536}
+{"it": 700, "gen_updates": 650, "s_per_it": 5.62, "mem_gib": 24.9, "c_acc": 0.8875, "c_loss": 0.575126, "d_gen": -2.657025, "d_teacher": 2.421408, "g_logit": -2.473467, "gnorm": 229.405529}
+{"it": 710, "gen_updates": 660, "s_per_it": 6.63, "mem_gib": 24.9, "c_acc": 0.871875, "c_loss": 0.574605, "d_gen": -2.347231, "d_teacher": 2.330401, "g_logit": -2.323787, "gnorm": 968.408723}
+{"it": 720, "gen_updates": 670, "s_per_it": 5.6, "mem_gib": 24.9, "c_acc": 0.95, "c_loss": 0.300756, "d_gen": -3.302272, "d_teacher": 2.674137, "g_logit": -3.417291, "gnorm": 67.733886}
+{"it": 730, "gen_updates": 680, "s_per_it": 5.64, "mem_gib": 24.9, "c_acc": 0.9375, "c_loss": 0.221562, "d_gen": -3.521256, "d_teacher": 3.014115, "g_logit": -3.507561, "gnorm": 100.00077}
+{"it": 740, "gen_updates": 690, "s_per_it": 5.75, "mem_gib": 24.9, "c_acc": 0.971875, "c_loss": 0.138255, "d_gen": -3.753518, "d_teacher": 3.362483, "g_logit": -3.768429, "gnorm": 553.137035}
+{"it": 750, "gen_updates": 700, "s_per_it": 5.58, "mem_gib": 24.9, "c_acc": 0.96875, "c_loss": 0.169519, "d_gen": -3.683156, "d_teacher": 3.372113, "g_logit": -3.977836, "gnorm": 65.564101}
+{"it": 760, "gen_updates": 710, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.971875, "c_loss": 0.13752, "d_gen": -3.875997, "d_teacher": 3.534403, "g_logit": -3.927255, "gnorm": 141.72445}
+{"it": 770, "gen_updates": 720, "s_per_it": 5.71, "mem_gib": 24.9, "c_acc": 0.95625, "c_loss": 0.233508, "d_gen": -3.549302, "d_teacher": 3.582107, "g_logit": -3.492271, "gnorm": 81.573391}
+{"it": 780, "gen_updates": 730, "s_per_it": 5.81, "mem_gib": 24.9, "c_acc": 0.95625, "c_loss": 0.179923, "d_gen": -3.765295, "d_teacher": 3.225188, "g_logit": -3.776723, "gnorm": 31.5222}
+{"it": 790, "gen_updates": 740, "s_per_it": 5.59, "mem_gib": 24.9, "c_acc": 0.978125, "c_loss": 0.097448, "d_gen": -4.034114, "d_teacher": 3.69661, "g_logit": -3.898952, "gnorm": 39.593995}
+{"it": 800, "gen_updates": 750, "s_per_it": 5.68, "mem_gib": 24.9, "c_acc": 0.978125, "c_loss": 0.102428, "d_gen": -4.094866, "d_teacher": 3.774902, "g_logit": -3.896441, "gnorm": 24.271654}
+{"it": 810, "gen_updates": 760, "s_per_it": 7.66, "mem_gib": 24.9, "c_acc": 0.9625, "c_loss": 0.202403, "d_gen": -3.854326, "d_teacher": 3.584347, "g_logit": -3.978992, "gnorm": 40.056661}
+{"it": 820, "gen_updates": 770, "s_per_it": 5.59, "mem_gib": 24.9, "c_acc": 0.971875, "c_loss": 0.156311, "d_gen": -3.827758, "d_teacher": 3.670029, "g_logit": -3.863046, "gnorm": 44.031802}
+{"it": 830, "gen_updates": 780, "s_per_it": 5.62, "mem_gib": 24.9, "c_acc": 0.98125, "c_loss": 0.102383, "d_gen": -4.115957, "d_teacher": 3.809848, "g_logit": -4.091031, "gnorm": 45.254895}
+{"it": 840, "gen_updates": 790, "s_per_it": 5.73, "mem_gib": 24.9, "c_acc": 0.971875, "c_loss": 0.12304, "d_gen": -4.100542, "d_teacher": 3.828115, "g_logit": -4.164267, "gnorm": 27.769859}
+{"it": 850, "gen_updates": 800, "s_per_it": 5.67, "mem_gib": 24.9, "c_acc": 0.778125, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 860, "gen_updates": 810, "s_per_it": 5.48, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 870, "gen_updates": 820, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 880, "gen_updates": 830, "s_per_it": 5.58, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 890, "gen_updates": 840, "s_per_it": 5.49, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 900, "gen_updates": 850, "s_per_it": 5.65, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 910, "gen_updates": 860, "s_per_it": 6.28, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 920, "gen_updates": 870, "s_per_it": 5.56, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 930, "gen_updates": 880, "s_per_it": 5.54, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 940, "gen_updates": 890, "s_per_it": 5.53, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 950, "gen_updates": 900, "s_per_it": 5.5, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 960, "gen_updates": 910, "s_per_it": 5.52, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 970, "gen_updates": 920, "s_per_it": 5.5, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 980, "gen_updates": 930, "s_per_it": 5.55, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 990, "gen_updates": 940, "s_per_it": 5.57, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 1000, "gen_updates": 950, "s_per_it": 5.5, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 1010, "gen_updates": 960, "s_per_it": 7.44, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}
+{"it": 1020, "gen_updates": 970, "s_per_it": 5.56, "mem_gib": 24.9, "c_acc": 0.0, "c_loss": NaN, "d_gen": NaN, "d_teacher": NaN, "g_logit": NaN, "gnorm": NaN}

logs/probe_adm_discriminator.jsonl

@@ -0,0 +1,15 @@
+{"it": 10, "gen_updates": 0, "s_per_it": 2.75, "mem_gib": 20.7, "fake": 0.27561}
+{"it": 20, "gen_updates": 0, "s_per_it": 2.64, "mem_gib": 20.7, "fake": 0.375196}
+{"it": 30, "gen_updates": 10, "s_per_it": 4.93, "mem_gib": 25.1, "adv_ratio": 35.733486, "d_acc": 0.500449, "d_fake": -0.10222, "d_loss": 2.126553, "d_real": -0.101114, "dmd_rms": 0.221341, "fake": 0.353202, "gnorm": 1.763772, "lam": 0.0, "ramp": 0.0}
+{"it": 40, "gen_updates": 20, "s_per_it": 9.97, "mem_gib": 25.3, "adv_ratio": 48.805241, "d_acc": 0.499809, "d_fake": 0.092451, "d_loss": 2.002986, "d_real": 0.095818, "dmd_rms": 0.325674, "fake": 0.416026, "gnorm": 1.199596, "lam": 0.0, "ramp": 0.0}
+{"it": 50, "gen_updates": 30, "s_per_it": 10.45, "mem_gib": 25.4, "adv_ratio": 9.468557, "d_acc": 0.503584, "d_fake": -0.081498, "d_loss": 2.012508, "d_real": -0.073537, "dmd_rms": 0.32834, "fake": 0.505973, "gnorm": 1.627691, "lam": 0.0, "ramp": 0.0}
+{"it": 60, "gen_updates": 40, "s_per_it": 10.57, "mem_gib": 25.4, "adv_ratio": 16.740142, "d_acc": 0.501583, "d_fake": 0.143709, "d_loss": 1.987252, "d_real": 0.157508, "dmd_rms": 0.274366, "fake": 0.463597, "gnorm": 1.094926, "lam": 0.0, "ramp": 0.0}
+{"it": 70, "gen_updates": 50, "s_per_it": 10.33, "mem_gib": 25.4, "adv_ratio": 17.995773, "d_acc": 0.502987, "d_fake": 0.218125, "d_loss": 1.988724, "d_real": 0.229557, "dmd_rms": 0.260837, "fake": 0.4254, "gnorm": 1.340356, "lam": 0.0, "ramp": 0.0}
+{"it": 80, "gen_updates": 60, "s_per_it": 10.4, "mem_gib": 25.4, "adv_ratio": 21.950003, "d_acc": 0.502617, "d_fake": 0.195303, "d_loss": 1.999953, "d_real": 0.200418, "dmd_rms": 0.2464, "fake": 0.405491, "gnorm": 2.077201, "lam": 0.0, "ramp": 0.0}
+{"it": 90, "gen_updates": 70, "s_per_it": 10.41, "mem_gib": 25.4, "adv_ratio": 4.992844, "d_acc": 0.512019, "d_fake": 0.056383, "d_loss": 2.005259, "d_real": 0.089869, "dmd_rms": 0.353305, "fake": 0.438201, "gnorm": 1.893705, "lam": 0.0, "ramp": 0.0}
+{"it": 100, "gen_updates": 80, "s_per_it": 10.52, "mem_gib": 25.4, "adv_ratio": 22.842478, "d_acc": 0.500302, "d_fake": 0.098687, "d_loss": 1.999297, "d_real": 0.099608, "dmd_rms": 0.232808, "fake": 0.345837, "gnorm": 1.049859, "lam": 0.0, "ramp": 0.0}
+{"it": 110, "gen_updates": 90, "s_per_it": 10.42, "mem_gib": 25.4, "adv_ratio": 14.325914, "d_acc": 0.504625, "d_fake": 0.015475, "d_loss": 1.990486, "d_real": 0.027477, "dmd_rms": 0.245574, "fake": 0.39951, "gnorm": 1.243515, "lam": 0.0, "ramp": 0.0}
+{"it": 120, "gen_updates": 100, "s_per_it": 10.54, "mem_gib": 25.4, "adv_ratio": 14.00046, "d_acc": 0.50106, "d_fake": 0.017234, "d_loss": 1.999404, "d_real": 0.018786, "dmd_rms": 0.253277, "fake": 0.425688, "gnorm": 1.833253, "lam": 0.0, "ramp": 0.0}
+{"it": 130, "gen_updates": 110, "s_per_it": 10.33, "mem_gib": 25.4, "adv_ratio": 8.808989, "d_acc": 0.506074, "d_fake": -0.007883, "d_loss": 1.982148, "d_real": 0.010178, "dmd_rms": 0.319683, "fake": 0.470164, "gnorm": 1.987517, "lam": 0.0, "ramp": 0.0}
+{"it": 140, "gen_updates": 120, "s_per_it": 10.55, "mem_gib": 25.4, "adv_ratio": 30.816719, "d_acc": 0.500643, "d_fake": 0.018101, "d_loss": 1.999332, "d_real": 0.018775, "dmd_rms": 0.23059, "fake": 0.44002, "gnorm": 1.406316, "lam": 0.0, "ramp": 0.0}
+{"it": 150, "gen_updates": 130, "s_per_it": 10.44, "mem_gib": 25.4, "adv_ratio": 21.776939, "d_acc": 0.50403, "d_fake": 0.049467, "d_loss": 2.006303, "d_real": 0.067547, "dmd_rms": 0.255741, "fake": 0.51779, "gnorm": 1.035783, "lam": 0.0, "ramp": 0.0}

logs/probe_teacher_discriminator.jsonl

@@ -0,0 +1,8 @@
+{"it": 10, "gen_updates": 0, "s_per_it": 2.76, "mem_gib": 20.7, "fake": 0.2758}
+{"it": 20, "gen_updates": 0, "s_per_it": 2.67, "mem_gib": 20.7, "fake": 0.374674}
+{"it": 30, "gen_updates": 10, "s_per_it": 4.96, "mem_gib": 25.5, "adv_ratio": 30.275976, "d_acc": 0.555563, "d_fake": -0.104033, "d_loss": 2.012383, "d_real": 0.087962, "dmd_rms": 0.205062, "fake": 0.32308, "gan_frac": 1.0, "gnorm": 0.77355, "lam": 0.0, "ramp": 0.0}
+{"it": 40, "gen_updates": 20, "s_per_it": 4.96, "mem_gib": 25.5, "adv_ratio": 99.557544, "d_acc": 0.58669, "d_fake": -0.229194, "d_loss": 1.836997, "d_real": 0.012343, "dmd_rms": 0.305676, "fake": 0.440187, "gan_frac": 1.0, "gnorm": 1.42406, "lam": 0.0, "ramp": 0.0}
+{"it": 50, "gen_updates": 30, "s_per_it": 4.8, "mem_gib": 25.5, "adv_ratio": 39.060537, "d_acc": 0.646695, "d_fake": -0.243651, "d_loss": 1.674095, "d_real": 0.232983, "dmd_rms": 0.284247, "fake": 0.443884, "gan_frac": 1.0, "gnorm": 3.348034, "lam": 0.0, "ramp": 0.0}
+{"it": 60, "gen_updates": 40, "s_per_it": 4.85, "mem_gib": 25.5, "adv_ratio": 48.76557, "d_acc": 0.583147, "d_fake": -0.154269, "d_loss": 1.790229, "d_real": 0.145679, "dmd_rms": 0.281149, "fake": 0.425242, "gan_frac": 1.0, "gnorm": 2.247003, "lam": 0.0, "ramp": 0.0}
+{"it": 70, "gen_updates": 50, "s_per_it": 4.77, "mem_gib": 25.5, "adv_ratio": 59.619794, "d_acc": 0.581655, "d_fake": -0.218971, "d_loss": 1.795289, "d_real": 0.126522, "dmd_rms": 0.310085, "fake": 0.388189, "gan_frac": 1.0, "gnorm": 1.043796, "lam": 0.0, "ramp": 0.0}
+{"it": 80, "gen_updates": 60, "s_per_it": 4.86, "mem_gib": 25.5, "adv_ratio": 68.792467, "d_acc": 0.682936, "d_fake": -0.367275, "d_loss": 1.531309, "d_real": 0.305816, "dmd_rms": 0.224524, "fake": 0.357255, "gan_frac": 1.0, "gnorm": 1.901103, "lam": 0.0, "ramp": 0.0}

lora.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a4c8e2d6244a0fcdc2822b87b3435226ae57c289aef99fc2ac54edeab9e09bb
+size 1189176832

time_embedding.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ebabc91e2a9e68c00e948dd80b05235a826e808817b5b2600b441457570d244
+size 41017896