stage_b.yaml

experiment_id: sotediffusion-wr3_3b-stage_b
model_version: 3B
dtype: float32
use_fsdp: False

batch_size: 8
grad_accum_steps: 4
updates: 1024000
backup_every: 2048
save_every: 1024
warmup_updates: 128

lr: 8.0e-6
optimizer_type: AdamW
adaptive_loss_weight: False
stochastic_rounding: False

image_size: 1024
multi_aspect_ratio: [1/1, 1/2, 1/3, 2/3, 3/4, 1/5, 2/5, 3/5, 4/5, 1/6, 5/6, 9/16]
shift: 4

checkpoint_path:  /home/ubuntu/out/
output_path: /home/ubuntu/out/
webdataset_path: file:/home/ubuntu/newest_best.tar

effnet_checkpoint_path: /home/ubuntu/models/wuerstchen3/effnet_encoder.safetensors
stage_a_checkpoint_path: /home/ubuntu/models/wuerstchen3/stage_a.safetensors
generator_checkpoint_path: /home/ubuntu/models/wuerstchen3/generator-020480.safetensors