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