peal.generators.rae_pipeline¶

Two-stage RAEv2 training pipeline behind the RAEDiffusionAutoencoder generator.

python -m peal.generators.rae_pipeline –config <generator yaml> –stage {stage1,stats,cache,stage2,all}

Every stage is idempotent: it looks for its own final artifact and returns immediately when that exists, so a supervisor can simply re-run “all” after a crash. RAEv2’s own scripts resume from the highest-epoch checkpoint in the experiment’s checkpoints/ directory, so a retry costs at most one epoch.

Layout under the generator’s base_path (default $PEAL_RUNS/imagenet/rae_clip):

config.yaml copy of the generator config (for –continue_training) stage1/<stage1_experiment>/ RAEv2 stage-1 run (checkpoints/ep-.pt, log, src/) stage1_assets/decoder.pt EMA decoder extracted from the final stage-1 checkpoint stage1_assets/stats.pt per-channel latent mean/var over the train set stage2/<stage2_experiment>/ RAEv2 stage-2 run (checkpoints/ep-.pt, log, src/) <cache_dir> (node-local /tmp) premixed stage-2 latent + CLS cache, metadata.json marks completion

Nothing here imports torch, so it can run as the driver process for torchrun.

Module attributes

RAEV2_REPO

RAEv2 is CC BY-NC 4.0, which cannot be relicensed under PEAL's LGPL-3.0.

Functions

`default_raev2_dir`	Where RAEv2 is expected, first match wins:
`expand_path`	Expand `<PEAL_RAEV2>`, `<PEAL_BASE>`, `$PEAL_RUNS`, `$PEAL_DATA` and `~` in a path.
`load_generator_config`	Load an RAE generator yaml and fill in the pipeline defaults.
`main`	CLI entry point: `--config <generator yaml> --stage <stage>`.
`require_raev2`	Return `path` if RAEv2 is installed there, else raise an actionable error.

Classes

RAEPipeline

Drive RAEv2's stage-1, stats, cache and stage-2 scripts as subprocesses.