ENV:
  BACKEND: nccl
META:
  VERSION: 'SD1.5'
  DESCRIPTION: "Stable Diffusion v1.5"
  IS_DEFAULT: False
  IS_SHARE: True
  INFERENCE_PARAS:
    INFERENCE_BATCH_SIZE: 1
    INFERENCE_PREFIX: ""
    DEFAULT_SAMPLER: "ddim"
    DEFAULT_SAMPLE_STEPS: 40
    INFERENCE_N_PROMPT: ""
    RESOLUTION: [512, 512]
  PARAS:
    -
      TRAIN_BATCH_SIZE: 2
      TRAIN_PREFIX: ""
      TRAIN_N_PROMPT: ""
      RESOLUTION: [512, 512]
      MEMORY: 29000
      EPOCHS: 50
      SAVE_INTERVAL: 25
      EPSEC: 0.818
      LEARNING_RATE: 0.0001
      IS_DEFAULT: False
      TUNER: FULL
    -
      TRAIN_BATCH_SIZE: 2
      TRAIN_PREFIX: ""
      TRAIN_N_PROMPT: ""
      RESOLUTION: [512, 512]
      MEMORY: 29000
      EPOCHS: 50
      SAVE_INTERVAL: 25
      EPSEC: 0.818
      LEARNING_RATE: 0.0001
      IS_DEFAULT: False
      TUNER: LORA
    -
      TRAIN_BATCH_SIZE: 2
      TRAIN_PREFIX: ""
      TRAIN_N_PROMPT: ""
      RESOLUTION: [512, 512]
      MEMORY: 29000
      EPOCHS: 200
      SAVE_INTERVAL: 25
      EPSEC: 0.818
      LEARNING_RATE: 0.0001
      IS_DEFAULT: False
      TUNER: SCE

    -
      TRAIN_BATCH_SIZE: 2
      TRAIN_PREFIX: ""
      TRAIN_N_PROMPT: ""
      RESOLUTION: [512, 512]
      MEMORY: 29000
      EPOCHS: 200
      SAVE_INTERVAL: 25
      EPSEC: 0.818
      LEARNING_RATE: 0.0001
      IS_DEFAULT: True
      TUNER: TEXT_SCE

    -
      TRAIN_BATCH_SIZE: 4
      TRAIN_PREFIX: ""
      TRAIN_N_PROMPT: ""
      RESOLUTION: [512, 512]
      MEMORY: 29000
      EPOCHS: 50
      SAVE_INTERVAL: 25
      EPSEC: 0.818
      LEARNING_RATE: 0.0001
      IS_DEFAULT: False
      TUNER: TEXT_LORA

  TUNERS:
    LORA:
      -
        NAME: SwiftLoRA
        R: 256
        LORA_ALPHA: 256
        LORA_DROPOUT: 0.0
        BIAS: "none"
        TARGET_MODULES: "model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2)$"
    TEXT_LORA:
      -
        NAME: SwiftLoRA
        R: 256
        LORA_ALPHA: 256
        LORA_DROPOUT: 0.0
        BIAS: "none"
        TARGET_MODULES: "(cond_stage_model.*(q_proj|k_proj|v_proj|out_proj|mlp.fc1|mlp.fc2))|(model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2))$"
    SCE:
      -
        NAME: SwiftSCETuning
        DIMS: [1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320]
        DOWN_RATIO: 1.0
        TARGET_MODULES: model.lsc_identity\.\d+$
        TUNER_MODE: identity
    TEXT_SCE:
      -
        NAME: SwiftSCETuning
        DIMS: [ 1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320 ]
        DOWN_RATIO: 1.0
        TARGET_MODULES: model.lsc_identity\.\d+$
        TUNER_MODE: identity
      -
        NAME: SwiftLoRA
        R: 256
        LORA_ALPHA: 256
        LORA_DROPOUT: 0.0
        BIAS: "none"
        TARGET_MODULES: "cond_stage_model.*(q_proj|k_proj|v_proj|out_proj|mlp.fc1|mlp.fc2)$"

  MODIFY_PARAS:
    TEXT_LORA:
      TRAIN:
        SOLVER.MODEL.COND_STAGE_MODEL.USE_GRAD: True
    TEXT_SCE:
      TRAIN:
        SOLVER.MODEL.COND_STAGE_MODEL.USE_GRAD: True

SOLVER:
  NAME: LatentDiffusionSolver
  RESUME_FROM:
  LOAD_MODEL_ONLY: True
  USE_FSDP: False
  SHARDING_STRATEGY:
  USE_AMP: True
  DTYPE: float16
  CHANNELS_LAST: True
  MAX_STEPS: 1000
  MAX_EPOCHS: -1
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: -1
  RESCALE_LR: False
  #
  WORK_DIR:
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
  #
  TUNER:
  #
  MODEL:
    NAME: LatentDiffusion
    PARAMETERIZATION: eps
    TIMESTEPS: 1000
    MIN_SNR_GAMMA:
    ZERO_TERMINAL_SNR: False
    PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-v1-5@v1-5-pruned-emaonly.safetensors
    IGNORE_KEYS: [ ]
    SCALE_FACTOR: 0.18215
    SIZE_FACTOR: 8
    # DEFAULT_N_PROMPT: 'lowres, error, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'
    DEFAULT_N_PROMPT:
    SCHEDULE_ARGS:
      "NAME": "scaled_linear"
      "BETA_MIN": 0.00085
      "BETA_MAX": 0.012
    USE_EMA: False
    #
    DIFFUSION_MODEL:
      NAME: DiffusionUNet
      IN_CHANNELS: 4
      OUT_CHANNELS: 4
      MODEL_CHANNELS: 320
      NUM_HEADS: 8
      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: 768
      DISABLE_MIDDLE_SELF_ATTN: False
      USE_LINEAR_IN_TRANSFORMER: False
      IGNORE_KEYS: []
    #
    FIRST_STAGE_MODEL:
      NAME: AutoencoderKL
      EMBED_DIM: 4
      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: ClipTokenizer
      PRETRAINED_PATH: ms://AI-ModelScope/clip-vit-large-patch14
      LENGTH: 77
      CLEAN: True
    #
    COND_STAGE_MODEL:
      NAME: FrozenCLIPEmbedder
      FREEZE: True
      USE_GRAD: False
      LAYER: last
      PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
    #
    LOSS:
      NAME: ReconstructLoss
      LOSS_TYPE: l2
  #
  SAMPLE_ARGS:
    SAMPLER: ddim
    SAMPLE_STEPS: 50
    SEED: 2023
    GUIDE_SCALE: 7.5
    GUIDE_RESCALE:
    DISCRETIZATION: trailing
    RUN_TRAIN_N: False
  #
  OPTIMIZER:
    NAME: AdamW
    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
    AMSGRAD: False
  #
  TRAIN_DATA:
    NAME: ImageTextPairMSDataset
    MODE: train
    MS_DATASET_NAME: style_custom_dataset
    MS_DATASET_NAMESPACE: damo
    MS_DATASET_SUBNAME: 3D
    PROMPT_PREFIX: ""
    MS_DATASET_SPLIT: train
    MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
    REPLACE_STYLE: False
    PIN_MEMORY: True
    BATCH_SIZE: 1
    NUM_WORKERS: 4
    SAMPLER:
      NAME: LoopSampler
    TRANSFORMS:
      - NAME: LoadImageFromFile
        RGB_ORDER: RGB
        BACKEND: pillow
      - NAME: FlexibleResize
        INTERPOLATION: bilinear
        SIZE: [ 512, 512 ]
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
      - NAME: FlexibleCenterCrop
        SIZE: [ 512, 512 ]
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
      - NAME: ImageToTensor
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
      - NAME: Normalize
        MEAN: [ 0.5,  0.5,  0.5 ]
        STD: [ 0.5,  0.5,  0.5 ]
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'image' ]
        BACKEND: torchvision
      - NAME: Select
        KEYS: [ 'image', 'prompt' ]
        META_KEYS: [ 'data_key' ]
  #
  EVAL_DATA:
    NAME: Text2ImageDataset
    MODE: eval
    PROMPT_FILE:
    PROMPT_DATA: [ "a boy wearing a jacket", "a dog running on the lawn" ]
    IMAGE_SIZE: [ 512, 512 ]
    FIELDS: [ "prompt" ]
    DELIMITER: '#;#'
    PROMPT_PREFIX: ''
    PIN_MEMORY: True
    BATCH_SIZE: 1
    NUM_WORKERS: 4
    TRANSFORMS:
      - NAME: Select
        KEYS: [ 'index', 'prompt' ]
        META_KEYS: [ 'image_size' ]
  #
  TRAIN_HOOKS:
    -
      NAME: BackwardHook
      PRIORITY: 0
    -
      NAME: LogHook
      LOG_INTERVAL: 10
      SHOW_GPU_MEM: True
    -
      NAME: TensorboardLogHook
    -
      NAME: CheckpointHook
      INTERVAL: 10000
      PRIORITY: 200
      SAVE_LAST: True
      SAVE_NAME_PREFIX: 'step'
      DISABLE_SNAPSHOT: True
  #
  EVAL_HOOKS:
    -
      NAME: ProbeDataHook
      PROB_INTERVAL: 100
      SAVE_LAST: True
      SAVE_NAME_PREFIX: 'step'
      SAVE_PROBE_PREFIX: 'image'