config/models/ace_0.6b_1024.yaml

NAME: ACE_0.6B_1024_REFINER
IS_DEFAULT: True
DEFAULT_PARAS:
  PARAS:
  #
  INPUT:
    INPUT_IMAGE:
    INPUT_MASK:
    TASK:
    PROMPT: ""
    NEGATIVE_PROMPT: ""
    OUTPUT_HEIGHT: 1024
    OUTPUT_WIDTH: 1024
    SAMPLER: ddim
    SAMPLE_STEPS: 50
    GUIDE_SCALE: 4.5
    GUIDE_RESCALE: 0.5
    SEED: -1
    TAR_INDEX: 0
    REFINER_SCALE: 0.2
    USE_ACE: True
    #REFINER_PROMPT: "High Resolution, Sharpness, Clarity, Detail Enhancement, Noise Reduction, HD, 4k, Image Restoration, HDR"
    REFINER_PROMPT: "High Resolution, Sharpness, Clarity, Detail Enhancement, Noise Reduction, HD, 4k, Image Restoration, HDR"
  OUTPUT:
    LATENT:
    IMAGES:
    SEED:
  MODULES_PARAS:
    FIRST_STAGE_MODEL:
      FUNCTION:
        - NAME: encode
          DTYPE: float16
          INPUT: ["IMAGE"]
        - NAME: decode
          DTYPE: float16
          INPUT: ["LATENT"]
    #
    DIFFUSION_MODEL:
      FUNCTION:
        - NAME: forward
          DTYPE: float16
          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE"]
    #
    COND_STAGE_MODEL:
      FUNCTION:
        - NAME: encode_list_of_list
          DTYPE: bfloat16
          INPUT: ["PROMPT"]
#
MODEL:
  NAME: LatentDiffusionACE
  PRETRAINED_MODEL:
  IGNORE_KEYS: [ ]
  SCALE_FACTOR: 0.18215
  SIZE_FACTOR: 8
  DECODER_BIAS: 0.5
  DEFAULT_N_PROMPT: ""
  TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
  USE_TEXT_POS_EMBEDDINGS: True
  #
  DIFFUSION:
    NAME: BaseDiffusion
    PREDICTION_TYPE: eps
    MIN_SNR_GAMMA:
    NOISE_SCHEDULER:
      NAME: LinearScheduler
      NUM_TIMESTEPS: 1000
      BETA_MIN: 0.0001
      BETA_MAX: 0.02
  #
  DIFFUSION_MODEL:
    NAME: ACE
    PRETRAINED_MODEL: hf://iic/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
    IGNORE_KEYS: [ ]
    PATCH_SIZE: 2
    IN_CHANNELS: 4
    HIDDEN_SIZE: 1152
    DEPTH: 28
    NUM_HEADS: 16
    MLP_RATIO: 4.0
    PRED_SIGMA: True
    DROP_PATH: 0.0
    WINDOW_DIZE: 0
    Y_CHANNELS: 4096
    MAX_SEQ_LEN: 4096
    QK_NORM: True
    USE_GRAD_CHECKPOINT: True
    ATTENTION_BACKEND: flash_attn
  #
  FIRST_STAGE_MODEL:
    NAME: AutoencoderKL
    EMBED_DIM: 4
    PRETRAINED_MODEL: hf://iic/ACE-0.6B-1024px@models/vae/vae.bin
    IGNORE_KEYS: []
    #
    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
  #
  COND_STAGE_MODEL:
    NAME: T5EmbedderHF
    PRETRAINED_MODEL: hf://iic/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
    TOKENIZER_PATH: hf://iic/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
    LENGTH: 120
    T5_DTYPE: bfloat16
    ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
    CLEAN: whitespace
    USE_GRAD: False

ACE_PROMPT: [
  "A cute cartoon rabbit holding a whiteboard that says 'ACE Refiner', standing in a sunny meadow filled with flowers, with a big smile and bright colors.",
  "A beautiful young woman with long flowing hair, wearing a summer dress, holding a whiteboard that reads 'ACE Refiner' while sitting on a park bench surrounded by cherry blossoms.",
  "An adorable cartoon cat wearing oversized glasses, holding a whiteboard that says 'ACE Refiner', perched on a stack of colorful books in a cozy library setting.",
  "A charming girl with pigtails, wearing a cute school uniform, enthusiastically holding a whiteboard that has 'ACE Refiner' written on it, in a bright and cheerful classroom full of educational posters.",
  "A friendly cartoon dog with floppy ears, sitting in front of a doghouse, proudly holding a whiteboard that says 'ACE Refiner', with a playful expression and a blue sky in the background.",
  "A cute anime girl with big expressive eyes, dressed in a colorful outfit, holding a whiteboard that reads 'ACE Refiner' in a fantastical landscape filled with mythical creatures.",
  "A vibrant cartoon fox holding a whiteboard that says 'ACE Refiner', standing on a rock by a sparkling stream, surrounded by lush greenery and butterflies.",
  "A stylish young woman in a business outfit, smiling as she holds a whiteboard written with 'ACE Refiner', in a modern office filled with plants and natural light.",
  "A cute cartoon unicorn holding a sparkling whiteboard that says 'ACE Refiner', frolicking in a magical forest, with rainbows and stars in the background.",
  "A happy family, consisting of a cute little girl and her playful puppy, holding a whiteboard that says 'ACE Refiner', together in their backyard on a sunny day."
]
REFINER_MODEL:
  NAME: ""
  IS_DEFAULT: False
  DEFAULT_PARAS:
    PARAS:
      RESOLUTIONS: [ [ 1024, 1024 ] ]
    INPUT:
      INPUT_IMAGE:
      INPUT_MASK:
      TASK:
      PROMPT: ""
      NEGATIVE_PROMPT: ""
      OUTPUT_HEIGHT: 1024
      OUTPUT_WIDTH: 1024
      SAMPLER: flow_euler
      SAMPLE_STEPS: 30
      GUIDE_SCALE: 3.5
      GUIDE_RESCALE:
    OUTPUT:
      LATENT:
      IMAGES:
      SEED:
    MODULES_PARAS:
      FIRST_STAGE_MODEL:
        FUNCTION:
          - NAME: encode
            DTYPE: bfloat16
            INPUT: [ "IMAGE" ]
          - NAME: decode
            DTYPE: bfloat16
            INPUT: [ "LATENT" ]
        PARAS:
          SCALE_FACTOR: 1.5305
          SHIFT_FACTOR: 0.0609
          SIZE_FACTOR: 8
      DIFFUSION_MODEL:
        FUNCTION:
          - NAME: forward
            DTYPE: bfloat16
            INPUT: [ "SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE" ]
      COND_STAGE_MODEL:
        FUNCTION:
          - NAME: encode
            DTYPE: bfloat16
            INPUT: [ "PROMPT" ]

  MODEL:
    DIFFUSION:
      NAME: DiffusionFluxRF
      PREDICTION_TYPE: raw
      NOISE_SCHEDULER:
        NAME: FlowMatchSigmaScheduler
        WEIGHTING_SCHEME: logit_normal
        SHIFT: 3.0
        LOGIT_MEAN: 0.0
        LOGIT_STD: 1.0
        MODE_SCALE: 1.29
    DIFFUSION_MODEL:
      NAME: FluxMR
      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@flux1-dev.safetensors
      IN_CHANNELS: 64
      OUT_CHANNELS: 64
      HIDDEN_SIZE: 3072
      NUM_HEADS: 24
      AXES_DIM: [ 16, 56, 56 ]
      THETA: 10000
      VEC_IN_DIM: 768
      GUIDANCE_EMBED: True
      CONTEXT_IN_DIM: 4096
      MLP_RATIO: 4.0
      QKV_BIAS: True
      DEPTH: 19
      DEPTH_SINGLE_BLOCKS: 38
      USE_GRAD_CHECKPOINT: True
      ATTN_BACKEND: flash_attn
    #
    FIRST_STAGE_MODEL:
      NAME: AutoencoderKLFlux
      EMBED_DIM: 16
      PRETRAINED_MODEL:  ms://AI-ModelScope/FLUX.1-dev@ae.safetensors
      IGNORE_KEYS: [ ]
      BATCH_SIZE: 8
      USE_CONV: False
      SCALE_FACTOR: 0.3611
      SHIFT_FACTOR: 0.1159
      #
      ENCODER:
        NAME: Encoder
        USE_CHECKPOINT: False
        CH: 128
        OUT_CH: 3
        NUM_RES_BLOCKS: 2
        IN_CHANNELS: 3
        ATTN_RESOLUTIONS: [ ]
        CH_MULT: [ 1, 2, 4, 4 ]
        Z_CHANNELS: 16
        DOUBLE_Z: True
        DROPOUT: 0.0
        RESAMP_WITH_CONV: True
      #
      DECODER:
        NAME: Decoder
        USE_CHECKPOINT: False
        CH: 128
        OUT_CH: 3
        NUM_RES_BLOCKS: 2
        IN_CHANNELS: 3
        ATTN_RESOLUTIONS: [ ]
        CH_MULT: [ 1, 2, 4, 4 ]
        Z_CHANNELS: 16
        DROPOUT: 0.0
        RESAMP_WITH_CONV: True
        GIVE_PRE_END: False
        TANH_OUT: False
    #
    COND_STAGE_MODEL:
      NAME: T5PlusClipFluxEmbedder
      T5_MODEL:
        NAME: HFEmbedder
        HF_MODEL_CLS: T5EncoderModel
        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder_2/
        HF_TOKENIZER_CLS: T5Tokenizer
        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer_2/
        MAX_LENGTH: 512
        OUTPUT_KEY: last_hidden_state
        D_TYPE: bfloat16
        BATCH_INFER: False
        CLEAN: whitespace
      CLIP_MODEL:
        NAME: HFEmbedder
        HF_MODEL_CLS: CLIPTextModel
        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder/
        HF_TOKENIZER_CLS: CLIPTokenizer
        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer/
        MAX_LENGTH: 77
        OUTPUT_KEY: pooler_output
        D_TYPE: bfloat16
        BATCH_INFER: True
        CLEAN: whitespace