peal.generators.rae_pipeline.RAEPipeline

class peal.generators.rae_pipeline.RAEPipeline

Bases: object

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

Each run_* method checks its own completion marker first, so the pipeline can be restarted after a crash. Training scripts run under torch.distributed.run with one process per visible GPU (or $RAE_NPROC), with cwd set to the RAEv2 checkout and its src and .deps directories on PYTHONPATH. Output goes to <base_path>/{stage1,stats,cache,stage2}.log.

Parameters:

• cfg (dict) – Config from load_generator_config(); uses base_path, raev2_dir, stage1_config, stage2_config, the experiment names, precisions, keep_checkpoints and the cache_* keys.

• logger (callable, optional) – Receives progress strings. Default print.

Variables:

• stage2_exp (stage1_exp,) – The RAEv2 experiment directories.

• stats_path (decoder_path,) – Stage-1 assets consumed by stage 2.

• stage2_yaml (stage1_yaml,) – The parsed RAEv2 stage configs.

Methods

cache_done	Whether the stage-2 latent cache has its metadata.json marker.
prune_checkpoints	Delete all but the newest keep epoch checkpoints of an experiment.
resolved_stage_config	Path of a stage config with PEAL placeholders already expanded.
run	Run one stage, or all four in order.
run_cache	Encode the train set once (latent + CLIP embedding) into premixed shards.
run_stage1	Train the RAEv2 stage-1 autoencoder unless already finished.
run_stage2	Train the RAEv2 stage-2 diffusion model on the latent cache.
run_stats	Extract the EMA decoder and per-channel latent statistics of stage 1.
save_config_copy	Copy the generator yaml to <base_path>/config.yaml (if different).
stage1_done	Whether the final stage-1 checkpoint exists.
stage1_final_ckpt	Path of the stage-1 checkpoint for the configured final epoch.
stage2_done	Whether the final stage-2 checkpoint exists.
stage2_final_ckpt	Path of the stage-2 checkpoint for the configured final epoch.
stats_done	Whether both decoder.pt and stats.pt exist.

__init__(cfg, logger=<built-in function print>)

Resolve the run layout from cfg and parse both RAEv2 stage configs.

resolved_stage_config(stage)

Path of a stage config with PEAL placeholders already expanded.

RAEv2’s train scripts receive --config <path> and parse it with OmegaConf, which knows nothing about <PEAL_RAEV2> or <PEAL_BASE>; a placeholder would reach it verbatim and surface as Incorrect path_or_model_id: '<PEAL_RAEV2>/configs/decoder/ViTXL'. The expanded copy is written once per process under <base>/resolved and reused.

Parameters:

stage (str) – "stage1" or "stage2".

Returns:

Path of the expanded copy, or of the original when it held no placeholder.

Return type:

str

prune_checkpoints(exp_dir, keep=None)

Delete all but the newest keep epoch checkpoints of an experiment.

ep-last.pt is a hard link and survives on its own. RAEv2 writes about 10 GB per epoch and never prunes.

Parameters:

• exp_dir (str) – RAEv2 experiment directory containing checkpoints/.

• keep (int, optional) – How many to keep; defaults to cfg["keep_checkpoints"]. A value of 0 or less prunes nothing.

save_config_copy(src_path)

Copy the generator yaml to <base_path>/config.yaml (if different).

stage1_final_ckpt()

Path of the stage-1 checkpoint for the configured final epoch.

stage1_done()

Whether the final stage-1 checkpoint exists.

run_stage1()

Train the RAEv2 stage-1 autoencoder unless already finished.

Prunes old checkpoints first, then runs src/train_stage1.py with stage1_config, stage1_dir and stage1_precision.

Raises:

RuntimeError – When the script exits 0 without writing the final checkpoint.

stats_done()

Whether both decoder.pt and stats.pt exist.

run_stats()

Extract the EMA decoder and per-channel latent statistics of stage 1.

Runs scripts/stage1/extract_decoder.py and scripts/stage1/compute_encoder_stats.py (the latter under torchrun over stats_num_samples train images), skipping whichever asset is already present.

Raises:

RuntimeError – When stage 1 has not finished, or the assets are still missing after the scripts exit 0.

cache_done()

Whether the stage-2 latent cache has its metadata.json marker.

run_cache()

Encode the train set once (latent + CLIP embedding) into premixed shards.

The shard-batch sampler serves every micro-batch from a single shard, so the source order is globally permuted first (cache_permutation_seed); without that each shard would hold one ImageFolder class. A cache directory that already holds rank* files but no metadata.json is treated as incomplete and rebuilt with --overwrite.

Raises:

RuntimeError – When the stage-1 assets are missing or the build exits 0 without writing metadata.json.

stage2_final_ckpt()

Path of the stage-2 checkpoint for the configured final epoch.

stage2_done()

Whether the final stage-2 checkpoint exists.

run_stage2()

Train the RAEv2 stage-2 diffusion model on the latent cache.

Verifies that the stage-2 yaml’s pretrained_decoder_path and normalization_stat_path point at this pipeline’s assets, prunes old checkpoints and runs src/train.py under torchrun.

Raises:

RuntimeError – When prerequisites are missing, the stage-2 config points at other assets, or the final checkpoint is missing after a clean exit.

run(stage='all')

Run one stage, or all four in order.

Parameters:

stage ({"stage1", "stats", "cache", "stage2", "all"}, optional) – Which stage to run. Default "all".