# 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.