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.