# Conditional diffusion model configuration (shared across encoders/areas).
# Single-channel 64x64 DDPM conditioned on a 50-d response vector.

unet:
  image_size: 64
  in_channels: 1
  out_channels: 1
  block_channels: [64, 128, 256, 512]
  layers_per_block: 2
  cross_attention_dim: 256

conditioning:
  embedding_dim: 256
  mlp_layers: [50, 256, 256, 256]   # 50 neurons → 256-d token
  mlp_activation: silu
  cfg_dropout_rate: 0.1             # classifier-free guidance dropout
  observation_noise_sigma: 0.22     # σ for ξ added to r at training time

scheduler:
  num_train_timesteps: 1000
  beta_schedule: linear
  beta_start: 1.0e-4
  beta_end: 0.02
  prediction_type: epsilon

train:
  data: tiny_imagenet
  data_split: train
  n_train_images: 100000
  optimizer: adamw
  lr: 1.0e-4
  weight_decay: 1.0e-4
  scheduler: cosine
  warmup_steps: 500
  batch_size: 128
  epochs: 200                       # 200-ep checkpoint used for figures
  precision: bfloat16
  gradient_clip: 1.0
  seed: 42