# SpatialDiffuseSlot — Semanticist DiT-L/ViT (tok_L init) + OUR 85 multi-res
# spatial tokens + OUR spatial-align attention mask. (2026-07-15 설계 확정본)
# Phase 1 (this config): DiT trunk FROZEN (freeze_dit) — encoder/pool/cond-embedder
#   adapt first with LONG warmup. Phase 2: resume with freeze_dit=false + low blr.
trainer:
  target: semanticist.engine.diffusion_trainer.DiffusionTrainer
  params:
    num_epoch: 400
    valid_size: 64
    blr: 2.5e-5              # Semanticist recipe (batch 256 -> lr 2.5e-5)
    cosine_lr: True
    warmup_epochs: 100       # 25% warmup (Semanticist recipe, 우리 기존 2%의 교훈)
    batch_size: 32           # per-gpu; 4 GPU -> eff 128 (large-batch regime 회피)
    num_workers: 8
    pin_memory: True
    grad_accum_steps: 2      # eff 256 == Semanticist
    precision: 'bf16'
    max_grad_norm: 3.0
    save_every: 10000
    sample_every: 5000
    fid_every: 50000
    result_folder: "./output/tokenizer/models_l_spatial_phase2"
    log_dir: "./output/tokenizer/models_l_spatial_phase2/logs"
    cfg: 3.0
    compile: False
    enable_ema: True
    model:
      target: semanticist.stage1.spatial_diffuse_slot.SpatialDiffuseSlot
      params:
        encoder: 'vit_base_patch16'
        enc_img_size: 256
        enc_causal: True
        num_slots: 256               # KEEP 256: tok_L encoder loads 100% (slots unused)
        slot_dim: 16
        norm_slots: True
        enable_nest: False
        enable_nest_after: -1
        dit_model: 'DiT-L-2'
        vae: 'xwen99/mar-vae-kl16'
        num_sampling_steps: '250'
        use_repa: True
        repa_encoder_depth: 8
        repa_loss_weight: 1.0
        eval_fid: True
        fid_stats: 'fid_stats/adm_in256_stats.npz'
        # --- SpatialDiffuseSlot 전용 ---
        level_sizes: [8, 4, 2, 1]    # 85 multi-res cond tokens
        pool_depth: 2                # 우리 L_repa attn-pool depth
        freeze_dit: False            # Phase-2: DiT UNFROZEN (encoder와 동일 warmup lr)
        init_from: /NHNHOME/WORKSPACE/0226010398_A/sr_diffusion/clevr_sudoku/semanticist/output/tokenizer/models_l_spatial/models/step380000
        ckpt_path: null              # trainer가 무조건 참조 (resume용; init은 init_from)
    dataset:
      target: semanticist.utils.datasets.ImageNet
      params:
        root: ./dataset/imagenet/
        split: train
        aug: 'randcrop'
        img_size: 256
    test_dataset:
      target: semanticist.utils.datasets.ImageNet
      params:
        root: ./dataset/imagenet/
        split: val
        aug: 'centercrop'
        img_size: 256