unet/config.json

{
  "_target": "hsigene.HSIGeneUNet",
  "image_size": 32,
  "in_channels": 4,
  "model_channels": 320,
  "out_channels": 4,
  "num_res_blocks": 2,
  "attention_resolutions": [
    4,
    2,
    1
  ],
  "channel_mult": [
    1,
    2,
    4,
    4
  ],
  "use_checkpoint": true,
  "num_heads": 8,
  "use_spatial_transformer": true,
  "transformer_depth": 1,
  "context_dim": 768,
  "legacy": false
}