Merge pull request #52 from modelscope/v1.1.0_dev

update v1.1.0
update
2024-10-23 10:19:29 +08:00 · 2024-10-23 10:15:31 +08:00 · 2024-10-21 11:59:55 +08:00 · 2024-10-21 11:59:41 +08:00 · 2024-10-21 11:58:43 +08:00 · 2024-10-21 11:36:02 +08:00
216 changed files with 27296 additions and 1765 deletions
@@ -18,3 +18,4 @@ scepter.egg-info
 #MANIFEST.in
 *resources
 *.ipynb_checkpoints*
+*.vscode
@@ -0,0 +1,163 @@
+{
+  "last_node_id": 4,
+  "last_link_id": 3,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "PreviewImage",
+      "pos": {
+        "0": 937,
+        "1": 243
+      },
+      "size": {
+        "0": 284.03680419921875,
+        "1": 246
+      },
+      "flags": {
+        "collapsed": false
+      },
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 4,
+      "type": "ParameterNode",
+      "pos": {
+        "0": 0,
+        "1": 244
+      },
+      "size": {
+        "0": 311.9849853515625,
+        "1": 236.4765625
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ParameterNode"
+      },
+      "widgets_values": [
+        "ddim",
+        50,
+        5,
+        0.5,
+        "trailing",
+        1024,
+        1024,
+        2024
+      ]
+    },
+    {
+      "id": 1,
+      "type": "ModelNode",
+      "pos": {
+        "0": 430,
+        "1": 243
+      },
+      "size": {
+        "0": 400,
+        "1": 234
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "parameters",
+          "type": "CONDITIONING",
+          "link": 3,
+          "shape": 7
+        },
+        {
+          "name": "mantras",
+          "type": "CONDITIONING",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "tuners",
+          "type": "CONDITIONING",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "controls",
+          "type": "CONDITIONING",
+          "link": null,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            1
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ModelNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0",
+        "ModelScope",
+        "a cat",
+        ""
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      3,
+      4,
+      0,
+      1,
+      0,
+      "CONDITIONING"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.9229599817706423,
+      "offset": [
+        128.91466506592707,
+        43.07074577975898
+      ]
+    }
+  },
+  "version": 0.4
+}
@@ -0,0 +1,202 @@
+{
+  "last_node_id": 5,
+  "last_link_id": 4,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "PreviewImage",
+      "pos": {
+        "0": 937,
+        "1": 243
+      },
+      "size": {
+        "0": 284.03680419921875,
+        "1": 246
+      },
+      "flags": {
+        "collapsed": false
+      },
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 4,
+      "type": "ParameterNode",
+      "pos": {
+        "0": 0,
+        "1": 244
+      },
+      "size": {
+        "0": 311.9849853515625,
+        "1": 236.4765625
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ParameterNode"
+      },
+      "widgets_values": [
+        "ddim",
+        50,
+        5,
+        0.5,
+        "trailing",
+        1024,
+        1024,
+        2024
+      ]
+    },
+    {
+      "id": 1,
+      "type": "ModelNode",
+      "pos": {
+        "0": 430,
+        "1": 243
+      },
+      "size": {
+        "0": 400,
+        "1": 234
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "parameters",
+          "type": "CONDITIONING",
+          "link": 3,
+          "shape": 7
+        },
+        {
+          "name": "mantras",
+          "type": "CONDITIONING",
+          "link": 4,
+          "shape": 7
+        },
+        {
+          "name": "tuners",
+          "type": "CONDITIONING",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "controls",
+          "type": "CONDITIONING",
+          "link": null,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            1
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ModelNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0",
+        "ModelScope",
+        "a cat",
+        ""
+      ]
+    },
+    {
+      "id": 5,
+      "type": "MantrasNode",
+      "pos": {
+        "0": 3,
+        "1": 547
+      },
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "MantrasNode"
+      },
+      "widgets_values": [
+        "Flat 2D Art"
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      3,
+      4,
+      0,
+      1,
+      0,
+      "CONDITIONING"
+    ],
+    [
+      4,
+      5,
+      0,
+      1,
+      1,
+      "CONDITIONING"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.0152559799477068,
+      "offset": [
+        58.35814123566077,
+        -65.12749162217047
+      ]
+    }
+  },
+  "version": 0.4
+}
@@ -0,0 +1,242 @@
+{
+  "last_node_id": 6,
+  "last_link_id": 5,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "PreviewImage",
+      "pos": {
+        "0": 937,
+        "1": 243
+      },
+      "size": {
+        "0": 284.03680419921875,
+        "1": 246
+      },
+      "flags": {
+        "collapsed": false
+      },
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 4,
+      "type": "ParameterNode",
+      "pos": {
+        "0": 0,
+        "1": 244
+      },
+      "size": {
+        "0": 311.9849853515625,
+        "1": 236.4765625
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ParameterNode"
+      },
+      "widgets_values": [
+        "ddim",
+        50,
+        5,
+        0.5,
+        "trailing",
+        1024,
+        1024,
+        2024
+      ]
+    },
+    {
+      "id": 5,
+      "type": "MantrasNode",
+      "pos": {
+        "0": 3,
+        "1": 547
+      },
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "MantrasNode"
+      },
+      "widgets_values": [
+        "Flat 2D Art"
+      ]
+    },
+    {
+      "id": 1,
+      "type": "ModelNode",
+      "pos": {
+        "0": 430,
+        "1": 243
+      },
+      "size": {
+        "0": 400,
+        "1": 234
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "parameters",
+          "type": "CONDITIONING",
+          "link": 3,
+          "shape": 7
+        },
+        {
+          "name": "mantras",
+          "type": "CONDITIONING",
+          "link": 4,
+          "shape": 7
+        },
+        {
+          "name": "tuners",
+          "type": "CONDITIONING",
+          "link": 5,
+          "shape": 7
+        },
+        {
+          "name": "controls",
+          "type": "CONDITIONING",
+          "link": null,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            1
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ModelNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0",
+        "ModelScope",
+        "a cat",
+        ""
+      ]
+    },
+    {
+      "id": 6,
+      "type": "TunerNode",
+      "pos": {
+        "0": 8,
+        "1": 680
+      },
+      "size": {
+        "0": 315,
+        "1": 82
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            5
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "TunerNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0_Pencil Sketch Drawing",
+        1
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      3,
+      4,
+      0,
+      1,
+      0,
+      "CONDITIONING"
+    ],
+    [
+      4,
+      5,
+      0,
+      1,
+      1,
+      "CONDITIONING"
+    ],
+    [
+      5,
+      6,
+      0,
+      1,
+      2,
+      "CONDITIONING"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.6934334949441366,
+      "offset": [
+        450.0923948318967,
+        30.783653239467935
+      ]
+    }
+  },
+  "version": 0.4
+}
@@ -0,0 +1,366 @@
+{
+  "last_node_id": 11,
+  "last_link_id": 9,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "ModelNode",
+      "pos": {
+        "0": 439,
+        "1": 123
+      },
+      "size": {
+        "0": 400,
+        "1": 234
+      },
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "parameters",
+          "type": "CONDITIONING",
+          "link": 3,
+          "shape": 7
+        },
+        {
+          "name": "mantras",
+          "type": "CONDITIONING",
+          "link": 4,
+          "shape": 7
+        },
+        {
+          "name": "tuners",
+          "type": "CONDITIONING",
+          "link": 5,
+          "shape": 7
+        },
+        {
+          "name": "controls",
+          "type": "CONDITIONING",
+          "link": 6,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            1
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ModelNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0",
+        "ModelScope",
+        "a cat",
+        ""
+      ]
+    },
+    {
+      "id": 5,
+      "type": "MantrasNode",
+      "pos": {
+        "0": 13,
+        "1": 336
+      },
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "MantrasNode"
+      },
+      "widgets_values": [
+        "Flat 2D Art"
+      ]
+    },
+    {
+      "id": 6,
+      "type": "TunerNode",
+      "pos": {
+        "0": 12,
+        "1": 455
+      },
+      "size": {
+        "0": 315,
+        "1": 82
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            5
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "TunerNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0_Pencil Sketch Drawing",
+        1
+      ]
+    },
+    {
+      "id": 11,
+      "type": "NoteNode",
+      "pos": {
+        "0": 845,
+        "1": 506
+      },
+      "size": {
+        "0": 398.3059387207031,
+        "1": 210.83267211914062
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "NoteNode"
+      },
+      "widgets_values": [
+        "This is a sample for quickly setting up ComfyUI using the Scepter open-source library:\n\n1) The first run involves downloading the models. By default, we will automatically pull models from ModelScope. The initial download may take some time, and you can also adjust the model source to change the model address.\n\n2) Currently, it supports various base models, basic settings for some inference hyperparameters, mantra settings, tuning model settings, and conditional generation node.\n\n3) In the future, we will gradually integrate interesting features to enrich the use cases.\n\n"
+      ]
+    },
+    {
+      "id": 8,
+      "type": "LoadImage",
+      "pos": {
+        "0": 471,
+        "1": 451
+      },
+      "size": {
+        "0": 315,
+        "1": 314
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            7
+          ]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "cat.jpg",
+        "image"
+      ]
+    },
+    {
+      "id": 7,
+      "type": "ControlNode",
+      "pos": {
+        "0": 15,
+        "1": 600
+      },
+      "size": {
+        "0": 330,
+        "1": 198
+      },
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "source_image",
+          "type": "IMAGE",
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            6
+          ]
+        },
+        {
+          "name": "Control Image",
+          "type": "IMAGE",
+          "links": [],
+          "slot_index": 1
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ControlNode"
+      },
+      "widgets_values": [
+        "SD_XL1.0_color",
+        "Color",
+        "CenterCrop",
+        1,
+        1024,
+        1024
+      ]
+    },
+    {
+      "id": 2,
+      "type": "PreviewImage",
+      "pos": {
+        "0": 959,
+        "1": 121
+      },
+      "size": {
+        "0": 284.03680419921875,
+        "1": 246
+      },
+      "flags": {
+        "collapsed": false
+      },
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 4,
+      "type": "ParameterNode",
+      "pos": {
+        "0": 13,
+        "1": 43
+      },
+      "size": {
+        "0": 311.9849853515625,
+        "1": 236.4765625
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "Result",
+          "type": "CONDITIONING",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ParameterNode"
+      },
+      "widgets_values": [
+        "ddim",
+        50,
+        5,
+        0.5,
+        "trailing",
+        1024,
+        1024,
+        2024
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      3,
+      4,
+      0,
+      1,
+      0,
+      "CONDITIONING"
+    ],
+    [
+      4,
+      5,
+      0,
+      1,
+      1,
+      "CONDITIONING"
+    ],
+    [
+      5,
+      6,
+      0,
+      1,
+      2,
+      "CONDITIONING"
+    ],
+    [
+      6,
+      7,
+      0,
+      1,
+      3,
+      "CONDITIONING"
+    ],
+    [
+      7,
+      8,
+      0,
+      7,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.7627768444385667,
+      "offset": [
+        254.62860058494985,
+        87.05734194144193
+      ]
+    }
+  },
+  "version": 0.4
+}
@@ -56,7 +56,7 @@ Given an original image, image editing aims to generate an image that align with
 - More models will be released in the future.

 ## Run StyleBooth
- Code implementation: See model configuration and code based on [🪄SCEPTER](https://github.com/modelscope/scepter).
+- Code implementation: See model configuration and code based on [🪄SCEPTER](https://github.com/modelscope/scepter/blob/main/docs/en/tasks/stylebooth.md).

 - Demo: Try [🖥️SCEPTER Studio](https://github.com/modelscope/scepter/tree/main?tab=readme-ov-file#%EF%B8%8F-scepter-studio).

@@ -114,3 +114,17 @@ class DiffusionInferenceTest(unittest.TestCase):
 if __name__ == '__main__':
    unittest.main()
 ```
+
+## StyleTuner and De-StyleTuner.
+
+### Base I2I Model.
+
+For style and de-style tuning, we use a private high-resolution I2I model trained with [InstructPix2Pix dataset](https://instruct-pix2pix.eecs.berkeley.edu/) as base model. However, one can try the same tunning process using this [yaml](https://github.com/modelscope/scepter/blob/main/scepter/methods/edit/edit_512_lora.yaml) based on any other I2I model, such as StyleBooth (shown in this yaml), [InstructPix2Pix](https://github.com/timothybrooks/instruct-pix2pix) or [MagicBrush](https://github.com/OSU-NLP-Group/MagicBrush).
+
+### Training Data.
+
+Please check the zips for correct format: [De-Text](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fdetext.zip), [Image2Hed](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fhed_pair.zip), [Image2Depth](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fimage2depth.zip), [Depth2Image](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fdepth2image.zip).
+
+### Launch.
+
+See [code](https://github.com/modelscope/scepter/blob/cbbef4a2da5b66fc33b9f8ece7f2fb4aac9d6e3c/tests/tools/test_train.py#L220) for more information.
@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import shutil
+import subprocess
+import sys
+
+if sys.argv[0] == 'install.py':
+    sys.path.append('.')   # for portable version
+
+source_folder = os.path.join(os.path.dirname(__file__), "scepter/workflow")
+current_dir = os.path.dirname(__file__)
+destination_folder = os.path.join(os.path.dirname(current_dir), "ComfyUI-Scepter")
+
+if not os.path.exists(destination_folder):
+    shutil.copytree(source_folder, destination_folder)
+    print(f"{os.path.abspath(source_folder)} copy to {os.path.abspath(destination_folder)} success!")
+else:
+    print(f"{os.path.abspath(destination_folder)} exist.")
+
+# pip install scepter
+subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'scepter'])
@@ -18,6 +18,9 @@ SCEPTER offers 3 core components:


 ## 🎉 News
+- [2024.10]: Support for inference and tuning with [FLUX](https://huggingface.co/black-forest-labs/FLUX.1-dev), as well as for building [ComfyUI](https://github.com/comfyanonymous/ComfyUI) workflows using this framework.
+- [🔥2024.09]: We introduce **ACE**, an **A**ll-round **C**reator and **E**ditor adept at executing a diverse array of image editing tasks tailored to your specifications. Built upon the cutting-edge Diffusion Transformer architecture, ACE has been extensively trained on a comprehensive dataset to seamlessly interpret and execute any natural language instruction. For further information, please consult the [project page](https://ali-vilab.github.io/ace-page/).
+- [2024.07]: Support the inference and training of open-source generative models based on the [DiT](https://arxiv.org/abs/2212.09748) architecture, such as [SD3](https://arxiv.org/pdf/2403.03206) and [PixArt](https://arxiv.org/abs/2310.00426).
 - [2024.05]: Introducing SCEPTER v1, supporting customized image edit tasks! Simply provide 10 image pairs, SCEPTER will tune an edit tuner for your own Image-to-Image tasks, like `Clay Style`, `De-Text`, `Segmentation`, etc.
 - [2024.04]: New [StyleBooth](https://ali-vilab.github.io/stylebooth-page/) demo on SCEPTER Studio for`Text-Based Style Editing`.
 - [2024.03]: We optimize the training UI and checkpoint management. New [LAR-Gen](https://arxiv.org/abs/2403.19534) model has been added on SCEPTER Studio, supporting `zoom-out`, `virtual try on`, `inpainting`.
@@ -30,45 +33,83 @@ SCEPTER offers 3 core components:

 ## 🖼 Gallery for Recent Works

-### Edit Tuners
+### <img src="https://github.com/ali-vilab/ace-page/raw/main/static/images/logo.png?raw=true" height=20> <img src="https://github.com/ali-vilab/ace-page/raw/main/static/images/icon.png?raw=true" height=20>

-Simply provide 10 image pairs, SCEPTER will tune an edit tuner for your own Image-to-Image tasks, like `Clay Style`, `De-Text`, `Segmentation`, etc.
-Try our official few-shot datasets: [De-Text](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fdetext.zip), [Image2Hed](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fhed_pair.zip), [Image2Depth](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fimage2depth.zip), [Depth2Image](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets%2Fdepth2image.zip).
+ACE is a unified foundational model framework that supports a wide range of visual generation tasks. By defining CU for unifying multi-modal inputs across different tasks and incorporating long-
+context CU, we introduce historical contextual information into visual generation tasks, paving
+the way for ChatGPT-like dialog systems in visual generation.

+[![Watch the demo](https://github.com/ali-vilab/ace-page/raw/main/static/images/teaser.jpg)](https://ali-vilab.github.io/ace-page/)
+
+### FLUX Tuners

 <table><tbody>
  <tr>
-    <th align="center" colspan="4">Clay Style<br>Prompt: "Convert this image into clay style"</th>
+    <th align="center" colspan="3">Yarn Style</th>
+    <th align="center" colspan="3">Soft Watercolor Style</th>
  </tr>
  <tr>
-    <td><img src="asset/images/edit_tuner/vermeer.jpeg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/clay_vermeer.jpeg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/cat_512.jpg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/clay_cat.jpeg" width="300"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_2_1.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_2_2.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_2_3.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_1_1.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_1_2.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_1_3.webp" width="200"></td>
  </tr>
  <tr>
-    <th align="center" colspan="2">De-Text<br>Prompt: "Remove the texts"</th>
-    <th align="center" colspan="2">Image2Hed<br>Prompt: "Convert to an edge map"</th>
+    <th align="center" colspan="3">Travel Style</th>
+    <th align="center" colspan="3">WuKong Style</th>
  </tr>
  <tr>
-    <td><img src="asset/images/edit_tuner/text.jpg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/detext.jpeg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/cat_512.jpg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/hed.jpeg" width="300"></td>
-  </tr>
-    <tr>
-    <th align="center" colspan="2">Image2Depth<br>Prompt: "Calculate the depth map"</th>
-    <th align="center" colspan="2">Depth2Image<br>Prompt: "Convert depth map into color image"</th>
-  </tr>
-  <tr>
-    <td><img src="asset/images/edit_tuner/house.jpg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/image2depth.jpeg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/depth.jpg" width="300"></td>
-    <td><img src="asset/images/edit_tuner/depth2image.jpeg" width="300"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_3_1.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_3_2.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_3_3.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_4_1.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_4_2.webp" width="200"></td>
+    <td><img src="asset/images/flux_tuner/flux_tuner_4_3.webp" width="200"></td>
  </tr>
 </tbody>
 </table>
-Note: Left image is input and right image is output.
+
+### ComfyUI Workflow
+
+![Workflow](asset/workflow/workflow.jpg)
+
+<table><tbody>
+  <tr>
+    <th align="center" colspan="4">Example Workflow Case</th>
+  </tr>
+  <tr>
+    <th align="center" colspan="1">Base</th>
+    <th align="center" colspan="1">+Mantra</th>
+    <th align="center" colspan="1">+Tuner</th>
+    <th align="center" colspan="1">+Control</th>
+  </tr>
+  <tr>
+    <td>
+      <a href="asset/workflow/sdxl_base.json" target="_blank">
+        <img src="asset/workflow/sdxl_base.jpg" width="200">
+      </a>
+    </td>
+    <td>
+      <a href="asset/workflow/sdxl_base_mantra.json" target="_blank">
+        <img src="asset/workflow/sdxl_base_mantra.jpg" width="200">
+      </a>
+    </td>
+    <td>
+      <a href="asset/workflow/sdxl_base_mantra_tuner.json" target="_blank">
+        <img src="asset/workflow/sdxl_base_mantra_tuner.jpg" width="200">
+      </a>
+    </td>
+    <td>
+      <a href="asset/workflow/sdxl_base_mantra_tuner_control.json" target="_blank">
+        <img src="asset/workflow/sdxl_base_mantra_tuner_control.jpg" width="200">
+      </a>
+    </td>
+  </tr>
+</tbody>
+</table>
+

 ## 🛠️ Installation

@@ -102,16 +143,18 @@ pip install scepter

 ### Currently supported approaches

-|            Tasks             |                   Methods                    | Links                                                                                                                                                                                                                                                       |
-|:----------------------------:|:--------------------------------------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-|   Text-to-image generation   |                   SD v1.5                    | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/runwayml/stable-diffusion-v1-5)                                                                                                         |
-|   Text-to-image generation   |                   SD v2.1                    | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/runwayml/stable-diffusion-v1-5)                                                                                                         |
-|   Text-to-image generation   |                    SD-XL                     | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)                                                                                               |
-|       Efficient Tuning       |                     LoRA                     | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=LoRA&color=red&logo=arxiv)](https://arxiv.org/abs/2106.09685)                                                                                                                         |
-|       Efficient Tuning       |            Res-Tuning(NeurIPS23)             | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=Res-Tuing&color=red&logo=arxiv)](https://arxiv.org/abs/2310.19859) [![Page link](https://img.shields.io/badge/Page-ResTuning-Gree)](https://res-tuning.github.io/)                    |
-| Controllable image synthesis | [🌟SCEdit(CVPR24)](docs/en/tasks/scedit.md)  | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=SCEdit&color=red&logo=arxiv)](https://arxiv.org/abs/2312.11392)   [![Page link](https://img.shields.io/badge/Page-SCEdit-Gree)](https://scedit.github.io/)                            |
-|        Image editing         |     [🌟LAR-Gen](docs/en/tasks/largen.md)     | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=LARGen&color=red&logo=arxiv)](https://arxiv.org/abs/2403.19534)   [![Page link](https://img.shields.io/badge/Page-LARGen-Gree)](https://ali-vilab.github.io/largen-page/)             |
-|        Image editing         |  [🌟StyleBooth](docs/en/tasks/stylebooth.md) | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=StyleBooth&color=red&logo=arxiv)](https://arxiv.org/abs/2404.12154)   [![Page link](https://img.shields.io/badge/Page-StyleBooth-Gree)](https://ali-vilab.github.io/stylebooth-page/) |
+|            Tasks             |                    Methods                     | Links                                                                                                                                                                                                                                                       |
+|:----------------------------:|:----------------------------------------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+|   Text-to-image Generation   |                    SD v1.5                     | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/runwayml/stable-diffusion-v1-5)                                                                                                         |
+|   Text-to-image Generation   |                    SD v2.1                     | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/runwayml/stable-diffusion-v1-5)                                                                                                         |
+|   Text-to-image Generation   |                     SD-XL                      | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)                                                                                               |
+|   Text-to-image Generation   |                      FLUX                      | [![Hugging Face Repo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repo-blue)](https://huggingface.co/black-forest-labs/FLUX.1-dev)                                                                                               |
+|       Efficient Tuning       |                      LoRA                      | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=LoRA&color=red&logo=arxiv)](https://arxiv.org/abs/2106.09685)                                                                                                                         |
+|       Efficient Tuning       |             Res-Tuning(NeurIPS23)              | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=Res-Tuing&color=red&logo=arxiv)](https://arxiv.org/abs/2310.19859) [![Page link](https://img.shields.io/badge/Page-ResTuning-Gree)](https://res-tuning.github.io/)                    |
+| Controllable Image Synthesis |  [🌟SCEdit(CVPR24)](docs/en/tasks/scedit.md)   | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=SCEdit&color=red&logo=arxiv)](https://arxiv.org/abs/2312.11392)   [![Page link](https://img.shields.io/badge/Page-SCEdit-Gree)](https://scedit.github.io/)                            |
+|        Image Editing         |      [🌟LAR-Gen](docs/en/tasks/largen.md)      | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=LARGen&color=red&logo=arxiv)](https://arxiv.org/abs/2403.19534)   [![Page link](https://img.shields.io/badge/Page-LARGen-Gree)](https://ali-vilab.github.io/largen-page/)             |
+|        Image Editing         |  [🌟StyleBooth](docs/en/tasks/stylebooth.md)   | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=StyleBooth&color=red&logo=arxiv)](https://arxiv.org/abs/2404.12154)   [![Page link](https://img.shields.io/badge/Page-StyleBooth-Gree)](https://ali-vilab.github.io/stylebooth-page/) |
+| Image Generation and Editing | [🌟ACE](https://ali-vilab.github.io/ace-page/) | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=ACE&color=red&logo=arxiv)](https://arxiv.org/abs/2410.00086)   [![Page link](https://img.shields.io/badge/Page-ACE-Gree)](https://ali-vilab.github.io/ace-page/)                      |


 ## 🖥️ SCEPTER Studio
@@ -145,6 +188,22 @@ Therefore, subsequent startups will become much faster (about one minute) as dow
 We deploy a work studio on Modelscope that includes only the inference tab, please refer to [ms_scepter_studio](https://www.modelscope.cn/studios/iic/scepter_studio/summary) and [hf_scepter_studio](https://huggingface.co/spaces/modelscope/scepter_studio)


+
+## ⚙️️ ComfyUI Workflow
+
+### Launch
+
+Manually install by moving custom_nodes to ComfyUI.
+```shell
+cd path/to/scepter
+pip install -e .
+cp -r path/to/scepter/workflow/ path/to/ComfyUI/custom_nodes/ComfyUI-Scepter
+cd path/to/ComfyUI
+python main.py
+```
+Alternatively, we will support ComfyUI Manager shortly.
+
+
 ## 🔍 Learn More

 - [Alibaba TongYi Vision Intelligence Lab](https://github.com/ali-vilab)
@@ -177,4 +236,4 @@ This project is licensed under the [Apache License (Version 2.0)](https://github


 ## Acknowledgement
-Thanks to [Stability-AI](https://github.com/Stability-AI), [SWIFT library](https://github.com/modelscope/swift/) and [Fooocus](https://github.com/lllyasviel/Fooocus) for their awesome work.
+Thanks to [Stability-AI](https://github.com/Stability-AI), [SWIFT library](https://github.com/modelscope/swift/), [Fooocus](https://github.com/lllyasviel/Fooocus) and [ComfyUI](https://github.com/comfyanonymous/ComfyUI) for their awesome work.
@@ -1,8 +1,9 @@
 albumentations
+beautifulsoup4
 bezier
 einops
 modelscope
-ms-swift>=2.0.1
+ms-swift
 numpy
 open_clip_torch
 opencv-python
@@ -13,3 +14,4 @@ pyyaml>=5.3.1
 scikit-image
 torchsde
 transformers
+scikit-learn
@@ -1,5 +1,5 @@
 bitsandbytes
-gradio>=3.47.1,<4.0.0
+gradio
 imagehash
 psutil
 tiktoken
@@ -20,7 +20,7 @@ SOLVER:
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -11,7 +11,7 @@ SOLVER:
  # NUM_FOLDS DESCRIPTION: Num folds for training. TYPE: int default: 0
  NUM_FOLDS: 1
  # WORK_DIR DESCRIPTION: Save dir of the training log or model. TYPE: str default: ''
-  WORK_DIR: ./exp12/
+  WORK_DIR: ./cache/save_data/example/
  LOG_FILE: std_log.txt
  # EVAL_INTERVAL DESCRIPTION: Eval the model interval. TYPE: int default: 1
  EVAL_INTERVAL: 1
@@ -102,7 +102,7 @@ SOLVER:
    # DATASET DESCRIPTION: the public dataset name TYPE: str default: 'cifar10'
    DATASET: cifar10
    # DATA_ROOT DESCRIPTION: the download data save path TYPE: str default: ''
-    DATA_ROOT: ./local_data/cifar10
+    DATA_ROOT: ./cache/cache_data/cifar10
    # MODE DESCRIPTION: test TYPE: str default: test
    MODE: test
    # PIN_MEMORY DESCRIPTION: pin_memory for data loader TYPE: bool default: False
@@ -0,0 +1,303 @@
+ENV:
+  BACKEND: nccl
+  SEED: 166666
+SOLVER:
+  # NAME DESCRIPTION:  TYPE:  default: 'LatentUfitSolver'
+  NAME: LatentDiffusionSolver
+  # MAX_STEPS DESCRIPTION: The total steps for training. TYPE: int default: 100000
+  MAX_STEPS: 100000
+  # USE_AMP DESCRIPTION: Use amp to surpport mix precision or not, default is False. TYPE: bool default: False
+  USE_AMP: True
+  # DTYPE DESCRIPTION: The precision for training. TYPE: str default: 'float32'
+  DTYPE: bfloat16
+  # USE_FAIRSCALE DESCRIPTION: Use fairscale as the backend of ddp, default False. TYPE: bool default: False
+  USE_FAIRSCALE: False
+  # USE_FSDP DESCRIPTION: Use fsdp as the backend of ddp, default False. TYPE: bool default: False
+  USE_FSDP: True
+  # LOAD_MODEL_ONLY DESCRIPTION: Only load the model rather than the optimizer and schedule, default is False. TYPE: bool default: False
+  LOAD_MODEL_ONLY: False
+  # RESUME_FROM DESCRIPTION: Resume from some state of training! TYPE: str default: ''
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_flux_dev_1024_lora
+  LOG_FILE: std_log.txt
+  # EVAL_INTERVAL DESCRIPTION: Eval the model interval. TYPE: int default: 1
+  EVAL_INTERVAL: 100
+  # LOG_TRAIN_NUM DESCRIPTION: The number samples used to log in training phase. TYPE: int default: -1
+  LOG_TRAIN_NUM: 16
+  # FSDP_REDUCE_DTYPE DESCRIPTION: The dtype of reduce in FSDP. TYPE: str default: 'float16'
+  FSDP_REDUCE_DTYPE: float32
+  # FSDP_BUFFER_DTYPE DESCRIPTION: The dtype of buffer in FSDP. TYPE: str default: 'float16'
+  FSDP_BUFFER_DTYPE: float32
+  # FSDP_SHARD_MODULES DESCRIPTION: The modules to be sharded in FSDP. TYPE: list default: ['model']
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.t5_model' ] #
+  SAVE_MODULES: [ 'model'] #
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 4
+      LORA_ALPHA: 4
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "(model.double_blocks.*(.qkv|.proj|.img_mod.lin|.txt_mod.lin))|(model.single_blocks.*(.linear1|.linear2|.modulation.lin))$"
+  ##
+  MODEL:
+    NAME: LatentDiffusionFlux
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      # NAME DESCRIPTION:  TYPE:  default: 'DiffusionFluxRF'
+      NAME: DiffusionFluxRF
+      PREDICTION_TYPE: raw
+      # NOISE_SCHEDULER DESCRIPTION:  TYPE:  default: ''
+      NOISE_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchSigmaScheduler'
+        NAME: FlowMatchSigmaScheduler
+        # WEIGHTING_SCHEME DESCRIPTION: The weighting scheme for sampling timesteps, choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']. TYPE: str default: 'logit_normal'
+        WEIGHTING_SCHEME: logit_normal
+        SHIFT: 3.0
+        # LOGIT_MEAN DESCRIPTION: The mean of the logit distribution for sampling timesteps. TYPE: float default: 0.0
+        LOGIT_MEAN: 0.0
+        # LOGIT_STD DESCRIPTION: The standard deviation of the logit distribution for sampling timesteps. TYPE: float default: 1.0
+        LOGIT_STD: 1.0
+        # MODE_SCALE DESCRIPTION: The scale factor for the mode of the logit distribution for sampling timesteps. TYPE: float default: 1.29
+        MODE_SCALE: 1.29
+      SAMPLER_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchFluxShiftScheduler'
+        NAME: FlowMatchFluxShiftScheduler
+        # SHIFT DESCRIPTION: Use timestamp shift or not, default is True. TYPE: bool default: True
+        SHIFT: False
+        # SIGMOID_SCALE DESCRIPTION: The scale of sigmoid function for sampling timesteps. TYPE: int default: 1
+        SIGMOID_SCALE: 1
+        # BASE_SHIFT DESCRIPTION: The base shift factor for the timestamp. TYPE: float default: 0.5
+        BASE_SHIFT: 0.5
+        # MAX_SHIFT DESCRIPTION: The max shift factor for the timestamp. TYPE: float default: 1.15
+        MAX_SHIFT: 1.15
+        #
+    DIFFUSION_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'Flux'
+      NAME: Flux
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@flux1-dev.safetensors
+      # IN_CHANNELS DESCRIPTION: model's input channels. TYPE: int default: 64
+      IN_CHANNELS: 64
+      # HIDDEN_SIZE DESCRIPTION: model's hidden size. TYPE: int default: 1024
+      HIDDEN_SIZE: 3072
+      # NUM_HEADS DESCRIPTION: number of heads in the transformer. TYPE: int default: 16
+      NUM_HEADS: 24
+      # AXES_DIM DESCRIPTION: dimensions of the axes of the positional encoding. TYPE: list default: [16, 56, 56]
+      AXES_DIM: [ 16, 56, 56 ]
+      # THETA DESCRIPTION: theta for positional encoding. TYPE: int default: 10000
+      THETA: 10000
+      # VEC_IN_DIM DESCRIPTION: dimension of the vector input. TYPE: int default: 768
+      VEC_IN_DIM: 768
+      # GUIDANCE_EMBED DESCRIPTION: whether to use guidance embedding. TYPE: bool default: False
+      GUIDANCE_EMBED: True
+      # CONTEXT_IN_DIM DESCRIPTION: dimension of the context input. TYPE: int default: 4096
+      CONTEXT_IN_DIM: 4096
+      # MLP_RATIO DESCRIPTION: ratio of mlp hidden size to hidden size. TYPE: float default: 4.0
+      MLP_RATIO: 4.0
+      # QKV_BIAS DESCRIPTION: whether to use bias in qkv projection. TYPE: bool default: True
+      QKV_BIAS: True
+      # DEPTH DESCRIPTION: number of transformer blocks. TYPE: int default: 19
+      DEPTH: 19
+      # DEPTH_SINGLE_BLOCKS DESCRIPTION: number of transformer blocks in the single stream block. TYPE: int default: 38
+      DEPTH_SINGLE_BLOCKS: 38
+      USE_GRAD_CHECKPOINT: True
+
+    #
+    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: True
+        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: True
+        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 DESCRIPTION:  TYPE:  default: 'T5PlusClipFluxEmbedder'
+      NAME: T5PlusClipFluxEmbedder
+      # T5_MODEL DESCRIPTION:  TYPE:  default: ''
+      T5_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: T5EncoderModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder_2/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: T5Tokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer_2/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 512
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: last_hidden_state
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: False
+        CLEAN: whitespace
+      # CLIP_MODEL DESCRIPTION:  TYPE:  default: ''
+      CLIP_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: CLIPTextModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: CLIPTokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 77
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: pooler_output
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: True
+        CLEAN: whitespace
+        USE_GRAD_CHECKPOINT: True
+  #
+  SAMPLE_ARGS:
+    SAMPLE_STEPS: 50
+    SAMPLER: flow_eluer
+    SEED: 2024
+    IMAGE_SIZE: [ 1024, 1024 ]
+    GUIDE_SCALE: 3.5
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 4e-4
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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: [ 1024, 1024 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 2
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+#      GRADIENT_CLIP: 1.0
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+    -
+      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
+      PRIORITY: 0
@@ -0,0 +1,302 @@
+ENV:
+  BACKEND: nccl
+  SEED: 166666
+SOLVER:
+  # NAME DESCRIPTION:  TYPE:  default: 'LatentUfitSolver'
+  NAME: LatentDiffusionSolver
+  # MAX_STEPS DESCRIPTION: The total steps for training. TYPE: int default: 100000
+  MAX_STEPS: 100000
+  # USE_AMP DESCRIPTION: Use amp to surpport mix precision or not, default is False. TYPE: bool default: False
+  USE_AMP: True
+  # DTYPE DESCRIPTION: The precision for training. TYPE: str default: 'float32'
+  DTYPE: bfloat16
+  # USE_FAIRSCALE DESCRIPTION: Use fairscale as the backend of ddp, default False. TYPE: bool default: False
+  USE_FAIRSCALE: False
+  # USE_FSDP DESCRIPTION: Use fsdp as the backend of ddp, default False. TYPE: bool default: False
+  USE_FSDP: True
+  # LOAD_MODEL_ONLY DESCRIPTION: Only load the model rather than the optimizer and schedule, default is False. TYPE: bool default: False
+  LOAD_MODEL_ONLY: False
+  # RESUME_FROM DESCRIPTION: Resume from some state of training! TYPE: str default: ''
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_flux_schnell_1024_lora
+  LOG_FILE: std_log.txt
+  # EVAL_INTERVAL DESCRIPTION: Eval the model interval. TYPE: int default: 1
+  EVAL_INTERVAL: 100
+  # LOG_TRAIN_NUM DESCRIPTION: The number samples used to log in training phase. TYPE: int default: -1
+  LOG_TRAIN_NUM: 16
+  # FSDP_REDUCE_DTYPE DESCRIPTION: The dtype of reduce in FSDP. TYPE: str default: 'float16'
+  FSDP_REDUCE_DTYPE: float32
+  # FSDP_BUFFER_DTYPE DESCRIPTION: The dtype of buffer in FSDP. TYPE: str default: 'float16'
+  FSDP_BUFFER_DTYPE: float32
+  # FSDP_SHARD_MODULES DESCRIPTION: The modules to be sharded in FSDP. TYPE: list default: ['model']
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.t5_model' ] #
+  SAVE_MODULES: [ 'model'] #
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 4
+      LORA_ALPHA: 4
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "(model.double_blocks.*(.qkv|.proj|.img_mod.lin|.txt_mod.lin))|(model.single_blocks.*(.linear1|.linear2|.modulation.lin))$"
+  ##
+  MODEL:
+    NAME: LatentDiffusionFlux
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      # NAME DESCRIPTION:  TYPE:  default: 'DiffusionFluxRF'
+      NAME: DiffusionFluxRF
+      PREDICTION_TYPE: raw
+      # NOISE_SCHEDULER DESCRIPTION:  TYPE:  default: ''
+      NOISE_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchSigmaScheduler'
+        NAME: FlowMatchSigmaScheduler
+        # WEIGHTING_SCHEME DESCRIPTION: The weighting scheme for sampling timesteps, choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']. TYPE: str default: 'logit_normal'
+        WEIGHTING_SCHEME: logit_normal
+        SHIFT: 3.0
+        # LOGIT_MEAN DESCRIPTION: The mean of the logit distribution for sampling timesteps. TYPE: float default: 0.0
+        LOGIT_MEAN: 0.0
+        # LOGIT_STD DESCRIPTION: The standard deviation of the logit distribution for sampling timesteps. TYPE: float default: 1.0
+        LOGIT_STD: 1.0
+        # MODE_SCALE DESCRIPTION: The scale factor for the mode of the logit distribution for sampling timesteps. TYPE: float default: 1.29
+        MODE_SCALE: 1.29
+      SAMPLER_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchFluxShiftScheduler'
+        NAME: FlowMatchFluxShiftScheduler
+        # SHIFT DESCRIPTION: Use timestamp shift or not, default is True. TYPE: bool default: True
+        SHIFT: False
+        # SIGMOID_SCALE DESCRIPTION: The scale of sigmoid function for sampling timesteps. TYPE: int default: 1
+        SIGMOID_SCALE: 1
+        # BASE_SHIFT DESCRIPTION: The base shift factor for the timestamp. TYPE: float default: 0.5
+        BASE_SHIFT: 0.5
+        # MAX_SHIFT DESCRIPTION: The max shift factor for the timestamp. TYPE: float default: 1.15
+        MAX_SHIFT: 1.15
+        #
+    DIFFUSION_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'Flux'
+      NAME: Flux
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-schnell@flux1-schnell.safetensors
+      # IN_CHANNELS DESCRIPTION: model's input channels. TYPE: int default: 64
+      IN_CHANNELS: 64
+      # HIDDEN_SIZE DESCRIPTION: model's hidden size. TYPE: int default: 1024
+      HIDDEN_SIZE: 3072
+      # NUM_HEADS DESCRIPTION: number of heads in the transformer. TYPE: int default: 16
+      NUM_HEADS: 24
+      # AXES_DIM DESCRIPTION: dimensions of the axes of the positional encoding. TYPE: list default: [16, 56, 56]
+      AXES_DIM: [ 16, 56, 56 ]
+      # THETA DESCRIPTION: theta for positional encoding. TYPE: int default: 10000
+      THETA: 10000
+      # VEC_IN_DIM DESCRIPTION: dimension of the vector input. TYPE: int default: 768
+      VEC_IN_DIM: 768
+      # GUIDANCE_EMBED DESCRIPTION: whether to use guidance embedding. TYPE: bool default: False
+      GUIDANCE_EMBED: False
+      # CONTEXT_IN_DIM DESCRIPTION: dimension of the context input. TYPE: int default: 4096
+      CONTEXT_IN_DIM: 4096
+      # MLP_RATIO DESCRIPTION: ratio of mlp hidden size to hidden size. TYPE: float default: 4.0
+      MLP_RATIO: 4.0
+      # QKV_BIAS DESCRIPTION: whether to use bias in qkv projection. TYPE: bool default: True
+      QKV_BIAS: True
+      # DEPTH DESCRIPTION: number of transformer blocks. TYPE: int default: 19
+      DEPTH: 19
+      # DEPTH_SINGLE_BLOCKS DESCRIPTION: number of transformer blocks in the single stream block. TYPE: int default: 38
+      DEPTH_SINGLE_BLOCKS: 38
+      USE_GRAD_CHECKPOINT: True
+
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLFlux
+      EMBED_DIM: 16
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-schnell@ae.safetensors
+      IGNORE_KEYS: [ ]
+      BATCH_SIZE: 8
+      USE_CONV: False
+      SCALE_FACTOR: 0.3611
+      SHIFT_FACTOR: 0.1159
+      #
+      ENCODER:
+        NAME: Encoder
+        USE_CHECKPOINT: True
+        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: True
+        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 DESCRIPTION:  TYPE:  default: 'T5PlusClipFluxEmbedder'
+      NAME: T5PlusClipFluxEmbedder
+      # T5_MODEL DESCRIPTION:  TYPE:  default: ''
+      T5_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: T5EncoderModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-schnell@text_encoder_2/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: T5Tokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-schnell@tokenizer_2/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 256
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: last_hidden_state
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: False
+        CLEAN: whitespace
+      # CLIP_MODEL DESCRIPTION:  TYPE:  default: ''
+      CLIP_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: CLIPTextModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-schnell@text_encoder/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: CLIPTokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-schnell@tokenizer/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 77
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: pooler_output
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: True
+        CLEAN: whitespace
+  #
+  SAMPLE_ARGS:
+    SAMPLE_STEPS: 4
+    SAMPLER: flow_eluer
+    SEED: 2024
+    IMAGE_SIZE: [ 1024, 1024 ]
+    GUIDE_SCALE: 3.5
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 4e-4
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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: [ 1024, 1024 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 2
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+#      GRADIENT_CLIP: 1.0
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+    -
+      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
+      PRIORITY: 0
@@ -0,0 +1,223 @@
+ENV:
+  BACKEND: nccl
+#
+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: 100
+  RESCALE_LR: False
+  #
+  WORK_DIR: ./cache/save_data/dit_pixart_alpha_1024_lora
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 128
+      LORA_ALPHA: 128
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.q|.k|.v|.o|mlp.fc1|mlp.fc2)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionPixart
+    PARAMETERIZATION: eps
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 0.18215
+    SIZE_FACTOR: 8
+    DECODER_BIAS: 0.5
+    DEFAULT_N_PROMPT:
+    SCHEDULE_ARGS:
+      "NAME": "linear"
+      "BETA_MIN": 0.0001
+      "BETA_MAX": 0.02
+    USE_EMA: False
+    LOAD_REFINER: False
+    #
+    DIFFUSION_MODEL:
+      NAME: PixArt
+      PRETRAINED_MODEL: ms://AI-ModelScope/PixArt-alpha@PixArt-XL-2-1024-MS.pth
+      INPUT_SIZE: 128
+      PATCH_SIZE: 2
+      IN_CHANNELS: 4
+      HIDDEN_SIZE: 1152
+      DEPTH: 28
+      NUM_HEADS: 16
+      MLP_RATIO: 4.0
+      CLASS_DROPOUT_PROB: 0.1
+      PRED_SIGMA: True
+      DROP_PATH: 0.0
+      WINDOW_DIZE: 0
+      USE_REL_POS: False
+      CAPTION_CHANNELS: 4096
+      LEWEI_SCALE: 2
+      MODEL_MAX_LENGTH: 120
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-base@512-base-ema.safetensors
+      EMBED_DIM: 4
+      IGNORE_KEYS: [ ]
+      BATCH_SIZE: 1
+      #
+      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: ms://AI-ModelScope/PixArt-alpha@t5-v1_1-xxl/
+      TOKENIZER_PATH: ms://AI-ModelScope/PixArt-alpha@t5-v1_1-xxl/
+      LENGTH: 120
+      CLEAN: heavy
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 20
+    SEED: 2024
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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: [ 1024, 1024 ]
+    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'
@@ -0,0 +1,233 @@
+ENV:
+  BACKEND: nccl
+#
+SOLVER:
+  NAME: LatentDiffusionSolver
+  RESUME_FROM:
+  LOAD_MODEL_ONLY: True
+  USE_FSDP: False
+  SHARDING_STRATEGY:
+  USE_AMP: True
+  DTYPE: float16
+  CHANNELS_LAST: True
+  MAX_STEPS: 500
+  MAX_EPOCHS: -1
+  NUM_FOLDS: 1
+  ACCU_STEP: 1
+  EVAL_INTERVAL: 50
+  RESCALE_LR: False
+  #
+  WORK_DIR: ./cache/save_data/dit_sd3_1024
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    - NAME: "ModelscopeFs"
+      TEMP_DIR: ./cache/cache_data
+  #
+  MODEL:
+    NAME: LatentDiffusionSD3
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 1.5305
+    SHIFT_FACTOR: 0.0609
+    DEFAULT_N_PROMPT:
+    SCHEDULE_ARGS:
+      "NAME": "shifted"
+      "SHIFT": 3
+    USE_EMA: False
+    T_WEIGHT: uniform
+    #
+    DIFFUSION_MODEL:
+      NAME: MMDiT
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+      IGNORE_KEYS: '^first_stage_model.'
+      IN_CHANNELS: 16
+      PATCH_SIZE: 2
+      OUT_CHANNELS: 16
+      DEPTH: 24
+      INPUT_SIZE:
+      ADM_IN_CHANNELS: 2048
+      CONTEXT_EMBEDDER_CONFIG: { 'target': 'torch.nn.Linear', 'params': { 'in_features': 4096, 'out_features': 1536 } }
+      NUM_PATCHES: 36864
+      POS_EMBED_MAX_SIZE: 192
+      POS_EMBED_SCALING_FACTOR:
+      USE_CHECKPOINT: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+      EMBED_DIM: 16
+      IGNORE_KEYS: '^model.diffusion_model.'
+      BATCH_SIZE: 1
+      USE_CONV: False
+      #
+      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: 16
+        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: 16
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+        GIVE_PRE_END: False
+        TANH_OUT: False
+    #
+    COND_STAGE_MODEL:
+      NAME: SD3TextEmbedder
+      P_ZERO: 0.0
+      CLIP_L:
+        NAME: FrozenCLIPEmbedder2
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer
+        MAX_LENGTH: 77
+        FREEZE: True
+        LAYER: penultimate
+        RETURN_POOLED: True
+        USE_FINAL_LAYER_NORM: False
+        IS_TRAINABLE: False
+      CLIP_G:
+        NAME: FrozenCLIPEmbedder2
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_2
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_2
+        MAX_LENGTH: 77
+        FREEZE: True
+        LAYER: penultimate
+        RETURN_POOLED: True
+        USE_FINAL_LAYER_NORM: False
+        IS_TRAINABLE: False
+      T5_XXL:
+        NAME: T5EmbedderHF
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_3
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_3
+        LENGTH: 256
+        CLEAN: whitespace
+        USE_GRAD: False
+        T5_DTYPE: float16
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: euler
+    SAMPLE_STEPS: 28
+    SEED: 1749023094
+    GUIDE_SCALE: 5.0
+    GUIDE_RESCALE: 0.0
+    DISCRETIZATION: trailing
+    RUN_TRAIN_N: False
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 1e-5
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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 cat holds a blackboard that writes \"hello world\"", "a dog running on the lawn" ]
+    IMAGE_SIZE: [ 1024, 1024 ]
+    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: 10000
+    -
+      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: 50
+      SAVE_LAST: True
+      SAVE_NAME_PREFIX: 'step'
+      SAVE_PROBE_PREFIX: 'image'
@@ -0,0 +1,241 @@
+ENV:
+  BACKEND: nccl
+#
+SOLVER:
+  NAME: LatentDiffusionSolver
+  RESUME_FROM:
+  LOAD_MODEL_ONLY: True
+  USE_FSDP: False
+  SHARDING_STRATEGY:
+  USE_AMP: True
+  DTYPE: float16
+  CHANNELS_LAST: True
+  MAX_STEPS: 500
+  MAX_EPOCHS: -1
+  NUM_FOLDS: 1
+  ACCU_STEP: 1
+  EVAL_INTERVAL: 50
+  RESCALE_LR: False
+  #
+  WORK_DIR: ./cache/save_data/dit_sd3_1024_lora
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    - NAME: "ModelscopeFs"
+      TEMP_DIR: ./cache/cache_data
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 128
+      LORA_ALPHA: 128
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.attn.qkv|.attn.proj|mlp.fc1|mlp.fc2)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionSD3
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 1.5305
+    SHIFT_FACTOR: 0.0609
+    DEFAULT_N_PROMPT:
+    SCHEDULE_ARGS:
+      "NAME": "shifted"
+      "SHIFT": 3
+    USE_EMA: False
+    T_WEIGHT: uniform
+    #
+    DIFFUSION_MODEL:
+      NAME: MMDiT
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+      IGNORE_KEYS: '^first_stage_model.'
+      IN_CHANNELS: 16
+      PATCH_SIZE: 2
+      OUT_CHANNELS: 16
+      DEPTH: 24
+      INPUT_SIZE:
+      ADM_IN_CHANNELS: 2048
+      CONTEXT_EMBEDDER_CONFIG: { 'target': 'torch.nn.Linear', 'params': { 'in_features': 4096, 'out_features': 1536 } }
+      NUM_PATCHES: 36864
+      POS_EMBED_MAX_SIZE: 192
+      POS_EMBED_SCALING_FACTOR:
+      USE_CHECKPOINT: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+      EMBED_DIM: 16
+      IGNORE_KEYS: '^model.diffusion_model.'
+      BATCH_SIZE: 1
+      USE_CONV: False
+      #
+      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: 16
+        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: 16
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+        GIVE_PRE_END: False
+        TANH_OUT: False
+    #
+    COND_STAGE_MODEL:
+      NAME: SD3TextEmbedder
+      P_ZERO: 0.0
+      CLIP_L:
+        NAME: FrozenCLIPEmbedder2
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer
+        MAX_LENGTH: 77
+        FREEZE: True
+        LAYER: penultimate
+        RETURN_POOLED: True
+        USE_FINAL_LAYER_NORM: False
+        IS_TRAINABLE: False
+      CLIP_G:
+        NAME: FrozenCLIPEmbedder2
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_2
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_2
+        MAX_LENGTH: 77
+        FREEZE: True
+        LAYER: penultimate
+        RETURN_POOLED: True
+        USE_FINAL_LAYER_NORM: False
+        IS_TRAINABLE: False
+      T5_XXL:
+        NAME: T5EmbedderHF
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_3
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_3
+        LENGTH: 256
+        CLEAN: whitespace
+        USE_GRAD: False
+        T5_DTYPE: float16
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: euler
+    SAMPLE_STEPS: 28
+    SEED: 1749023094
+    GUIDE_SCALE: 5.0
+    GUIDE_RESCALE: 0.0
+    DISCRETIZATION: trailing
+    RUN_TRAIN_N: False
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 5e-5
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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 cat holds a blackboard that writes \"hello world\"", "a dog running on the lawn" ]
+    IMAGE_SIZE: [ 1024, 1024 ]
+    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: 10000
+    -
+      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: 50
+      SAVE_LAST: True
+      SAVE_NAME_PREFIX: 'step'
+      SAVE_PROBE_PREFIX: 'image'
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_full
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  MODEL:
    NAME: LatentDiffusion
@@ -124,7 +125,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.0064
+    LEARNING_RATE: 0.00001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -190,7 +191,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_lora
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  TUNER:
    -
      NAME: SwiftLoRA
@@ -132,7 +133,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -198,7 +199,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_textlora
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  TUNER:
    -
      NAME: SwiftLoRA
@@ -140,7 +141,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -206,7 +207,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_512_full
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  MODEL:
    NAME: LatentDiffusion
@@ -120,7 +121,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.0064
+    LEARNING_RATE: 0.00001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -186,7 +187,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_512_lora
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -129,7 +130,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.0064
+    LEARNING_RATE: 0.00001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -195,7 +196,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_768_full
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  MODEL:
    NAME: LatentDiffusion
@@ -120,7 +121,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.0064
+    LEARNING_RATE: 0.00001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -186,7 +187,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_768_lora
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -129,7 +130,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -195,7 +196,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_full
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  MODEL:
    NAME: LatentDiffusionXL
@@ -238,7 +239,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.0064
+    LEARNING_RATE: 0.00001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -306,7 +307,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      - NAME: Select
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_lora
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -247,7 +248,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -315,7 +316,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      - NAME: Select
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_textlora
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -255,7 +256,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -323,7 +324,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      - NAME: Select
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_sce_ctr_hed
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -141,7 +142,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -167,13 +168,13 @@ SOLVER:
        RGB_ORDER: RGB
        BACKEND: pillow
      - NAME: Resize
-        SIZE: 768
+        SIZE: 512
        INTERPOLATION: bilinear
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
      - NAME: CenterCrop
-        SIZE: 768
+        SIZE: 512
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
@@ -215,13 +216,13 @@ SOLVER:
        RGB_ORDER: RGB
        BACKEND: pillow
      - NAME: Resize
-        SIZE: 768
+        SIZE: 512
        INTERPOLATION: bilinear
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
      - NAME: CenterCrop
-        SIZE: 768
+        SIZE: 512
        INPUT_KEY: [ 'img' ]
        OUTPUT_KEY: [ 'img' ]
        BACKEND: pillow
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_768_sce_ctr_canny
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -140,7 +141,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_768_sce_ctr_pose
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -140,7 +141,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_canny
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -254,7 +255,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_color
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -255,7 +256,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_color_datatxt
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -255,7 +256,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_depth
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -255,7 +256,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_sce_t2i
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -134,7 +135,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -200,7 +201,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_sce_t2i_swift
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -132,7 +133,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -198,7 +199,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd15_512_textsce_t2i_swift
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -140,7 +141,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -206,7 +207,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_768_sce_t2i
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -130,7 +131,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -196,7 +197,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sd21_768_sce_t2i_swift
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -128,7 +129,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -194,7 +195,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      -
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_t2i
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -247,7 +248,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -315,7 +316,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      - NAME: Select
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_t2i_datatxt
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
    FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
@@ -247,7 +248,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_sce_t2i_swift
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -245,7 +246,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -313,7 +314,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      - NAME: Select
@@ -14,13 +14,14 @@ SOLVER:
  NUM_FOLDS: 1
  ACCU_STEP: 1
  EVAL_INTERVAL: 100
+  RESCALE_LR: False
  #
  WORK_DIR: ./cache/save_data/sdxl_1024_textsce_t2i_swift
  LOG_FILE: std_log.txt
  #
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  TUNER:
    -
@@ -253,7 +254,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -321,7 +322,7 @@ SOLVER:
    NUM_WORKERS: 4
    FILE_SYSTEM:
      NAME: "ModelscopeFs"
-      TEMP_DIR: "./cache/data"
+      TEMP_DIR: "./cache/cache_data"
    #
    TRANSFORMS:
      - NAME: Select
@@ -0,0 +1,201 @@
+NAME: FLUX1.0_DEV
+IS_DEFAULT: False
+DEFAULT_PARAS:
+  PARAS:
+    RESOLUTIONS: [[1024, 1024]]
+  INPUT:
+    IMAGE:
+    ORIGINAL_SIZE_AS_TUPLE: [1024, 1024]
+    TARGET_SIZE_AS_TUPLE: [1024, 1024]
+    PROMPT: ""
+    NEGATIVE_PROMPT:
+      DEFAULT: ""
+      VISIBLE: False
+    PROMPT_PREFIX: ""
+    SAMPLE:
+      VALUES: ["flow_eluer"]
+      DEFAULT: "flow_eluer"
+    SAMPLE_STEPS: 50
+    GUIDE_SCALE: 3.5
+    GUIDE_RESCALE:
+      DEFAULT: 0.0
+      VISIBLE: False
+    DISCRETIZATION:
+      VALUES: []
+      DEFAULT:
+      VISIBLE: False
+  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:
+  NAME: LatentDiffusionFlux
+  PARAMETERIZATION: rf
+  TIMESTEPS: 1000
+  MIN_SNR_GAMMA:
+  ZERO_TERMINAL_SNR: False
+  PRETRAINED_MODEL:
+  IGNORE_KEYS: [ ]
+  DEFAULT_N_PROMPT:
+  USE_EMA: False
+  EVAL_EMA: False
+  DIFFUSION:
+    # NAME DESCRIPTION:  TYPE:  default: 'DiffusionFluxRF'
+    NAME: DiffusionFluxRF
+    PREDICTION_TYPE: raw
+    # NOISE_SCHEDULER DESCRIPTION:  TYPE:  default: ''
+    NOISE_SCHEDULER:
+      # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchSigmaScheduler'
+      NAME: FlowMatchSigmaScheduler
+      # WEIGHTING_SCHEME DESCRIPTION: The weighting scheme for sampling timesteps, choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']. TYPE: str default: 'logit_normal'
+      WEIGHTING_SCHEME: logit_normal
+      SHIFT: 3.0
+      # LOGIT_MEAN DESCRIPTION: The mean of the logit distribution for sampling timesteps. TYPE: float default: 0.0
+      LOGIT_MEAN: 0.0
+      # LOGIT_STD DESCRIPTION: The standard deviation of the logit distribution for sampling timesteps. TYPE: float default: 1.0
+      LOGIT_STD: 1.0
+      # MODE_SCALE DESCRIPTION: The scale factor for the mode of the logit distribution for sampling timesteps. TYPE: float default: 1.29
+      MODE_SCALE: 1.29
+      #
+  DIFFUSION_MODEL:
+    # NAME DESCRIPTION:  TYPE:  default: 'Flux'
+    NAME: Flux
+    PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@flux1-dev.safetensors
+    # IN_CHANNELS DESCRIPTION: model's input channels. TYPE: int default: 64
+    IN_CHANNELS: 64
+    # HIDDEN_SIZE DESCRIPTION: model's hidden size. TYPE: int default: 1024
+    HIDDEN_SIZE: 3072
+    # NUM_HEADS DESCRIPTION: number of heads in the transformer. TYPE: int default: 16
+    NUM_HEADS: 24
+    # AXES_DIM DESCRIPTION: dimensions of the axes of the positional encoding. TYPE: list default: [16, 56, 56]
+    AXES_DIM: [ 16, 56, 56 ]
+    # THETA DESCRIPTION: theta for positional encoding. TYPE: int default: 10000
+    THETA: 10000
+    # VEC_IN_DIM DESCRIPTION: dimension of the vector input. TYPE: int default: 768
+    VEC_IN_DIM: 768
+    # GUIDANCE_EMBED DESCRIPTION: whether to use guidance embedding. TYPE: bool default: False
+    GUIDANCE_EMBED: True
+    # CONTEXT_IN_DIM DESCRIPTION: dimension of the context input. TYPE: int default: 4096
+    CONTEXT_IN_DIM: 4096
+    # MLP_RATIO DESCRIPTION: ratio of mlp hidden size to hidden size. TYPE: float default: 4.0
+    MLP_RATIO: 4.0
+    # QKV_BIAS DESCRIPTION: whether to use bias in qkv projection. TYPE: bool default: True
+    QKV_BIAS: True
+    # DEPTH DESCRIPTION: number of transformer blocks. TYPE: int default: 19
+    DEPTH: 19
+    # DEPTH_SINGLE_BLOCKS DESCRIPTION: number of transformer blocks in the single stream block. TYPE: int default: 38
+    DEPTH_SINGLE_BLOCKS: 38
+
+  #
+  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: True
+      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: True
+      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 DESCRIPTION:  TYPE:  default: 'T5PlusClipFluxEmbedder'
+    NAME: T5PlusClipFluxEmbedder
+    # T5_MODEL DESCRIPTION:  TYPE:  default: ''
+    T5_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+      NAME: HFEmbedder
+      # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_MODEL_CLS: T5EncoderModel
+      # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+      MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder_2/
+      # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_TOKENIZER_CLS: T5Tokenizer
+      # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+      TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer_2/
+      # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+      MAX_LENGTH: 512
+      # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+      OUTPUT_KEY: last_hidden_state
+      # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+      D_TYPE: bfloat16
+      # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+      BATCH_INFER: False
+      CLEAN: whitespace
+    # CLIP_MODEL DESCRIPTION:  TYPE:  default: ''
+    CLIP_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+      NAME: HFEmbedder
+      # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_MODEL_CLS: CLIPTextModel
+      # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+      MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder/
+      # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_TOKENIZER_CLS: CLIPTokenizer
+      # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+      TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer/
+      # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+      MAX_LENGTH: 77
+      # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+      OUTPUT_KEY: pooler_output
+      # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+      D_TYPE: bfloat16
+      # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+      BATCH_INFER: True
+      CLEAN: whitespace
@@ -0,0 +1,212 @@
+NAME: FLUX1.0_SCHNELL
+IS_DEFAULT: False
+DEFAULT_PARAS:
+  PARAS:
+    RESOLUTIONS: [[1024, 1024]]
+  INPUT:
+    IMAGE:
+    ORIGINAL_SIZE_AS_TUPLE: [1024, 1024]
+    TARGET_SIZE_AS_TUPLE: [1024, 1024]
+    PROMPT: ""
+    NEGATIVE_PROMPT:
+      DEFAULT: ""
+      VISIBLE: False
+    PROMPT_PREFIX: ""
+    SAMPLE:
+      VALUES: ["flow_eluer"]
+      DEFAULT: "flow_eluer"
+    SAMPLE_STEPS: 4
+    GUIDE_SCALE: 3.5
+    GUIDE_RESCALE:
+      DEFAULT: 0.0
+      VISIBLE: False
+    DISCRETIZATION:
+      VALUES: []
+      DEFAULT:
+      VISIBLE: False
+  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:
+  NAME: LatentDiffusionFlux
+  PARAMETERIZATION: rf
+  TIMESTEPS: 1000
+  MIN_SNR_GAMMA:
+  ZERO_TERMINAL_SNR: False
+  PRETRAINED_MODEL:
+  IGNORE_KEYS: [ ]
+  DEFAULT_N_PROMPT:
+  USE_EMA: False
+  EVAL_EMA: False
+  DIFFUSION:
+    # NAME DESCRIPTION:  TYPE:  default: 'DiffusionFluxRF'
+    NAME: DiffusionFluxRF
+    PREDICTION_TYPE: raw
+    # NOISE_SCHEDULER DESCRIPTION:  TYPE:  default: ''
+    NOISE_SCHEDULER:
+      # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchSigmaScheduler'
+      NAME: FlowMatchSigmaScheduler
+      # WEIGHTING_SCHEME DESCRIPTION: The weighting scheme for sampling timesteps, choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']. TYPE: str default: 'logit_normal'
+      WEIGHTING_SCHEME: logit_normal
+      SHIFT: 3.0
+      # LOGIT_MEAN DESCRIPTION: The mean of the logit distribution for sampling timesteps. TYPE: float default: 0.0
+      LOGIT_MEAN: 0.0
+      # LOGIT_STD DESCRIPTION: The standard deviation of the logit distribution for sampling timesteps. TYPE: float default: 1.0
+      LOGIT_STD: 1.0
+      # MODE_SCALE DESCRIPTION: The scale factor for the mode of the logit distribution for sampling timesteps. TYPE: float default: 1.29
+      MODE_SCALE: 1.29
+    SAMPLER_SCHEDULER:
+      # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchFluxShiftScheduler'
+      NAME: FlowMatchFluxShiftScheduler
+      # SHIFT DESCRIPTION: Use timestamp shift or not, default is True. TYPE: bool default: True
+      SHIFT: False
+      # SIGMOID_SCALE DESCRIPTION: The scale of sigmoid function for sampling timesteps. TYPE: int default: 1
+      SIGMOID_SCALE: 1
+      # BASE_SHIFT DESCRIPTION: The base shift factor for the timestamp. TYPE: float default: 0.5
+      BASE_SHIFT: 0.5
+      # MAX_SHIFT DESCRIPTION: The max shift factor for the timestamp. TYPE: float default: 1.15
+      MAX_SHIFT: 1.15
+      #
+  DIFFUSION_MODEL:
+    # NAME DESCRIPTION:  TYPE:  default: 'Flux'
+    NAME: Flux
+    PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-schnell@flux1-schnell.safetensors
+    # IN_CHANNELS DESCRIPTION: model's input channels. TYPE: int default: 64
+    IN_CHANNELS: 64
+    # HIDDEN_SIZE DESCRIPTION: model's hidden size. TYPE: int default: 1024
+    HIDDEN_SIZE: 3072
+    # NUM_HEADS DESCRIPTION: number of heads in the transformer. TYPE: int default: 16
+    NUM_HEADS: 24
+    # AXES_DIM DESCRIPTION: dimensions of the axes of the positional encoding. TYPE: list default: [16, 56, 56]
+    AXES_DIM: [ 16, 56, 56 ]
+    # THETA DESCRIPTION: theta for positional encoding. TYPE: int default: 10000
+    THETA: 10000
+    # VEC_IN_DIM DESCRIPTION: dimension of the vector input. TYPE: int default: 768
+    VEC_IN_DIM: 768
+    # GUIDANCE_EMBED DESCRIPTION: whether to use guidance embedding. TYPE: bool default: False
+    GUIDANCE_EMBED: False
+    # CONTEXT_IN_DIM DESCRIPTION: dimension of the context input. TYPE: int default: 4096
+    CONTEXT_IN_DIM: 4096
+    # MLP_RATIO DESCRIPTION: ratio of mlp hidden size to hidden size. TYPE: float default: 4.0
+    MLP_RATIO: 4.0
+    # QKV_BIAS DESCRIPTION: whether to use bias in qkv projection. TYPE: bool default: True
+    QKV_BIAS: True
+    # DEPTH DESCRIPTION: number of transformer blocks. TYPE: int default: 19
+    DEPTH: 19
+    # DEPTH_SINGLE_BLOCKS DESCRIPTION: number of transformer blocks in the single stream block. TYPE: int default: 38
+    DEPTH_SINGLE_BLOCKS: 38
+
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKLFlux
+    EMBED_DIM: 16
+    PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-schnell@ae.safetensors
+    IGNORE_KEYS: [ ]
+    BATCH_SIZE: 8
+    USE_CONV: False
+    SCALE_FACTOR: 0.3611
+    SHIFT_FACTOR: 0.1159
+    #
+    ENCODER:
+      NAME: Encoder
+      USE_CHECKPOINT: True
+      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: True
+      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 DESCRIPTION:  TYPE:  default: 'T5PlusClipFluxEmbedder'
+    NAME: T5PlusClipFluxEmbedder
+    # T5_MODEL DESCRIPTION:  TYPE:  default: ''
+    T5_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+      NAME: HFEmbedder
+      # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_MODEL_CLS: T5EncoderModel
+      # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+      MODEL_PATH: ms://AI-ModelScope/FLUX.1-schnell@text_encoder_2/
+      # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_TOKENIZER_CLS: T5Tokenizer
+      # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+      TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-schnell@tokenizer_2/
+      # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+      MAX_LENGTH: 256
+      # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+      OUTPUT_KEY: last_hidden_state
+      # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+      D_TYPE: bfloat16
+      # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+      BATCH_INFER: False
+      CLEAN: whitespace
+    # CLIP_MODEL DESCRIPTION:  TYPE:  default: ''
+    CLIP_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+      NAME: HFEmbedder
+      # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_MODEL_CLS: CLIPTextModel
+      # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+      MODEL_PATH: ms://AI-ModelScope/FLUX.1-schnell@text_encoder/
+      # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+      HF_TOKENIZER_CLS: CLIPTokenizer
+      # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+      TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-schnell@tokenizer/
+      # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+      MAX_LENGTH: 77
+      # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+      OUTPUT_KEY: pooler_output
+      # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+      D_TYPE: bfloat16
+      # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+      BATCH_INFER: True
+      CLEAN: whitespace
@@ -0,0 +1,120 @@
+NAME: PIXART_ALPHA
+IS_DEFAULT: False
+DEFAULT_PARAS:
+  PARAS:
+    RESOLUTIONS: [[1024, 1024]]
+  INPUT:
+    IMAGE:
+    ORIGINAL_SIZE_AS_TUPLE: [1024, 1024]
+    TARGET_SIZE_AS_TUPLE: [1024, 1024]
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    PROMPT_PREFIX: ""
+    SAMPLE: ddim
+    SAMPLE_STEPS: 20
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+    DISCRETIZATION: trailing
+  OUTPUT:
+    LATENT:
+    IMAGES:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        -
+          NAME: encode
+          DTYPE: float32
+          INPUT: ["IMAGE"]
+        -
+          NAME: decode
+          DTYPE: float32
+          INPUT: ["LATENT"]
+      PARAS:
+        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
+          DTYPE: float32
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  PRETRAINED_MODEL:
+  DECODER_BIAS: 0.5
+  SCHEDULE:
+    PARAMETERIZATION: "eps"
+    TIMESTEPS: 1000
+    ZERO_TERMINAL_SNR: False
+    SCHEDULE_ARGS:
+      "NAME": "linear"
+      "BETA_MIN": 0.0001
+      "BETA_MAX": 0.02
+  #
+  DIFFUSION_MODEL:
+    NAME: PixArt
+    PRETRAINED_MODEL: ms://AI-ModelScope/PixArt-alpha@PixArt-XL-2-1024-MS.pth
+    INPUT_SIZE: 128
+    PATCH_SIZE: 2
+    IN_CHANNELS: 4
+    HIDDEN_SIZE: 1152
+    DEPTH: 28
+    NUM_HEADS: 16
+    MLP_RATIO: 4.0
+    CLASS_DROPOUT_PROB: 0.1
+    PRED_SIGMA: True
+    DROP_PATH: 0.0
+    WINDOW_DIZE: 0
+    USE_REL_POS: False
+    CAPTION_CHANNELS: 4096
+    LEWEI_SCALE: 2
+    MODEL_MAX_LENGTH: 120
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKL
+    PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-base@512-base-ema.safetensors
+    EMBED_DIM: 4
+    IGNORE_KEYS: [ ]
+    BATCH_SIZE: 1
+    #
+    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: ms://AI-ModelScope/PixArt-alpha@t5-v1_1-xxl/
+    TOKENIZER_PATH: ms://AI-ModelScope/PixArt-alpha@t5-v1_1-xxl/
+    LENGTH: 120
+    CLEAN: heavy
+    USE_GRAD: False
@@ -0,0 +1,148 @@
+NAME: SD3
+IS_DEFAULT: False
+DEFAULT_PARAS:
+  PARAS:
+    RESOLUTIONS: [[1024, 1024]]
+  INPUT:
+    IMAGE:
+    ORIGINAL_SIZE_AS_TUPLE: [1024, 1024]
+    TARGET_SIZE_AS_TUPLE: [1024, 1024]
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    PROMPT_PREFIX: ""
+    SAMPLE: euler
+    SAMPLE_STEPS: 28
+    GUIDE_SCALE: 5.0
+    GUIDE_RESCALE: 0.0
+    DISCRETIZATION: trailing
+  OUTPUT:
+    LATENT:
+    IMAGES:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        -
+          NAME: encode
+          DTYPE: float32
+          INPUT: ["IMAGE"]
+        -
+          NAME: decode
+          DTYPE: float32
+          INPUT: ["LATENT"]
+      PARAS:
+        SCALE_FACTOR: 1.5305
+        SHIFT_FACTOR: 0.0609
+        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
+          DTYPE: float32
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  PRETRAINED_MODEL:
+  SCHEDULE:
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 1.5305
+    SHIFT_FACTOR: 0.0609
+    DEFAULT_N_PROMPT:
+    SCHEDULE_ARGS:
+      "NAME": "shifted"
+      "SHIFT": 3
+    T_WEIGHT: uniform
+  #
+  DIFFUSION_MODEL:
+    NAME: MMDiT
+    PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+    IGNORE_KEYS: '^first_stage_model.'
+    IN_CHANNELS: 16
+    PATCH_SIZE: 2
+    OUT_CHANNELS: 16
+    DEPTH: 24
+    INPUT_SIZE:
+    ADM_IN_CHANNELS: 2048
+    CONTEXT_EMBEDDER_CONFIG: { 'target': 'torch.nn.Linear', 'params': { 'in_features': 4096, 'out_features': 1536 } }
+    NUM_PATCHES: 36864
+    POS_EMBED_MAX_SIZE: 192
+    POS_EMBED_SCALING_FACTOR:
+    USE_CHECKPOINT: True
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKL
+    PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+    EMBED_DIM: 16
+    IGNORE_KEYS: '^model.diffusion_model.'
+    BATCH_SIZE: 1
+    USE_CONV: False
+    #
+    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: 16
+      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: 16
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+      GIVE_PRE_END: False
+      TANH_OUT: False
+  #
+  COND_STAGE_MODEL:
+    NAME: SD3TextEmbedder
+    P_ZERO: 0.0
+    CLIP_L:
+      NAME: FrozenCLIPEmbedder2
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder
+      TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer
+      MAX_LENGTH: 77
+      FREEZE: True
+      LAYER: penultimate
+      RETURN_POOLED: True
+      USE_FINAL_LAYER_NORM: False
+      IS_TRAINABLE: False
+    CLIP_G:
+      NAME: FrozenCLIPEmbedder2
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_2
+      TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_2
+      MAX_LENGTH: 77
+      FREEZE: True
+      LAYER: penultimate
+      RETURN_POOLED: True
+      USE_FINAL_LAYER_NORM: False
+      IS_TRAINABLE: False
+    T5_XXL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_3
+      TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_3
+      LENGTH: 256
+      CLEAN: whitespace
+      USE_GRAD: False
+      T5_DTYPE: float16
@@ -6,65 +6,81 @@ DIFFUSION_PARAS:
             'dpm2_karras', 'dpm2_ancestral_karras', 'dpmpp_2s_ancestral_karras', 'dpmpp_2m_karras',
             'dpmpp_sde_karras', 'dpmpp_2m_sde_karras']
    DEFAULT: 'dpmpp_2s_ancestral'
+    VISIBLE: True
  NEGATIVE_PROMPT:
    DEFAULT:
+    VISIBLE: True
  PROMPT_PREFIX:
    DEFAULT:
+    VISIBLE: True
  SAMPLES:
    MIN: 1
    MAX: 4
    DEFAULT: 1
+    VISIBLE: True
  SAMPLE_STEPS:
    MIN: 1
    MAX: 100
    DEFAULT: 30
+    VISIBLE: True
  GUIDE_SCALE:
    MIN: 0
    MAX: 10
    DEFAULT: 5.0
+    VISIBLE: True
  GUIDE_RESCALE:
    MIN: 0
    MAX: 1.0
    DEFAULT: 0.5
+    VISIBLE: True
  DISCRETIZATION:
    VALUES: ["trailing", "leading", "linspace"]
    DEFAULT: "linspace"
+    VISIBLE: True
  REFINE_SAMPLERS:
    VALUES: [ 'ddim', 'euler', 'euler_ancestral', 'heun', 'dpm2',
              'dpm2_ancestral', 'dpmpp_2m', 'dpmpp_sde', 'dpmpp_2m_sde', 'dpmpp_2s_ancestral',
              'dpm2_karras', 'dpm2_ancestral_karras', 'dpmpp_2s_ancestral_karras', 'dpmpp_2m_karras',
              'dpmpp_sde_karras', 'dpmpp_2m_sde_karras' ]
    DEFAULT: 'dpmpp_2s_ancestral'
+    VISIBLE: True
  REFINE_SAMPLE_STEPS:
    MIN: 0
    MAX: 100
    DEFAULT: 30
+    VISIBLE: True
  REFINE_GUIDE_SCALE:
    MIN: 0
    MAX: 10
    DEFAULT: 5.0
+    VISIBLE: True
  REFINE_GUIDE_RESCALE:
    MIN: 0
    MAX: 1.0
    DEFAULT: 0.5
+    VISIBLE: True
  REFINE_DISCRETIZATION:
    VALUES: [ "trailing", "leading", "linspace" ]
    DEFAULT: "linspace"
+    VISIBLE: True
  AESTHETIC_SCORE:
    MIN: 0.0
    MAX: 10.0
    DEFAULT: 6.0
+    VISIBLE: True
  NEGATIVE_AESTHETIC_SCORE:
    MIN: 0.0
    MAX: 10.0
    DEFAULT: 2.5
+    VISIBLE: True
  REFINE_STRENGTH:
    MIN: 0
    MAX: 1.0
    DEFAULT: 0.15
  RESOLUTIONS:
-    VALUES: [ [512, 512], [768, 768],
-              [704, 1408], [704, 1344], [768, 1344],
+    VALUES: [
+             [512, 512], [768, 768],
+             [704, 1408], [704, 1344], [768, 1344],
             [720, 1280],
             [768, 1280], [832, 1216], [832, 1152],
             [896, 1152], [896, 1088], [960, 1088],
@@ -74,8 +90,13 @@ DIFFUSION_PARAS:
             [1280, 768],
             [1344, 768], [1344, 704], [1408, 704],
             [1472, 704], [1536, 640], [1600, 640],
-             [1664, 576], [1728, 576]]
+             [1664, 576], [1728, 576],
+             [2048, 2048], [2048, 1920], [1920, 2048],
+             [1536, 2560], [2560, 1536], [2560, 1440],
+             [2560, 1440]
+            ]
    DEFAULT: [1024, 1024]
+    VISIBLE: True
 EXTENSION_PARAS:
  MANTRA_BOOK: scepter/methods/studio/extensions/mantra_book/mantra_book.yaml
  OFFICIAL_TUNERS: scepter/methods/studio/extensions/tuners/official_tuners.yaml
@@ -2,11 +2,10 @@ WORK_DIR: datasets
 EXPORT_DIR: export_datasets
 FILE_SYSTEM:
  -
-    # NAME DESCRIPTION:  TYPE:  default: ''
    NAME: LocalFs
    AUTO_CLEAN: False
-
 PROCESSORS:
+  # Caption processor
  - NAME: BlipImageBase
    TYPE: caption
    MODEL_PATH: ms://cubeai/blip-image-captioning-base
@@ -15,6 +14,81 @@ PROCESSORS:
    PARAS:
      - LANGUAGE_NAME: English
        LANGUAGE_ZH_NAME: 英语
+  - NAME: InternVL15
+    TYPE: caption
+    MODEL_PATH: ms://AI-ModelScope/InternVL-Chat-V1-5
+    DEVICE: "gpu"
+    MEMORY: 49968
+    PARAS:
+      - PROMPT: 用中文描述这张图片
+        LANGUAGE_NAME: Chinese
+        LANGUAGE_ZH_NAME: 中文
+      - PROMPT: Generate the caption in English
+        LANGUAGE_NAME: English
+        LANGUAGE_ZH_NAME: 英语
+  - NAME: QWVLQuantize
+    TYPE: caption
+    DEVICE: "gpu"
+    MEMORY: 7885
+    MODEL_PATH: ms://qwen/Qwen-VL:v1.0.3
+    PARAS:
+      - PROMPT: 用中文描述这张图片
+        LANGUAGE_NAME: Chinese
+        LANGUAGE_ZH_NAME: 中文
+        MAX_NEW_TOKENS:
+          VALUE: 1024
+          MAX: 2048
+          STEP: 128
+          MIN: 256
+        MIN_NEW_TOKENS:
+          VALUE: 16
+          MAX: 1024
+          STEP: 16
+          MIN: 0
+        NUM_BEAMS:
+          VALUE: 1
+          MAX: 12
+          STEP: 1
+          MIN: 1
+        REPETITION_PENALTY:
+          VALUE: 1.0
+          MAX: 100.0
+          STEP: 1.0
+          MIN: 1.0
+        TEMPERATURE:
+          VALUE: 1.0
+          MAX: 100.0
+          STEP: 1.0
+          MIN: 1.0
+      - PROMPT: Generate the caption in English
+        LANGUAGE_NAME: English
+        LANGUAGE_ZH_NAME: 英语
+        MAX_NEW_TOKENS:
+          VALUE: 1024
+          MAX: 2048
+          STEP: 128
+          MIN: 256
+        MIN_NEW_TOKENS:
+          VALUE: 16
+          MAX: 1024
+          STEP: 16
+          MIN: 0
+        NUM_BEAMS:
+          VALUE: 1
+          MAX: 12
+          STEP: 1
+          MIN: 1
+        REPETITION_PENALTY:
+          VALUE: 1.0
+          MAX: 100.0
+          STEP: 1.0
+          MIN: 1.0
+        TEMPERATURE:
+          VALUE: 1.0
+          MAX: 100.0
+          STEP: 1.0
+          MIN: 1.0
+
  - NAME: QWVL
    TYPE: caption
    MODEL_PATH: ms://qwen/Qwen-VL:v1.0.3
@@ -77,75 +151,14 @@ PROCESSORS:
          MAX: 100.0
          STEP: 1.0
          MIN: 1.0
-  -
-    NAME: QWVLQuantize
-    TYPE: caption
-    DEVICE: "gpu"
-    MEMORY: 7885
-    MODEL_PATH: ms://qwen/Qwen-VL:v1.0.3
-    PARAS:
-      - PROMPT: 用中文描述这张图片
-        LANGUAGE_NAME: Chinese
-        LANGUAGE_ZH_NAME: 中文
-        MAX_NEW_TOKENS:
-          VALUE: 1024
-          MAX: 2048
-          STEP: 128
-          MIN: 256
-        MIN_NEW_TOKENS:
-          VALUE: 16
-          MAX: 1024
-          STEP: 16
-          MIN: 0
-        NUM_BEAMS:
-          VALUE: 1
-          MAX: 12
-          STEP: 1
-          MIN: 1
-        REPETITION_PENALTY:
-          VALUE: 1.0
-          MAX: 100.0
-          STEP: 1.0
-          MIN: 1.0
-        TEMPERATURE:
-          VALUE: 1.0
-          MAX: 100.0
-          STEP: 1.0
-          MIN: 1.0
-      - PROMPT: Generate the caption in English
-        LANGUAGE_NAME: English
-        LANGUAGE_ZH_NAME: 英语
-        MAX_NEW_TOKENS:
-          VALUE: 1024
-          MAX: 2048
-          STEP: 128
-          MIN: 256
-        MIN_NEW_TOKENS:
-          VALUE: 16
-          MAX: 1024
-          STEP: 16
-          MIN: 0
-        NUM_BEAMS:
-          VALUE: 1
-          MAX: 12
-          STEP: 1
-          MIN: 1
-        REPETITION_PENALTY:
-          VALUE: 1.0
-          MAX: 100.0
-          STEP: 1.0
-          MIN: 1.0
-        TEMPERATURE:
-          VALUE: 1.0
-          MAX: 100.0
-          STEP: 1.0
-          MIN: 1.0
-  -
-    NAME: CenterCrop
+
+  # Simple processor
+  - NAME: CenterCrop
    TYPE: image
    DEVICE: "cpu"
    MEMORY: 10
    PARAS:
+      CAPTION_INTERACTIVE: False
      HEIGHT_RATIO:
        VALUE: 1
        MAX: 20
@@ -156,18 +169,284 @@ PROCESSORS:
        MAX: 20
        STEP: 1
        MIN: 1
-#  - NAME: PaddingCrop
-#    TYPE: image
-#    DEVICE: "cpu"
-#    MEMORY: 10
-#    PARAS:
-#      HEIGHT_RATIO:
-#        VALUE: 3
-#        MAX: 25
-#        STEP: 1
-#        MIN: 1
-#      WIDTH_RATIO:
-#        VALUE: 4
-#        MAX: 20
-#        STEP: 1
-#        MIN: 1
+  - NAME: ChangeSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      CAPTION_INTERACTIVE: False
+      SRC_IMAGE_INTERACTIVE: True
+      SRC_IMAGE_MASK_INTERACTIVE: True
+      TARGET_IMAGE_INTERACTIVE: True
+      PREVIEW_BTN_VISIBLE: False
+  - NAME: MaskEditSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      SRC_IMAGE_INTERACTIVE: True
+      SRC_IMAGE_TOOL: sketch
+      CAPTION_INTERACTIVE: False
+  - NAME: SourceMaskSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      SRC_IMAGE_INTERACTIVE: True
+      SRC_IMAGE_TOOL: sketch
+      CAPTION_INTERACTIVE: False
+  - NAME: MaskSwapEditSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      PREVIEW_BTN_VISIBLE: False
+      SRC_IMAGE_INTERACTIVE: True
+      SRC_IMAGE_TOOL: sketch
+      CAPTION_INTERACTIVE: False
+  - NAME: SwapMaskSwapEditSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      PREVIEW_BTN_VISIBLE: False
+      SRC_IMAGE_INTERACTIVE: True
+      SRC_IMAGE_TOOL: sketch
+      CAPTION_INTERACTIVE: False
+  - NAME: SwapSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      PREVIEW_BTN_VISIBLE: True
+      SRC_IMAGE_INTERACTIVE: True
+      TARGET_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: DefaultMaskSample
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      PREVIEW_BTN_VISIBLE: True
+      SRC_IMAGE_INTERACTIVE: False
+      TARGET_IMAGE_INTERACTIVE: False
+      CAPTION_INTERACTIVE: False
+  #  - NAME: PaddingCrop
+  #    TYPE: image
+  #    DEVICE: "cpu"
+  #    MEMORY: 10
+  #    PARAS:
+  #      HEIGHT_RATIO:
+  #        VALUE: 3
+  #        MAX: 25
+  #        STEP: 1
+  #        MIN: 1
+  #      WIDTH_RATIO:
+  #        VALUE: 4
+  #        MAX: 20
+  #        STEP: 1
+  #        MIN: 1
+  # Annotator processor
+  - NAME: CannyExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "CannyAnnotator"
+      LOW_THRESHOLD: 100
+      HIGH_THRESHOLD: 200
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: ColorExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "ColorAnnotator"
+      RATIO: 64
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: InfoDrawContourExtractor
+    TYPE: image
+    DEVICE: "gpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "InfoDrawContourAnnotator"
+      INPUT_NC: 3
+      OUTPUT_NC: 1
+      N_RESIDUAL_BLOCKS: 3
+      SIGMOID: True
+      PRETRAINED_MODEL: "ms://iic/scepter_annotator@annotator/ckpts/informative_drawing_contour_style.pth"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: DegradationExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "DegradationAnnotator"
+      RANDOM_DEGRADATION: True
+      PARAMS:
+        gaussian_noise: { }
+        resize: { 'scale': [ 0.4, 0.8 ] }
+        jpeg: { 'jpeg_level': [ 25, 75 ] }
+        gaussian_blur: { 'kernel_size': [ 7, 9, 11, 13, 15 ], 'sigma': [ 0.9, 1.8 ] }
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: MidasExtractor
+    TYPE: image
+    DEVICE: "gpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "MidasDetector"
+      PRETRAINED_MODEL: "ms://iic/scepter_annotator@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: DoodleExtractor
+    TYPE: image
+    DEVICE: "gpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "DoodleAnnotator"
+      PROCESSOR_TYPE: "pidinet_sketch"
+      PROCESSOR_CFG:
+        - NAME: "PiDiAnnotator"
+          PRETRAINED_MODEL: "ms://iic/scepter_annotator@annotator/ckpts/table5_pidinet.pth"
+        - NAME: "SketchAnnotator"
+          PRETRAINED_MODEL: "ms://iic/scepter_annotator@annotator/ckpts/sketch_simplification_gan.pth"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: GrayExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "GrayAnnotator"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: InpaintingExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "InpaintingAnnotator"
+      RETURN_MASK: False
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+
+  - NAME: InpaintingSourceExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: OpenposeExtractor
+    TYPE: image
+    DEVICE: "gpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "OpenposeAnnotator"
+      BODY_MODEL_PATH: "ms://iic/scepter_annotator@annotator/ckpts/body_pose_model.pth"
+      HAND_MODEL_PATH: "ms://iic/scepter_annotator@annotator/ckpts/hand_pose_model.pth"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: OutpaintingExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "OutpaintingAnnotator"
+      RETURN_MASK: False
+      KEEP_PADDING_RATIO: 1
+      RANDOM_CFG:
+        DIRECTION_RANGE: [ 'left', 'right', 'up', 'down' ]
+        RATIO_RANGE: [ 0.1, 0.7 ]
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+      USE_MASK_VISIBLE: True
+      PREVIEW_BTN_VISIBLE: False
+  - NAME: OutpaintingResize
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "OutpaintingResize"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+      PREVIEW_BTN_VISIBLE: False
+  - NAME: InfoDrawAnimeAnnotator
+    TYPE: image
+    DEVICE: "gpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "InfoDrawAnimeAnnotator"
+      INPUT_NC: 3
+      OUTPUT_NC: 1
+      N_RESIDUAL_BLOCKS: 3
+      SIGMOID: True
+      PRETRAINED_MODEL: "ms://iic/scepter_annotator@annotator/ckpts/informative_drawing_anime_style.pth"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: ESAMExtractor
+    TYPE: image
+    DEVICE: "gpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "ESAMAnnotator"
+      PRETRAINED_MODEL: "ms://iic/scepter_annotator@annotator/ckpts/efficient_sam_vits.pt"
+      SAVE_MODE: 'P'
+      GRID_SIZE: 32
+      USE_DOMINANT_COLOR: True
+      RETURN_MASK: False
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+  - NAME: InvertExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "InvertAnnotator"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
+
+  - NAME: LamaExtractor
+    TYPE: image
+    DEVICE: "cpu"
+    MEMORY: 10
+    MODEL:
+      NAME: "LamaAnnotator"
+      PRETRAINED_MODEL: "ms:///iic/cv_fft_inpainting_lama/"
+    PARAS:
+      SRC_IMAGE_TOOL: sketch
+      SRC_IMAGE_INTERACTIVE: True
+      CAPTION_INTERACTIVE: False
@@ -0,0 +1,343 @@
+ENV:
+  BACKEND: nccl
+META:
+  VERSION: 'FLUX1.0_DEV'
+  DESCRIPTION: "flux 1.0 dev"
+  IS_DEFAULT: False
+  IS_SHARE: True
+  INFERENCE_PARAS:
+    INFERENCE_BATCH_SIZE: 1
+    INFERENCE_PREFIX: ""
+    DEFAULT_SAMPLER: "flow_euler"
+    DEFAULT_SAMPLE_STEPS: 50
+    INFERENCE_N_PROMPT: ""
+    RESOLUTION: [1024, 1024]
+  PARAS:
+    -
+      TRAIN_BATCH_SIZE: 1
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      MEMORY: 89000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 4e-4
+      IS_DEFAULT: False
+      TUNER: FULL
+    -
+      TRAIN_BATCH_SIZE: 1
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      MEMORY: 89000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 4e-4
+      IS_DEFAULT: True
+      TUNER: LORA
+  #
+  TUNERS:
+    LORA:
+      -
+        NAME: SwiftLoRA
+        R: 4
+        LORA_ALPHA: 4
+        LORA_DROPOUT: 0.0
+        BIAS: "none"
+        TARGET_MODULES: "(model.double_blocks.*(.qkv|.proj|.img_mod.lin|.txt_mod.lin))|(model.single_blocks.*(.linear1|.linear2|.modulation.lin))$"
+#
+SOLVER:
+  NAME: LatentDiffusionSolver
+  # MAX_STEPS DESCRIPTION: The total steps for training. TYPE: int default: 100000
+  MAX_STEPS: 100000
+  # USE_AMP DESCRIPTION: Use amp to surpport mix precision or not, default is False. TYPE: bool default: False
+  USE_AMP: True
+  # DTYPE DESCRIPTION: The precision for training. TYPE: str default: 'float32'
+  DTYPE: bfloat16
+  # USE_FAIRSCALE DESCRIPTION: Use fairscale as the backend of ddp, default False. TYPE: bool default: False
+  USE_FAIRSCALE: False
+  # USE_FSDP DESCRIPTION: Use fsdp as the backend of ddp, default False. TYPE: bool default: False
+  USE_FSDP: True
+  # LOAD_MODEL_ONLY DESCRIPTION: Only load the model rather than the optimizer and schedule, default is False. TYPE: bool default: False
+  LOAD_MODEL_ONLY: False
+  # RESUME_FROM DESCRIPTION: Resume from some state of training! TYPE: str default: ''
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_flux_dev_1024_lora
+  LOG_FILE: std_log.txt
+  # EVAL_INTERVAL DESCRIPTION: Eval the model interval. TYPE: int default: 1
+  EVAL_INTERVAL: 100
+  # LOG_TRAIN_NUM DESCRIPTION: The number samples used to log in training phase. TYPE: int default: -1
+  LOG_TRAIN_NUM: 16
+  # FSDP_REDUCE_DTYPE DESCRIPTION: The dtype of reduce in FSDP. TYPE: str default: 'float16'
+  FSDP_REDUCE_DTYPE: float32
+  # FSDP_BUFFER_DTYPE DESCRIPTION: The dtype of buffer in FSDP. TYPE: str default: 'float16'
+  FSDP_BUFFER_DTYPE: float32
+  # FSDP_SHARD_MODULES DESCRIPTION: The modules to be sharded in FSDP. TYPE: list default: ['model']
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.t5_model' ] #
+  SAVE_MODULES: [ 'model'] #
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  #
+  TUNER:
+  #
+  MODEL:
+    NAME: LatentDiffusionFlux
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      # NAME DESCRIPTION:  TYPE:  default: 'DiffusionFluxRF'
+      NAME: DiffusionFluxRF
+      PREDICTION_TYPE: raw
+      # NOISE_SCHEDULER DESCRIPTION:  TYPE:  default: ''
+      NOISE_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchSigmaScheduler'
+        NAME: FlowMatchSigmaScheduler
+        # WEIGHTING_SCHEME DESCRIPTION: The weighting scheme for sampling timesteps, choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']. TYPE: str default: 'logit_normal'
+        WEIGHTING_SCHEME: logit_normal
+        SHIFT: 3.0
+        # LOGIT_MEAN DESCRIPTION: The mean of the logit distribution for sampling timesteps. TYPE: float default: 0.0
+        LOGIT_MEAN: 0.0
+        # LOGIT_STD DESCRIPTION: The standard deviation of the logit distribution for sampling timesteps. TYPE: float default: 1.0
+        LOGIT_STD: 1.0
+        # MODE_SCALE DESCRIPTION: The scale factor for the mode of the logit distribution for sampling timesteps. TYPE: float default: 1.29
+        MODE_SCALE: 1.29
+      SAMPLER_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchFluxShiftScheduler'
+        NAME: FlowMatchFluxShiftScheduler
+        # SHIFT DESCRIPTION: Use timestamp shift or not, default is True. TYPE: bool default: True
+        SHIFT: False
+        # SIGMOID_SCALE DESCRIPTION: The scale of sigmoid function for sampling timesteps. TYPE: int default: 1
+        SIGMOID_SCALE: 1
+        # BASE_SHIFT DESCRIPTION: The base shift factor for the timestamp. TYPE: float default: 0.5
+        BASE_SHIFT: 0.5
+        # MAX_SHIFT DESCRIPTION: The max shift factor for the timestamp. TYPE: float default: 1.15
+        MAX_SHIFT: 1.15
+        #
+    DIFFUSION_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'Flux'
+      NAME: Flux
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@flux1-dev.safetensors
+      # IN_CHANNELS DESCRIPTION: model's input channels. TYPE: int default: 64
+      IN_CHANNELS: 64
+      # HIDDEN_SIZE DESCRIPTION: model's hidden size. TYPE: int default: 1024
+      HIDDEN_SIZE: 3072
+      # NUM_HEADS DESCRIPTION: number of heads in the transformer. TYPE: int default: 16
+      NUM_HEADS: 24
+      # AXES_DIM DESCRIPTION: dimensions of the axes of the positional encoding. TYPE: list default: [16, 56, 56]
+      AXES_DIM: [ 16, 56, 56 ]
+      # THETA DESCRIPTION: theta for positional encoding. TYPE: int default: 10000
+      THETA: 10000
+      # VEC_IN_DIM DESCRIPTION: dimension of the vector input. TYPE: int default: 768
+      VEC_IN_DIM: 768
+      # GUIDANCE_EMBED DESCRIPTION: whether to use guidance embedding. TYPE: bool default: False
+      GUIDANCE_EMBED: False
+      # CONTEXT_IN_DIM DESCRIPTION: dimension of the context input. TYPE: int default: 4096
+      CONTEXT_IN_DIM: 4096
+      # MLP_RATIO DESCRIPTION: ratio of mlp hidden size to hidden size. TYPE: float default: 4.0
+      MLP_RATIO: 4.0
+      # QKV_BIAS DESCRIPTION: whether to use bias in qkv projection. TYPE: bool default: True
+      QKV_BIAS: True
+      # DEPTH DESCRIPTION: number of transformer blocks. TYPE: int default: 19
+      DEPTH: 19
+      # DEPTH_SINGLE_BLOCKS DESCRIPTION: number of transformer blocks in the single stream block. TYPE: int default: 38
+      DEPTH_SINGLE_BLOCKS: 38
+      USE_GRAD_CHECKPOINT: True
+
+    #
+    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: True
+        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: True
+        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 DESCRIPTION:  TYPE:  default: 'T5PlusClipFluxEmbedder'
+      NAME: T5PlusClipFluxEmbedder
+      # T5_MODEL DESCRIPTION:  TYPE:  default: ''
+      T5_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: T5EncoderModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder_2/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: T5Tokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer_2/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 512
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: last_hidden_state
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: False
+        CLEAN: whitespace
+      # CLIP_MODEL DESCRIPTION:  TYPE:  default: ''
+      CLIP_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: CLIPTextModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: CLIPTokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 77
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: pooler_output
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: True
+        CLEAN: whitespace
+  #
+  SAMPLE_ARGS:
+    SAMPLE_STEPS: 50
+    SAMPLER: flow_eluer
+    SEED: 2024
+    IMAGE_SIZE: [ 1024, 1024 ]
+    SHIFT: True
+    GUIDE_SCALE: 3.5
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 4e-4
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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 cat holds a blackboard that writes \"hello world\"", "a dog running on the lawn" ]
+    IMAGE_SIZE: [ 1024, 1024 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 2
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+#      GRADIENT_CLIP: 1.0
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+    - 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'
@@ -0,0 +1,342 @@
+ENV:
+  BACKEND: nccl
+META:
+  VERSION: 'FLUX1.0_SCHNELL'
+  DESCRIPTION: "flux 1.0 schnell"
+  IS_DEFAULT: False
+  IS_SHARE: True
+  INFERENCE_PARAS:
+    INFERENCE_BATCH_SIZE: 1
+    INFERENCE_PREFIX: ""
+    DEFAULT_SAMPLER: "flow_euler"
+    DEFAULT_SAMPLE_STEPS: 4
+    INFERENCE_N_PROMPT: ""
+    RESOLUTION: [1024, 1024]
+  PARAS:
+    -
+      TRAIN_BATCH_SIZE: 1
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      MEMORY: 89000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 4e-4
+      IS_DEFAULT: False
+      TUNER: FULL
+    -
+      TRAIN_BATCH_SIZE: 1
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      MEMORY: 89000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 4e-4
+      IS_DEFAULT: True
+      TUNER: LORA
+  #
+  TUNERS:
+    LORA:
+      -
+        NAME: SwiftLoRA
+        R: 4
+        LORA_ALPHA: 4
+        LORA_DROPOUT: 0.0
+        BIAS: "none"
+        TARGET_MODULES: "(model.double_blocks.*(.qkv|.proj|.img_mod.lin|.txt_mod.lin))|(model.single_blocks.*(.linear1|.linear2|.modulation.lin))$"
+#
+SOLVER:
+  NAME: LatentDiffusionSolver
+  # MAX_STEPS DESCRIPTION: The total steps for training. TYPE: int default: 100000
+  MAX_STEPS: 100000
+  # USE_AMP DESCRIPTION: Use amp to surpport mix precision or not, default is False. TYPE: bool default: False
+  USE_AMP: True
+  # DTYPE DESCRIPTION: The precision for training. TYPE: str default: 'float32'
+  DTYPE: bfloat16
+  # USE_FAIRSCALE DESCRIPTION: Use fairscale as the backend of ddp, default False. TYPE: bool default: False
+  USE_FAIRSCALE: False
+  # USE_FSDP DESCRIPTION: Use fsdp as the backend of ddp, default False. TYPE: bool default: False
+  USE_FSDP: True
+  # LOAD_MODEL_ONLY DESCRIPTION: Only load the model rather than the optimizer and schedule, default is False. TYPE: bool default: False
+  LOAD_MODEL_ONLY: False
+  # RESUME_FROM DESCRIPTION: Resume from some state of training! TYPE: str default: ''
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_flux_schnell_1024_lora
+  LOG_FILE: std_log.txt
+  # EVAL_INTERVAL DESCRIPTION: Eval the model interval. TYPE: int default: 1
+  EVAL_INTERVAL: 100
+  # LOG_TRAIN_NUM DESCRIPTION: The number samples used to log in training phase. TYPE: int default: -1
+  LOG_TRAIN_NUM: 16
+  # FSDP_REDUCE_DTYPE DESCRIPTION: The dtype of reduce in FSDP. TYPE: str default: 'float16'
+  FSDP_REDUCE_DTYPE: float32
+  # FSDP_BUFFER_DTYPE DESCRIPTION: The dtype of buffer in FSDP. TYPE: str default: 'float16'
+  FSDP_BUFFER_DTYPE: float32
+  # FSDP_SHARD_MODULES DESCRIPTION: The modules to be sharded in FSDP. TYPE: list default: ['model']
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.t5_model' ] #
+  SAVE_MODULES: [ 'model'] #
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  #
+  TUNER:
+  #
+  MODEL:
+    NAME: LatentDiffusionFlux
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      # NAME DESCRIPTION:  TYPE:  default: 'DiffusionFluxRF'
+      NAME: DiffusionFluxRF
+      PREDICTION_TYPE: raw
+      # NOISE_SCHEDULER DESCRIPTION:  TYPE:  default: ''
+      NOISE_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchSigmaScheduler'
+        NAME: FlowMatchSigmaScheduler
+        # WEIGHTING_SCHEME DESCRIPTION: The weighting scheme for sampling timesteps, choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']. TYPE: str default: 'logit_normal'
+        WEIGHTING_SCHEME: logit_normal
+        SHIFT: 3.0
+        # LOGIT_MEAN DESCRIPTION: The mean of the logit distribution for sampling timesteps. TYPE: float default: 0.0
+        LOGIT_MEAN: 0.0
+        # LOGIT_STD DESCRIPTION: The standard deviation of the logit distribution for sampling timesteps. TYPE: float default: 1.0
+        LOGIT_STD: 1.0
+        # MODE_SCALE DESCRIPTION: The scale factor for the mode of the logit distribution for sampling timesteps. TYPE: float default: 1.29
+        MODE_SCALE: 1.29
+      SAMPLER_SCHEDULER:
+        # NAME DESCRIPTION:  TYPE:  default: 'FlowMatchFluxShiftScheduler'
+        NAME: FlowMatchFluxShiftScheduler
+        # SHIFT DESCRIPTION: Use timestamp shift or not, default is True. TYPE: bool default: True
+        SHIFT: False
+        # SIGMOID_SCALE DESCRIPTION: The scale of sigmoid function for sampling timesteps. TYPE: int default: 1
+        SIGMOID_SCALE: 1
+        # BASE_SHIFT DESCRIPTION: The base shift factor for the timestamp. TYPE: float default: 0.5
+        BASE_SHIFT: 0.5
+        # MAX_SHIFT DESCRIPTION: The max shift factor for the timestamp. TYPE: float default: 1.15
+        MAX_SHIFT: 1.15
+        #
+    DIFFUSION_MODEL:
+      # NAME DESCRIPTION:  TYPE:  default: 'Flux'
+      NAME: Flux
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-schnell@flux1-schnell.safetensors
+      # IN_CHANNELS DESCRIPTION: model's input channels. TYPE: int default: 64
+      IN_CHANNELS: 64
+      # HIDDEN_SIZE DESCRIPTION: model's hidden size. TYPE: int default: 1024
+      HIDDEN_SIZE: 3072
+      # NUM_HEADS DESCRIPTION: number of heads in the transformer. TYPE: int default: 16
+      NUM_HEADS: 24
+      # AXES_DIM DESCRIPTION: dimensions of the axes of the positional encoding. TYPE: list default: [16, 56, 56]
+      AXES_DIM: [ 16, 56, 56 ]
+      # THETA DESCRIPTION: theta for positional encoding. TYPE: int default: 10000
+      THETA: 10000
+      # VEC_IN_DIM DESCRIPTION: dimension of the vector input. TYPE: int default: 768
+      VEC_IN_DIM: 768
+      # GUIDANCE_EMBED DESCRIPTION: whether to use guidance embedding. TYPE: bool default: False
+      GUIDANCE_EMBED: False
+      # CONTEXT_IN_DIM DESCRIPTION: dimension of the context input. TYPE: int default: 4096
+      CONTEXT_IN_DIM: 4096
+      # MLP_RATIO DESCRIPTION: ratio of mlp hidden size to hidden size. TYPE: float default: 4.0
+      MLP_RATIO: 4.0
+      # QKV_BIAS DESCRIPTION: whether to use bias in qkv projection. TYPE: bool default: True
+      QKV_BIAS: True
+      # DEPTH DESCRIPTION: number of transformer blocks. TYPE: int default: 19
+      DEPTH: 19
+      # DEPTH_SINGLE_BLOCKS DESCRIPTION: number of transformer blocks in the single stream block. TYPE: int default: 38
+      DEPTH_SINGLE_BLOCKS: 38
+      USE_GRAD_CHECKPOINT: True
+
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLFlux
+      EMBED_DIM: 16
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-schnell@ae.safetensors
+      IGNORE_KEYS: [ ]
+      BATCH_SIZE: 8
+      USE_CONV: False
+      SCALE_FACTOR: 0.3611
+      SHIFT_FACTOR: 0.1159
+      #
+      ENCODER:
+        NAME: Encoder
+        USE_CHECKPOINT: True
+        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: True
+        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 DESCRIPTION:  TYPE:  default: 'T5PlusClipFluxEmbedder'
+      NAME: T5PlusClipFluxEmbedder
+      # T5_MODEL DESCRIPTION:  TYPE:  default: ''
+      T5_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: T5EncoderModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-schnell@text_encoder_2/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: T5Tokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-schnell@tokenizer_2/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 256
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: last_hidden_state
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: False
+        CLEAN: whitespace
+      # CLIP_MODEL DESCRIPTION:  TYPE:  default: ''
+      CLIP_MODEL:
+        # NAME DESCRIPTION:  TYPE:  default: 'HFEmbedder'
+        NAME: HFEmbedder
+        # HF_MODEL_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_MODEL_CLS: CLIPTextModel
+        # MODEL_PATH DESCRIPTION: model folder path TYPE: NoneType default: None
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-schnell@text_encoder/
+        # HF_TOKENIZER_CLS DESCRIPTION: huggingface cls in transfomer TYPE: NoneType default: None
+        HF_TOKENIZER_CLS: CLIPTokenizer
+        # TOKENIZER_PATH DESCRIPTION: tokenizer folder path TYPE: NoneType default: None
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-schnell@tokenizer/
+        # MAX_LENGTH DESCRIPTION: max length of input TYPE: int default: 77
+        MAX_LENGTH: 77
+        # OUTPUT_KEY DESCRIPTION: output key TYPE: str default: 'last_hidden_state'
+        OUTPUT_KEY: pooler_output
+        # D_TYPE DESCRIPTION: dtype TYPE: str default: 'bfloat16'
+        D_TYPE: bfloat16
+        # BATCH_INFER DESCRIPTION: batch infer TYPE: bool default: False
+        BATCH_INFER: True
+        CLEAN: whitespace
+  #
+  SAMPLE_ARGS:
+    SAMPLE_STEPS: 4
+    SAMPLER: flow_eluer
+    SEED: 2024
+    IMAGE_SIZE: [ 1024, 1024 ]
+    GUIDE_SCALE: 3.5
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 4e-4
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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 cat holds a blackboard that writes \"hello world\"", "a dog running on the lawn" ]
+    IMAGE_SIZE: [ 1024, 1024 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 2
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+#      GRADIENT_CLIP: 1.0
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+    - 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'
@@ -0,0 +1,264 @@
+ENV:
+  BACKEND: nccl
+META:
+  VERSION: 'PIXART_ALPHA'
+  DESCRIPTION: "PIXART ALPHA"
+  IS_DEFAULT: False
+  IS_SHARE: True
+  INFERENCE_PARAS:
+    INFERENCE_BATCH_SIZE: 1
+    INFERENCE_PREFIX: ""
+    DEFAULT_SAMPLER: "ddim"
+    DEFAULT_SAMPLE_STEPS: 20
+    INFERENCE_N_PROMPT: ""
+    RESOLUTION: [1024, 1024]
+  PARAS:
+    -
+      TRAIN_BATCH_SIZE: 2
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      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: [1024, 1024]
+      MEMORY: 29000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 0.0001
+      IS_DEFAULT: False
+      TUNER: LORA
+  #
+  TUNERS:
+    LORA:
+      -
+        NAME: SwiftLoRA
+        R: 256
+        LORA_ALPHA: 256
+        LORA_DROPOUT: 0.0
+        BIAS: "none"
+        TARGET_MODULES: "model.*(.q|.k|.v|.o|mlp.fc1|mlp.fc2)$"
+#
+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: cache/scepter_ui/self_train/dit/pixart_alpha_pro
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  #
+  TUNER:
+  #
+  MODEL:
+    NAME: LatentDiffusionPixart
+    PARAMETERIZATION: eps
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 0.18215
+    SIZE_FACTOR: 8
+    DECODER_BIAS: 0.5
+    DEFAULT_N_PROMPT:
+    SCHEDULE_ARGS:
+      "NAME": "linear"
+      "BETA_MIN": 0.0001
+      "BETA_MAX": 0.02
+    USE_EMA: False
+    LOAD_REFINER: False
+    #
+    DIFFUSION_MODEL:
+      NAME: PixArt
+      PRETRAINED_MODEL: ms://AI-ModelScope/PixArt-alpha@PixArt-XL-2-1024-MS.pth
+      INPUT_SIZE: 128
+      PATCH_SIZE: 2
+      IN_CHANNELS: 4
+      HIDDEN_SIZE: 1152
+      DEPTH: 28
+      NUM_HEADS: 16
+      MLP_RATIO: 4.0
+      CLASS_DROPOUT_PROB: 0.1
+      PRED_SIGMA: True
+      DROP_PATH: 0.0
+      WINDOW_DIZE: 0
+      USE_REL_POS: False
+      CAPTION_CHANNELS: 4096
+      LEWEI_SCALE: 2
+      MODEL_MAX_LENGTH: 120
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-base@512-base-ema.safetensors
+      EMBED_DIM: 4
+      IGNORE_KEYS: [ ]
+      BATCH_SIZE: 1
+      #
+      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: ms://AI-ModelScope/PixArt-alpha@t5-v1_1-xxl/
+      TOKENIZER_PATH: ms://AI-ModelScope/PixArt-alpha@t5-v1_1-xxl/
+      LENGTH: 120
+      CLEAN: heavy
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 20
+    SEED: 2024
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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: [ 1024, 1024 ]
+    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'
@@ -0,0 +1,284 @@
+ENV:
+  BACKEND: nccl
+META:
+  VERSION: 'SD3'
+  DESCRIPTION: "SD3 ALPHA"
+  IS_DEFAULT: False
+  IS_SHARE: True
+  INFERENCE_PARAS:
+    INFERENCE_BATCH_SIZE: 1
+    INFERENCE_PREFIX: ""
+    DEFAULT_SAMPLER: "euler"
+    DEFAULT_SAMPLE_STEPS: 28
+    INFERENCE_N_PROMPT: ""
+    RESOLUTION: [1024, 1024]
+  PARAS:
+    -
+      TRAIN_BATCH_SIZE: 2
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      MEMORY: 29000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 1e-5
+      IS_DEFAULT: False
+      TUNER: FULL
+    -
+      TRAIN_BATCH_SIZE: 2
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [1024, 1024]
+      MEMORY: 29000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 5e-5
+      IS_DEFAULT: True
+      TUNER: LORA
+  #
+  TUNERS:
+    LORA:
+      -
+        NAME: SwiftLoRA
+        R: 128
+        LORA_ALPHA: 128
+        LORA_DROPOUT: 0.0
+        BIAS: "none"
+        TARGET_MODULES: "model.*(.attn.qkv|.attn.proj|mlp.fc1|mlp.fc2)$"
+#
+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: cache/scepter_ui/self_train/dit/sd3
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    - NAME: "ModelscopeFs"
+      TEMP_DIR: "./cache/cache_data"
+  #
+  FREEZE:
+  #
+  TUNER:
+  #
+  MODEL:
+    NAME: LatentDiffusionSD3
+    PARAMETERIZATION: rf
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA:
+    ZERO_TERMINAL_SNR: False
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 1.5305
+    SHIFT_FACTOR: 0.0609
+    DEFAULT_N_PROMPT:
+    SCHEDULE_ARGS:
+      "NAME": "shifted"
+      "SHIFT": 3
+    USE_EMA: False
+    T_WEIGHT: uniform
+    #
+    DIFFUSION_MODEL:
+      NAME: MMDiT
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+      IGNORE_KEYS: '^first_stage_model.'
+      IN_CHANNELS: 16
+      PATCH_SIZE: 2
+      OUT_CHANNELS: 16
+      DEPTH: 24
+      INPUT_SIZE:
+      ADM_IN_CHANNELS: 2048
+      CONTEXT_EMBEDDER_CONFIG: { 'target': 'torch.nn.Linear', 'params': { 'in_features': 4096, 'out_features': 1536 } }
+      NUM_PATCHES: 36864
+      POS_EMBED_MAX_SIZE: 192
+      POS_EMBED_SCALING_FACTOR:
+      USE_CHECKPOINT: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium@sd3_medium.safetensors
+      EMBED_DIM: 16
+      IGNORE_KEYS: '^model.diffusion_model.'
+      BATCH_SIZE: 1
+      USE_CONV: False
+      #
+      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: 16
+        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: 16
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+        GIVE_PRE_END: False
+        TANH_OUT: False
+    #
+    COND_STAGE_MODEL:
+      NAME: SD3TextEmbedder
+      P_ZERO: 0.0
+      CLIP_L:
+        NAME: FrozenCLIPEmbedder2
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer
+        MAX_LENGTH: 77
+        FREEZE: True
+        LAYER: penultimate
+        RETURN_POOLED: True
+        USE_FINAL_LAYER_NORM: False
+        IS_TRAINABLE: False
+      CLIP_G:
+        NAME: FrozenCLIPEmbedder2
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_2
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_2
+        MAX_LENGTH: 77
+        FREEZE: True
+        LAYER: penultimate
+        RETURN_POOLED: True
+        USE_FINAL_LAYER_NORM: False
+        IS_TRAINABLE: False
+      T5_XXL:
+        NAME: T5EmbedderHF
+        PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@text_encoder_3
+        TOKENIZER_PATH: ms://AI-ModelScope/stable-diffusion-3-medium-diffusers@tokenizer_3
+        LENGTH: 256
+        CLEAN: whitespace
+        USE_GRAD: False
+        T5_DTYPE: float16
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: euler
+    SAMPLE_STEPS: 28
+    SEED: 1749023094
+    GUIDE_SCALE: 5.0
+    GUIDE_RESCALE: 0.0
+    DISCRETIZATION: trailing
+    RUN_TRAIN_N: False
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 5e-5
+    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: [ 1024, 1024 ]
+        INPUT_KEY: [ 'img' ]
+        OUTPUT_KEY: [ 'img' ]
+        BACKEND: pillow
+      - NAME: FlexibleCenterCrop
+        SIZE: [ 1024, 1024 ]
+        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 cat holds a blackboard that writes \"hello world\"", "a dog running on the lawn" ]
+    IMAGE_SIZE: [ 1024, 1024 ]
+    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: 1
+    -
+      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'
@@ -142,13 +142,14 @@ SOLVER:
  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/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
  #
@@ -258,7 +259,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -267,7 +268,7 @@ SOLVER:
  TRAIN_DATA:
    NAME: ImageTextPairMSDataset
    MODE: train
-    MS_DATASET_NAME: cache/save_data/delogo/
+    MS_DATASET_NAME: ./cache/cache_data/delogo/
    MS_DATASET_NAMESPACE: ""
    MS_DATASET_SPLIT: "train"
    MS_DATASET_SUBNAME: ""
@@ -147,6 +147,7 @@ SOLVER:
  MAX_EPOCHS: -1
  # NUM_FOLDS DESCRIPTION: Num folds for training. TYPE: int default: 1
  NUM_FOLDS: 1
+  RESCALE_LR: False
  #
  EVAL_INTERVAL: -1
  # WORK_DIR DESCRIPTION: Save dir of the training log or model. TYPE: str default: ''
@@ -155,7 +156,7 @@ SOLVER:
  LOG_FILE: std_log.txt
  FILE_SYSTEM:
    NAME: "ModelscopeFs"
-    TEMP_DIR: "./cache/data"
+    TEMP_DIR: "./cache/cache_data"
  TUNER:
  # MODEL DESCRIPTION:  TYPE:  default: ''
  MODEL:
@@ -538,7 +539,7 @@ SOLVER:
  OPTIMIZER:
    # NAME DESCRIPTION:  TYPE:  default: ''
    NAME: AdamW
-    LEARNING_RATE: 0.0064
+    LEARNING_RATE: 0.00001
    EPS: 1e-8
    AMSGRAD: False
  #
@@ -137,13 +137,14 @@ SOLVER:
  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/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
  #
@@ -249,7 +250,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -80,13 +80,14 @@ SOLVER:
  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/data"
+    TEMP_DIR: "./cache/cache_data"
  #
  FREEZE:
  #
@@ -191,7 +192,7 @@ SOLVER:
  #
  OPTIMIZER:
    NAME: AdamW
-    LEARNING_RATE: 0.064
+    LEARNING_RATE: 0.0001
    BETAS: [ 0.9, 0.999 ]
    EPS: 1e-8
    WEIGHT_DECAY: 1e-2
@@ -13,3 +13,5 @@ BASE_MODEL_VERSION:
    TUNER_TYPE: [ 'LORA', 'SCE', 'FULL' ]
  - BASE_MODEL: 'EDIT'
    TUNER_TYPE: [ 'LORA', 'SCE', 'FULL' ]
+  - BASE_MODEL: 'FLUX1.0_DEV'
+    TUNER_TYPE: [ 'LORA' ]
@@ -3,9 +3,21 @@
 from scepter.modules.annotator.base_annotator import GeneralAnnotator
 from scepter.modules.annotator.canny import CannyAnnotator
 from scepter.modules.annotator.color import ColorAnnotator
+from scepter.modules.annotator.degradation import DegradationAnnotator
+from scepter.modules.annotator.doodle import DoodleAnnotator
+from scepter.modules.annotator.gray import GrayAnnotator
 from scepter.modules.annotator.hed import HedAnnotator
 from scepter.modules.annotator.identity import IdentityAnnotator
+from scepter.modules.annotator.informative_drawing import (
+    InfoDrawAnimeAnnotator, InfoDrawContourAnnotator,
+    InfoDrawOpenSketchAnnotator)
+from scepter.modules.annotator.inpainting import InpaintingAnnotator
 from scepter.modules.annotator.invert import InvertAnnotator
 from scepter.modules.annotator.midas_op import MidasDetector
 from scepter.modules.annotator.mlsd_op import MLSDdetector
 from scepter.modules.annotator.openpose import OpenposeAnnotator
+from scepter.modules.annotator.outpainting import OutpaintingAnnotator, OutpaintingResize
+from scepter.modules.annotator.pidinet import PiDiAnnotator
+from scepter.modules.annotator.segmentation import ESAMAnnotator
+from scepter.modules.annotator.sketch import SketchAnnotator
+from scepter.modules.annotator.lama import LamaAnnotator
@@ -0,0 +1,135 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import math
+import random
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+
+from PIL import Image
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import Config, dict_to_yaml
+
+
+def gaussian_noise_op(im, v):
+    from basicsr.data.degradations import random_add_gaussian_noise
+    noise_level = v.get('noise_level', [10, 20])
+    out = random_add_gaussian_noise(
+        im,
+        sigma_range=noise_level,
+        clip=True,
+        rounds=False,
+        gray_prob=0.4,
+    )
+    out = np.clip(out, 0.0, 1.0)
+    return out
+
+
+def resize_op(im, v):
+    scale = v.get('scale', [0.5, 0.8])
+    h, w = im.shape[:2]
+    scale = random.uniform(scale[0], scale[1])
+    h_, w_ = int(h * scale), int(w * scale)
+    mode = v.get('mode', 'nearest')
+    if mode == 'nearest':
+        interpolation = cv2.INTER_NEAREST
+    elif mode == 'bilinear':
+        interpolation = cv2.INTER_LINEAR
+    elif mode == 'bicubic':
+        interpolation = cv2.INTER_CUBIC
+    else:
+        interpolation = cv2.INTER_NEAREST
+    im = cv2.resize(im, (w_, h_), interpolation=interpolation)
+    out = cv2.resize(im, (w, h), interpolation=interpolation)
+    out = np.clip(out, 0.0, 1.0)
+    return out
+
+
+def jpeg_op(im, v):
+    from basicsr.data.degradations import add_jpg_compression
+    jpeg_level = v.get('jpeg_level', [50, 75])
+    v = int(random.uniform(jpeg_level[0], jpeg_level[1]))
+    out = add_jpg_compression(im, v)
+    out = np.clip(out, 0.0, 1.0)
+    return out
+
+
+def gaussian_blur_op(im, v):
+    from basicsr.data.degradations import random_mixed_kernels
+    kernel_range = v.get('kernel_size', [7, 9])
+    kernel_size = random.choice(kernel_range)
+    kernel_size = min(int(kernel_size) // 2 * 2 + 1, 21)
+    blur_sigma = v.get('sigma', [0.9, 1.0])
+    kernel = random_mixed_kernels(
+        ('iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso',
+         'plateau_aniso'), (0.45, 0.25, 0.12, 0.03, 0.12, 0.03),
+        kernel_size,
+        blur_sigma,
+        blur_sigma, [-math.pi, math.pi], [0.5, 2.0], [1, 1.5],
+        noise_range=None)
+
+    pad_size = (21 - kernel_size) // 2
+    kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
+    out = cv2.filter2D(im, -1, kernel)
+    out = np.clip(out, 0.0, 1.0)
+    return out
+
+
+@ANNOTATORS.register_class()
+class DegradationAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        self.params = cfg.get('PARAMS', {
+            'gaussian_noise': {},
+            'resize': {},
+            'jpeg': {},
+            'gaussian_blur': {},
+        })
+        if not isinstance(self.params, dict):
+            self.params = Config.get_dict(self.params)
+        self.random_degradation = cfg.get('RANDOM_DEGRADATION', False)
+
+    def forward(self, image):
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        elif isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        elif isinstance(image, np.ndarray):
+            image = image.copy()
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+        if np.max(image) > 1.0:
+            image = (image / 255.).astype(np.float32)
+
+        degradation_list = list(self.params.keys())
+        if self.random_degradation:
+            random.shuffle(degradation_list)
+
+        for degradation_type in degradation_list:
+            if degradation_type == 'gaussian_noise':
+                image = gaussian_noise_op(image, self.params[degradation_type])
+            elif degradation_type == 'resize':
+                image = resize_op(image, self.params[degradation_type])
+            elif degradation_type == 'jpeg':
+                image = jpeg_op(image, self.params[degradation_type])
+            elif degradation_type == 'gaussian_blur':
+                image = gaussian_blur_op(image, self.params[degradation_type])
+            else:
+                raise NotImplementedError(
+                    f'ERROR: degradation_type: {degradation_type} is invalid.')
+        image = (image * 255.0).astype(np.uint8)
+
+        assert len(image.shape) < 4
+        return image
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            DegradationAnnotator.para_dict,
+                            set_name=True)
@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+
+import math
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as TT
+from einops import rearrange
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+
+
+@ANNOTATORS.register_class()
+class DoodleAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        self.processor_type = cfg.get('PROCESSOR_TYPE', 'pidinet_sketch')
+        processor_cfg = cfg.get('PROCESSOR_CFG', None)
+        if self.processor_type == 'pidinet_sketch':
+            self.pidinet_ins = ANNOTATORS.build(processor_cfg[0])
+            self.sketch_ins = ANNOTATORS.build(processor_cfg[1])
+        else:
+            raise 'Unsurpport PROCESSOR for DoodleAnnotator'
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    @torch.autocast('cuda', enabled=False)
+    def forward(self, image):
+        if self.processor_type == 'pidinet_sketch':
+            pidinet_res = self.pidinet_ins(image)
+            sketch_res = self.sketch_ins(pidinet_res)
+            doodle_res = sketch_res
+        else:
+            raise 'Unsurpport PROCESSOR for DoodleAnnotator'
+        return doodle_res
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            DoodleAnnotator.para_dict,
+                            set_name=True)
@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+
+
+@ANNOTATORS.register_class()
+class GrayAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+
+    def forward(self, image):
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        elif isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        elif isinstance(image, np.ndarray):
+            image = image.copy()
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+        gray_map = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        return gray_map[..., None].repeat(3, axis=2)
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            GrayAnnotator.para_dict,
+                            set_name=True)
@@ -0,0 +1,178 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision
+import torchvision.transforms as transforms
+from einops import rearrange
+from PIL import Image
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+from torchvision.transforms import InterpolationMode
+
+norm_layer = nn.InstanceNorm2d
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+
+        conv_block = [
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            norm_layer(in_features),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            norm_layer(in_features)
+        ]
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+
+class ContourInference(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
+        super(ContourInference, self).__init__()
+
+        # Initial convolution block
+        model0 = [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(input_nc, 64, 7),
+            norm_layer(64),
+            nn.ReLU(inplace=True)
+        ]
+        self.model0 = nn.Sequential(*model0)
+
+        # Downsampling
+        model1 = []
+        in_features = 64
+        out_features = in_features * 2
+        for _ in range(2):
+            model1 += [
+                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                norm_layer(out_features),
+                nn.ReLU(inplace=True)
+            ]
+            in_features = out_features
+            out_features = in_features * 2
+        self.model1 = nn.Sequential(*model1)
+
+        model2 = []
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model2 += [ResidualBlock(in_features)]
+        self.model2 = nn.Sequential(*model2)
+
+        # Upsampling
+        model3 = []
+        out_features = in_features // 2
+        for _ in range(2):
+            model3 += [
+                nn.ConvTranspose2d(in_features,
+                                   out_features,
+                                   3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(out_features),
+                nn.ReLU(inplace=True)
+            ]
+            in_features = out_features
+            out_features = in_features // 2
+        self.model3 = nn.Sequential(*model3)
+
+        # Output layer
+        model4 = [nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, 7)]
+        if sigmoid:
+            model4 += [nn.Sigmoid()]
+
+        self.model4 = nn.Sequential(*model4)
+
+    def forward(self, x, cond=None):
+        out = self.model0(x)
+        out = self.model1(out)
+        out = self.model2(out)
+        out = self.model3(out)
+        out = self.model4(out)
+
+        return out
+
+
+@ANNOTATORS.register_class()
+class InfoDrawContourAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        input_nc = cfg.get('INPUT_NC', 3)
+        output_nc = cfg.get('OUTPUT_NC', 1)
+        n_residual_blocks = cfg.get('N_RESIDUAL_BLOCKS', 3)
+        sigmoid = cfg.get('SIGMOID', True)
+        pretrained_model = cfg.get('PRETRAINED_MODEL', None)
+
+        self.model = ContourInference(input_nc, output_nc, n_residual_blocks,
+                                      sigmoid)
+        with FS.get_from(pretrained_model, wait_finish=True) as local_path:
+            self.model.load_state_dict(torch.load(local_path))
+        self.model = self.model.eval().requires_grad_(False).to(we.device_id)
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    @torch.autocast('cuda', enabled=False)
+    def forward(self, image):
+        is_batch = False if len(image.shape) == 3 else True
+        if isinstance(image, torch.Tensor):
+            if len(image.shape) == 3:
+                image = rearrange(image, 'h w c -> 1 c h w')
+                B, C, H, W = image.shape
+            elif len(image.shape) == 4:
+                B, C, H, W = image.shape
+            else:
+                raise "Unsurpport input image's shape"
+        elif isinstance(image, np.ndarray):
+            image = torch.from_numpy(image.copy()).float()
+            if len(image.shape) == 3:
+                image = rearrange(image, 'h w c -> 1 c h w')
+                B, C, H, W = image.shape
+            elif len(image.shape) == 4:
+                B, C, H, W = image.shape
+            else:
+                raise "Unsurpport input image's shape"
+        else:
+            raise "Unsurpport input image's type"
+
+        image = image.float().div(255).to(we.device_id)
+        contour_map = self.model(image)
+        contour_map = (contour_map.squeeze(dim=1) * 255.0).clip(
+            0, 255).cpu().numpy().astype(np.uint8)
+        contour_map = contour_map[..., None].repeat(3, -1)
+        if not is_batch:
+            contour_map = contour_map.squeeze()
+        return contour_map
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            InfoDrawContourAnnotator.para_dict,
+                            set_name=True)
+
+
+@ANNOTATORS.register_class()
+class InfoDrawAnimeAnnotator(InfoDrawContourAnnotator):
+    pass
+
+
+@ANNOTATORS.register_class()
+class InfoDrawOpenSketchAnnotator(InfoDrawContourAnnotator):
+    pass
@@ -0,0 +1,271 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import math
+import random
+from abc import ABCMeta
+from enum import Enum
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import Config, dict_to_yaml
+
+
+def invert_image(im):
+    im_arr = np.array(im)
+    mask_1 = im_arr == 0
+    mask_2 = im_arr == 255
+    im_arr[mask_1] = 255
+    im_arr[mask_2] = 0
+    new_im = im_arr
+    return new_im
+
+
+class DrawMethod(Enum):
+    LINE = 'line'
+    CIRCLE = 'circle'
+    SQUARE = 'square'
+
+
+def make_random_irregular_mask(shape,
+                               max_angle=4,
+                               max_len=60,
+                               max_width=20,
+                               min_times=0,
+                               max_times=10,
+                               draw_method=DrawMethod.LINE):
+    draw_method = DrawMethod(draw_method)
+
+    height, width = shape
+    mask = np.zeros((height, width), np.float32)
+    times = np.random.randint(min_times, max_times + 1)
+    for i in range(times):
+        start_x = np.random.randint(width)
+        start_y = np.random.randint(height)
+        for j in range(1 + np.random.randint(5)):
+            angle = 0.01 + np.random.randint(max_angle)
+            if i % 2 == 0:
+                angle = 2 * 3.1415926 - angle
+            length = 10 + np.random.randint(max_len)
+            brush_w = 5 + np.random.randint(max_width)
+            end_x = np.clip(
+                (start_x + length * np.sin(angle)).astype(np.int32), 0, width)
+            end_y = np.clip(
+                (start_y + length * np.cos(angle)).astype(np.int32), 0, height)
+            if draw_method == DrawMethod.LINE:
+                cv2.line(mask, (start_x, start_y), (end_x, end_y), 1.0,
+                         brush_w)
+            elif draw_method == DrawMethod.CIRCLE:
+                cv2.circle(mask, (start_x, start_y),
+                           radius=brush_w,
+                           color=1.,
+                           thickness=-1)
+            elif draw_method == DrawMethod.SQUARE:
+                radius = brush_w // 2
+                mask[start_y - radius:start_y + radius,
+                     start_x - radius:start_x + radius] = 1
+            start_x, start_y = end_x, end_y
+    return mask[None, ...]
+
+
+class RandomIrregularMaskGenerator:
+    def __init__(self,
+                 max_angle=4,
+                 max_len=60,
+                 max_width=20,
+                 min_times=0,
+                 max_times=10,
+                 ramp_kwargs=None,
+                 draw_method=DrawMethod.LINE):
+        self.max_angle = max_angle
+        self.max_len = max_len
+        self.max_width = max_width
+        self.min_times = min_times
+        self.max_times = max_times
+        self.draw_method = draw_method
+        self.ramp = None
+        # self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None
+
+    def __call__(self, img, iter_i=None, raw_image=None):
+        coef = self.ramp(iter_i) if (self.ramp is not None) and (
+            iter_i is not None) else 1
+        cur_max_len = int(max(1, self.max_len * coef))
+        cur_max_width = int(max(1, self.max_width * coef))
+        cur_max_times = int(self.min_times + 1 +
+                            (self.max_times - self.min_times) * coef)
+        return make_random_irregular_mask(img.shape[1:],
+                                          max_angle=self.max_angle,
+                                          max_len=cur_max_len,
+                                          max_width=cur_max_width,
+                                          min_times=self.min_times,
+                                          max_times=cur_max_times,
+                                          draw_method=self.draw_method)
+
+
+def make_random_rectangle_mask(shape,
+                               margin=10,
+                               bbox_min_size=30,
+                               bbox_max_size=100,
+                               min_times=0,
+                               max_times=3):
+    height, width = shape
+    mask = np.zeros((height, width), np.float32)
+    bbox_max_size = min(bbox_max_size, height - margin * 2, width - margin * 2)
+    times = np.random.randint(min_times, max_times + 1)
+    for i in range(times):
+        box_width = np.random.randint(bbox_min_size, bbox_max_size)
+        box_height = np.random.randint(bbox_min_size, bbox_max_size)
+        start_x = np.random.randint(margin, width - margin - box_width + 1)
+        start_y = np.random.randint(margin, height - margin - box_height + 1)
+        mask[start_y:start_y + box_height, start_x:start_x + box_width] = 1
+    return mask[None, ...]
+
+
+class RandomRectangleMaskGenerator:
+    def __init__(self,
+                 margin=10,
+                 bbox_min_size=30,
+                 bbox_max_size=100,
+                 min_times=0,
+                 max_times=3,
+                 ramp_kwargs=None):
+        self.margin = margin
+        self.bbox_min_size = bbox_min_size
+        self.bbox_max_size = bbox_max_size
+        self.min_times = min_times
+        self.max_times = max_times
+        self.ramp = None
+        # self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None
+
+    def __call__(self, img, iter_i=None, raw_image=None):
+        coef = self.ramp(iter_i) if (self.ramp is not None) and (
+            iter_i is not None) else 1
+        cur_bbox_max_size = int(self.bbox_min_size + 1 +
+                                (self.bbox_max_size - self.bbox_min_size) *
+                                coef)
+        cur_max_times = int(self.min_times +
+                            (self.max_times - self.min_times) * coef)
+        return make_random_rectangle_mask(img.shape[1:],
+                                          margin=self.margin,
+                                          bbox_min_size=self.bbox_min_size,
+                                          bbox_max_size=cur_bbox_max_size,
+                                          min_times=self.min_times,
+                                          max_times=cur_max_times)
+
+
+class MixedMaskGenerator:
+    def __init__(self,
+                 irregular_proba=1 / 3,
+                 irregular_kwargs=None,
+                 box_proba=1 / 3,
+                 box_kwargs=None,
+                 invert_proba=0):
+        self.probas = []
+        self.gens = []
+
+        if irregular_proba > 0:
+            self.probas.append(irregular_proba)
+            if irregular_kwargs is None:
+                irregular_kwargs = {}
+            else:
+                irregular_kwargs = dict(irregular_kwargs)
+            irregular_kwargs['draw_method'] = DrawMethod.LINE
+            self.gens.append(RandomIrregularMaskGenerator(**irregular_kwargs))
+
+        if box_proba > 0:
+            self.probas.append(box_proba)
+            if box_kwargs is None:
+                box_kwargs = {}
+            self.gens.append(RandomRectangleMaskGenerator(**box_kwargs))
+
+        self.probas = np.array(self.probas, dtype='float32')
+        self.probas /= self.probas.sum()
+        self.invert_proba = invert_proba
+
+    def __call__(self, img, iter_i=None, raw_image=None):
+        kind = np.random.choice(len(self.probas), p=self.probas)
+        gen = self.gens[kind]
+        result = gen(img, iter_i=iter_i, raw_image=raw_image)
+        if self.invert_proba > 0 and random.random() < self.invert_proba:
+            result = 1 - result
+        return result
+
+
+@ANNOTATORS.register_class()
+class InpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        self.mask_cfg = cfg.get(
+            'MASK_CFG', {
+                'irregular_proba': 0.5,
+                'irregular_kwargs': {
+                    'min_times': 4,
+                    'max_times': 10,
+                    'max_width': 100,
+                    'max_angle': 4,
+                    'max_len': 200
+                },
+                'box_proba': 0.5,
+                'box_kwargs': {
+                    'margin': 0,
+                    'bbox_min_size': 30,
+                    'bbox_max_size': 150,
+                    'max_times': 5,
+                    'min_times': 1
+                }
+            })
+        self.mask_cfg = Config.get_dict(self.mask_cfg) if not isinstance(
+            self.mask_cfg, dict) else self.mask_cfg
+        self.mask_generator = MixedMaskGenerator(**self.mask_cfg)
+        self.return_mask = cfg.get('RETURN_MASK', False)
+        self.return_invert = cfg.get('RETURN_INVERT', True)
+        self.mask_color = cfg.get('MASK_COLOR', 0)
+
+    def forward(self, image, mask=None, return_mask=None, return_invert=None, mask_color=None):
+        return_mask = return_mask if return_mask is not None else self.return_mask
+        return_invert = return_invert if return_invert is not None else self.return_invert
+        mask_color = mask_color if mask_color is not None else self.mask_color
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        elif isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        elif isinstance(image, np.ndarray):
+            image = image.copy()
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+
+        if mask is not None:
+            if mask_color:
+                image[np.array(mask) == 255] = mask_color
+            else:
+                image[np.array(mask) == 255] = 0
+        else:
+            img = np.transpose(image, (2, 0, 1))
+            mask = self.mask_generator(img)
+            mask = (np.transpose(mask, (1, 2, 0)).squeeze(-1) * 255).astype(np.uint8)
+            if return_invert:
+                mask = invert_image(mask)
+            colored_mask = np.zeros_like(image)
+            if mask_color: colored_mask[:] = mask_color
+            image = np.where(mask[:, :, np.newaxis] == 255, colored_mask, image)
+
+        if return_mask:
+            ret_data = {
+                'image': np.array(image),
+                'mask': np.array(mask)
+            }
+        else:
+            ret_data = np.array(image)
+        return ret_data
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            InpaintingAnnotator.para_dict,
+                            set_name=True)
@@ -0,0 +1,98 @@
+from abc import ABCMeta
+
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+
+def dilate_mask(mask, dilate_factor=15):
+    mask = mask.astype(np.uint8)
+    mask = cv2.dilate(
+        mask,
+        np.ones((dilate_factor, dilate_factor), np.uint8),
+        iterations=1
+    )
+    return mask
+
+@ANNOTATORS.register_class()
+class LamaAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        from modelscope.pipelines.builder import PIPELINES
+        from modelscope.pipelines.cv import ImageInpaintingPipeline
+        from modelscope.pipelines import pipeline
+        from modelscope.utils.constant import Tasks
+        from modelscope.metainfo import Pipelines
+        from modelscope.models.cv.image_inpainting.refinement import refine_predict
+        from torch.utils.data._utils.collate import default_collate
+
+        @PIPELINES.register_module(Tasks.image_inpainting, module_name=Pipelines.image_inpainting + "-v2")
+        class ImageInpaintingPipelineV2(ImageInpaintingPipeline):
+            def perform_inference(self, data):
+                px_budget = 9000000
+                batch = default_collate([data])
+                if self.refine:
+                    assert 'unpad_to_size' in batch, 'Unpadded size is required for the refinement'
+                    assert 'cuda' in str(self.device), 'GPU is required for refinement'
+                    gpu_ids = str(self.device).split(':')[-1]
+                    cur_res = refine_predict(
+                        batch,
+                        self.infer_model,
+                        gpu_ids=gpu_ids,
+                        modulo=self.pad_out_to_modulo,
+                        n_iters=15,
+                        lr=0.002,
+                        min_side=512,
+                        max_scales=3,
+                        px_budget=px_budget)
+                    cur_res = cur_res[0].permute(1, 2, 0).detach().cpu().numpy()
+                else:
+                    with torch.no_grad():
+                        batch = self.move_to_device(batch, self.device)
+                        batch['mask'] = (batch['mask'] > 0) * 1
+                        batch = self.infer_model(batch)
+                        cur_res = batch['inpainted'][0].permute(
+                            1, 2, 0).detach().cpu().numpy()
+                        unpad_to_size = batch.get('unpad_to_size', None)
+                        if unpad_to_size is not None:
+                            orig_height, orig_width = unpad_to_size
+                            cur_res = cur_res[:orig_height, :orig_width]
+
+                cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
+                cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR)
+                return cur_res
+
+        lama_model_dir = FS.get_dir_to_local_dir(cfg.PRETRAINED_MODEL)
+        self.lama_model = pipeline(Tasks.image_inpainting, model=lama_model_dir,
+                              pipeline_name=Pipelines.image_inpainting + "-v2", refine=True,
+                              device="cuda:{}".format(we.device_id))
+    def forward(self, image, mask):
+        mask = dilate_mask(mask, dilate_factor=19)
+        input_mask = Image.fromarray(mask)
+        mask_expanded = np.tile(np.expand_dims(mask, axis=-1), (1, 1, 3))
+        input_image_np = np.array(image)
+        input_image_np[mask_expanded == 255] = 0
+        input_image = Image.fromarray(input_image_np)
+        input = {
+            'img': input_image,
+            'mask': input_mask,
+        }
+        result = self.lama_model(input)
+        output_img = result['output_img']
+        return output_img[..., ::-1]
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            LamaAnnotator.para_dict,
+                            set_name=True)
+
+
+
@@ -9,9 +9,10 @@ import os
 from abc import ABCMeta
 from collections import OrderedDict

+import numpy as np
+
 import cv2
 import matplotlib
-import numpy as np
 import torch
 import torch.nn as nn
 from PIL import Image
@@ -478,7 +479,7 @@ class Hand(object):
            map_ori = heatmap_avg[:, :, part]
            one_heatmap = gaussian_filter(map_ori, sigma=3)
            binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
-            # 全部小于阈值
+
            if np.sum(binary) == 0:
                all_peaks.append([0, 0])
                continue
@@ -0,0 +1,189 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import math
+import random
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+
+
+@ANNOTATORS.register_class()
+class OutpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        self.mask_blur = cfg.get('MASK_BLUR', 0)
+        self.random_cfg = cfg.get('RANDOM_CFG', None)
+        self.return_mask = cfg.get('RETURN_MASK', False)
+        self.return_source = cfg.get('RETURN_SOURCE', True)
+        self.keep_padding_ratio = cfg.get('KEEP_PADDING_RATIO', 64)
+        self.mask_color = cfg.get('MASK_COLOR', 0)
+
+    def get_box(self, mask):
+        locs = np.where(mask == 255)
+        if len(locs) < 1 or locs[0].shape[0] < 1 or locs[1].shape[0] < 1:
+            return None
+        left, right = np.min(locs[1]), np.max(locs[1])
+        top, bottom = np.min(locs[0]), np.max(locs[0])
+        return [left, top, right, bottom]
+
+    def forward(self,
+                image,
+                ratio=0.3,
+                mask=None,
+                direction=['left', 'right', 'up', 'down'],
+                return_mask=None,
+                return_source=None,
+                mask_color=None):
+        return_mask = return_mask if return_mask is not None else self.return_mask
+        return_source = return_source if return_source is not None else self.return_source
+        mask_color = mask_color if mask_color is not None else self.mask_color
+        if isinstance(image, Image.Image):
+            image = image
+        elif isinstance(image, torch.Tensor):
+            image = Image.fromarray(image.detach().cpu().numpy())
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image.copy())
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+
+        if self.random_cfg:
+            direction_range = self.random_cfg.get(
+                'DIRECTION_RANGE', ['left', 'right', 'up', 'down'])
+            ratio_range = self.random_cfg.get('RATIO_RANGE', [0.0, 1.0])
+            direction = random.sample(
+                direction_range,
+                random.choice(list(range(1,
+                                         len(direction_range) + 1))))
+            ratio = random.uniform(ratio_range[0], ratio_range[1])
+
+        if mask is None:
+            init_image = image
+            src_width, src_height = init_image.width, init_image.height
+            left = int(ratio * src_width) if 'left' in direction else 0
+            right = int(ratio * src_width) if 'right' in direction else 0
+            up = int(ratio * src_height) if 'up' in direction else 0
+            down = int(ratio * src_height) if 'down' in direction else 0
+            # print(direction, ratio, left, right, up, down)
+            tar_width = math.ceil(
+                (src_width + left + right) /
+                self.keep_padding_ratio) * self.keep_padding_ratio
+            tar_height = math.ceil(
+                (src_height + up + down) /
+                self.keep_padding_ratio) * self.keep_padding_ratio
+            if left > 0:
+                left = left * (tar_width - src_width) // (left + right)
+            if right > 0:
+                right = tar_width - src_width - left
+            if up > 0:
+                up = up * (tar_height - src_height) // (up + down)
+            if down > 0:
+                down = tar_height - src_height - up
+            if mask_color is not None:
+                img = Image.new('RGB', (tar_width, tar_height), color=mask_color)
+            else:
+                img = Image.new('RGB', (tar_width, tar_height))
+            img.paste(init_image, (left, up))
+            mask = Image.new('L', (img.width, img.height), 'white')
+            draw = ImageDraw.Draw(mask)
+
+            draw.rectangle(
+                (left + (self.mask_blur * 2 if left > 0 else 0), up +
+                 (self.mask_blur * 2 if up > 0 else 0), mask.width - right -
+                 (self.mask_blur * 2 if right > 0 else 0), mask.height - down -
+                 (self.mask_blur * 2 if down > 0 else 0)),
+                fill='black')
+        else:
+            bbox = self.get_box(np.array(mask))
+            if bbox is None:
+                img = image
+                mask = mask
+                init_image = image
+            else:
+                mask = Image.new('L', (image.width, image.height), 'white')
+                mask_zero = Image.new('L', (bbox[2]-bbox[0], bbox[3]-bbox[1]), 'black')
+                mask.paste(mask_zero, (bbox[0], bbox[1]))
+                crop_image = image.crop(bbox)
+                init_image = Image.new('RGB', (image.width, image.height), 'black')
+                init_image.paste(crop_image, (bbox[0], bbox[1]))
+                img = image
+        if return_mask:
+            if return_source:
+                ret_data = {'src_image': np.array(init_image), 'image': np.array(img), 'mask': np.array(mask)}
+            else:
+                ret_data = {'image': np.array(img), 'mask': np.array(mask)}
+        else:
+            if return_source:
+                ret_data = {'src_image': np.array(init_image), 'image': np.array(img)}
+            else:
+                ret_data = np.array(img)
+        return ret_data
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            OutpaintingAnnotator.para_dict,
+                            set_name=True)
+
+@ANNOTATORS.register_class()
+class OutpaintingResize(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+
+    def get_box(self, mask):
+        locs = np.where(mask == 0)
+        if len(locs) < 1 or locs[0].shape[0] < 1 or locs[1].shape[0] < 1:
+            return None
+        left, right = np.min(locs[1]), np.max(locs[1])
+        top, bottom = np.min(locs[0]), np.max(locs[0])
+        return [left, top, right, bottom]
+
+    def forward(self,
+                image,
+                target_image,
+                mask=None
+                ):
+        if isinstance(image, Image.Image):
+            image = image
+        elif isinstance(image, torch.Tensor):
+            image = Image.fromarray(image.detach().cpu().numpy())
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image.copy())
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+
+        if isinstance(target_image, Image.Image):
+            target_image = target_image
+        elif isinstance(target_image, torch.Tensor):
+            target_image = Image.fromarray(target_image.detach().cpu().numpy())
+        elif isinstance(target_image, np.ndarray):
+            target_image = Image.fromarray(target_image.copy())
+        else:
+            raise f'Unsurpport datatype{type(target_image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+
+        bbox = self.get_box(np.array(mask))
+        if bbox is None:
+            init_image = image
+        else:
+            paste_img = image.resize((bbox[2]-bbox[0], bbox[3]-bbox[1]))
+            init_image = Image.new('RGB', (target_image.width, target_image.height), 'black')
+            init_image.paste(paste_img, (bbox[0], bbox[1]))
+        ret_data = {'src_image': np.array(init_image)}
+        return ret_data
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            OutpaintingResize.para_dict,
+                            set_name=True)
@@ -0,0 +1,935 @@
+# -*- coding: utf-8 -*-
+
+import math
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as TT
+from einops import rearrange
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+
+CONFIGS = {
+    'baseline': {
+        'layer0': 'cv',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'cv',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'cv',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'c-v15': {
+        'layer0': 'cd',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'cv',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'cv',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'a-v15': {
+        'layer0': 'ad',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'cv',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'cv',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'r-v15': {
+        'layer0': 'rd',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'cv',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'cv',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cv',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'cvvv4': {
+        'layer0': 'cd',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'cd',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'cd',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'cd',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'avvv4': {
+        'layer0': 'ad',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'ad',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'ad',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'ad',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'rvvv4': {
+        'layer0': 'rd',
+        'layer1': 'cv',
+        'layer2': 'cv',
+        'layer3': 'cv',
+        'layer4': 'rd',
+        'layer5': 'cv',
+        'layer6': 'cv',
+        'layer7': 'cv',
+        'layer8': 'rd',
+        'layer9': 'cv',
+        'layer10': 'cv',
+        'layer11': 'cv',
+        'layer12': 'rd',
+        'layer13': 'cv',
+        'layer14': 'cv',
+        'layer15': 'cv',
+    },
+    'cccv4': {
+        'layer0': 'cd',
+        'layer1': 'cd',
+        'layer2': 'cd',
+        'layer3': 'cv',
+        'layer4': 'cd',
+        'layer5': 'cd',
+        'layer6': 'cd',
+        'layer7': 'cv',
+        'layer8': 'cd',
+        'layer9': 'cd',
+        'layer10': 'cd',
+        'layer11': 'cv',
+        'layer12': 'cd',
+        'layer13': 'cd',
+        'layer14': 'cd',
+        'layer15': 'cv',
+    },
+    'aaav4': {
+        'layer0': 'ad',
+        'layer1': 'ad',
+        'layer2': 'ad',
+        'layer3': 'cv',
+        'layer4': 'ad',
+        'layer5': 'ad',
+        'layer6': 'ad',
+        'layer7': 'cv',
+        'layer8': 'ad',
+        'layer9': 'ad',
+        'layer10': 'ad',
+        'layer11': 'cv',
+        'layer12': 'ad',
+        'layer13': 'ad',
+        'layer14': 'ad',
+        'layer15': 'cv',
+    },
+    'rrrv4': {
+        'layer0': 'rd',
+        'layer1': 'rd',
+        'layer2': 'rd',
+        'layer3': 'cv',
+        'layer4': 'rd',
+        'layer5': 'rd',
+        'layer6': 'rd',
+        'layer7': 'cv',
+        'layer8': 'rd',
+        'layer9': 'rd',
+        'layer10': 'rd',
+        'layer11': 'cv',
+        'layer12': 'rd',
+        'layer13': 'rd',
+        'layer14': 'rd',
+        'layer15': 'cv',
+    },
+    'c16': {
+        'layer0': 'cd',
+        'layer1': 'cd',
+        'layer2': 'cd',
+        'layer3': 'cd',
+        'layer4': 'cd',
+        'layer5': 'cd',
+        'layer6': 'cd',
+        'layer7': 'cd',
+        'layer8': 'cd',
+        'layer9': 'cd',
+        'layer10': 'cd',
+        'layer11': 'cd',
+        'layer12': 'cd',
+        'layer13': 'cd',
+        'layer14': 'cd',
+        'layer15': 'cd',
+    },
+    'a16': {
+        'layer0': 'ad',
+        'layer1': 'ad',
+        'layer2': 'ad',
+        'layer3': 'ad',
+        'layer4': 'ad',
+        'layer5': 'ad',
+        'layer6': 'ad',
+        'layer7': 'ad',
+        'layer8': 'ad',
+        'layer9': 'ad',
+        'layer10': 'ad',
+        'layer11': 'ad',
+        'layer12': 'ad',
+        'layer13': 'ad',
+        'layer14': 'ad',
+        'layer15': 'ad',
+    },
+    'r16': {
+        'layer0': 'rd',
+        'layer1': 'rd',
+        'layer2': 'rd',
+        'layer3': 'rd',
+        'layer4': 'rd',
+        'layer5': 'rd',
+        'layer6': 'rd',
+        'layer7': 'rd',
+        'layer8': 'rd',
+        'layer9': 'rd',
+        'layer10': 'rd',
+        'layer11': 'rd',
+        'layer12': 'rd',
+        'layer13': 'rd',
+        'layer14': 'rd',
+        'layer15': 'rd',
+    },
+    'carv4': {
+        'layer0': 'cd',
+        'layer1': 'ad',
+        'layer2': 'rd',
+        'layer3': 'cv',
+        'layer4': 'cd',
+        'layer5': 'ad',
+        'layer6': 'rd',
+        'layer7': 'cv',
+        'layer8': 'cd',
+        'layer9': 'ad',
+        'layer10': 'rd',
+        'layer11': 'cv',
+        'layer12': 'cd',
+        'layer13': 'ad',
+        'layer14': 'rd',
+        'layer15': 'cv'
+    }
+}
+
+
+def create_conv_func(op_type):
+    assert op_type in ['cv', 'cd', 'ad',
+                       'rd'], 'unknown op type: %s' % str(op_type)
+    if op_type == 'cv':
+        return F.conv2d
+    if op_type == 'cd':
+
+        def func(x,
+                 weights,
+                 bias=None,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1):
+            assert dilation in [1,
+                                2], 'dilation for cd_conv should be in 1 or 2'
+            assert weights.size(2) == 3 and weights.size(3) == 3, \
+                'kernel size for cd_conv should be 3x3'
+            assert padding == dilation, 'padding for cd_conv set wrong'
+
+            weights_c = weights.sum(dim=[2, 3], keepdim=True)
+            yc = F.conv2d(x,
+                          weights_c,
+                          stride=stride,
+                          padding=0,
+                          groups=groups)
+            y = F.conv2d(x,
+                         weights,
+                         bias,
+                         stride=stride,
+                         padding=padding,
+                         dilation=dilation,
+                         groups=groups)
+            return y - yc
+
+        return func
+    elif op_type == 'ad':
+
+        def func(x,
+                 weights,
+                 bias=None,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1):
+            assert dilation in [1,
+                                2], 'dilation for ad_conv should be in 1 or 2'
+            assert weights.size(2) == 3 and weights.size(3) == 3, \
+                'kernel size for ad_conv should be 3x3'
+            assert padding == dilation, 'padding for ad_conv set wrong'
+
+            shape = weights.shape
+            weights = weights.view(shape[0], shape[1], -1)
+            # clock-wise
+            weights_conv = (
+                weights -
+                weights[:, :, [3, 0, 1, 6, 4, 2, 7, 8, 5]]).view(shape)
+            y = F.conv2d(x,
+                         weights_conv,
+                         bias,
+                         stride=stride,
+                         padding=padding,
+                         dilation=dilation,
+                         groups=groups)
+            return y
+
+        return func
+    elif op_type == 'rd':
+
+        def func(x,
+                 weights,
+                 bias=None,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1):
+            assert dilation in [1,
+                                2], 'dilation for rd_conv should be in 1 or 2'
+            assert weights.size(2) == 3 and weights.size(3) == 3, \
+                'kernel size for rd_conv should be 3x3'
+            padding = 2 * dilation
+
+            shape = weights.shape
+            if weights.is_cuda:
+                buffer = torch.cuda.FloatTensor(shape[0], shape[1],
+                                                5 * 5).fill_(0)
+            else:
+                buffer = torch.zeros(shape[0], shape[1], 5 * 5)
+            weights = weights.view(shape[0], shape[1], -1)
+            buffer[:, :, [0, 2, 4, 10, 14, 20, 22, 24]] = weights[:, :, 1:]
+            buffer[:, :, [6, 7, 8, 11, 13, 16, 17, 18]] = -weights[:, :, 1:]
+            buffer[:, :, 12] = 0
+            buffer = buffer.view(shape[0], shape[1], 5, 5)
+            y = F.conv2d(x,
+                         buffer,
+                         bias,
+                         stride=stride,
+                         padding=padding,
+                         dilation=dilation,
+                         groups=groups)
+            return y
+
+        return func
+    else:
+        print('impossible to be here unless you force that', flush=True)
+        return None
+
+
+def config_model(model):
+    model_options = list(CONFIGS.keys())
+    assert model in model_options, \
+        'unrecognized model, please choose from %s' % str(model_options)
+
+    pdcs = []
+    for i in range(16):
+        layer_name = 'layer%d' % i
+        op = CONFIGS[model][layer_name]
+        pdcs.append(create_conv_func(op))
+    return pdcs
+
+
+def config_model_converted(model):
+    model_options = list(CONFIGS.keys())
+    assert model in model_options, \
+        'unrecognized model, please choose from %s' % str(model_options)
+
+    pdcs = []
+    for i in range(16):
+        layer_name = 'layer%d' % i
+        op = CONFIGS[model][layer_name]
+        pdcs.append(op)
+    return pdcs
+
+
+def convert_pdc(op, weight):
+    if op == 'cv':
+        return weight
+    elif op == 'cd':
+        shape = weight.shape
+        weight_c = weight.sum(dim=[2, 3])
+        weight = weight.view(shape[0], shape[1], -1)
+        weight[:, :, 4] = weight[:, :, 4] - weight_c
+        weight = weight.view(shape)
+        return weight
+    elif op == 'ad':
+        shape = weight.shape
+        weight = weight.view(shape[0], shape[1], -1)
+        weight_conv = (weight -
+                       weight[:, :, [3, 0, 1, 6, 4, 2, 7, 8, 5]]).view(shape)
+        return weight_conv
+    elif op == 'rd':
+        shape = weight.shape
+        buffer = torch.zeros(shape[0], shape[1], 5 * 5, device=weight.device)
+        weight = weight.view(shape[0], shape[1], -1)
+        buffer[:, :, [0, 2, 4, 10, 14, 20, 22, 24]] = weight[:, :, 1:]
+        buffer[:, :, [6, 7, 8, 11, 13, 16, 17, 18]] = -weight[:, :, 1:]
+        buffer = buffer.view(shape[0], shape[1], 5, 5)
+        return buffer
+    raise ValueError('wrong op {}'.format(str(op)))
+
+
+def convert_pidinet(state_dict, config):
+    pdcs = config_model_converted(config)
+    new_dict = {}
+    for pname, p in state_dict.items():
+        if 'init_block.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[0], p)
+        elif 'block1_1.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[1], p)
+        elif 'block1_2.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[2], p)
+        elif 'block1_3.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[3], p)
+        elif 'block2_1.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[4], p)
+        elif 'block2_2.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[5], p)
+        elif 'block2_3.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[6], p)
+        elif 'block2_4.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[7], p)
+        elif 'block3_1.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[8], p)
+        elif 'block3_2.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[9], p)
+        elif 'block3_3.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[10], p)
+        elif 'block3_4.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[11], p)
+        elif 'block4_1.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[12], p)
+        elif 'block4_2.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[13], p)
+        elif 'block4_3.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[14], p)
+        elif 'block4_4.conv1.weight' in pname:
+            new_dict[pname] = convert_pdc(pdcs[15], p)
+        else:
+            new_dict[pname] = p
+    return new_dict
+
+
+class Conv2d(nn.Module):
+    def __init__(self,
+                 pdc,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=False):
+        super().__init__()
+        if in_channels % groups != 0:
+            raise ValueError('in_channels must be divisible by groups')
+        if out_channels % groups != 0:
+            raise ValueError('out_channels must be divisible by groups')
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels // groups, kernel_size,
+                         kernel_size))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+        self.pdc = pdc
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in)
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input):
+        return self.pdc(input, self.weight, self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+
+class CSAM(nn.Module):
+    """
+    Compact Spatial Attention Module
+    """
+    def __init__(self, channels):
+        super().__init__()
+
+        mid_channels = 4
+        self.relu1 = nn.ReLU()
+        self.conv1 = nn.Conv2d(channels,
+                               mid_channels,
+                               kernel_size=1,
+                               padding=0)
+        self.conv2 = nn.Conv2d(mid_channels,
+                               1,
+                               kernel_size=3,
+                               padding=1,
+                               bias=False)
+        self.sigmoid = nn.Sigmoid()
+        nn.init.constant_(self.conv1.bias, 0)
+
+    def forward(self, x):
+        y = self.relu1(x)
+        y = self.conv1(y)
+        y = self.conv2(y)
+        y = self.sigmoid(y)
+
+        return x * y
+
+
+class CDCM(nn.Module):
+    """
+    Compact Dilation Convolution based Module
+    """
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        self.relu1 = nn.ReLU()
+        self.conv1 = nn.Conv2d(in_channels,
+                               out_channels,
+                               kernel_size=1,
+                               padding=0)
+        self.conv2_1 = nn.Conv2d(out_channels,
+                                 out_channels,
+                                 kernel_size=3,
+                                 dilation=5,
+                                 padding=5,
+                                 bias=False)
+        self.conv2_2 = nn.Conv2d(out_channels,
+                                 out_channels,
+                                 kernel_size=3,
+                                 dilation=7,
+                                 padding=7,
+                                 bias=False)
+        self.conv2_3 = nn.Conv2d(out_channels,
+                                 out_channels,
+                                 kernel_size=3,
+                                 dilation=9,
+                                 padding=9,
+                                 bias=False)
+        self.conv2_4 = nn.Conv2d(out_channels,
+                                 out_channels,
+                                 kernel_size=3,
+                                 dilation=11,
+                                 padding=11,
+                                 bias=False)
+        nn.init.constant_(self.conv1.bias, 0)
+
+    def forward(self, x):
+        x = self.relu1(x)
+        x = self.conv1(x)
+        x1 = self.conv2_1(x)
+        x2 = self.conv2_2(x)
+        x3 = self.conv2_3(x)
+        x4 = self.conv2_4(x)
+        return x1 + x2 + x3 + x4
+
+
+class MapReduce(nn.Module):
+    """
+    Reduce feature maps into a single edge map
+    """
+    def __init__(self, channels):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, 1, kernel_size=1, padding=0)
+        nn.init.constant_(self.conv.bias, 0)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class PDCBlock(nn.Module):
+    def __init__(self, pdc, inplane, ouplane, stride=1):
+        super().__init__()
+        self.stride = stride
+
+        self.stride = stride
+        if self.stride > 1:
+            self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.shortcut = nn.Conv2d(inplane,
+                                      ouplane,
+                                      kernel_size=1,
+                                      padding=0)
+        self.conv1 = Conv2d(pdc,
+                            inplane,
+                            inplane,
+                            kernel_size=3,
+                            padding=1,
+                            groups=inplane,
+                            bias=False)
+        self.relu2 = nn.ReLU()
+        self.conv2 = nn.Conv2d(inplane,
+                               ouplane,
+                               kernel_size=1,
+                               padding=0,
+                               bias=False)
+
+    def forward(self, x):
+        if self.stride > 1:
+            x = self.pool(x)
+        y = self.conv1(x)
+        y = self.relu2(y)
+        y = self.conv2(y)
+        if self.stride > 1:
+            x = self.shortcut(x)
+        y = y + x
+        return y
+
+
+class PDCBlock_converted(nn.Module):
+    """
+    CPDC, APDC can be converted to vanilla 3x3 convolution
+    RPDC can be converted to vanilla 5x5 convolution
+    """
+    def __init__(self, pdc, inplane, ouplane, stride=1):
+        super().__init__()
+        self.stride = stride
+
+        if self.stride > 1:
+            self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.shortcut = nn.Conv2d(inplane,
+                                      ouplane,
+                                      kernel_size=1,
+                                      padding=0)
+        if pdc == 'rd':
+            self.conv1 = nn.Conv2d(inplane,
+                                   inplane,
+                                   kernel_size=5,
+                                   padding=2,
+                                   groups=inplane,
+                                   bias=False)
+        else:
+            self.conv1 = nn.Conv2d(inplane,
+                                   inplane,
+                                   kernel_size=3,
+                                   padding=1,
+                                   groups=inplane,
+                                   bias=False)
+        self.relu2 = nn.ReLU()
+        self.conv2 = nn.Conv2d(inplane,
+                               ouplane,
+                               kernel_size=1,
+                               padding=0,
+                               bias=False)
+
+    def forward(self, x):
+        if self.stride > 1:
+            x = self.pool(x)
+        y = self.conv1(x)
+        y = self.relu2(y)
+        y = self.conv2(y)
+        if self.stride > 1:
+            x = self.shortcut(x)
+        y = y + x
+        return y
+
+
+class PiDiNet(nn.Module):
+    def __init__(self,
+                 inplane,
+                 pdcs,
+                 dil=None,
+                 sa=False,
+                 convert=False,
+                 mean=[0.485, 0.456, 0.406],
+                 std=[0.229, 0.224, 0.225]):
+        super().__init__()
+        self.sa = sa
+        if dil is not None:
+            assert isinstance(dil, int), 'dil should be an int'
+        self.dil = dil
+        self.mean = mean
+        self.std = std
+
+        self.fuseplanes = []
+
+        self.inplane = inplane
+        if convert:
+            if pdcs[0] == 'rd':
+                init_kernel_size = 5
+                init_padding = 2
+            else:
+                init_kernel_size = 3
+                init_padding = 1
+            self.init_block = nn.Conv2d(3,
+                                        self.inplane,
+                                        kernel_size=init_kernel_size,
+                                        padding=init_padding,
+                                        bias=False)
+            block_class = PDCBlock_converted
+        else:
+            self.init_block = Conv2d(pdcs[0],
+                                     3,
+                                     self.inplane,
+                                     kernel_size=3,
+                                     padding=1)
+            block_class = PDCBlock
+
+        self.block1_1 = block_class(pdcs[1], self.inplane, self.inplane)
+        self.block1_2 = block_class(pdcs[2], self.inplane, self.inplane)
+        self.block1_3 = block_class(pdcs[3], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane)  # C
+
+        inplane = self.inplane
+        self.inplane = self.inplane * 2
+        self.block2_1 = block_class(pdcs[4], inplane, self.inplane, stride=2)
+        self.block2_2 = block_class(pdcs[5], self.inplane, self.inplane)
+        self.block2_3 = block_class(pdcs[6], self.inplane, self.inplane)
+        self.block2_4 = block_class(pdcs[7], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane)  # 2C
+
+        inplane = self.inplane
+        self.inplane = self.inplane * 2
+        self.block3_1 = block_class(pdcs[8], inplane, self.inplane, stride=2)
+        self.block3_2 = block_class(pdcs[9], self.inplane, self.inplane)
+        self.block3_3 = block_class(pdcs[10], self.inplane, self.inplane)
+        self.block3_4 = block_class(pdcs[11], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane)  # 4C
+
+        self.block4_1 = block_class(pdcs[12],
+                                    self.inplane,
+                                    self.inplane,
+                                    stride=2)
+        self.block4_2 = block_class(pdcs[13], self.inplane, self.inplane)
+        self.block4_3 = block_class(pdcs[14], self.inplane, self.inplane)
+        self.block4_4 = block_class(pdcs[15], self.inplane, self.inplane)
+        self.fuseplanes.append(self.inplane)  # 4C
+
+        self.conv_reduces = nn.ModuleList()
+        if self.sa and self.dil is not None:
+            self.attentions = nn.ModuleList()
+            self.dilations = nn.ModuleList()
+            for i in range(4):
+                self.dilations.append(CDCM(self.fuseplanes[i], self.dil))
+                self.attentions.append(CSAM(self.dil))
+                self.conv_reduces.append(MapReduce(self.dil))
+        elif self.sa:
+            self.attentions = nn.ModuleList()
+            for i in range(4):
+                self.attentions.append(CSAM(self.fuseplanes[i]))
+                self.conv_reduces.append(MapReduce(self.fuseplanes[i]))
+        elif self.dil is not None:
+            self.dilations = nn.ModuleList()
+            for i in range(4):
+                self.dilations.append(CDCM(self.fuseplanes[i], self.dil))
+                self.conv_reduces.append(MapReduce(self.dil))
+        else:
+            for i in range(4):
+                self.conv_reduces.append(MapReduce(self.fuseplanes[i]))
+
+        self.classifier = nn.Conv2d(4, 1, kernel_size=1)  # has bias
+        nn.init.constant_(self.classifier.weight, 0.25)
+        nn.init.constant_(self.classifier.bias, 0)
+
+    def get_weights(self):
+        conv_weights = []
+        bn_weights = []
+        relu_weights = []
+        for pname, p in self.named_parameters():
+            if 'bn' in pname:
+                bn_weights.append(p)
+            elif 'relu' in pname:
+                relu_weights.append(p)
+            else:
+                conv_weights.append(p)
+
+        return conv_weights, bn_weights, relu_weights
+
+    def forward(self, x):
+        """x: [B, 3, H, W] within range [0, 1].
+        """
+        x = (x - x.new_tensor(self.mean).view(1, -1, 1, 1)) / \
+            x.new_tensor(self.std).view(1, -1, 1, 1)
+        h, w = x.size()[2:]
+
+        x = self.init_block(x)
+
+        x1 = self.block1_1(x)
+        x1 = self.block1_2(x1)
+        x1 = self.block1_3(x1)
+
+        x2 = self.block2_1(x1)
+        x2 = self.block2_2(x2)
+        x2 = self.block2_3(x2)
+        x2 = self.block2_4(x2)
+
+        x3 = self.block3_1(x2)
+        x3 = self.block3_2(x3)
+        x3 = self.block3_3(x3)
+        x3 = self.block3_4(x3)
+
+        x4 = self.block4_1(x3)
+        x4 = self.block4_2(x4)
+        x4 = self.block4_3(x4)
+        x4 = self.block4_4(x4)
+
+        x_fuses = []
+        if self.sa and self.dil is not None:
+            for i, xi in enumerate([x1, x2, x3, x4]):
+                x_fuses.append(self.attentions[i](self.dilations[i](xi)))
+        elif self.sa:
+            for i, xi in enumerate([x1, x2, x3, x4]):
+                x_fuses.append(self.attentions[i](xi))
+        elif self.dil is not None:
+            for i, xi in enumerate([x1, x2, x3, x4]):
+                x_fuses.append(self.dilations[i](xi))
+        else:
+            x_fuses = [x1, x2, x3, x4]
+
+        e1 = self.conv_reduces[0](x_fuses[0])
+        e1 = F.interpolate(e1, (h, w), mode='bilinear', align_corners=False)
+
+        e2 = self.conv_reduces[1](x_fuses[1])
+        e2 = F.interpolate(e2, (h, w), mode='bilinear', align_corners=False)
+
+        e3 = self.conv_reduces[2](x_fuses[2])
+        e3 = F.interpolate(e3, (h, w), mode='bilinear', align_corners=False)
+
+        e4 = self.conv_reduces[3](x_fuses[3])
+        e4 = F.interpolate(e4, (h, w), mode='bilinear', align_corners=False)
+
+        outputs = [e1, e2, e3, e4]
+        output = self.classifier(torch.cat(outputs, dim=1))
+
+        outputs.append(output)
+        outputs = [torch.sigmoid(r) for r in outputs]
+        return outputs[-1]
+
+
+@ANNOTATORS.register_class()
+class PiDiAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        pretrained_model = cfg.get('PRETRAINED_MODEL', None)
+        vanilla_cnn = cfg.get('VANILLA_CNN', True)
+        pdcs = config_model_converted(
+            'carv4') if vanilla_cnn else config_model('carv4')
+        self.model = PiDiNet(60, pdcs, dil=24, sa=True,
+                             convert=vanilla_cnn).eval()
+        if pretrained_model:
+            with FS.get_from(pretrained_model, wait_finish=True) as local_path:
+                state = torch.load(local_path,
+                                   map_location='cpu')['state_dict']
+                if vanilla_cnn:
+                    state = convert_pidinet(state, 'carv4')
+                state = {
+                    k[len('module.'):] if k.startswith('module.') else k: v
+                    for k, v in state.items()
+                }
+                self.model.load_state_dict(state)
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    @torch.autocast('cuda', enabled=False)
+    def forward(self, image, return_grayscale=False):
+        is_batch = False if len(image.shape) == 3 else True
+        if isinstance(image, torch.Tensor):
+            if len(image.shape) == 3:
+                image = rearrange(image, 'h w c -> 1 c h w')
+            elif len(image.shape) == 4:
+                image = rearrange(image, 'b h w c -> b c h w')
+            else:
+                raise "Unsurpport input image's shape"
+        elif isinstance(image, np.ndarray):
+            image = torch.from_numpy(image.copy()).float()
+            if len(image.shape) == 3:
+                image = rearrange(image, 'h w c -> 1 c h w')
+            elif len(image.shape) == 4:
+                image = rearrange(image, 'b h w c -> b c h w')
+            else:
+                raise "Unsurpport input image's shape"
+        else:
+            raise "Unsurpport input image's type"
+        image = image.float().div(255)
+        image = image.to(we.device_id)
+        edge = self.model(image)
+        edge = edge.squeeze(dim=1)
+        edge = 255 - (edge * 255.0).clip(0, 255)  # return white background
+        edge = edge.cpu().numpy()
+        edge = edge.astype(np.uint8)
+        if not is_batch:
+            edge = edge.squeeze()
+        if not return_grayscale:
+            edge = edge[..., None].repeat(3, -1)
+        return edge
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            PiDiAnnotator.para_dict,
+                            set_name=True)
@@ -0,0 +1,382 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import math
+import random
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from scipy import ndimage
+from pycocotools import mask as mask_utils
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+from sklearn.cluster import KMeans
+from torchvision.ops.boxes import batched_nms
+
+
+def find_dominant_color(image, k=1):
+    pixels = image.reshape((-1, 3))
+    mask = (pixels != [0, 0, 0]).all(axis=1)
+    pixels = pixels[mask]
+    try:
+        kmeans = KMeans(n_clusters=k, n_init='auto')
+        kmeans.fit(pixels)
+        dominant_color = kmeans.cluster_centers_.astype(int)[0]
+    except:
+        dominant_color = np.array([255, 255, 255])
+    return dominant_color
+
+
+def cv2_resize_crop(image, resize_size, crop_size):
+    resize_height, resize_width = resize_size
+    crop_height, crop_width = crop_size
+
+    resized_image = cv2.resize(image, (resize_width, resize_height))
+
+    center_x, center_y = resize_width // 2, resize_height // 2
+    crop_start_x = max(center_x - crop_width // 2, 0)
+    crop_start_y = max(center_y - crop_height // 2, 0)
+    crop_end_x = crop_start_x + crop_width
+    crop_end_y = crop_start_y + crop_height
+
+    crop_end_x = min(crop_end_x, resize_width)
+    crop_end_y = min(crop_end_y, resize_height)
+
+    center_cropped_image = resized_image[crop_start_y:crop_end_y,
+                                         crop_start_x:crop_end_x]
+
+    return center_cropped_image
+
+
+@ANNOTATORS.register_class()
+class ESAMAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        try:
+            from efficient_sam.efficient_sam import build_efficient_sam
+            from segment_anything.utils.amg import (
+                batched_mask_to_box,
+                calculate_stability_score,
+                mask_to_rle_pytorch,
+                remove_small_regions,
+                rle_to_mask,
+            )
+        except:
+            raise NotImplementedError(
+                f'Please install efficient_sam and segment_anything modules.')
+
+        pretrained_model = cfg.get('PRETRAINED_MODEL', None)
+        if pretrained_model:
+            with FS.get_from(pretrained_model, wait_finish=True) as local_path:
+                self.efficient_sam_module = build_efficient_sam(
+                    encoder_patch_embed_dim=384,
+                    encoder_num_heads=6,
+                    checkpoint=local_path).eval().to(we.device_id)
+        self.GRID_SIZE = cfg.get('GRID_SIZE', 16)
+        self.save_mode = cfg.get('SAVE_MODE', 'P')
+        self.use_dominant_color = cfg.get('USE_DOMINANT_COLOR', False)
+        self.return_mask = cfg.get('RETURN_MASK', False)
+
+    @torch.no_grad()
+    def get_predictions_given_embeddings_and_queries(self, img, points,
+                                                     point_labels, model):
+        from segment_anything.utils.amg import calculate_stability_score
+        predicted_masks, predicted_iou = [], []
+        bs = 128
+        num = int(float(self.GRID_SIZE * self.GRID_SIZE) /
+                  bs) if self.GRID_SIZE * self.GRID_SIZE % bs == 0 else int(
+                      float(self.GRID_SIZE * self.GRID_SIZE) / bs) + 1
+        for i in range(num):
+            predicted_mask_item, predicted_iou_item = model(
+                img[None, ...], points[:, i * bs:(i + 1) * bs, ...],
+                point_labels[:, i * bs:(i + 1) * bs, :])
+            predicted_masks.append(predicted_mask_item)
+            predicted_iou.append(predicted_iou_item)
+            torch.cuda.empty_cache()
+        # # predicted:  torch.Size([1, 1024, 3]) torch.Size([1, 1024, 3, 512, 512])
+        # print('predicted: ', predicted_iou.size(), predicted_masks.size())
+        predicted_masks = torch.cat(predicted_masks, dim=1)
+        predicted_iou = torch.cat(predicted_iou, dim=1)
+        sorted_ids = torch.argsort(predicted_iou, dim=-1, descending=True)
+        predicted_iou_scores = torch.take_along_dim(predicted_iou,
+                                                    sorted_ids,
+                                                    dim=2)
+        predicted_masks = torch.take_along_dim(predicted_masks,
+                                               sorted_ids[..., None, None],
+                                               dim=2)
+        predicted_masks = predicted_masks[0]
+        iou = predicted_iou_scores[0, :, 0]
+        index_iou = iou > 0.7
+        iou_ = iou[index_iou]
+        masks = predicted_masks[index_iou]
+        score = calculate_stability_score(masks, 0.0, 1.0)
+        score = score[:, 0]
+        index = score > 0.9
+        masks = masks[index]
+        iou_ = iou_[index]
+        masks = torch.ge(masks, 0.0)
+        return masks, iou_
+
+    def singel_mask_to_rle(self, mask):
+        rle = mask_utils.encode(
+            np.array(mask[:, :, None], order='F', dtype='uint8'))[0]
+        rle['counts'] = rle['counts'].decode('utf-8')
+        return rle
+
+    def process_small_region(self, rles):
+        from segment_anything.utils.amg import rle_to_mask, remove_small_regions, \
+            batched_mask_to_box, mask_to_rle_pytorch
+        new_masks = []
+        scores = []
+        min_area = 100
+        nms_thresh = 0.7
+        for rle in rles:
+            mask = rle_to_mask(rle[0])
+
+            mask, changed = remove_small_regions(mask, min_area, mode='holes')
+            unchanged = not changed
+            mask, changed = remove_small_regions(mask,
+                                                 min_area,
+                                                 mode='islands')
+            unchanged = unchanged and not changed
+
+            new_masks.append(torch.as_tensor(mask).unsqueeze(0))
+            # Give score=0 to changed masks and score=1 to unchanged masks
+            # so NMS will prefer ones that didn't need postprocessing
+            scores.append(float(unchanged))
+
+        # Recalculate boxes and remove any new duplicates
+        masks = torch.cat(new_masks, dim=0).to(we.device_id)
+        boxes = batched_mask_to_box(masks)
+        keep_by_nms = batched_nms(
+            boxes.float(),
+            torch.as_tensor(scores).to(we.device_id),
+            torch.zeros_like(boxes[:, 0]),  # categories
+            iou_threshold=nms_thresh,
+        )
+
+        # Only recalculate RLEs for masks that have changed
+        for i_mask in keep_by_nms:
+            if scores[i_mask] == 0.0:
+                mask_torch = masks[i_mask].unsqueeze(0)
+                rles[i_mask] = mask_to_rle_pytorch(mask_torch)
+        masks = [rle_to_mask(rles[i][0]) for i in keep_by_nms]
+        return masks
+
+    def run_everything_ours(self, img_tensor, model):
+        from segment_anything.utils.amg import mask_to_rle_pytorch
+        img_tensor = img_tensor.squeeze(0)
+        _, original_image_h, original_image_w = img_tensor.shape
+        xy = []
+        for i in range(self.GRID_SIZE):
+            curr_x = 0.5 + i / self.GRID_SIZE * original_image_w
+            for j in range(self.GRID_SIZE):
+                curr_y = 0.5 + j / self.GRID_SIZE * original_image_h
+                xy.append([curr_x, curr_y])
+
+        xy = torch.from_numpy(np.array(xy))
+        points = xy
+        num_pts = xy.shape[0]
+        point_labels = torch.ones(num_pts, 1)
+        with torch.no_grad():
+            predicted_masks, predicted_iou = self.get_predictions_given_embeddings_and_queries(
+                img_tensor,
+                points.reshape(1, num_pts, 1, 2).to(we.device_id),
+                point_labels.reshape(1, num_pts, 1).to(we.device_id),
+                model,
+            )
+        # print('predicted_masks: ', predicted_masks[0][0:1].dtype, predicted_masks[0][0:1].device)
+        rle = [mask_to_rle_pytorch(m[0:1]) for m in predicted_masks]
+        # transform to numpy
+        size, counts = [], []
+        for rle_item in rle:
+            size.append(rle_item[0]['size'])
+            counts += rle_item[0]['counts']
+            counts += '#'
+        predicted_masks = self.process_small_region(rle)
+        return predicted_masks
+
+    def forward(self, image, return_mask=None):
+        return_mask = return_mask if return_mask is not None else self.return_mask
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        elif isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        elif isinstance(image, np.ndarray):
+            image = image.copy()
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+        h, w = image.shape[:2]
+        max_rate = max(float(w) / 1024.0, float(h) / 1024.0)
+        w_ori = int(float(w) / max_rate)
+        h_ori = int(float(h) / max_rate)
+        # image = T.ToTensor()(T.Resize((h_ori, w_ori))(Image.fromarray(image)))
+        # image_pad = T.Pad((0, 0, 1024 - w_ori, 1024 - h_ori))(image)
+        image_pad = T.Pad((0, 0, 1024 - w_ori, 1024 - h_ori))(T.Resize(
+            (h_ori, w_ori))(Image.fromarray(image)))
+        input_image = T.ToTensor()(image_pad)
+        input_image = input_image.unsqueeze(0).to(we.device_id)
+
+        mask_efficient_sam_vits = self.run_everything_ours(
+            input_image, self.efficient_sam_module)
+        annos = []
+        mask_efficient_sam_vits = sorted(list(mask_efficient_sam_vits),
+                                         key=lambda m: int(m.sum()),
+                                         reverse=True)
+        mask_efficient_sam_vits = mask_efficient_sam_vits[:256]
+        for mask in mask_efficient_sam_vits:
+            mask_item = mask_utils.encode(
+                np.array(mask[:, :, None], order='F', dtype='uint8'))[0]
+            mask_item['counts'] = mask_item['counts'].decode('utf-8')
+            mask_area = int(mask.sum())
+            annos.append({'mask': mask_item, 'mask_area': mask_area})
+
+        annos = sorted(annos, key=lambda x: x['mask_area'], reverse=True)
+        seg_img = None
+        dominant_palette = []
+        image_pad_np = np.array(image_pad)
+        for idx, anno in enumerate(annos):
+            color = idx
+            if idx > 255:
+                break
+            mask = np.array(mask_utils.decode(anno['mask'])).astype(np.uint8)
+            h, w = mask.shape[:2]
+            if seg_img is None:
+                seg_img = np.ones((h, w, 3)) * 255
+            if self.use_dominant_color:
+                masked_image = cv2.bitwise_and(image_pad_np,
+                                               image_pad_np,
+                                               mask=mask)
+                dominant_color = find_dominant_color(masked_image).tolist()
+                dominant_palette.append(dominant_color)
+            seg_img[mask.astype(bool)] = [color, color, color]
+        seg_img = Image.fromarray(seg_img.astype(np.uint8)).convert('L')
+
+        resize_rate = max(float(h_ori) / 1024.0, float(w_ori) / 1024.0)
+        h_new = int(float(h_ori) / resize_rate)
+        w_new = int(float(w_ori) / resize_rate)
+        seg_img = seg_img.crop((0, 0, w_new, h_new))
+        if self.save_mode == 'P':
+            palette = []
+            for i in range(256):
+                if not self.use_dominant_color:
+                    palette_item = [random.randint(0, 255) for _ in range(3)]
+                else:
+                    palette_item = dominant_palette[i] if i < len(
+                        dominant_palette) else [255, 255, 255]
+                palette += palette_item
+            seg_img = seg_img.convert('P')
+            seg_img.putpalette(palette)
+            seg_rgb_img = seg_img.convert('RGB')
+            if return_mask:
+                return {
+                    'image': np.array(seg_rgb_img),
+                    'mask': np.array(seg_img)
+                }
+            else:
+                return np.array(seg_rgb_img)
+        else:
+            return np.array(seg_img)
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            ESAMAnnotator.para_dict,
+                            set_name=True)
+
+
+
+
+@ANNOTATORS.register_class()
+class SAMAnnotatorDraw(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        from segment_anything import sam_model_registry, SamPredictor
+        from segment_anything.utils.transforms import ResizeLongestSide
+
+        self.transform = ResizeLongestSide(1024)
+        self.task_type = cfg.get('TASK_TYPE', 'input_box')
+        self.sam_model = cfg.get('SAM_MODEL', 'vit_b')
+        pretrained_model = cfg.get('PRETRAINED_MODEL', 'sam_vit_b_01ec64.pth')
+
+        if pretrained_model:
+            with FS.get_from(pretrained_model, wait_finish=True) as local_path:
+                seg_model = sam_model_registry[self.sam_model](checkpoint=local_path).eval().to(we.device_id)
+                self.sam_predictor = SamPredictor(seg_model)
+
+    def forward(self, image, input_box=None, mask=None, task_type=None, multimask_output=False):
+        task_type = task_type if task_type is not None else self.task_type
+
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        elif isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        elif isinstance(image, np.ndarray):
+            image = image.copy()
+        else:
+            raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+
+        if mask is not None:
+            if isinstance(mask, Image.Image):
+                mask = np.array(mask)
+            elif isinstance(mask, torch.Tensor):
+                mask = mask.detach().cpu().numpy()
+            elif isinstance(mask, np.ndarray):
+                mask = mask.copy()
+            else:
+                raise f'Unsurpport datatype{type(mask)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
+
+        original_size = image.shape[:2]
+        if task_type == 'mask_point':
+            scribble = mask.transpose(2, 1, 0)[0]
+            labeled_array, num_features = ndimage.label(scribble >= 255)
+            centers = ndimage.center_of_mass(scribble, labeled_array, range(1, num_features + 1))
+            point_coords = np.array(centers)
+            point_labels = np.array([1] * len(centers))
+            sample = {'point_coords': point_coords, 'point_labels': point_labels}
+
+        elif task_type == 'mask_box':
+            scribble = mask.transpose(2, 1, 0)[0]
+            labeled_array, num_features = ndimage.label(scribble >= 255)
+            centers = ndimage.center_of_mass(scribble, labeled_array, range(1, num_features + 1))
+            centers = np.array(centers)
+            ### (x1, y1, x2, y2)
+            x_min = centers[:, 0].min()
+            x_max = centers[:, 0].max()
+            y_min = centers[:, 1].min()
+            y_max = centers[:, 1].max()
+            bbox = np.array([x_min, y_min, x_max, y_max])
+            sample = {'box': bbox}
+
+        elif task_type == 'input_box':
+            if isinstance(input_box, list):
+                input_box = np.array(input_box)
+            sample = {'box': input_box}
+
+        self.sam_predictor.set_image(image)
+        masks, scores, logits = self.sam_predictor.predict(**sample, multimask_output=True)
+        index = np.argmax(scores)
+
+        ret_data = {
+            "mask": (masks[index]* 255).astype(np.uint8),
+            "score": scores[index]
+        }
+        return ret_data
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            SAMAnnotatorDraw.para_dict,
+                            set_name=True)
@@ -0,0 +1,160 @@
+# -*- coding: utf-8 -*-
+
+import math
+from abc import ABCMeta
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as TT
+from einops import rearrange
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+
+
+class SketchNet(nn.Module):
+    def __init__(self, mean, std):
+        assert isinstance(mean, float) and isinstance(std, float)
+        super().__init__()
+        self.mean = mean
+        self.std = std
+
+        # layers
+        self.layers = nn.Sequential(nn.Conv2d(1, 48, 5, 2, 2),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(48, 128, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(128, 128, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(128, 128, 3, 2, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(128, 256, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(256, 256, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(256, 256, 3, 2, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(256, 512, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(512, 1024, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(1024, 1024, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(1024, 1024, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(1024, 1024, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(1024, 512, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(512, 256, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.ConvTranspose2d(256, 256, 4, 2, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(256, 256, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(256, 128, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.ConvTranspose2d(128, 128, 4, 2, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(128, 128, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(128, 48, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.ConvTranspose2d(48, 48, 4, 2, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(48, 24, 3, 1, 1),
+                                    nn.ReLU(inplace=True),
+                                    nn.Conv2d(24, 1, 3, 1, 1), nn.Sigmoid())
+
+    def forward(self, x):
+        """x: [B, 1, H, W] within range [0, 1]. Sketch pixels in dark color.
+        """
+        x = (x - self.mean) / self.std
+        return self.layers(x)
+
+
+@ANNOTATORS.register_class()
+class SketchAnnotator(BaseAnnotator, metaclass=ABCMeta):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        pretrained_model = cfg.get('PRETRAINED_MODEL', None)
+        self.model = SketchNet(mean=0.9664114577640158,
+                               std=0.0858381272736797).eval()
+        if pretrained_model:
+            with FS.get_from(pretrained_model, wait_finish=True) as local_path:
+                state = torch.load(local_path, map_location='cpu')
+                self.model.load_state_dict(state)
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    @torch.autocast('cuda', enabled=False)
+    def forward(self, image):
+        is_batch = False if len(image.shape) == 3 else True
+        if isinstance(image, torch.Tensor):
+            if len(image.shape) == 3:
+                if torch.equal(image[:, :, 0], image[:, :, 1]) and torch.equal(
+                        image[:, :, 1], image[:, :, 2]):
+                    image = image[:, :, 0]
+                else:
+                    raise "Unsurpport input image's shape and each channel is different"
+            elif len(image.shape) == 4:
+                if (torch.equal(image[:, :, :, 0], image[:, :, :, 1])
+                        and torch.equal(image[:, :, :, 1], image[:, :, :, 2])):
+                    image = image[:, :, :, 0]
+                else:
+                    raise "Unsurpport input image's shape and each channel is different"
+            if len(image.shape) == 2:
+                image = rearrange(image, 'h w -> 1 h w')
+                B, H, W = image.shape
+            elif len(image.shape) == 3:
+                B, H, W = image.shape
+            else:
+                raise "Unsurpport input image's shape"
+        elif isinstance(image, np.ndarray):
+            image = image.copy()
+            if len(image.shape) == 3:
+                if np.array_equal(image[:, :, 0],
+                                  image[:, :, 1]) and np.array_equal(
+                                      image[:, :, 1], image[:, :, 2]):
+                    image = image[:, :, 0]
+                else:
+                    raise "Unsurpport input image's shape and each channel is different"
+            elif len(image.shape) == 4:
+                if (np.array_equal(image[:, :, :, 0], image[:, :, :, 1]) and
+                        np.array_equal(image[:, :, :, 1], image[:, :, :, 2])):
+                    image = image[:, :, :, 0]
+                else:
+                    raise "Unsurpport input image's shape and each channel is different"
+            image = torch.from_numpy(image).float()
+            if len(image.shape) == 2:
+                image = rearrange(image, 'h w -> 1 1 h w')
+            elif len(image.shape) == 3:
+                image = rearrange(image, 'b h w -> b 1 h w')
+            else:
+                raise "Unsurpport input image's shape"
+        else:
+            raise "Unsurpport input image's type"
+        image = image.float().div(255)
+        image = image.to(we.device_id)
+        edge = self.model(image)
+        edge = edge.squeeze(dim=1)
+        edge = (edge * 255.0).clip(0, 255)
+        edge = edge.cpu().numpy()
+        edge = edge.astype(np.uint8)
+        if not is_batch:
+            edge = edge.squeeze()
+        return edge[..., None].repeat(3, -1)
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('ANNOTATORS',
+                            __class__.__name__,
+                            SketchAnnotator.para_dict,
+                            set_name=True)
@@ -82,6 +82,8 @@ class BaseDataset(Dataset, metaclass=ABCMeta):
                                               overwrite=False)
        self.worker_id = worker_id
        self.logger = self.worker_logger
+        self.local_we["seed"] += (worker_id + we.rank)
+        self.seed = self.local_we["seed"]
        we.set_env(self.local_we)

    @abstractmethod
@@ -258,6 +258,8 @@ class DataObject(object):
            if sampler_name == 'MixtureOfSamplers':
                subsampler_configs = self.data_sampler_config.get(
                    'SUB_SAMPLERS', [])
+                keep_order = self.data_sampler_config.get(
+                    'KEEP_ORDER', False)
                subsamplers = list()
                subsampler_probs = list()
                for ssconfig in subsampler_configs:
@@ -277,7 +279,7 @@ class DataObject(object):
                    subsampler_probs.append(prob)
                self.batch_sampler = MixtureOfSamplers(subsamplers,
                                                       subsampler_probs, rank,
-                                                       seed)
+                                                       seed, keep_order = keep_order)
            elif sampler_name == 'MultiLevelBatchSampler':
                self.batch_sampler = self._instantiate_multi_level_batch_sampler(
                    self.data_sampler_config, self.batch_size, rank, seed)
@@ -523,12 +523,15 @@ class MultiLevelBatchSampler(BaseSampler):
 class MixtureOfSamplers(BaseSampler):
    para_dict = {'SUB_SAMPLERS': []}

-    def __init__(self, samplers, probabilities, rank=0, seed=8888):
+    def __init__(self, samplers, probabilities, rank=0, seed=8888, keep_order = False):
        self.samplers = samplers
        self.iterators = [iter(u) for u in samplers]
        self.probabilities = probabilities
        self.seed = seed
-        self.rng = np.random.default_rng(seed + rank)
+        if keep_order:
+            self.rng = np.random.default_rng(seed)
+        else:
+            self.rng = np.random.default_rng(seed + rank)

    def __iter__(self):
        while True:
@@ -13,12 +13,6 @@ from scepter.modules.utils.config import Config
 from scepter.modules.utils.distribute import we
 from scepter.modules.utils.file_system import FS

-try:
-    from swift import SwiftModel
-except Exception:
-    warnings.warn('Import swift failed, please check it.')
-
-
 class ControlInference():
    def __init__(self, logger=None):
        self.logger = logger
@@ -26,6 +20,11 @@ class ControlInference():

    # @classmethod
    def unregister_controllers(self, control_model_ins, diffusion_model):
+        try:
+            from swift import SwiftModel
+        except Exception:
+            warnings.warn('Import swift failed, please check it.')
+
        self.logger.info('Unloading control model')
        if isinstance(diffusion_model['model'], SwiftModel):
            if (hasattr(diffusion_model['model'].base_model, 'control_blocks')
@@ -42,6 +41,11 @@ class ControlInference():

    # @classmethod
    def register_controllers(self, control_model_ins, diffusion_model):
+        try:
+            from swift import SwiftModel
+        except Exception:
+            warnings.warn('Import swift failed, please check it.')
+
        self.logger.info('Loading control model')
        if control_model_ins is None or control_model_ins == '':
            self.unregister_controllers(control_model_ins, diffusion_model)
@@ -33,10 +33,12 @@ class DiffusionInference():
    '''
    def __init__(self, logger=None):
        self.logger = logger
+        self.is_redefine_paras = True
        self.loaded_model = {}
        self.loaded_model_name = [
            'diffusion_model', 'first_stage_model', 'cond_stage_model'
        ]
+        self.diffusion_insclass = GaussianDiffusion
        self.tuner_infer = TunerInference(self.logger)
        self.control_infer = ControlInference(self.logger)

@@ -45,7 +47,8 @@ class DiffusionInference():
        self.is_default = cfg.get('IS_DEFAULT', False)
        module_paras = self.load_default(cfg.get('DEFAULT_PARAS', None))
        assert cfg.have('MODEL')
-        cfg.MODEL = self.redefine_paras(cfg.MODEL)
+        if self.is_redefine_paras:
+            cfg.MODEL = self.redefine_paras(cfg.MODEL)
        self.diffusion = self.load_schedule(cfg.MODEL.SCHEDULE)
        self.diffusion_model = self.infer_model(
            cfg.MODEL.DIFFUSION_MODEL, module_paras.get(
@@ -81,14 +84,16 @@ class DiffusionInference():

    def redefine_paras(self, cfg):
        if cfg.get('PRETRAINED_MODEL', None):
-            assert FS.isfile(cfg.PRETRAINED_MODEL)
            with FS.get_from(cfg.PRETRAINED_MODEL,
                             wait_finish=True) as local_path:
                if local_path.endswith('safetensors'):
                    from safetensors.torch import load_file as load_safetensors
                    sd = load_safetensors(local_path)
                else:
-                    sd = torch.load(local_path, map_location='cpu')
+                    if 'weights_only' in torch.load.__code__.co_varnames:
+                        sd = torch.load(local_path, map_location='cpu', weights_only=True)
+                    else:
+                        sd = torch.load(local_path, map_location='cpu')
                first_stage_model_path = os.path.join(
                    os.path.dirname(local_path), 'first_stage_model.pth')
                cond_stage_model_path = os.path.join(
@@ -308,7 +313,7 @@ class DiffusionInference():
        module_paras = {}
        if cfg is not None:
            self.paras = cfg.PARAS
-            self.input = {k.lower(): v for k, v in cfg.INPUT.items()}
+            self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict)) else v for k, v in cfg.INPUT.items()}
            self.output = {k.lower(): v for k, v in cfg.OUTPUT.items()}
            module_paras = cfg.MODULES_PARAS
        return module_paras
@@ -316,8 +321,8 @@ class DiffusionInference():
    def load_schedule(self, cfg):
        parameterization = cfg.get('PARAMETERIZATION', 'eps')
        assert parameterization in [
-            'eps', 'x0', 'v'
-        ], 'currently only supporting "eps" and "x0" and "v"'
+            'eps', 'x0', 'v', 'rf'
+        ], 'currently only supporting "eps" and "x0" and "v" and "rf"'
        num_timesteps = cfg.get('TIMESTEPS', 1000)

        schedule_args = {
@@ -336,8 +341,8 @@ class DiffusionInference():
                                n=num_timesteps,
                                zero_terminal_snr=zero_terminal_snr,
                                **schedule_args)
-        diffusion = GaussianDiffusion(sigmas=sigmas,
-                                      prediction_type=parameterization)
+        diffusion = self.diffusion_insclass(sigmas=sigmas,
+                                            prediction_type=parameterization)
        return diffusion

    def get_batch(self, value_dict, num_samples=1):
@@ -0,0 +1,216 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import copy
+import math
+import random
+import torch
+from scepter.modules.utils.distribute import we
+from .control_inference import ControlInference
+from .diffusion_inference import DiffusionInference, get_model
+from .tuner_inference import TunerInference
+from scepter.modules.model.registry import DIFFUSIONS, TOKENIZERS
+
+
+class FluxInference(DiffusionInference):
+    def __init__(self, logger=None):
+        self.logger = logger
+        self.is_redefine_paras = False
+        self.loaded_model = {}
+        self.loaded_model_name = [
+            'diffusion_model', 'first_stage_model', 'cond_stage_model'
+        ]
+        self.tuner_infer = TunerInference(self.logger)
+        self.control_infer = ControlInference(self.logger)
+
+    def init_from_cfg(self, cfg):
+        self.name = cfg.NAME
+        self.is_default = cfg.get('IS_DEFAULT', False)
+        module_paras = self.load_default(cfg.get('DEFAULT_PARAS', None))
+        assert cfg.have('MODEL')
+        if self.is_redefine_paras:
+            cfg.MODEL = self.redefine_paras(cfg.MODEL)
+        self.diffusion_model = self.infer_model(
+            cfg.MODEL.DIFFUSION_MODEL, module_paras.get(
+                'DIFFUSION_MODEL',
+                None)) if cfg.MODEL.have('DIFFUSION_MODEL') else None
+        self.first_stage_model = self.infer_model(
+            cfg.MODEL.FIRST_STAGE_MODEL,
+            module_paras.get(
+                'FIRST_STAGE_MODEL',
+                None)) if cfg.MODEL.have('FIRST_STAGE_MODEL') else None
+        self.cond_stage_model = self.infer_model(
+            cfg.MODEL.COND_STAGE_MODEL,
+            module_paras.get(
+                'COND_STAGE_MODEL',
+                None)) if cfg.MODEL.have('COND_STAGE_MODEL') else None
+        self.refiner_cond_model = self.infer_model(
+            cfg.MODEL.REFINER_COND_MODEL,
+            module_paras.get(
+                'REFINER_COND_MODEL',
+                None)) if cfg.MODEL.have('REFINER_COND_MODEL') else None
+        self.refiner_diffusion_model = self.infer_model(
+            cfg.MODEL.REFINER_MODEL, module_paras.get(
+                'REFINER_MODEL',
+                None)) if cfg.MODEL.have('REFINER_MODEL') else None
+        self.tokenizer = TOKENIZERS.build(
+            cfg.MODEL.TOKENIZER,
+            logger=self.logger) if cfg.MODEL.have('TOKENIZER') else None
+
+        if self.tokenizer is not None:
+            self.cond_stage_model['cfg'].KWARGS = {
+                'vocab_size': self.tokenizer.vocab_size
+            }
+        self.diffusion = DIFFUSIONS.build(cfg.MODEL.DIFFUSION, logger=self.logger) \
+            if cfg.MODEL.have('DIFFUSION') else None
+        assert self.diffusion is not None
+
+    @torch.no_grad()
+    def encode_first_stage(self, x, **kwargs):
+        _, dtype = self.get_function_info(self.first_stage_model, 'encode')
+        with torch.autocast('cuda',
+                            enabled= dtype in ('float16', 'bfloat16'),
+                            dtype=getattr(torch, dtype)):
+            z = get_model(self.first_stage_model).encode(x)
+            if isinstance(z, (tuple, list)):
+                z = z[0]
+            return z
+
+    @torch.no_grad()
+    def decode_first_stage(self, z):
+        _, dtype = self.get_function_info(self.first_stage_model, 'decode')
+        with torch.autocast('cuda',
+                            enabled=dtype in ('float16', 'bfloat16'),
+                            dtype=getattr(torch, dtype)):
+            return get_model(self.first_stage_model).decode(z)
+
+    @torch.no_grad()
+    def __call__(self,
+                 input,
+                 num_samples=1,
+                 cat_uc=True,
+                 tuner_model=None,
+                 control_model=None,
+                 **kwargs):
+
+        value_input = copy.deepcopy(self.input)
+        value_input.update(input)
+        print(value_input)
+        height, width = value_input['target_size_as_tuple']
+        value_output = copy.deepcopy(self.output)
+        # register tuner
+        if tuner_model is not None and tuner_model != '' and len(
+                tuner_model) > 0:
+            if not isinstance(tuner_model, list):
+                tuner_model = [tuner_model]
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            self.tuner_infer.register_tuner(tuner_model, self.diffusion_model,
+                                            self.cond_stage_model)
+            self.dynamic_unload(self.diffusion_model,
+                                'diffusion_model',
+                                skip_loaded=True)
+            self.dynamic_unload(self.cond_stage_model,
+                                'cond_stage_model',
+                                skip_loaded=True)
+
+        # cond stage
+        self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+        function_name, dtype = self.get_function_info(self.cond_stage_model)
+        with torch.autocast('cuda',
+                            enabled=dtype == 'float16',
+                            dtype=getattr(torch, dtype)):
+            ctx = getattr(get_model(self.cond_stage_model),
+                                  function_name)(value_input['prompt'])
+        self.dynamic_unload(self.cond_stage_model,
+                            'cond_stage_model',
+                            skip_loaded=True)
+
+        # get noise
+        seed = kwargs.pop('seed', -1)
+        g = torch.Generator(device=we.device_id)
+        seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
+        g.manual_seed(seed)
+        if 'seed' in value_output:
+            value_output['seed'] = seed
+        for sample_id in range(num_samples):
+            if self.diffusion_model is not None:
+                noise = torch.randn(
+                    num_samples,
+                    16,
+                    # allow for packing
+                    2 * math.ceil(height / 16),
+                    2 * math.ceil(width / 16),
+                    device=we.device_id,
+                    dtype=getattr(torch, dtype),
+                    generator=torch.Generator(device=we.device_id).manual_seed(seed),
+                )
+                self.dynamic_load(self.diffusion_model, 'diffusion_model')
+                # UNet use input n_prompt
+                function_name, dtype = self.get_function_info(
+                    self.diffusion_model)
+                with torch.autocast('cuda',
+                                    enabled= dtype in ('float16', 'bfloat16'),
+                                    dtype=getattr(torch, dtype)):
+                    solver_sample = value_input.get('sample', 'flow_eluer')
+                    sample_steps = value_input.get('sample_steps', 20)
+                    guide_scale = value_input.get('guide_scale', 3.5)
+                    if guide_scale is not None:
+                        guide_scale = torch.full((noise.shape[0],), guide_scale, device=noise.device,
+                                                 dtype=noise.dtype)
+                    else:
+                        guide_scale = None
+
+                    latent = self.diffusion.sample(
+                        noise=noise,
+                        sampler=solver_sample,
+                        model=get_model(self.diffusion_model),
+                        model_kwargs={"cond": ctx, "guidance": guide_scale},
+                        steps=sample_steps,
+                        show_progress=True,
+                        guide_scale=guide_scale,
+                        return_intermediate=None,
+                        **kwargs)
+
+                self.dynamic_unload(self.diffusion_model,
+                                    'diffusion_model',
+                                    skip_loaded=True)
+
+            if 'latent' in value_output:
+                if value_output['latent'] is None or (
+                        isinstance(value_output['latent'], list)
+                        and len(value_output['latent']) < 1):
+                    value_output['latent'] = []
+                value_output['latent'].append(latent)
+
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            x_samples = self.decode_first_stage(latent).float()
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=True)
+            images = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+            if 'images' in value_output:
+                if value_output['images'] is None or (
+                        isinstance(value_output['images'], list)
+                        and len(value_output['images']) < 1):
+                    value_output['images'] = []
+                value_output['images'].append(images)
+
+        for k, v in value_output.items():
+            if isinstance(v, list):
+                value_output[k] = torch.cat(v, dim=0)
+            if isinstance(v, torch.Tensor):
+                value_output[k] = v.cpu()
+
+        # unregister tuner
+        if tuner_model is not None and tuner_model != '' and len(
+                tuner_model) > 0:
+            self.tuner_infer.unregister_tuner(tuner_model,
+                                              self.diffusion_model,
+                                              self.cond_stage_model)
+
+        # unregister control
+        if control_model is not None and control_model != '':
+            self.control_infer.unregister_controllers(control_model,
+                                                      self.diffusion_model)
+
+        return value_output
@@ -8,16 +8,12 @@ from collections import OrderedDict
 import gradio as gr
 import torch
 import torchvision.transforms.functional as TF
+from scepter.modules.model.network.diffusion.diffusion import GaussianDiffusion
 from scepter.modules.model.utils.data_utils import crop_back
 from scepter.modules.utils.distribute import we
 from scepter.modules.utils.file_system import FS

-from .diffusion_inference import DiffusionInference
-
-
-def get_model(model_tuple):
-    assert 'model' in model_tuple
-    return model_tuple['model']
+from .diffusion_inference import DiffusionInference, get_model


 class LargenInference(DiffusionInference):
@@ -28,10 +24,12 @@ class LargenInference(DiffusionInference):
    '''
    def __init__(self, logger=None):
        self.logger = logger
+        self.is_redefine_paras = True
        self.loaded_model = {}
        self.loaded_model_name = [
            'diffusion_model', 'first_stage_model', 'cond_stage_model'
        ]
+        self.diffusion_insclass = GaussianDiffusion

    def redefine_paras(self, cfg):
        if cfg.get('PRETRAINED_MODEL', None):
@@ -42,8 +40,10 @@ class LargenInference(DiffusionInference):
                    from safetensors.torch import load_file as load_safetensors
                    sd = load_safetensors(local_path)
                else:
-                    sd = torch.load(local_path, map_location='cpu')
-
+                    if 'weights_only' in torch.load.__code__.co_varnames:
+                        sd = torch.load(local_path, map_location='cpu', weights_only=True)
+                    else:
+                        sd = torch.load(local_path, map_location='cpu')
                if 'model' in sd:
                    sd = sd['model']

@@ -0,0 +1,211 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import copy
+import os.path
+import random
+from collections import OrderedDict
+
+import torch
+import torch.nn.functional as F
+from PIL.Image import Image
+from scepter.modules.model.network.diffusion.diffusion import GaussianDiffusion
+from scepter.modules.model.network.diffusion.schedules import noise_schedule
+from scepter.modules.model.registry import (BACKBONES, EMBEDDERS, MODELS,
+                                            TOKENIZERS)
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+from scepter.studio.utils.env import get_available_memory
+
+from .control_inference import ControlInference
+from .diffusion_inference import DiffusionInference, get_model
+from .tuner_inference import TunerInference
+
+
+class PixArtInference(DiffusionInference):
+    '''
+        define vae, unet, text-encoder, tuner, refiner components
+        support to load the components dynamicly.
+        create and load model when run this model at the first time.
+    '''
+    def __init__(self, logger=None):
+        self.logger = logger
+        self.is_redefine_paras = False
+        self.loaded_model = {}
+        self.loaded_model_name = [
+            'diffusion_model', 'first_stage_model', 'cond_stage_model'
+        ]
+        self.diffusion_insclass = GaussianDiffusion
+        self.tuner_infer = TunerInference(self.logger)
+        self.control_infer = ControlInference(self.logger)
+
+    @torch.no_grad()
+    def __call__(self,
+                 input,
+                 num_samples=1,
+                 intermediate_callback=None,
+                 refine_strength=0,
+                 img_to_img_strength=0,
+                 cat_uc=True,
+                 tuner_model=None,
+                 control_model=None,
+                 **kwargs):
+
+        value_input = copy.deepcopy(self.input)
+        value_input.update(input)
+        print(value_input)
+        height, width = value_input['target_size_as_tuple']
+        value_output = copy.deepcopy(self.output)
+
+        # register tuner
+        if tuner_model is not None and tuner_model != '' and len(
+                tuner_model) > 0:
+            if not isinstance(tuner_model, list):
+                tuner_model = [tuner_model]
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            self.tuner_infer.register_tuner(tuner_model, self.diffusion_model,
+                                            self.cond_stage_model)
+            self.dynamic_unload(self.diffusion_model,
+                                'diffusion_model',
+                                skip_loaded=True)
+            self.dynamic_unload(self.cond_stage_model,
+                                'cond_stage_model',
+                                skip_loaded=True)
+
+        # cond stage
+        self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+        function_name, dtype = self.get_function_info(self.cond_stage_model)
+        with torch.autocast('cuda',
+                            enabled=dtype == 'float16',
+                            dtype=getattr(torch, dtype)):
+            context, null_context = {}, {}
+            cont_mask = None
+            cont, cont_mask = getattr(get_model(self.cond_stage_model),
+                                      function_name)(value_input['prompt'],
+                                                     return_mask=True)
+            context['crossattn'] = cont.float()
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+            null_context['crossattn'] = get_model(
+                self.diffusion_model).y_embedder.y_embedding[None].repeat(
+                    num_samples, 1, 1)
+            self.dynamic_unload(self.diffusion_model, 'diffusion_model')
+        self.dynamic_unload(self.cond_stage_model,
+                            'cond_stage_model',
+                            skip_loaded=True)
+
+        # get noise
+        seed = kwargs.pop('seed', -1)
+        g = torch.Generator(device=we.device_id)
+        seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
+        g.manual_seed(seed)
+        if 'seed' in value_output:
+            value_output['seed'] = seed
+        for sample_id in range(num_samples):
+            if self.diffusion_model is not None:
+                noise = torch.empty(
+                    1,
+                    4,
+                    height // self.first_stage_model['paras']['size_factor'],
+                    width // self.first_stage_model['paras']['size_factor'],
+                    device=we.device_id).normal_(generator=g)
+
+                self.dynamic_load(self.diffusion_model, 'diffusion_model')
+                # UNet use input n_prompt
+                function_name, dtype = self.get_function_info(
+                    self.diffusion_model)
+                with torch.autocast('cuda',
+                                    enabled=dtype == 'float16',
+                                    dtype=getattr(torch, dtype)):
+                    latent = self.diffusion.sample(
+                        solver=value_input.get('sample', 'ddim'),
+                        noise=noise,
+                        model=get_model(self.diffusion_model),
+                        model_kwargs=[{
+                            'cond': context,
+                            'mask': cont_mask,
+                            'data_info': {
+                                'img_hw':
+                                torch.tensor([[height, width]],
+                                             dtype=torch.float,
+                                             device=we.device_id).repeat(
+                                                 num_samples, 1),
+                                'aspect_ratio':
+                                torch.tensor([[1.]],
+                                             device=we.device_id).repeat(
+                                                 num_samples, 1)
+                            }
+                        }, {
+                            'cond': null_context,
+                            'mask': cont_mask,
+                            'data_info': {
+                                'img_hw':
+                                torch.tensor([[height, width]],
+                                             dtype=torch.float,
+                                             device=we.device_id).repeat(
+                                                 num_samples, 1),
+                                'aspect_ratio':
+                                torch.tensor([[1.]],
+                                             device=we.device_id).repeat(
+                                                 num_samples, 1)
+                            }
+                        }],
+                        cat_uc=False,
+                        steps=value_input.get('sample_steps', 50),
+                        guide_scale=value_input.get('guide_scale', 7.5),
+                        guide_rescale=value_input.get('guide_rescale', 0.5),
+                        discretization=value_input.get('discretization',
+                                                       'trailing'),
+                        show_progress=True,
+                        seed=seed,
+                        condition_fn=None,
+                        clamp=None,
+                        percentile=None,
+                        t_max=None,
+                        t_min=None,
+                        discard_penultimate_step=None,
+                        return_intermediate=None,
+                        **kwargs)
+
+                self.dynamic_unload(self.diffusion_model,
+                                    'diffusion_model',
+                                    skip_loaded=True)
+
+            if 'latent' in value_output:
+                if value_output['latent'] is None or (
+                        isinstance(value_output['latent'], list)
+                        and len(value_output['latent']) < 1):
+                    value_output['latent'] = []
+                value_output['latent'].append(latent)
+
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            x_samples = self.decode_first_stage(latent).float()
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=True)
+            images = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+            if 'images' in value_output:
+                if value_output['images'] is None or (
+                        isinstance(value_output['images'], list)
+                        and len(value_output['images']) < 1):
+                    value_output['images'] = []
+                value_output['images'].append(images)
+
+        for k, v in value_output.items():
+            if isinstance(v, list):
+                value_output[k] = torch.cat(v, dim=0)
+            if isinstance(v, torch.Tensor):
+                value_output[k] = v.cpu()
+
+        # unregister tuner
+        if tuner_model is not None and tuner_model != '' and len(
+                tuner_model) > 0:
+            self.tuner_infer.unregister_tuner(tuner_model,
+                                              self.diffusion_model,
+                                              self.cond_stage_model)
+
+        # unregister control
+        if control_model is not None and control_model != '':
+            self.control_infer.unregister_controllers(control_model,
+                                                      self.diffusion_model)
+
+        return value_output
@@ -0,0 +1,189 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import copy
+import random
+
+import gradio as gr
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from scepter.modules.model.network.diffusion.diffusion import \
+    GaussianDiffusionRF
+from scepter.modules.utils.distribute import we
+
+from .control_inference import ControlInference
+from .diffusion_inference import DiffusionInference, get_model
+from .tuner_inference import TunerInference
+
+
+class SD3Inference(DiffusionInference):
+    def __init__(self, logger=None):
+        self.logger = logger
+        self.is_redefine_paras = False
+        self.loaded_model = {}
+        self.loaded_model_name = [
+            'diffusion_model', 'first_stage_model', 'cond_stage_model'
+        ]
+        self.diffusion_insclass = GaussianDiffusionRF
+        self.tuner_infer = TunerInference(self.logger)
+        self.control_infer = ControlInference(self.logger)
+
+    @torch.no_grad()
+    def decode_first_stage(self, z):
+        _, dtype = self.get_function_info(self.first_stage_model, 'decode')
+        # with torch.autocast('cuda', enabled=dtype == 'float16', dtype=getattr(torch, dtype)):
+        z = 1. / self.first_stage_model['paras'][
+            'scale_factor'] * z + self.first_stage_model['paras'][
+                'shift_factor']
+        return get_model(self.first_stage_model).decode(z)
+
+    @torch.no_grad()
+    def __call__(self,
+                 input,
+                 num_samples=1,
+                 cat_uc=True,
+                 tuner_model=None,
+                 control_model=None,
+                 **kwargs):
+
+        value_input = copy.deepcopy(self.input)
+        value_input.update(input)
+        print(value_input)
+        height, width = value_input['target_size_as_tuple']
+        value_output = copy.deepcopy(self.output)
+
+        # register tuner
+        if tuner_model is not None and tuner_model != '' and len(
+                tuner_model) > 0:
+            if not isinstance(tuner_model, list):
+                tuner_model = [tuner_model]
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            self.tuner_infer.register_tuner(tuner_model, self.diffusion_model,
+                                            self.cond_stage_model)
+            self.dynamic_unload(self.diffusion_model,
+                                'diffusion_model',
+                                skip_loaded=True)
+            self.dynamic_unload(self.cond_stage_model,
+                                'cond_stage_model',
+                                skip_loaded=True)
+
+        # cond stage
+        self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+        function_name, dtype = self.get_function_info(self.cond_stage_model)
+        context, null_context = {}, {}
+        with torch.autocast('cuda',
+                            enabled=dtype == 'float16',
+                            dtype=getattr(torch, dtype)):
+            ctx, pooled = getattr(get_model(self.cond_stage_model),
+                                  function_name)(value_input['prompt'])
+            null_ctx, null_pooled = getattr(get_model(self.cond_stage_model),
+                                            function_name)([''] * num_samples)
+            context['crossattn'] = ctx.float()
+            context['y'] = pooled.float()
+            null_context['crossattn'] = null_ctx.float()
+            null_context['y'] = null_pooled.float()
+        self.dynamic_unload(self.cond_stage_model,
+                            'cond_stage_model',
+                            skip_loaded=True)
+
+        # get noise
+        seed = kwargs.pop('seed', -1)
+        g = torch.Generator(device=we.device_id)
+        seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
+        g.manual_seed(seed)
+        if 'seed' in value_output:
+            value_output['seed'] = seed
+        for sample_id in range(num_samples):
+            if self.diffusion_model is not None:
+                noise = torch.empty(
+                    1,
+                    16,
+                    height // self.first_stage_model['paras']['size_factor'],
+                    width // self.first_stage_model['paras']['size_factor'],
+                    device=we.device_id).normal_(generator=g)
+
+                self.dynamic_load(self.diffusion_model, 'diffusion_model')
+                # UNet use input n_prompt
+                function_name, dtype = self.get_function_info(
+                    self.diffusion_model)
+                with torch.autocast('cuda',
+                                    enabled=dtype == 'float16',
+                                    dtype=getattr(torch, dtype)):
+                    solver_sample = value_input.get('sample', 'ddim')
+                    sample_steps = value_input.get('sample_steps', 50)
+                    guide_scale = value_input.get('guide_scale', 7.5)
+                    guide_rescale = value_input.get('guide_rescale', 0.5)
+                    discretization = value_input.get('discretization',
+                                                     'trailing')
+                    latent = self.diffusion.sample(
+                        solver=solver_sample,
+                        noise=noise,
+                        model=get_model(self.diffusion_model),
+                        model_kwargs=[{
+                            'cond': context
+                        }, {
+                            'cond': null_context
+                        }]
+                        if guide_scale is not None and guide_scale > 0 else {
+                            'cond': context,
+                        },
+                        cat_uc=False,
+                        steps=sample_steps,
+                        guide_scale=guide_scale,
+                        guide_rescale=guide_rescale,
+                        discretization=discretization,
+                        show_progress=True,
+                        seed=seed,
+                        condition_fn=None,
+                        clamp=None,
+                        percentile=None,
+                        t_max=None,
+                        t_min=None,
+                        discard_penultimate_step=None,
+                        return_intermediate=None,
+                        **kwargs)
+
+                self.dynamic_unload(self.diffusion_model,
+                                    'diffusion_model',
+                                    skip_loaded=True)
+
+            if 'latent' in value_output:
+                if value_output['latent'] is None or (
+                        isinstance(value_output['latent'], list)
+                        and len(value_output['latent']) < 1):
+                    value_output['latent'] = []
+                value_output['latent'].append(latent)
+
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            x_samples = self.decode_first_stage(latent).float()
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=True)
+            images = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+            if 'images' in value_output:
+                if value_output['images'] is None or (
+                        isinstance(value_output['images'], list)
+                        and len(value_output['images']) < 1):
+                    value_output['images'] = []
+                value_output['images'].append(images)
+
+        for k, v in value_output.items():
+            if isinstance(v, list):
+                value_output[k] = torch.cat(v, dim=0)
+            if isinstance(v, torch.Tensor):
+                value_output[k] = v.cpu()
+
+        # unregister tuner
+        if tuner_model is not None and tuner_model != '' and len(
+                tuner_model) > 0:
+            self.tuner_infer.unregister_tuner(tuner_model,
+                                              self.diffusion_model,
+                                              self.cond_stage_model)
+
+        # unregister control
+        if control_model is not None and control_model != '':
+            self.control_infer.unregister_controllers(control_model,
+                                                      self.diffusion_model)
+
+        return value_output
@@ -7,18 +7,14 @@ import gradio as gr
 import torch
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
+from scepter.modules.model.network.diffusion.diffusion import GaussianDiffusion
 from scepter.modules.utils.distribute import we

 from .control_inference import ControlInference
-from .diffusion_inference import DiffusionInference
+from .diffusion_inference import DiffusionInference, get_model
 from .tuner_inference import TunerInference


-def get_model(model_tuple):
-    assert 'model' in model_tuple
-    return model_tuple['model']
-
-
 class StyleboothInference(DiffusionInference):
    '''
        define vae, unet, text-encoder, tuner, refiner components
@@ -27,10 +23,12 @@ class StyleboothInference(DiffusionInference):
    '''
    def __init__(self, logger=None):
        self.logger = logger
+        self.is_redefine_paras = True
        self.loaded_model = {}
        self.loaded_model_name = [
            'diffusion_model', 'first_stage_model', 'cond_stage_model'
        ]
+        self.diffusion_insclass = GaussianDiffusion
        self.tuner_infer = TunerInference(self.logger)
        self.control_infer = ControlInference(self.logger)

@@ -13,10 +13,6 @@ try:
    from peft.utils import CONFIG_NAME, SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME
 except Exception as e:
    warnings.warn(f'Import peft error, please deal with this problem: {e}')
-try:
-    from swift import Swift, SwiftModel
-except Exception as e:
-    warnings.warn(f'Import swift error, please deal with this problem: {e}')


 class TunerInference():
@@ -27,6 +23,11 @@ class TunerInference():
    # @classmethod
    def unregister_tuner(self, tuner_model_list, diffusion_model,
                         cond_stage_model):
+        try:
+            from swift import SwiftModel
+        except Exception as e:
+            warnings.warn(f'Import swift error, please deal with this problem: {e}')
+
        self.logger.info('Unloading tuner model')
        if isinstance(diffusion_model['model'], SwiftModel):
            for adapter_name in diffusion_model['model'].adapters:
@@ -41,6 +42,11 @@ class TunerInference():
    # @classmethod
    def register_tuner(self, tuner_model_list, diffusion_model,
                       cond_stage_model):
+        try:
+            from swift import Swift
+        except Exception as e:
+            warnings.warn(f'Import swift error, please deal with this problem: {e}')
+
        self.logger.info('Loading tuner model')
        if len(tuner_model_list) < 1:
            self.unregister_tuner(tuner_model_list, diffusion_model,
@@ -137,7 +143,10 @@ class TunerInference():
                        state_dict = {}
                        is_bin_file = True
                        if os.path.isfile(bin_file):
-                            state_dict = torch.load(bin_file)
+                            if 'weights_only' in torch.load.__code__.co_varnames:
+                                state_dict = torch.load(bin_file, weights_only=True)
+                            else:
+                                state_dict = torch.load(bin_file)
                        elif os.path.isfile(safe_file):
                            is_bin_file = False
                            from safetensors.torch import \
@@ -2,4 +2,4 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.

 from scepter.modules.model import (backbone, embedder, head, loss, metric,
-                                   neck, network, tokenizer, tuner)
+                                   neck, network, tokenizer, tuner, diffusion)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
jiangzeyinzi	eac03e9856	Merge pull request #52 from modelscope/v1.1.0_dev update v1.1.0	2024-10-23 10:19:29 +08:00
jiangzeyinzi	379b94ab4f	update	2024-10-23 10:15:31 +08:00
jiangzeyinzi	342c6b8a15	Merge branch 'v1.1.0_dev' of https://github.com/modelscope/scepter into v1.1.0_dev	2024-10-21 11:59:55 +08:00
jiangzeyinzi	4ddcb08c7b	update	2024-10-21 11:59:41 +08:00
jiangzeyinzi	b2169a1597	Update readme.md	2024-10-21 11:58:43 +08:00
jiangzeyinzi	cffd54a02a	update	2024-10-21 11:36:02 +08:00
zeyinzi.jzyz	0bba2c319d	update v1.1.0	2024-10-21 00:35:53 +08:00
LouieStark	7d6451efad	update project page url	2024-10-01 23:45:22 +08:00
LouieStark	a5decd17aa	update readme	2024-09-30 14:22:14 +08:00
jiangzeyinzi	8a14866562	Merge pull request #46 from yaosheng216/patch-4 Update ldm_sce.py	2024-09-27 09:49:27 +08:00
jiangzeyinzi	5a94f6c4a2	Merge pull request #45 from yaosheng216/patch-3 Update sd15_512_sce_ctr_hed.yaml	2024-09-27 09:48:43 +08:00
Great	4ad0f6038c	Update ldm_sce.py	2024-09-27 09:39:22 +08:00
Great	4f56623627	Update sd15_512_sce_ctr_hed.yaml	2024-09-27 09:36:27 +08:00
jiangzeyinzi	30afc0a1de	Merge pull request #39 from yaosheng216/patch-2 Update readme.md	2024-07-18 17:51:53 +08:00
Great	aeabef75b4	Update readme.md	2024-07-18 17:40:28 +08:00
jiangzeyinzi	9376149415	Merge pull request #38 from modelscope/v1.0.3_dev v1.0.3	2024-07-18 17:31:44 +08:00
zeyinzi.jzyz	01fd8335af	v1.0.3 update	2024-07-18 14:12:42 +08:00
Zhen Han	7a9f90efb2	Update stylebooth.md	2024-07-14 20:10:10 +08:00
Zhen Han	cbbef4a2da	Update stylebooth.md	2024-07-14 19:53:11 +08:00
hanzhen.hz	b242449b25	v1.0.2	2024-06-05 16:12:05 +08:00
jiangzeyinzi	2fb8fd1872	Merge pull request #34 from modelscope/v1.0.0_dev v1.0.0 update	2024-06-03 17:49:47 +08:00