config.yaml

model: 
  transport:
    path_type: linear
    prediction: v
    weighting: lognormal
  network:
    target: nit.models.c2i.nit_model.NiT
    params:
      class_dropout_prob: 0.1
      num_classes: 1000
      depth: 28
      hidden_size: 1152
      patch_size: 1
      in_channels: 32
      num_heads: 16
      qk_norm: True
      encoder_depth: 8
      z_dim: 1280
      use_checkpoint: False
  # pretrained_vae:
  vae_dir: mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers
  slice_vae: False
  tile_vae: False
  # repa encoder
  enc_type: radio
  enc_dir: checkpoints/radio_v2.5-h.pth.tar
  proj_coeff: 1.0
  # ema
  use_ema: True
  ema_decay: 0.9999
  
data:
  data_type: improved_pack
  dataset:
    packed_json: datasets/imagenet1k/sampler_meta/dc-ae-f32c32-sana-1.1-diffusers_merge_LPFHP_16384.json
    jsonl_dir: datasets/imagenet1k/data_meta/dc-ae-f32c32-sana-1.1-diffusers_merge_meta.jsonl
    data_types: ['native-resolution', 'fixed-256x256', 'fixed-512x512']
    latent_dirs: [
      'datasets/imagenet1k/dc-ae-f32c32-sana-1.1-diffusers-native-resolution',
      'datasets/imagenet1k/dc-ae-f32c32-sana-1.1-diffusers-256x256',
      'datasets/imagenet1k/dc-ae-f32c32-sana-1.1-diffusers-512x512',
    ]
    image_dir: <Your imagenet1k directory>/train
  dataloader:
    num_workers: 4
    batch_size: 1  # Batch size (per device) for the training dataloader.

  
  
training:
  tracker: null
  tracker_kwargs: {'wandb': {'group': 'c2i'}}
  max_train_steps: 2000000
  checkpointing_steps: 2000
  checkpoints_total_limit: 2
  resume_from_checkpoint: latest
  learning_rate: 5.0e-5
  learning_rate_base_batch_size: 4
  scale_lr: True
  lr_scheduler: constant # "linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"]
  lr_warmup_steps: 0
  gradient_accumulation_steps: 1
  optimizer: 
    target: torch.optim.AdamW
    params:
      # betas: ${tuple:0.9, 0.999}
      betas: [0.9, 0.95]
      weight_decay: 1.0e-2
      eps: 1.0e-6
  max_grad_norm: 1.0
  proportion_empty_prompts: 0.0
  mixed_precision: bf16 # ["no", "fp16", "bf16"]
  allow_tf32: True 
  validation_steps: 500
  checkpoint_list: [200000, 500000, 100000, 150000]