NAME: SD2.1
IS_DEFAULT: False
DEFAULT_PARAS:
  PARAS:
    RESOLUTIONS: [[768, 768]]
  INPUT:
    IMAGE:
    PROMPT: ""
    NEGATIVE_PROMPT: ""
    PROMPT_PREFIX: ""
    TARGET_SIZE_AS_TUPLE: [768, 768]
    SAMPLE: ddim
    SAMPLE_STEPS: 50
    GUIDE_SCALE: 7.5
    GUIDE_RESCALE: 0.5
    DISCRETIZATION: trailing
  OUTPUT:
    LATENT:
    IMAGES:
    SEED:
  MODULES_PARAS:
    FIRST_STAGE_MODEL:
      FUNCTION:
        -
          NAME: encode
          DTYPE: float16
          INPUT: ["IMAGE"]
        -
          NAME: decode
          DTYPE: float16
          INPUT: ["LATENT"]
      PARAS:
        # SCALE_FACTOR DESCRIPTION: The vae embeding scale. TYPE: float default: 0.18215
        SCALE_FACTOR: 0.18215
        SIZE_FACTOR: 8
    DIFFUSION_MODEL:
      FUNCTION:
        -
          NAME: forward
          DTYPE: float16
          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE", "GUIDE_RESCALE", "DISCRETIZATION"]
    COND_STAGE_MODEL:
      FUNCTION:
        -
          NAME: encode_text
          DTYPE: float16
          INPUT: ["PROMPT", "NEGATIVE_PROMPT"]

MODEL:
  PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-1@v2-1_768-ema-pruned.safetensors
  SCHEDULE:
    PARAMETERIZATION: "v"
    TIMESTEPS: 1000
    ZERO_TERMINAL_SNR: False
    SCHEDULE_ARGS:
      # NAME DESCRIPTION:  TYPE:  default: ''
      NAME: "scaled_linear"
      BETA_MIN: 0.00085
      BETA_MAX: 0.0120
  #
  DIFFUSION_MODEL:
    NAME: DiffusionUNet
    IN_CHANNELS: 4
    OUT_CHANNELS: 4
    MODEL_CHANNELS: 320
    NUM_HEADS_CHANNELS: 64
    NUM_RES_BLOCKS: 2
    ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
    CHANNEL_MULT: [ 1, 2, 4, 4 ]
    CONV_RESAMPLE: True
    DIMS: 2
    USE_CHECKPOINT: False
    USE_SCALE_SHIFT_NORM: False
    RESBLOCK_UPDOWN: False
    USE_SPATIAL_TRANSFORMER: True
    TRANSFORMER_DEPTH: 1
    CONTEXT_DIM: 1024
    DISABLE_MIDDLE_SELF_ATTN: False
    USE_LINEAR_IN_TRANSFORMER: True
    PRETRAINED_MODEL:
  #
  FIRST_STAGE_MODEL:
    NAME: AutoencoderKL
    EMBED_DIM: 4
    PRETRAINED_MODEL:
    IGNORE_KEYS: [ ]
    BATCH_SIZE: 4
    #
    ENCODER:
      NAME: Encoder
      CH: 128
      OUT_CH: 3
      NUM_RES_BLOCKS: 2
      IN_CHANNELS: 3
      ATTN_RESOLUTIONS: [ ]
      CH_MULT: [ 1, 2, 4, 4 ]
      Z_CHANNELS: 4
      DOUBLE_Z: True
      DROPOUT: 0.0
      RESAMP_WITH_CONV: True
    #
    DECODER:
      NAME: Decoder
      CH: 128
      OUT_CH: 3
      NUM_RES_BLOCKS: 2
      IN_CHANNELS: 3
      ATTN_RESOLUTIONS: [ ]
      CH_MULT: [ 1, 2, 4, 4 ]
      Z_CHANNELS: 4
      DROPOUT: 0.0
      RESAMP_WITH_CONV: True
      GIVE_PRE_END: False
      TANH_OUT: False
  #
  TOKENIZER:
    NAME: OpenClipTokenizer
    LENGTH: 77
  #
  COND_STAGE_MODEL:
    NAME: FrozenOpenCLIPEmbedder
    ARCH: ViT-H-14
    PRETRAINED_MODEL:
    LAYER: penultimate