config.json

{
  "model_config": {
    "down_channels": [
      32,
      64,
      128,
      256
    ],
    "down_sample": [
      true,
      true,
      false
    ],
    "im_channels": 1,
    "im_size": 28,
    "mid_channels": [
      256,
      256,
      128
    ],
    "num_down_layers": 2,
    "num_heads": 4,
    "num_mid_layers": 2,
    "num_up_layers": 2,
    "time_emb_dim": 128
  }
}