configs/model/openai_unet.yaml

#########
# v1 2d #
#########

openai_unet_2d_v1:
  type: openai_unet_2d_next
  args:
    in_channels: 4
    out_channels: 4
    model_channels: 320
    attention_resolutions: [ 4, 2, 1 ]
    num_res_blocks: [ 2, 2, 2, 2 ]
    channel_mult: [ 1, 2, 4, 4 ]
    num_heads: 8
    context_dim: 768
    use_checkpoint: True
    parts: [global, data, context]

openai_unet_2d_v1_g:
  super_cfg: openai_unet_2d_v1
  args:
    parts: [global]

openai_unet_2d_v1_d:
  super_cfg: openai_unet_2d_v1
  args:
    parts: [data]

openai_unet_2d_v1_c:
  super_cfg: openai_unet_2d_v1
  args:
    parts: [context]

openai_unet_2d_v1_gd:
  super_cfg: openai_unet_2d_v1
  args:
    parts: [global, data]

openai_unet_2d_v1_gc:
  super_cfg: openai_unet_2d_v1
  args:
    parts: [global, context]

openai_unet_2d_v1_dc:
  super_cfg: openai_unet_2d_v1
  args:
    parts: [data, context]

#########
# v1 0d #
#########

openai_unet_0d_v1:
  type: openai_unet_0d_next
  args:
    input_channels: 768
    model_channels: 320
    output_channels: 768
    num_noattn_blocks: [ 2, 2, 2, 2 ]
    channel_mult: [ 1, 2, 4, 4 ]
    second_dim: [ 4, 4, 4, 4 ]
    with_attn: [true, true, true, false]
    num_heads: 8
    context_dim: 768
    use_checkpoint: True
    parts: [global, data, context]

openai_unet_0d_v1_g:
  super_cfg: openai_unet_0d_v1
  args:
    parts: [global]

openai_unet_0d_v1_d:
  super_cfg: openai_unet_0d_v1
  args:
    parts: [data]

openai_unet_0d_v1_c:
  super_cfg: openai_unet_0d_v1
  args:
    parts: [context]

openai_unet_0d_v1_gd:
  super_cfg: openai_unet_0d_v1
  args:
    parts: [global, data]

openai_unet_0d_v1_gc:
  super_cfg: openai_unet_0d_v1
  args:
    parts: [global, context]

openai_unet_0d_v1_dc:
  super_cfg: openai_unet_0d_v1
  args:
    parts: [data, context]