model_config.yml

img_size: 32
data_path: "./data3"

CLIP:
  batch_size: 256
  epochs: 23
  lr: 3.0e-4
  embed_dim: 256

  # Image encoder
  vision_layers: [1, 2, 2, 2]
  vision_width: 32
  vision_patch_size: null

  # Text encoder
  context_length: 33
  vocab_size: 49408
  transformer_width: 64
  transformer_heads: 16
  transformer_layers: 8

  model_path: "./models/clip4.pth"

Decoder:
  # Define hyperparameters
  diffusion_timesteps: 300
  noise_schedule: "cosine"
  batch_size: 64
  epochs: 1500
  lr: 0.001
  # grad_clip: 0.005
  grad_clip: null
  null_text_emb_rate: 0.5
  null_clip_emb_rate: 0.1
  guidance_scale: 3

  # UNet
  down_channels: [64, 128, 256]
  time_emb_dim: 32

  # UNet Transformer
  n_vocab: 49408
  context_length: 33
  transformer_width: 64
  transformer_layers: 6
  transformer_heads: 8

  # UNet attention block
  qkv_heads: 8

  normalize_clip_embeddings: True
  model_path: "./models/decoder4_cosine_ori.pth"

Prior:
  # Define hyperparameters
  diffusion_timesteps: 200
  noise_schedule: "cosine"
  batch_size: 256
  epochs: 100
  lr: 0.01
  normalize_clip_embeddings: True

  xf_layers: 8
  # dim_per_head: 64
  xf_heads: 16
  # ff_mult: 4

  model_path: "./models/prior4_cosine.pth"