Merge pull request #85 from yaosheng216/patch-17

Update pyproject.toml

.github/workflows/publish.yml

@@ -0,0 +1,24 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "scepter/workflow/pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    # if this is a forked repository. Skipping the workflow.
+    if: github.event.repository.fork == false 
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -9,13 +9,12 @@
 *.bin
 *.idea
 *.csv
+cache
 build
 dist
 dev
 scepter.egg-info
 .readthedocs.yml
-1.9
-#MANIFEST.in
 *resources
 *.ipynb_checkpoints*
 *.vscode

docs/en/tasks/ace.md

@@ -0,0 +1,184 @@
+<p align="center">
+
+  <h2 align="center"><img src="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/figures/icon.png" height=16> : All-round Creator and Editor Following <br> Instructions via Diffusion Transformer</h2>
+
+  <p align="center">
+    <a href="https://arxiv.org/abs/2410.00086"><img src='https://img.shields.io/badge/arXiv-ACE-red' alt='Paper PDF'></a>
+    <a href='https://ali-vilab.github.io/ace-page'><img src='https://img.shields.io/badge/Project_Page-ACE-blue' alt='Project Page'></a>
+    <a href='https://github.com/modelscope/scepter'><img src='https://img.shields.io/badge/Scepter-ACE-green'></a>
+    <a href='https://huggingface.co/spaces/scepter-studio/ACE-Chat'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Space-orange'></a>
+    <a href='https://huggingface.co/scepter-studio/ACE-0.6B-512px'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-orange'></a>
+    <a href='https://www.modelscope.cn/models/iic/ACE-0.6B-512px'><img src='https://img.shields.io/badge/ModelScope-Model-purple'></a>
+    <br>
+    <strong>Zhen Han*</strong>
+    ·
+    <strong>Zeyinzi Jiang*</strong>
+    ·
+    <strong>Yulin Pan*</strong>
+    ·
+    <strong>Jingfeng Zhang*</strong>
+    ·
+    <strong>Chaojie Mao*</strong>
+    <br>
+    <strong>Chenwei Xie</strong>
+    ·
+    <strong>Yu Liu</strong>
+    ·
+    <strong>Jingren Zhou</strong>
+    <br>
+    Tongyi Lab, Alibaba Group
+  </p>
+  <table align="center">
+    <tr>
+    <td>
+      <img src="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/figures/teaser.png">
+    </td>
+    </tr>
+  </table>
+
+
+## 🚀 Installation
+Install the necessary packages with `pip`: 
+```bash
+pip install -r requirements.txt
+```
+
+##  🔥 ACE Models
+|    **Model**     |                                                                                                                                                                                                            **Status**                                                                                                                                                                                                             | 
+|:----------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
+|  ACE-0.6B-512px  |          [![Demo link](https://img.shields.io/badge/Demo-ACE_Chat-purple)](https://huggingface.co/spaces/scepter-studio/ACE-Chat)<br>[![ModelScope link](https://img.shields.io/badge/ModelScope-Model-blue)](https://www.modelscope.cn/models/iic/ACE-0.6B-512px)  [![HuggingFace link](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-yellow)](https://huggingface.co/scepter-studio/ACE-0.6B-512px)          |
+| ACE-0.6B-1024px  | [![Demo link](https://img.shields.io/badge/Demo-ACE_Refiner_Chat-purple)](https://huggingface.co/spaces/scepter-studio/ACE-Refiner-Chat)<br>[![ModelScope link](https://img.shields.io/badge/ModelScope-Model-blue)](https://www.modelscope.cn/models/iic/ACE-0.6B-1024px)  [![HuggingFace link](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-yellow)](https://huggingface.co/scepter-studio/ACE-0.6B-1024px) |             |
+## 🖼 Model Performance Visualization
+
+The current model's parameters scale of ACE is 0.6B, which imposes certain limitations on the quality of image generation. [FLUX.1-Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev), on the other hand, 
+has a significant advantage in text-to-image generation quality. By using SDEdit, we can effectively leverage the generative capabilities of FLUX to further enhance the image results generated by ACE. Based on the above considerations, we have designed the ACE-Refiner pipeline, as shown in the diagram below.
+
+![ACE_REFINER](https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/ace_method/ace_refiner_process.webp)
+
+As shown in the figure below, when the strength 
+σ of the generated image is high, the generated image will suffer from fidelity loss compared to the original image. Conversely, lower 
+σ does not significantly improve the image quality. Therefore, users can make a trade-off between fidelity to the generated result and the image quality based on their own needs. 
+Users can set the value of "REFINER_SCALE" in the configuration file `config/inference_config/models/ace_0.6b_1024_refiner.yaml`. 
+We recommend that users use the advance options in the [webui-demo](#-chat-bot-) for effect verification.
+
+![ACE_REFINER_EXAMPLE](https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/ace_method/ace_refiner.webp)
+
+
+We compared the generation and editing performance of different models on several tasks, as shown as following. 
+![Samples](https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/ace_method/samples_compare.webp)
+
+
+## 🔥 Training
+
+We offer a demonstration training YAML that enables the end-to-end training of ACE using a toy dataset. For a comprehensive overview of the hyperparameter configurations, please consult `config/ace_0.6b_512_train.yaml`.
+
+### Prepare datasets
+
+Please find the dataset class located in `modules/data/dataset/dataset.py`, 
+designed to facilitate end-to-end training using an open-source toy dataset. 
+Download a dataset zip file from [modelscope](https://www.modelscope.cn/models/iic/scepter/resolve/master/datasets/hed_pair.zip), and then extract its contents into the `cache/datasets/` directory.
+
+Should you wish to prepare your own datasets, we recommend consulting `modules/data/dataset/dataset.py` for detailed guidance on the required data format.
+
+### Prepare initial weight
+The ACE checkpoint has been uploaded to both ModelScope and HuggingFace platforms:
+* [ModelScope](https://www.modelscope.cn/models/iic/ACE-0.6B-512px)
+* [HuggingFace](https://huggingface.co/scepter-studio/ACE-0.6B-512px)
+
+In the provided training YAML configuration, we have designated the Modelscope URL as the default checkpoint URL. Should you wish to transition to Hugging Face, you can effortlessly achieve this by modifying the PRETRAINED_MODEL value within the YAML file (replace the prefix "ms://iic" to "hf://scepter-studio").
+
+
+### Start training
+
+You can easily start training procedure by executing the following command:
+```bash
+# ACE-0.6B-512px
+PYTHONPATH=. python tools/run_train.py --cfg config/ace_0.6b_512_train.yaml
+# ACE-0.6B-1024px
+PYTHONPATH=. python tools/run_train.py --cfg config/ace_0.6b_1024_train.yaml
+```
+
+## 🚀 Inference
+
+We provide a simple inference demo that allows users to generate images from text descriptions.
+```bash
+ PYTHONPATH=. python tools/run_inference.py --cfg config/inference_config/models/ace_0.6b_512.yaml --instruction "make the boy cry, his eyes filled with tears" --seed 199999 --input_image examples/input_images/example0.webp
+```
+We recommend runing the examples for quick testing. Running the following command will run the example inference and the results will be saved in `examples/output_images/`.
+```bash
+PYTHONPATH=. python tools/run_inference.py --cfg config/inference_config/models/ace_0.6b_512.yaml
+```
+
+## 💬 Chat Bot 
+We have developed an chatbot UI utilizing Gradio, designed to transform user input in natural language into visually stunning images that align semantically with the provided instructions. Users can effortlessly initiate the chatbot app by executing the following command:
+```bash
+python chatbot/run_gradio.py --cfg chatbot/config/chatbot_ui.yaml --server_port 2024
+```
+
+<table align="center">
+  <tr>
+  <td>
+    <img src="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/videos/demo_chat.gif">
+  </td>
+  </tr>
+</table>
+
+## ⚙️️ ComfyUI Workflow
+
+![Workflow](https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_example.jpg)
+
+We support the use of ACE in the ComfyUI Workflow through the following methods:
+
+1) Automatic installation directly via the ComfyUI Manager by searching for the **ComfyUI-Scepter** node.
+2) Manually install by moving custom_nodes from Scepter to ComfyUI.
+```shell
+git clone https://github.com/modelscope/scepter.git
+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
+```
+
+**Note**: You can use the nodes by dragging the sample images below into ComfyUI. Additionally, our nodes can automatically pull models from ModelScope or HuggingFace by selecting the *model_source* field, or you can place the already downloaded models in a local path.
+
+<table><tbody>
+  <tr>
+    <th align="center" colspan="4">ACE Workflow Examples</th>
+  </tr>
+  <tr>
+    <th align="center" colspan="1">Control</th>
+    <th align="center" colspan="1">Semantic</th>
+    <th align="center" colspan="1">Element</th>
+  </tr>
+  <tr>
+    <td>
+      <a href="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_control.png" target="_blank">
+        <img src="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_control.png" width="200">
+      </a>
+    </td>
+    <td>
+      <a href="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_semantic.png" target="_blank">
+        <img src="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_semantic.png" width="200">
+      </a>
+    </td>
+    <td>
+      <a href="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_element.png" target="_blank">
+        <img src="https://raw.githubusercontent.com/ali-vilab/ACE/refs/heads/main/assets/comfyui/ace_element.png" width="200">
+      </a>
+    </td>
+  </tr>
+</tbody>
+</table>
+
+
+## 📝 Citation
+
+```bibtex
+@article{han2024ace,
+  title={ACE: All-round Creator and Editor Following Instructions via Diffusion Transformer},
+  author={Han, Zhen and Jiang, Zeyinzi and Pan, Yulin and Zhang, Jingfeng and Mao, Chaojie and Xie, Chenwei and Liu, Yu and Zhou, Jingren},
+  journal={arXiv preprint arXiv:2410.00086},
+  year={2024}
+}
+```

readme.md

@@ -14,12 +14,15 @@ SCEPTER integrates popular community-driven implementations as well as proprieta
 SCEPTER offers 3 core components:
 - [Generative training and inference framework](#tutorials)
 - [Easy implementation of popular approaches](#currently-supported-approaches)
-- [Interactive user interface: SCEPTER Studio](#launch)
+- [Interactive user interface: SCEPTER Studio & Comfy UI](#launch)
 
 
 ## 🎉 News
+- [🔥🔥🔥 2025.01]: We report ACE++, an instruction-based diffusion framework that tackles various image generation and editing tasks. The code and paper is available on [ACE++](https://ali-vilab.github.io/ACE_plus_page/).
+- [2024.11]: Supports video files, video annotation, caption translation in data management, and inference & training of the [CogVideoX](https://arxiv.org/abs/2408.06072).
+- [2024.10]: We are pleased to announce the release of the code for [ACE](https://arxiv.org/abs/2410.00086), supporting Customized Training / Comfy UI Workflow / gradio-based ChatBot Interface. 
 - [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.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`.
@@ -30,46 +33,65 @@ SCEPTER offers 3 core components:
 - [2023.12]: We propose [SCEdit](https://arxiv.org/abs/2312.11392), an efficient and controllable generation framework.
 - [2023.12]: We release [🪄SCEPTER](https://github.com/modelscope/scepter/) library.
 
+[//]: # (## 🖼 Gallery for Recent Works)
 
-## 🖼 Gallery for Recent Works
+[//]: # ()
+[//]: # (### FLUX 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>
+[//]: # ()
+[//]: # (<table><tbody>)
 
-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.
+[//]: # (  <tr>)
 
-[![Watch the demo](https://github.com/ali-vilab/ace-page/raw/main/static/images/teaser.jpg)](https://ali-vilab.github.io/ace-page/)
+[//]: # (    <th align="center" colspan="3">Yarn Style</th>)
 
-### FLUX Tuners
+[//]: # (    <th align="center" colspan="3">Soft Watercolor Style</th>)
 
-<table><tbody>
-  <tr>
-    <th align="center" colspan="3">Yarn Style</th>
-    <th align="center" colspan="3">Soft Watercolor Style</th>
-  </tr>
-  <tr>
-    <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="3">Travel Style</th>
-    <th align="center" colspan="3">WuKong Style</th>
-  </tr>
-  <tr>
-    <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>
+[//]: # (  </tr>)
+
+[//]: # (  <tr>)
+
+[//]: # (    <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="3">Travel Style</th>)
+
+[//]: # (    <th align="center" colspan="3">WuKong Style</th>)
+
+[//]: # (  </tr>)
+
+[//]: # (  <tr>)
+
+[//]: # (    <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>)
 
 ### ComfyUI Workflow
 
@@ -143,18 +165,19 @@ 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)                                                                                               |
-|   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/)                      |
+|            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/) [![Demo link](https://img.shields.io/badge/Demo-ACE-purple)](https://huggingface.co/spaces/scepter-studio/ACE-Chat) <br> [![ModelScope link](https://img.shields.io/badge/ModelScope-Model-blue)](https://www.modelscope.cn/models/iic/ACE-0.6B-512px)  [![HuggingFace link](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-yellow)](https://huggingface.co/scepter-studio/ACE-0.6B-512px)               |
+| Image Generation and Editing | [🌟ACE++](https://ali-vilab.github.io/ACE_plus_page/) | [![Arxiv   link](https://img.shields.io/static/v1?label=arXiv&message=ACEPlus&color=red&logo=arxiv)](https://arxiv.org/abs/2501.02487)   [![Page link](https://img.shields.io/badge/Page-ACE++-Gree)](https://ali-vilab.github.io/ACE_plus_page/) [![Demo link](https://img.shields.io/badge/Demo-ACE++-purple)](https://huggingface.co/spaces/scepter-studio/ACE-Plus) <br> [![ModelScope link](https://img.shields.io/badge/ModelScope-Model-blue)](https://www.modelscope.cn/models/iic/ACE_Plus/summary)  [![HuggingFace link](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-yellow)](https://huggingface.co/ali-vilab/ACE_Plus/tree/main) |
 
 
 ## 🖥️ SCEPTER Studio
@@ -191,18 +214,20 @@ We deploy a work studio on Modelscope that includes only the inference tab, plea
 
 ## ⚙️️ ComfyUI Workflow
 
-### Launch
+We support the use of all models in the ComfyUI Workflow through the following methods:
 
-Manually install by moving custom_nodes to ComfyUI.
+1) Automatic installation directly via the ComfyUI Manager by searching for the **ComfyUI-Scepter** node.
+2) Manually install by moving custom_nodes from Scepter to ComfyUI.
 ```shell
+git clone https://github.com/modelscope/scepter.git
 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.
 
+**Note**: You can use the nodes by dragging the sample images into ComfyUI. Additionally, our nodes can automatically pull models from ModelScope or HuggingFace by selecting the *model_source* field, or you can place the already downloaded models in a local path.
 
 ## 🔍 Learn More
 

requirements/framework.txt

@@ -1,8 +1,7 @@
 albumentations
 beautifulsoup4
-bezier
 einops
-modelscope
+modelscope[framework]
 ms-swift
 numpy
 open_clip_torch
@@ -12,6 +11,7 @@ oss2>=2.15.0
 pycocotools
 pyyaml>=5.3.1
 scikit-image
+scikit-learn
+sentencepiece
 torchsde
 transformers
-scikit-learn

requirements/recommended.txt

@@ -1,4 +1,5 @@
 git+https://github.com/cocodataset/panopticapi.git
-torch==2.0.1
-torchvision==0.15.2
-xformers==0.0.21
+torch==2.4.1
+torchvision==.19.1
+flash-attn==2.5.8
+xformers==0.0.28

requirements/scepter_studio.txt

@@ -1,5 +1,6 @@
 bitsandbytes
 gradio
+gradio_imageslider
 imagehash
 psutil
 tiktoken

scepter/__init__.py

@@ -1,18 +1,28 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-import os
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-import scepter
-from scepter.modules import data, model, opt, solver, transform, utils
-from scepter.tools.helper import get_module_list as module_list
-from scepter.tools.helper import \
-    get_module_object_config as configures_by_objects
-from scepter.tools.helper import get_module_objects as objects_by_module
-from scepter.version import __version__, version_info
 
-dirname = os.path.dirname(scepter.__file__)
+if TYPE_CHECKING:
+    from scepter.modules import data, model, opt, solver, transform, utils
+    from scepter.tools.helper import get_module_list as module_list
+    from scepter.tools.helper import \
+        get_module_object_config as configures_by_objects
+    from scepter.tools.helper import get_module_objects as objects_by_module
+    from scepter.version import __version__, version_info
+else:
+    _import_structure = {
+        'modules': ['data', 'model', 'opt', 'solver', 'transform', 'utils'],
+        'helper': ['get_module_list', 'get_module_object_config', 'get_module_objects'],
+        'version': ['__version__', 'version_info']
+    }
 
-__all__ = [
-    utils, transform, data, model, solver, version_info, opt, '__version__',
-    'dirname'
-]
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/methods/edit/dit_ace_0.6b_1024.yaml

@@ -0,0 +1,161 @@
+ENV:
+  BACKEND: nccl
+  SEED: 2024
+#
+SOLVER:
+  NAME: ACESolver
+  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/ace_0.6b_1024
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    - NAME: "HuggingfaceFs"
+      TEMP_DIR: ./cache/cache_data
+    - NAME: "LocalFs"
+      TEMP_DIR: ./cache/cache_data
+    - NAME: "ModelscopeFs"
+      TEMP_DIR: ./cache/cache_data
+
+  #
+  MODEL:
+    NAME: LatentDiffusionACE
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 0.18215
+    SIZE_FACTOR: 8
+    DECODER_BIAS: 0.5
+    DEFAULT_N_PROMPT:
+    USE_EMA: True
+    EVAL_EMA: False
+    TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+    USE_TEXT_POS_EMBEDDINGS: True
+    #
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: eps
+      MIN_SNR_GAMMA:
+      NOISE_SCHEDULER:
+        NAME: LinearScheduler
+        NUM_TIMESTEPS: 1000
+        BETA_MIN: 0.0001
+        BETA_MAX: 0.02
+    #
+    DIFFUSION_MODEL:
+      NAME: ACE
+      PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+      IGNORE_KEYS: [ ]
+      PATCH_SIZE: 2
+      IN_CHANNELS: 4
+      HIDDEN_SIZE: 1152
+      DEPTH: 28
+      NUM_HEADS: 16
+      MLP_RATIO: 4.0
+      PRED_SIGMA: True
+      DROP_PATH: 0.0
+      WINDOW_DIZE: 0
+      Y_CHANNELS: 4096
+      MAX_SEQ_LEN: 4096
+      QK_NORM: True
+      USE_GRAD_CHECKPOINT: True
+      ATTENTION_BACKEND: flash_attn
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      EMBED_DIM: 4
+      PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/vae/vae.bin
+      IGNORE_KEYS: []
+      #
+      ENCODER:
+        NAME: Encoder
+        CH: 128
+        OUT_CH: 3
+        NUM_RES_BLOCKS: 2
+        IN_CHANNELS: 3
+        ATTN_RESOLUTIONS: [ ]
+        CH_MULT: [ 1, 2, 4, 4 ]
+        Z_CHANNELS: 4
+        DOUBLE_Z: True
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+      #
+      DECODER:
+        NAME: Decoder
+        CH: 128
+        OUT_CH: 3
+        NUM_RES_BLOCKS: 2
+        IN_CHANNELS: 3
+        ATTN_RESOLUTIONS: [ ]
+        CH_MULT: [ 1, 2, 4, 4 ]
+        Z_CHANNELS: 4
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+        GIVE_PRE_END: False
+        TANH_OUT: False
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+      TOKENIZER_PATH: ms://iic/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+      LENGTH: 120
+      T5_DTYPE: bfloat16
+      ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+      CLEAN: whitespace
+      USE_GRAD: False
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 20
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 1e-7
+    EPS: 1e-10
+    WEIGHT_DECAY: 5e-4
+  #
+  TRAIN_DATA:
+    NAME: ImageTextPairMSDatasetForACE
+    MODE: train
+    MS_DATASET_NAME: cache/datasets/hed_pair
+    MS_DATASET_NAMESPACE: ""
+    MS_DATASET_SPLIT: "train"
+    MS_DATASET_SUBNAME: ""
+    PROMPT_PREFIX: ""
+    REPLACE_STYLE: False
+    MAX_SEQ_LEN: 4096
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 1
+    SAMPLER:
+      NAME: LoopSampler
+  #
+  TRAIN_HOOKS:
+    -
+      NAME: BackwardHook
+      PRIORITY: 0
+    -
+      NAME: LogHook
+      LOG_INTERVAL: 50
+    -
+      NAME: CheckpointHook
+      INTERVAL: 100
+    -
+      NAME: ProbeDataHook
+      PROB_INTERVAL: 100

scepter/methods/edit/dit_ace_0.6b_512.yaml

@@ -0,0 +1,161 @@
+ENV:
+  BACKEND: nccl
+  SEED: 2024
+#
+SOLVER:
+  NAME: ACESolver
+  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/ace_0.6b_512
+  LOG_FILE: std_log.txt
+  #
+  FILE_SYSTEM:
+    - NAME: "HuggingfaceFs"
+      TEMP_DIR: ./cache/cache_data
+    - NAME: "LocalFs"
+      TEMP_DIR: ./cache/cache_data
+    - NAME: "ModelscopeFs"
+      TEMP_DIR: ./cache/cache_data
+
+  #
+  MODEL:
+    NAME: LatentDiffusionACE
+    PRETRAINED_MODEL:
+    IGNORE_KEYS: [ ]
+    SCALE_FACTOR: 0.18215
+    SIZE_FACTOR: 8
+    DECODER_BIAS: 0.5
+    DEFAULT_N_PROMPT:
+    USE_EMA: True
+    EVAL_EMA: False
+    TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+    USE_TEXT_POS_EMBEDDINGS: True
+    #
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: eps
+      MIN_SNR_GAMMA:
+      NOISE_SCHEDULER:
+        NAME: LinearScheduler
+        NUM_TIMESTEPS: 1000
+        BETA_MIN: 0.0001
+        BETA_MAX: 0.02
+    #
+    DIFFUSION_MODEL:
+      NAME: ACE
+      PRETRAINED_MODEL: ms://iic/ACE-0.6B-512px@models/dit/ace_0.6b_512px.pth
+      IGNORE_KEYS: [ ]
+      PATCH_SIZE: 2
+      IN_CHANNELS: 4
+      HIDDEN_SIZE: 1152
+      DEPTH: 28
+      NUM_HEADS: 16
+      MLP_RATIO: 4.0
+      PRED_SIGMA: True
+      DROP_PATH: 0.0
+      WINDOW_DIZE: 0
+      Y_CHANNELS: 4096
+      MAX_SEQ_LEN: 1024
+      QK_NORM: True
+      USE_GRAD_CHECKPOINT: True
+      ATTENTION_BACKEND: flash_attn
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKL
+      EMBED_DIM: 4
+      PRETRAINED_MODEL: ms://iic/ACE-0.6B-512px@models/vae/vae.bin
+      IGNORE_KEYS: []
+      #
+      ENCODER:
+        NAME: Encoder
+        CH: 128
+        OUT_CH: 3
+        NUM_RES_BLOCKS: 2
+        IN_CHANNELS: 3
+        ATTN_RESOLUTIONS: [ ]
+        CH_MULT: [ 1, 2, 4, 4 ]
+        Z_CHANNELS: 4
+        DOUBLE_Z: True
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+      #
+      DECODER:
+        NAME: Decoder
+        CH: 128
+        OUT_CH: 3
+        NUM_RES_BLOCKS: 2
+        IN_CHANNELS: 3
+        ATTN_RESOLUTIONS: [ ]
+        CH_MULT: [ 1, 2, 4, 4 ]
+        Z_CHANNELS: 4
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+        GIVE_PRE_END: False
+        TANH_OUT: False
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://iic/ACE-0.6B-512px@models/text_encoder/t5-v1_1-xxl/
+      TOKENIZER_PATH: ms://iic/ACE-0.6B-512px@models/tokenizer/t5-v1_1-xxl
+      LENGTH: 120
+      T5_DTYPE: bfloat16
+      ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+      CLEAN: whitespace
+      USE_GRAD: False
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 20
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+  #
+  OPTIMIZER:
+    NAME: AdamW
+    LEARNING_RATE: 1e-7
+    EPS: 1e-10
+    WEIGHT_DECAY: 5e-4
+  #
+  TRAIN_DATA:
+    NAME: ImageTextPairMSDatasetForACE
+    MODE: train
+    MS_DATASET_NAME: cache/datasets/hed_pair
+    MS_DATASET_NAMESPACE: ""
+    MS_DATASET_SPLIT: "train"
+    MS_DATASET_SUBNAME: ""
+    PROMPT_PREFIX: ""
+    REPLACE_STYLE: False
+    MAX_SEQ_LEN: 1024
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 1
+    SAMPLER:
+      NAME: LoopSampler
+  #
+  TRAIN_HOOKS:
+    -
+      NAME: BackwardHook
+      PRIORITY: 0
+    -
+      NAME: LogHook
+      LOG_INTERVAL: 50
+    -
+      NAME: CheckpointHook
+      INTERVAL: 100
+    -
+      NAME: ProbeDataHook
+      PROB_INTERVAL: 100

scepter/methods/examples/generation/dit_cogvideox1.5_5b_i2v_lora.yaml

@@ -0,0 +1,277 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox1.5_5b_i2v_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  FPS: 16
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 64
+      LORA_ALPHA: 64
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 0.7  
+    NOISED_IMAGE_DROPOUT: 0.05
+    INVERT_SCALE_LATENTS: True
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      USE_DYNAMIC_CFG: False
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 1.0  
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: # 5b-I2V diff
+        - ms://ZhipuAI/CogVideoX1.5-5B-I2V@transformer/diffusion_pytorch_model-00001-of-00003.safetensors
+        - ms://ZhipuAI/CogVideoX1.5-5B-I2V@transformer/diffusion_pytorch_model-00002-of-00003.safetensors
+        - ms://ZhipuAI/CogVideoX1.5-5B-I2V@transformer/diffusion_pytorch_model-00003-of-00003.safetensors
+      NUM_ATTENTION_HEADS: 48  
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 32
+      LATENT_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      OFS_EMBED_DIM: 512 # v1.5 diff
+      NUM_LAYERS: 42  
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 300
+      SAMPLE_HEIGHT: 300
+      SAMPLE_FRAMES: 81
+      PATCH_SIZE: 2
+      PATCH_SIZE_T: 2  # v1.5 diff
+      PATCH_BIAS: False  # v1.5 diff
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 224
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: True 
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+      GRADIENT_CHECKPOINTING: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://ZhipuAI/CogVideoX1.5-5B-I2V@vae/diffusion_pytorch_model.safetensors 
+      SAMPLE_HEIGHT: 768
+      SAMPLE_WIDTH: 1360
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 224
+      CLEAN:
+      USE_GRAD: False
+      T5_DTYPE: bfloat16
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 81
+    IMAGE_SIZE: [768, 1360]
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDataset
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 0
+    NUM_FRAMES: 85
+    FPS: 16
+    HEIGHT: 768
+    WIDTH: 1360
+    PROMPT_PREFIX: 'DISNEY '
+    DATA_TYPE: 'i2v'
+    SAMPLER:
+      NAME: MixtureOfSamplers
+      SUB_SAMPLERS:
+        - NAME: MultiLevelBatchSampler
+          PROB: 1.0
+          FIELDS: [ "video_path", "prompt" ]
+          DELIMITER: '#;#'
+          PATH_PREFIX: cache/datasets/Disney-VideoGeneration-Dataset/
+          INDEX_FILE: cache/datasets/Disney-VideoGeneration-Dataset/index.jsonl
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ "video", "image", "prompt" ]
+        META_KEYS: [ ]
+  #
+#  EVAL_DATA:
+#    NAME: Text2ImageDataset
+#    MODE: eval
+#    PROMPT_FILE:
+#    PROMPT_DATA: [ "A cat running.#;#asset/images/edit_tuner/cat_512.jpg" ]
+#    FIELDS: [ "prompt", "img_path" ]
+#    DELIMITER: '#;#'
+#    PROMPT_PREFIX: ''
+#    PIN_MEMORY: True
+#    BATCH_SIZE: 1
+#    USE_NUM: 8
+#    NUM_WORKERS: 0
+#    IMAGE_SIZE: [768, 1360]
+#    TRANSFORMS:
+#      - NAME: LoadImageFromFileList
+#        FILE_KEYS: [ 'img_path' ]
+#        RGB_ORDER: RGB
+#        BACKEND: pillow
+#      - NAME: FlexibleResize
+#        INTERPOLATION: bilinear
+#        SIZE: [768, 1360]
+#        INPUT_KEY: [ 'img' ]
+#        OUTPUT_KEY: [ 'img' ]
+#        BACKEND: pillow
+#      - NAME: FlexibleCenterCrop
+#        SIZE: [768, 1360]
+#        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: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+  #
+#  EVAL_HOOKS:
+#    - NAME: ProbeDataHook
+#      PROB_INTERVAL: 100
+#      PRIORITY: 0

scepter/methods/examples/generation/dit_cogvideox1.5_5b_lora.yaml

@@ -0,0 +1,248 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox1.5_5b_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  FPS: 16
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 64
+      LORA_ALPHA: 64
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 0.7
+    INVERT_SCALE_LATENTS: True
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      USE_DYNAMIC_CFG: False
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 1.0  
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: 
+        - ms://ZhipuAI/CogVideoX1.5-5B@transformer/diffusion_pytorch_model-00001-of-00003.safetensors
+        - ms://ZhipuAI/CogVideoX1.5-5B@transformer/diffusion_pytorch_model-00002-of-00003.safetensors
+        - ms://ZhipuAI/CogVideoX1.5-5B@transformer/diffusion_pytorch_model-00003-of-00003.safetensors
+      NUM_ATTENTION_HEADS: 48  
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      NUM_LAYERS: 42  
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 300
+      SAMPLE_HEIGHT: 300
+      SAMPLE_FRAMES: 81
+      PATCH_SIZE: 2
+      PATCH_SIZE_T: 2  # v1.5 diff
+      PATCH_BIAS: False  # v1.5 diff
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 224
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: True 
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+      GRADIENT_CHECKPOINTING: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://ZhipuAI/CogVideoX1.5-5B@vae/diffusion_pytorch_model.safetensors 
+      SAMPLE_HEIGHT: 768
+      SAMPLE_WIDTH: 1360
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 224
+      CLEAN:
+      USE_GRAD: False
+      T5_DTYPE: bfloat16
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 81
+    IMAGE_SIZE: [768, 1360]
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDataset
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 0
+    NUM_FRAMES: 85
+    FPS: 16
+    HEIGHT: 768
+    WIDTH: 1360
+    PROMPT_PREFIX: 'DISNEY '
+    SAMPLER:
+      NAME: MixtureOfSamplers
+      SUB_SAMPLERS:
+        - NAME: MultiLevelBatchSampler
+          PROB: 1.0
+          FIELDS: [ "video_path", "prompt" ]
+          DELIMITER: '#;#'
+          PATH_PREFIX: cache/datasets/Disney-VideoGeneration-Dataset/
+          INDEX_FILE: cache/datasets/Disney-VideoGeneration-Dataset/index.jsonl
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'video', "prompt" ]
+        META_KEYS: [ ]
+  #
+  EVAL_DATA:
+    NAME: Text2ImageDataset
+    MODE: eval
+    PROMPT_FILE:
+    PROMPT_DATA: [ "A girl riding a bike." ]
+    IMAGE_SIZE: [ 768, 1360 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: 'DISNEY ' # ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    USE_NUM: 8
+    NUM_WORKERS: 0
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+  #
+  EVAL_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0

scepter/methods/examples/generation/dit_cogvideox_2b_lora.yaml

@@ -0,0 +1,239 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox_2b_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  FPS: 8
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 64
+      LORA_ALPHA: 64
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 1.15258426
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 3.0
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@transformer/diffusion_pytorch_model.safetensors
+      NUM_ATTENTION_HEADS: 30
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      NUM_LAYERS: 30
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 90
+      SAMPLE_HEIGHT: 60
+      SAMPLE_FRAMES: 49
+      PATCH_SIZE: 2
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 226
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: False
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+      GRADIENT_CHECKPOINTING: False
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@vae/diffusion_pytorch_model.safetensors
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: False
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: False
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: False
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 226
+      CLEAN:
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 49
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDataset
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 4
+    NUM_FRAMES: 49
+    FPS: 8
+    HEIGHT: 480
+    WIDTH: 720
+    PROMPT_PREFIX: 'DISNEY '
+    SAMPLER:
+      NAME: MixtureOfSamplers
+      SUB_SAMPLERS:
+        - NAME: MultiLevelBatchSampler
+          PROB: 1.0
+          FIELDS: [ "video_path", "prompt" ]
+          DELIMITER: '#;#'
+          PATH_PREFIX: cache/datasets/Disney-VideoGeneration-Dataset/
+          INDEX_FILE: cache/datasets/Disney-VideoGeneration-Dataset/index.jsonl
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'video', "prompt" ]
+        META_KEYS: [ ]
+  #
+  EVAL_DATA:
+    NAME: Text2ImageDataset
+    MODE: eval
+    PROMPT_FILE:
+    PROMPT_DATA: [ "A girl riding a bike.", "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance." ]
+    IMAGE_SIZE: [ 480, 720 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: 'DISNEY '
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    USE_NUM: 8
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+  #
+  EVAL_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0

scepter/methods/examples/generation/dit_cogvideox_5b_i2v_lora.yaml

@@ -0,0 +1,270 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox_5b_i2v_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  FPS: 8
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 64
+      LORA_ALPHA: 64
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 0.7  # 5b diff
+    NOISED_IMAGE_DROPOUT: 0.05
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 1.0  # 5b diff
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: # 5b-I2V diff
+        - ms://AI-ModelScope/CogVideoX-5b-I2V@transformer/diffusion_pytorch_model-00001-of-00003.safetensors
+        - ms://AI-ModelScope/CogVideoX-5b-I2V@transformer/diffusion_pytorch_model-00002-of-00003.safetensors
+        - ms://AI-ModelScope/CogVideoX-5b-I2V@transformer/diffusion_pytorch_model-00003-of-00003.safetensors
+      NUM_ATTENTION_HEADS: 48  # 5b diff
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 32  # 5b-I2V diff
+      LATENT_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      NUM_LAYERS: 42  # 5b diff
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 90
+      SAMPLE_HEIGHT: 60
+      SAMPLE_FRAMES: 49
+      PATCH_SIZE: 2
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 226
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: True # 5b diff
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: True  # 5b-I2V diff
+      GRADIENT_CHECKPOINTING: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-5b-I2V@vae/diffusion_pytorch_model.safetensors # 5b diff
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 226
+      CLEAN:
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 49
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDataset
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 0
+    NUM_FRAMES: 49
+    FPS: 8
+    HEIGHT: 480
+    WIDTH: 720
+    PROMPT_PREFIX: 'DISNEY '
+    DATA_TYPE: 'i2v'
+    SAMPLER:
+      NAME: MixtureOfSamplers
+      SUB_SAMPLERS:
+        - NAME: MultiLevelBatchSampler
+          PROB: 1.0
+          FIELDS: [ "video_path", "prompt" ]
+          DELIMITER: '#;#'
+          PATH_PREFIX: cache/datasets/Disney-VideoGeneration-Dataset/
+          INDEX_FILE: cache/datasets/Disney-VideoGeneration-Dataset/index.jsonl
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ "video", "image", "prompt" ]
+        META_KEYS: [ ]
+  #
+#  EVAL_DATA:
+#    NAME: Text2ImageDataset
+#    MODE: eval
+#    PROMPT_FILE:
+#    PROMPT_DATA: [ "A cat running.#;#asset/images/edit_tuner/cat_512.jpg" ]
+#    FIELDS: [ "prompt", "img_path" ]
+#    DELIMITER: '#;#'
+#    PROMPT_PREFIX: ''
+#    PIN_MEMORY: True
+#    BATCH_SIZE: 1
+#    USE_NUM: 8
+#    NUM_WORKERS: 0
+#    IMAGE_SIZE: [ 480, 720 ]
+#    TRANSFORMS:
+#      - NAME: LoadImageFromFileList
+#        FILE_KEYS: [ 'img_path' ]
+#        RGB_ORDER: RGB
+#        BACKEND: pillow
+#      - NAME: FlexibleResize
+#        INTERPOLATION: bilinear
+#        SIZE: [ 480, 720 ]
+#        INPUT_KEY: [ 'img' ]
+#        OUTPUT_KEY: [ 'img' ]
+#        BACKEND: pillow
+#      - NAME: FlexibleCenterCrop
+#        SIZE: [ 480, 720 ]
+#        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: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+  #
+#  EVAL_HOOKS:
+#    - NAME: ProbeDataHook
+#      PROB_INTERVAL: 100
+#      PRIORITY: 0

scepter/methods/examples/generation/dit_cogvideox_5b_lora.yaml

@@ -0,0 +1,277 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox_5b_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  FPS: 8
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+    - NAME: SwiftLoRA
+      R: 64
+      LORA_ALPHA: 64
+      LORA_DROPOUT: 0.0
+      BIAS: "none"
+      TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 0.7  # 5b diff
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 1.0  # 5b diff
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: # 5b diff
+        - ms://AI-ModelScope/CogVideoX-5b@transformer/diffusion_pytorch_model-00001-of-00002.safetensors
+        - ms://AI-ModelScope/CogVideoX-5b@transformer/diffusion_pytorch_model-00002-of-00002.safetensors
+      NUM_ATTENTION_HEADS: 48  # 5b diff
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      NUM_LAYERS: 42  # 5b diff
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 90
+      SAMPLE_HEIGHT: 60
+      SAMPLE_FRAMES: 49
+      PATCH_SIZE: 2
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 226
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: True # 5b diff
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+      GRADIENT_CHECKPOINTING: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-5b@vae/diffusion_pytorch_model.safetensors # 5b diff
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 226
+      CLEAN:
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 49
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDatasetOTF
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 4
+    NUM_FRAMES: 49
+    FPS: 8
+    HEIGHT: 480
+    WIDTH: 720
+    PROMPT_PREFIX: 'DISNEY '
+    DELIMITER: '#;#'
+    FIELDS: [ 'video_path', 'prompt' ]
+    PATH_PREFIX: cache/datasets/Disney-VideoGeneration-Dataset/
+    DATA_FILE: cache/datasets/Disney-VideoGeneration-Dataset/index.jsonl
+    SAMPLER:
+      NAME: LoopSampler
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'video', 'video_latent', "prompt" ]
+        META_KEYS: [ ]
+    MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-5b@vae/diffusion_pytorch_model.safetensors
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+  #
+  EVAL_DATA:
+    NAME: Text2ImageDataset
+    MODE: eval
+    PROMPT_FILE:
+    PROMPT_DATA: [ "A girl riding a bike.", "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance." ]
+    IMAGE_SIZE: [ 480, 720 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: 'DISNEY '
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    USE_NUM: 8
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+  #
+  EVAL_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0

scepter/methods/examples/generation/dit_flux1.0_dev_1024_lora.yaml

@@ -2,33 +2,22 @@ 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: ''
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
   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']
@@ -58,61 +47,36 @@ SOLVER:
     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
 
@@ -157,55 +121,34 @@ SOLVER:
         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
+    SAMPLER: flow_euler
     SEED: 2024
     IMAGE_SIZE: [ 1024, 1024 ]
     GUIDE_SCALE: 3.5

scepter/methods/examples/generation/dit_flux1.0_schnell_1024_lora.yaml

@@ -2,35 +2,24 @@ 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: ''
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
   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'] #
+  SAVE_MODULES: [ 'model']
   TRAIN_MODULES: ['model']
   #
   FILE_SYSTEM:
@@ -58,12 +47,9 @@ SOLVER:
     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
@@ -157,54 +143,33 @@ SOLVER:
         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
+    SAMPLER: flow_euler
     SEED: 2024
     IMAGE_SIZE: [ 1024, 1024 ]
     GUIDE_SCALE: 3.5

scepter/methods/studio/chatbot/chatbot.yaml

@@ -0,0 +1,25 @@
+WORK_DIR: chatbot
+FILE_SYSTEM:
+  - NAME: LocalFs
+    TEMP_DIR: ./cache/cache_data
+  - NAME: ModelscopeFs
+    TEMP_DIR: ./cache/cache_data
+  - NAME: HuggingfaceFs
+    TEMP_DIR: ./cache/cache_data
+#
+ENABLE_I2V: False
+SKIP_EXAMPLES: True
+#
+MODEL:
+  EDIT_MODEL:
+    MODEL_CFG_DIR: scepter/methods/studio/chatbot/models/
+  I2V:
+    MODEL_NAME: CogVideoX-5b-I2V
+    MODEL_DIR: ms://ZhipuAI/CogVideoX-5b-I2V/
+  CAPTIONER:
+    MODEL_NAME: InternVL2-2B
+    MODEL_DIR: ms://OpenGVLab/InternVL2-2B/
+    PROMPT: '<image>\nThis image is the first frame of a video. Based on this image, please imagine what changes may occur in the next few seconds of the video. Please output brief description, such as "a dog running" or "a person turns to left". No more than 30 words.'
+  ENHANCER:
+    MODEL_NAME: Meta-Llama-3.1-8B-Instruct
+    MODEL_DIR: ms://LLM-Research/Meta-Llama-3.1-8B-Instruct/

scepter/methods/studio/chatbot/models/ace_0.6b_1024.yaml

@@ -0,0 +1,128 @@
+NAME: ACE_0.6B_1024
+IS_DEFAULT: False
+USE_DYNAMIC_MODEL: True
+DEFAULT_PARAS:
+  PARAS:
+  #
+  INPUT:
+    INPUT_IMAGE:
+    INPUT_MASK:
+    TASK:
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    OUTPUT_HEIGHT: 1024
+    OUTPUT_WIDTH: 1024
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+    SEED: -1
+    TAR_INDEX: 0
+  OUTPUT:
+    LATENT:
+    IMAGES:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        - NAME: encode
+          DTYPE: float16
+          INPUT: ["IMAGE"]
+        - NAME: decode
+          DTYPE: float16
+          INPUT: ["LATENT"]
+    #
+    DIFFUSION_MODEL:
+      FUNCTION:
+        - NAME: forward
+          DTYPE: float16
+          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE"]
+    #
+    COND_STAGE_MODEL:
+      FUNCTION:
+        - NAME: encode_list_of_list
+          DTYPE: bfloat16
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  NAME: LatentDiffusionACE
+  PRETRAINED_MODEL:
+  IGNORE_KEYS: [ ]
+  SCALE_FACTOR: 0.18215
+  SIZE_FACTOR: 8
+  DECODER_BIAS: 0.5
+  DEFAULT_N_PROMPT: ""
+  TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+  USE_TEXT_POS_EMBEDDINGS: True
+  #
+  DIFFUSION:
+    NAME: BaseDiffusion
+    PREDICTION_TYPE: eps
+    MIN_SNR_GAMMA:
+    NOISE_SCHEDULER:
+      NAME: LinearScheduler
+      NUM_TIMESTEPS: 1000
+      BETA_MIN: 0.0001
+      BETA_MAX: 0.02
+  #
+  DIFFUSION_MODEL:
+    NAME: ACE
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+    IGNORE_KEYS: [ ]
+    PATCH_SIZE: 2
+    IN_CHANNELS: 4
+    HIDDEN_SIZE: 1152
+    DEPTH: 28
+    NUM_HEADS: 16
+    MLP_RATIO: 4.0
+    PRED_SIGMA: True
+    DROP_PATH: 0.0
+    WINDOW_DIZE: 0
+    Y_CHANNELS: 4096
+    MAX_SEQ_LEN: 4096
+    QK_NORM: True
+    USE_GRAD_CHECKPOINT: True
+    ATTENTION_BACKEND: flash_attn
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKL
+    EMBED_DIM: 4
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/vae/vae.bin
+    IGNORE_KEYS: []
+    #
+    ENCODER:
+      NAME: Encoder
+      CH: 128
+      OUT_CH: 3
+      NUM_RES_BLOCKS: 2
+      IN_CHANNELS: 3
+      ATTN_RESOLUTIONS: [ ]
+      CH_MULT: [ 1, 2, 4, 4 ]
+      Z_CHANNELS: 4
+      DOUBLE_Z: True
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+    #
+    DECODER:
+      NAME: Decoder
+      CH: 128
+      OUT_CH: 3
+      NUM_RES_BLOCKS: 2
+      IN_CHANNELS: 3
+      ATTN_RESOLUTIONS: [ ]
+      CH_MULT: [ 1, 2, 4, 4 ]
+      Z_CHANNELS: 4
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+      GIVE_PRE_END: False
+      TANH_OUT: False
+  #
+  COND_STAGE_MODEL:
+    NAME: T5EmbedderHF
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+    TOKENIZER_PATH: ms://iic/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+    LENGTH: 120
+    T5_DTYPE: bfloat16
+    ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+    CLEAN: whitespace
+    USE_GRAD: False

scepter/methods/studio/chatbot/models/ace_0.6b_1024_refiner.yaml

@@ -0,0 +1,284 @@
+NAME: ACE_0.6B_1024_REFINER
+IS_DEFAULT: False
+USE_DYNAMIC_MODEL: True
+DEFAULT_PARAS:
+  PARAS:
+  #
+  INPUT:
+    INPUT_IMAGE:
+    INPUT_MASK:
+    TASK:
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    OUTPUT_HEIGHT: 1024
+    OUTPUT_WIDTH: 1024
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+    SEED: -1
+    TAR_INDEX: 0
+    REFINER_SCALE: 0.2
+    USE_ACE: True
+    #REFINER_PROMPT: "High Resolution, Sharpness, Clarity, Detail Enhancement, Noise Reduction, HD, 4k, Image Restoration, HDR"
+    REFINER_PROMPT: "High Resolution, Sharpness, Clarity, Detail Enhancement, Noise Reduction, HD, 4k, Image Restoration, HDR"
+  OUTPUT:
+    LATENT:
+    IMAGES:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        - NAME: encode
+          DTYPE: float16
+          INPUT: ["IMAGE"]
+        - NAME: decode
+          DTYPE: float16
+          INPUT: ["LATENT"]
+    #
+    DIFFUSION_MODEL:
+      FUNCTION:
+        - NAME: forward
+          DTYPE: float16
+          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE"]
+    #
+    COND_STAGE_MODEL:
+      FUNCTION:
+        - NAME: encode_list_of_list
+          DTYPE: bfloat16
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  NAME: LatentDiffusionACE
+  PRETRAINED_MODEL:
+  IGNORE_KEYS: [ ]
+  SCALE_FACTOR: 0.18215
+  SIZE_FACTOR: 8
+  DECODER_BIAS: 0.5
+  DEFAULT_N_PROMPT: ""
+  TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+  USE_TEXT_POS_EMBEDDINGS: True
+  #
+  DIFFUSION:
+    NAME: BaseDiffusion
+    PREDICTION_TYPE: eps
+    MIN_SNR_GAMMA:
+    NOISE_SCHEDULER:
+      NAME: LinearScheduler
+      NUM_TIMESTEPS: 1000
+      BETA_MIN: 0.0001
+      BETA_MAX: 0.02
+  #
+  DIFFUSION_MODEL:
+    NAME: ACE
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+    IGNORE_KEYS: [ ]
+    PATCH_SIZE: 2
+    IN_CHANNELS: 4
+    HIDDEN_SIZE: 1152
+    DEPTH: 28
+    NUM_HEADS: 16
+    MLP_RATIO: 4.0
+    PRED_SIGMA: True
+    DROP_PATH: 0.0
+    WINDOW_DIZE: 0
+    Y_CHANNELS: 4096
+    MAX_SEQ_LEN: 4096
+    QK_NORM: True
+    USE_GRAD_CHECKPOINT: True
+    ATTENTION_BACKEND: flash_attn
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKL
+    EMBED_DIM: 4
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/vae/vae.bin
+    IGNORE_KEYS: []
+    #
+    ENCODER:
+      NAME: Encoder
+      CH: 128
+      OUT_CH: 3
+      NUM_RES_BLOCKS: 2
+      IN_CHANNELS: 3
+      ATTN_RESOLUTIONS: [ ]
+      CH_MULT: [ 1, 2, 4, 4 ]
+      Z_CHANNELS: 4
+      DOUBLE_Z: True
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+    #
+    DECODER:
+      NAME: Decoder
+      CH: 128
+      OUT_CH: 3
+      NUM_RES_BLOCKS: 2
+      IN_CHANNELS: 3
+      ATTN_RESOLUTIONS: [ ]
+      CH_MULT: [ 1, 2, 4, 4 ]
+      Z_CHANNELS: 4
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+      GIVE_PRE_END: False
+      TANH_OUT: False
+  #
+  COND_STAGE_MODEL:
+    NAME: T5EmbedderHF
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+    TOKENIZER_PATH: ms://iic/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+    LENGTH: 120
+    T5_DTYPE: bfloat16
+    ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+    CLEAN: whitespace
+    USE_GRAD: False
+
+ACE_PROMPT: [
+  "A cute cartoon rabbit holding a whiteboard that says 'ACE Refiner', standing in a sunny meadow filled with flowers, with a big smile and bright colors.",
+  "A beautiful young woman with long flowing hair, wearing a summer dress, holding a whiteboard that reads 'ACE Refiner' while sitting on a park bench surrounded by cherry blossoms.",
+  "An adorable cartoon cat wearing oversized glasses, holding a whiteboard that says 'ACE Refiner', perched on a stack of colorful books in a cozy library setting.",
+  "A charming girl with pigtails, wearing a cute school uniform, enthusiastically holding a whiteboard that has 'ACE Refiner' written on it, in a bright and cheerful classroom full of educational posters.",
+  "A friendly cartoon dog with floppy ears, sitting in front of a doghouse, proudly holding a whiteboard that says 'ACE Refiner', with a playful expression and a blue sky in the background.",
+  "A cute anime girl with big expressive eyes, dressed in a colorful outfit, holding a whiteboard that reads 'ACE Refiner' in a fantastical landscape filled with mythical creatures.",
+  "A vibrant cartoon fox holding a whiteboard that says 'ACE Refiner', standing on a rock by a sparkling stream, surrounded by lush greenery and butterflies.",
+  "A stylish young woman in a business outfit, smiling as she holds a whiteboard written with 'ACE Refiner', in a modern office filled with plants and natural light.",
+  "A cute cartoon unicorn holding a sparkling whiteboard that says 'ACE Refiner', frolicking in a magical forest, with rainbows and stars in the background.",
+  "A happy family, consisting of a cute little girl and her playful puppy, holding a whiteboard that says 'ACE Refiner', together in their backyard on a sunny day."
+]
+REFINER_MODEL:
+  NAME: ""
+  IS_DEFAULT: False
+  DEFAULT_PARAS:
+    PARAS:
+      RESOLUTIONS: [ [ 1024, 1024 ] ]
+    INPUT:
+      INPUT_IMAGE:
+      INPUT_MASK:
+      TASK:
+      PROMPT: ""
+      NEGATIVE_PROMPT: ""
+      OUTPUT_HEIGHT: 1024
+      OUTPUT_WIDTH: 1024
+      SAMPLER: flow_euler
+      SAMPLE_STEPS: 30
+      GUIDE_SCALE: 3.5
+      GUIDE_RESCALE:
+    OUTPUT:
+      LATENT:
+      IMAGES:
+      SEED:
+    MODULES_PARAS:
+      FIRST_STAGE_MODEL:
+        FUNCTION:
+          - NAME: encode
+            DTYPE: bfloat16
+            INPUT: [ "IMAGE" ]
+          - NAME: decode
+            DTYPE: bfloat16
+            INPUT: [ "LATENT" ]
+        PARAS:
+          SCALE_FACTOR: 1.5305
+          SHIFT_FACTOR: 0.0609
+          SIZE_FACTOR: 8
+      DIFFUSION_MODEL:
+        FUNCTION:
+          - NAME: forward
+            DTYPE: bfloat16
+            INPUT: [ "SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE" ]
+      COND_STAGE_MODEL:
+        FUNCTION:
+          - NAME: encode
+            DTYPE: bfloat16
+            INPUT: [ "PROMPT" ]
+
+  MODEL:
+    DIFFUSION:
+      NAME: DiffusionFluxRF
+      PREDICTION_TYPE: raw
+      NOISE_SCHEDULER:
+        NAME: FlowMatchSigmaScheduler
+        WEIGHTING_SCHEME: logit_normal
+        SHIFT: 3.0
+        LOGIT_MEAN: 0.0
+        LOGIT_STD: 1.0
+        MODE_SCALE: 1.29
+    DIFFUSION_MODEL:
+      NAME: FluxMR
+      PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@flux1-dev.safetensors
+      IN_CHANNELS: 64
+      OUT_CHANNELS: 64
+      HIDDEN_SIZE: 3072
+      NUM_HEADS: 24
+      AXES_DIM: [ 16, 56, 56 ]
+      THETA: 10000
+      VEC_IN_DIM: 768
+      GUIDANCE_EMBED: True
+      CONTEXT_IN_DIM: 4096
+      MLP_RATIO: 4.0
+      QKV_BIAS: True
+      DEPTH: 19
+      DEPTH_SINGLE_BLOCKS: 38
+      USE_GRAD_CHECKPOINT: True
+      ATTN_BACKEND: flash_attn
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLFlux
+      EMBED_DIM: 16
+      PRETRAINED_MODEL:  ms://AI-ModelScope/FLUX.1-dev@ae.safetensors
+      IGNORE_KEYS: [ ]
+      BATCH_SIZE: 8
+      USE_CONV: False
+      SCALE_FACTOR: 0.3611
+      SHIFT_FACTOR: 0.1159
+      #
+      ENCODER:
+        NAME: Encoder
+        USE_CHECKPOINT: False
+        CH: 128
+        OUT_CH: 3
+        NUM_RES_BLOCKS: 2
+        IN_CHANNELS: 3
+        ATTN_RESOLUTIONS: [ ]
+        CH_MULT: [ 1, 2, 4, 4 ]
+        Z_CHANNELS: 16
+        DOUBLE_Z: True
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+      #
+      DECODER:
+        NAME: Decoder
+        USE_CHECKPOINT: False
+        CH: 128
+        OUT_CH: 3
+        NUM_RES_BLOCKS: 2
+        IN_CHANNELS: 3
+        ATTN_RESOLUTIONS: [ ]
+        CH_MULT: [ 1, 2, 4, 4 ]
+        Z_CHANNELS: 16
+        DROPOUT: 0.0
+        RESAMP_WITH_CONV: True
+        GIVE_PRE_END: False
+        TANH_OUT: False
+    #
+    COND_STAGE_MODEL:
+      NAME: T5PlusClipFluxEmbedder
+      T5_MODEL:
+        NAME: HFEmbedder
+        HF_MODEL_CLS: T5EncoderModel
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder_2/
+        HF_TOKENIZER_CLS: T5Tokenizer
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer_2/
+        MAX_LENGTH: 512
+        OUTPUT_KEY: last_hidden_state
+        D_TYPE: bfloat16
+        BATCH_INFER: False
+        CLEAN: whitespace
+      CLIP_MODEL:
+        NAME: HFEmbedder
+        HF_MODEL_CLS: CLIPTextModel
+        MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder/
+        HF_TOKENIZER_CLS: CLIPTokenizer
+        TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer/
+        MAX_LENGTH: 77
+        OUTPUT_KEY: pooler_output
+        D_TYPE: bfloat16
+        BATCH_INFER: True
+        CLEAN: whitespace

scepter/methods/studio/chatbot/models/ace_0.6b_512.yaml

@@ -0,0 +1,128 @@
+NAME: ACE_0.6B_512
+IS_DEFAULT: True
+USE_DYNAMIC_MODEL: True
+DEFAULT_PARAS:
+  PARAS:
+  #
+  INPUT:
+    INPUT_IMAGE:
+    INPUT_MASK:
+    TASK:
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    OUTPUT_HEIGHT: 512
+    OUTPUT_WIDTH: 512
+    SAMPLER: ddim
+    SAMPLE_STEPS: 20
+    GUIDE_SCALE: 4.5
+    GUIDE_RESCALE: 0.5
+    SEED: -1
+    TAR_INDEX: 0
+  OUTPUT:
+    LATENT:
+    IMAGES:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        - NAME: encode
+          DTYPE: float16
+          INPUT: ["IMAGE"]
+        - NAME: decode
+          DTYPE: float16
+          INPUT: ["LATENT"]
+    #
+    DIFFUSION_MODEL:
+      FUNCTION:
+        - NAME: forward
+          DTYPE: float16
+          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE"]
+    #
+    COND_STAGE_MODEL:
+      FUNCTION:
+        - NAME: encode_list_of_list
+          DTYPE: bfloat16
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  NAME: LatentDiffusionACE
+  PRETRAINED_MODEL:
+  IGNORE_KEYS: [ ]
+  SCALE_FACTOR: 0.18215
+  SIZE_FACTOR: 8
+  DECODER_BIAS: 0.5
+  DEFAULT_N_PROMPT: ""
+  TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+  USE_TEXT_POS_EMBEDDINGS: True
+  #
+  DIFFUSION:
+    NAME: BaseDiffusion
+    PREDICTION_TYPE: eps
+    MIN_SNR_GAMMA:
+    NOISE_SCHEDULER:
+      NAME: LinearScheduler
+      NUM_TIMESTEPS: 1000
+      BETA_MIN: 0.0001
+      BETA_MAX: 0.02
+  #
+  DIFFUSION_MODEL:
+    NAME: ACE
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-512px@models/dit/ace_0.6b_512px.pth
+    IGNORE_KEYS: [ ]
+    PATCH_SIZE: 2
+    IN_CHANNELS: 4
+    HIDDEN_SIZE: 1152
+    DEPTH: 28
+    NUM_HEADS: 16
+    MLP_RATIO: 4.0
+    PRED_SIGMA: True
+    DROP_PATH: 0.0
+    WINDOW_DIZE: 0
+    Y_CHANNELS: 4096
+    MAX_SEQ_LEN: 1024
+    QK_NORM: True
+    USE_GRAD_CHECKPOINT: True
+    ATTENTION_BACKEND: flash_attn
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKL
+    EMBED_DIM: 4
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-512px@models/vae/vae.bin
+    IGNORE_KEYS: []
+    #
+    ENCODER:
+      NAME: Encoder
+      CH: 128
+      OUT_CH: 3
+      NUM_RES_BLOCKS: 2
+      IN_CHANNELS: 3
+      ATTN_RESOLUTIONS: [ ]
+      CH_MULT: [ 1, 2, 4, 4 ]
+      Z_CHANNELS: 4
+      DOUBLE_Z: True
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+    #
+    DECODER:
+      NAME: Decoder
+      CH: 128
+      OUT_CH: 3
+      NUM_RES_BLOCKS: 2
+      IN_CHANNELS: 3
+      ATTN_RESOLUTIONS: [ ]
+      CH_MULT: [ 1, 2, 4, 4 ]
+      Z_CHANNELS: 4
+      DROPOUT: 0.0
+      RESAMP_WITH_CONV: True
+      GIVE_PRE_END: False
+      TANH_OUT: False
+  #
+  COND_STAGE_MODEL:
+    NAME: T5EmbedderHF
+    PRETRAINED_MODEL: ms://iic/ACE-0.6B-512px@models/text_encoder/t5-v1_1-xxl/
+    TOKENIZER_PATH: ms://iic/ACE-0.6B-512px@models/tokenizer/t5-v1_1-xxl
+    LENGTH: 120
+    T5_DTYPE: bfloat16
+    ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+    CLEAN: whitespace
+    USE_GRAD: False

scepter/methods/studio/inference/dit/cogvideox_2b_pro.yaml

@@ -0,0 +1,152 @@
+NAME: COGVIDEOX_2B
+IS_DEFAULT: False
+DEFAULT_PARAS:
+  PARAS:
+    RESOLUTIONS: [[480, 720]]
+  INPUT:
+    IMAGE:
+    ORIGINAL_SIZE_AS_TUPLE: [480, 720]
+    TARGET_SIZE_AS_TUPLE: [480, 720]
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    PROMPT_PREFIX: ""
+    SAMPLE: ddim
+    SAMPLE_STEPS: 50
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    DISCRETIZATION: trailing
+    NUM_FRAMES:
+      DEFAULT: 49
+      VISIBLE: True
+    FPS:
+      DEFAULT: 8
+      VISIBLE: True
+  OUTPUT:
+    VIDEOS:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        -
+          NAME: decode
+          DTYPE: bfloat16
+          INPUT: ["LATENT"]
+      PARAS:
+          SCALING_FACTOR_IMAGE: 1.15258426
+    DIFFUSION_MODEL:
+      FUNCTION:
+        -
+          NAME: forward
+          DTYPE: bfloat16
+          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE", "GUIDE_RESCALE", "DISCRETIZATION", "NUM_FRAMES", "FPS"]
+      PARAS:
+        USE_ROTARY_POSITIONAL_EMBEDDINGS: False
+        PATCH_SIZE: 2
+        LATENT_CHANNELS: 16
+        SCALE_FACTOR_SPATIAL: 8
+        SCALE_FACTOR_TEMPORAL: 4
+        ATTENTION_HEAD_DIM: 64
+        SAMPLE_HEIGHT: 480
+        SAMPLE_WIDTH: 720
+    COND_STAGE_MODEL:
+      FUNCTION:
+        -
+          NAME: encode
+          DTYPE: bfloat16
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  PRETRAINED_MODEL:
+  DIFFUSION:
+    NAME: BaseDiffusion
+    PREDICTION_TYPE: v
+    NOISE_SCHEDULER:
+      NAME: ScaledLinearScheduler
+      BETA_MIN: 0.00085
+      BETA_MAX: 0.012
+      SNR_SHIFT_SCALE: 3.0
+      RESCALE_BETAS_ZERO_SNR: True
+    DIFFUSION_SAMPLERS:
+      NAME: DDIMSampler
+      DISCRETIZATION_TYPE: trailing
+      ETA: 0.0
+  #
+  DIFFUSION_MODEL:
+    NAME: CogVideoXTransformer3DModel
+    DTYPE: bfloat16
+    PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@transformer/diffusion_pytorch_model.safetensors
+    NUM_ATTENTION_HEADS: 30
+    ATTENTION_HEAD_DIM: 64
+    IN_CHANNELS: 16
+    OUT_CHANNELS: 16
+    FLIP_SIN_TO_COS: True
+    FREQ_SHIFT: 0
+    TIME_EMBED_DIM: 512
+    TEXT_EMBED_DIM: 4096
+    NUM_LAYERS: 30
+    DROPOUT: 0.0
+    ATTENTION_BIAS: True
+    SAMPLE_WIDTH: 90
+    SAMPLE_HEIGHT: 60
+    SAMPLE_FRAMES: 49
+    PATCH_SIZE: 2
+    TEMPORAL_COMPRESSION_RATIO: 4
+    MAX_TEXT_SEQ_LENGTH: 226
+    ACTIVATION_FN: "gelu-approximate"
+    TIMESTEP_ACTIVATION_FN: "silu"
+    NORM_ELEMENTWISE_AFFINE: True
+    NORM_EPS: 1e-5
+    SPATIAL_INTERPOLATION_SCALE: 1.875
+    TEMPORAL_INTERPOLATION_SCALE: 1.0
+    USE_ROTARY_POSITIONAL_EMBEDDINGS: False
+    USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+    GRADIENT_CHECKPOINTING: False
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKLCogVideoX
+    DTYPE: bfloat16
+    PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@vae/diffusion_pytorch_model.safetensors
+    SAMPLE_HEIGHT: 480
+    SAMPLE_WIDTH: 720
+    USE_QUANT_CONV: False
+    USE_POST_QUANT_CONV: False
+    USE_SLICING: True
+    USE_TILING: True
+    GRADIENT_CHECKPOINTING: False
+    ENCODER:
+      NAME: CogVideoXEncoder3D
+      IN_CHANNELS: 3
+      OUT_CHANNELS: 16
+      UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+      BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+      LAYERS_PER_BLOCK: 3
+      ACT_FN: "silu"
+      NORM_EPS: 1e-6
+      NORM_NUM_GROUPS: 32
+      DROPOUT: 0.0
+      PAD_MODE: "first"
+      TEMPORAL_COMPRESSION_RATIO: 4
+      GRADIENT_CHECKPOINTING: False
+    DECODER:
+      NAME: CogVideoXDecoder3D
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 3
+      UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+      BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+      LAYERS_PER_BLOCK: 3
+      ACT_FN: "silu"
+      NORM_EPS: 1e-6
+      NORM_NUM_GROUPS: 32
+      DROPOUT: 0.0
+      PAD_MODE: "first"
+      TEMPORAL_COMPRESSION_RATIO: 4
+      GRADIENT_CHECKPOINTING: False
+  #
+  COND_STAGE_MODEL:
+    NAME: T5EmbedderHF
+    PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+    TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+    LENGTH: 226
+    CLEAN:
+    USE_GRAD: False
+    T5_DTYPE: bfloat16

scepter/methods/studio/inference/dit/cogvideox_5b_pro.yaml

@@ -0,0 +1,154 @@
+NAME: COGVIDEOX_5B
+IS_DEFAULT: False
+DEFAULT_PARAS:
+  PARAS:
+    RESOLUTIONS: [[480, 720]]
+  INPUT:
+    IMAGE:
+    ORIGINAL_SIZE_AS_TUPLE: [480, 720]
+    TARGET_SIZE_AS_TUPLE: [480, 720]
+    PROMPT: ""
+    NEGATIVE_PROMPT: ""
+    PROMPT_PREFIX: ""
+    SAMPLE: ddim
+    SAMPLE_STEPS: 50
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    DISCRETIZATION: trailing
+    NUM_FRAMES:
+      DEFAULT: 49
+      VISIBLE: True
+    FPS:
+      DEFAULT: 8
+      VISIBLE: True
+  OUTPUT:
+    VIDEOS:
+    SEED:
+  MODULES_PARAS:
+    FIRST_STAGE_MODEL:
+      FUNCTION:
+        -
+          NAME: decode
+          DTYPE: bfloat16
+          INPUT: ["LATENT"]
+      PARAS:
+          SCALING_FACTOR_IMAGE: 0.7  # 5b diff
+    DIFFUSION_MODEL:
+      FUNCTION:
+        -
+          NAME: forward
+          DTYPE: bfloat16
+          INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE", "GUIDE_RESCALE", "DISCRETIZATION", "NUM_FRAMES", "FPS"]
+      PARAS:
+        USE_ROTARY_POSITIONAL_EMBEDDINGS: True
+        PATCH_SIZE: 2
+        LATENT_CHANNELS: 16
+        SCALE_FACTOR_SPATIAL: 8
+        SCALE_FACTOR_TEMPORAL: 4
+        ATTENTION_HEAD_DIM: 64
+        SAMPLE_HEIGHT: 480
+        SAMPLE_WIDTH: 720
+    COND_STAGE_MODEL:
+      FUNCTION:
+        -
+          NAME: encode
+          DTYPE: bfloat16
+          INPUT: ["PROMPT"]
+#
+MODEL:
+  PRETRAINED_MODEL:
+  DIFFUSION:
+    NAME: BaseDiffusion
+    PREDICTION_TYPE: v
+    NOISE_SCHEDULER:
+      NAME: ScaledLinearScheduler
+      BETA_MIN: 0.00085
+      BETA_MAX: 0.012
+      SNR_SHIFT_SCALE: 1.0  # 5b diff
+      RESCALE_BETAS_ZERO_SNR: True
+    DIFFUSION_SAMPLERS:
+      NAME: DDIMSampler
+      DISCRETIZATION_TYPE: trailing
+      ETA: 0.0
+  #
+  DIFFUSION_MODEL:
+    NAME: CogVideoXTransformer3DModel
+    DTYPE: bfloat16
+    PRETRAINED_MODEL: # 5b diff
+      - ms://AI-ModelScope/CogVideoX-5b@transformer/diffusion_pytorch_model-00001-of-00002.safetensors
+      - ms://AI-ModelScope/CogVideoX-5b@transformer/diffusion_pytorch_model-00002-of-00002.safetensors
+    NUM_ATTENTION_HEADS: 48  # 5b diff
+    ATTENTION_HEAD_DIM: 64
+    IN_CHANNELS: 16
+    OUT_CHANNELS: 16
+    FLIP_SIN_TO_COS: True
+    FREQ_SHIFT: 0
+    TIME_EMBED_DIM: 512
+    TEXT_EMBED_DIM: 4096
+    NUM_LAYERS: 42  # 5b diff
+    DROPOUT: 0.0
+    ATTENTION_BIAS: True
+    SAMPLE_WIDTH: 90
+    SAMPLE_HEIGHT: 60
+    SAMPLE_FRAMES: 49
+    PATCH_SIZE: 2
+    TEMPORAL_COMPRESSION_RATIO: 4
+    MAX_TEXT_SEQ_LENGTH: 226
+    ACTIVATION_FN: "gelu-approximate"
+    TIMESTEP_ACTIVATION_FN: "silu"
+    NORM_ELEMENTWISE_AFFINE: True
+    NORM_EPS: 1e-5
+    SPATIAL_INTERPOLATION_SCALE: 1.875
+    TEMPORAL_INTERPOLATION_SCALE: 1.0
+    USE_ROTARY_POSITIONAL_EMBEDDINGS: True # 5b diff
+    USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+    GRADIENT_CHECKPOINTING: True
+  #
+  FIRST_STAGE_MODEL:
+    NAME: AutoencoderKLCogVideoX
+    DTYPE: bfloat16
+    PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-5b@vae/diffusion_pytorch_model.safetensors # 5b diff
+    SAMPLE_HEIGHT: 480
+    SAMPLE_WIDTH: 720
+    USE_QUANT_CONV: False
+    USE_POST_QUANT_CONV: False
+    USE_SLICING: True
+    USE_TILING: True
+    GRADIENT_CHECKPOINTING: True
+    ENCODER:
+      NAME: CogVideoXEncoder3D
+      IN_CHANNELS: 3
+      OUT_CHANNELS: 16
+      UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+      BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+      LAYERS_PER_BLOCK: 3
+      ACT_FN: "silu"
+      NORM_EPS: 1e-6
+      NORM_NUM_GROUPS: 32
+      DROPOUT: 0.0
+      PAD_MODE: "first"
+      TEMPORAL_COMPRESSION_RATIO: 4
+      GRADIENT_CHECKPOINTING: True
+    DECODER:
+      NAME: CogVideoXDecoder3D
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 3
+      UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+      BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+      LAYERS_PER_BLOCK: 3
+      ACT_FN: "silu"
+      NORM_EPS: 1e-6
+      NORM_NUM_GROUPS: 32
+      DROPOUT: 0.0
+      PAD_MODE: "first"
+      TEMPORAL_COMPRESSION_RATIO: 4
+      GRADIENT_CHECKPOINTING: True
+  #
+  COND_STAGE_MODEL:
+    NAME: T5EmbedderHF
+    PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+    TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+    LENGTH: 226
+    CLEAN:
+    USE_GRAD: False
+    T5_DTYPE: bfloat16

scepter/methods/studio/inference/dit/flux1.0_dev_pro.yaml

@@ -13,8 +13,8 @@ DEFAULT_PARAS:
       VISIBLE: False
     PROMPT_PREFIX: ""
     SAMPLE:
-      VALUES: ["flow_eluer"]
-      DEFAULT: "flow_eluer"
+      VALUES: ["flow_euler"]
+      DEFAULT: "flow_euler"
     SAMPLE_STEPS: 50
     GUIDE_SCALE: 3.5
     GUIDE_RESCALE:

scepter/methods/studio/inference/dit/flux1.0_schnell_pro.yaml

@@ -13,8 +13,8 @@ DEFAULT_PARAS:
       VISIBLE: False
     PROMPT_PREFIX: ""
     SAMPLE:
-      VALUES: ["flow_eluer"]
-      DEFAULT: "flow_eluer"
+      VALUES: ["flow_euler"]
+      DEFAULT: "flow_euler"
     SAMPLE_STEPS: 4
     GUIDE_SCALE: 3.5
     GUIDE_RESCALE:

scepter/methods/studio/inference/inference.yaml

@@ -1,4 +1,5 @@
 WORK_DIR: "inference"
+SKIP_EXAMPLES: True
 DIFFUSION_PARAS:
   SAMPLE:
     VALUES: ['ddim', 'euler', 'euler_ancestral', 'heun', 'dpm2',
@@ -18,6 +19,16 @@ DIFFUSION_PARAS:
     MAX: 4
     DEFAULT: 1
     VISIBLE: True
+  NUM_FRAMES:
+    MIN: 1
+    MAX: 100
+    DEFAULT: 49
+    VISIBLE: False
+  FPS:
+    MIN: 1
+    MAX: 50
+    DEFAULT: 8
+    VISIBLE: False
   SAMPLE_STEPS:
     MIN: 1
     MAX: 100
@@ -93,7 +104,8 @@ DIFFUSION_PARAS:
              [1664, 576], [1728, 576],
              [2048, 2048], [2048, 1920], [1920, 2048],
              [1536, 2560], [2560, 1536], [2560, 1440],
-             [2560, 1440]
+             [2560, 1440],
+             [480, 720], [720, 480]
             ]
     DEFAULT: [1024, 1024]
     VISIBLE: True

scepter/methods/studio/preprocess/preprocess.yaml

@@ -450,3 +450,28 @@ PROCESSORS:
       SRC_IMAGE_TOOL: sketch
       SRC_IMAGE_INTERACTIVE: True
       CAPTION_INTERACTIVE: False
+
+VIDEO_PROCESSORS:
+  - NAME: CogVLM2Llama3Caption
+    TYPE: caption
+    MODEL_PATH: ms://ZhipuAI/cogvlm2-llama3-caption
+    DEVICE: "gpu"
+    MEMORY: 20000
+    PROMPT: Please describe this video in detail.
+    TEMPERATURE: 0.1
+    MAX_NEW_TOKENS: 2048
+    PAD_TOKEN_ID: 128002
+    TOP_K: 1
+    TOP_P: 0.1
+
+TRANSLATION_PROCESSORS:
+  - NAME: OpusMtZhEn
+    TYPE: caption
+    MODEL_PATH: ms://cubeai/trans-opus-mt-zh-en
+    DEVICE: "gpu"
+    MEMORY: 5000
+  - NAME: OpusMtEnZh
+    TYPE: caption
+    MODEL_PATH: ms://cubeai/trans-opus-mt-en-zh
+    DEVICE: "gpu"
+    MEMORY: 5000

scepter/methods/studio/scepter_ui.yaml

@@ -87,3 +87,7 @@ INTERFACE:
     NAME_EN: Inference
     IFID: inference
     CONFIG: scepter/methods/studio/inference/inference.yaml
+  - NAME: 对话式编辑
+    NAME_EN: ChatBot
+    IFID: chatbot
+    CONFIG: scepter/methods/studio/chatbot/chatbot.yaml

scepter/methods/studio/self_train/dit/cogvideox_2b_pro.yaml

@@ -0,0 +1,315 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+META:
+  VERSION: 'COGVIDEOX_2B'
+  DESCRIPTION: "cogvideox 2b"
+  IS_DEFAULT: False
+  IS_SHARE: True
+  INFERENCE_PARAS:
+    INFERENCE_BATCH_SIZE: 1
+    INFERENCE_PREFIX: ""
+    DEFAULT_SAMPLER: "ddim"
+    DEFAULT_SAMPLE_STEPS: 50
+    INFERENCE_N_PROMPT: ""
+    RESOLUTION: [ 480, 720 ]
+  PARAS:
+    - TRAIN_BATCH_SIZE: 1
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [ 480, 720 ]
+      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: [ 480, 720 ]
+      MEMORY: 89000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 4e-4
+      IS_DEFAULT: True
+      TUNER: LORA
+  #
+  TUNERS:
+    LORA:
+      - NAME: SwiftLoRA
+        R: 64
+        LORA_ALPHA: 64
+        LORA_DROPOUT: 0.0
+        BIAS: "none"
+        TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox_2b_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  FPS: 8
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 1.15258426
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 3.0
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@transformer/diffusion_pytorch_model.safetensors
+      NUM_ATTENTION_HEADS: 30
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      NUM_LAYERS: 30
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 90
+      SAMPLE_HEIGHT: 60
+      SAMPLE_FRAMES: 49
+      PATCH_SIZE: 2
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 226
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: False
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+      GRADIENT_CHECKPOINTING: False
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@vae/diffusion_pytorch_model.safetensors
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: False
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: False
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: False
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 226
+      CLEAN:
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 49
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDatasetOTF
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 4
+    PROMPT_PREFIX: ''
+    DELIMITER: '#;#'
+    FIELDS: [ 'video_path', 'width', 'height', 'prompt' ]
+    PATH_PREFIX:
+    DATA_FILE:
+    SAMPLER:
+      NAME: LoopSampler
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'video', 'video_latent', "prompt" ]
+        META_KEYS: [ ]
+    MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-2b@vae/diffusion_pytorch_model.safetensors
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+  #
+  EVAL_DATA:
+    NAME: Text2ImageDataset
+    MODE: eval
+    PROMPT_FILE:
+    PROMPT_DATA: [ "A girl riding a bike.", "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance." ]
+    IMAGE_SIZE: [ 480, 720 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+#    USE_NUM: 8
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+      SAVE_LAST: True
+      SAVE_NAME_PREFIX: 'step'
+      DISABLE_SNAPSHOT: True
+  #
+  EVAL_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+      SAVE_LAST: True
+      SAVE_NAME_PREFIX: 'step'
+      SAVE_PROBE_PREFIX: 'image'

scepter/methods/studio/self_train/dit/cogvideox_5b_pro.yaml

@@ -0,0 +1,317 @@
+ENV:
+  BACKEND: nccl
+  SEED: 42
+  TENSOR_PARALLEL_SIZE: 1
+  PIPELINE_PARALLEL_SIZE: 1
+  SYS_ENVS:
+    TORCH_CUDNN_V8_API_ENABLED: '1'
+    TOKENIZERS_PARALLELISM: 'false'
+    TF_CPP_MIN_LOG_LEVEL: '3'
+    PYTORCH_CUDA_ALLOC_CONF: 'expandable_segments:True'
+META:
+  VERSION: 'COGVIDEOX_5B'
+  DESCRIPTION: "cogvideox 5b"
+  IS_DEFAULT: False
+  IS_SHARE: True
+  INFERENCE_PARAS:
+    INFERENCE_BATCH_SIZE: 1
+    INFERENCE_PREFIX: ""
+    DEFAULT_SAMPLER: "ddim"
+    DEFAULT_SAMPLE_STEPS: 50
+    INFERENCE_N_PROMPT: ""
+    RESOLUTION: [ 480, 720 ]
+  PARAS:
+    - TRAIN_BATCH_SIZE: 1
+      TRAIN_PREFIX: ""
+      TRAIN_N_PROMPT: ""
+      RESOLUTION: [ 480, 720 ]
+      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: [ 480, 720 ]
+      MEMORY: 89000
+      EPOCHS: 50
+      SAVE_INTERVAL: 25
+      EPSEC: 0.818
+      LEARNING_RATE: 4e-4
+      IS_DEFAULT: True
+      TUNER: LORA
+  #
+  TUNERS:
+    LORA:
+      - NAME: SwiftLoRA
+        R: 64
+        LORA_ALPHA: 64
+        LORA_DROPOUT: 0.0
+        BIAS: "none"
+        TARGET_MODULES: "model.*(.to_k|.to_q|.to_v|.to_out.0)$"
+#
+SOLVER:
+  NAME: LatentDiffusionVideoSolver
+  MAX_STEPS: 2000
+  USE_AMP: True
+  DTYPE: bfloat16
+  USE_FAIRSCALE: False
+  USE_FSDP: True
+  LOAD_MODEL_ONLY: False
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
+  RESUME_FROM:
+  WORK_DIR: ./cache/save_data/dit_cogvideox_5b_lora
+  LOG_FILE: std_log.txt
+  EVAL_INTERVAL: 100
+  LOG_TRAIN_NUM: 4
+  FPS: 8
+  SHARDING_STRATEGY: full_shard
+  FSDP_REDUCE_DTYPE: float32
+  FSDP_BUFFER_DTYPE: float32
+  FSDP_SHARD_MODULES: [ 'model', 'cond_stage_model.model']
+  SAVE_MODULES: [ 'model', 'cond_stage_model.model']
+  TRAIN_MODULES: ['model']
+  #
+  FILE_SYSTEM:
+    NAME: "ModelscopeFs"
+    TEMP_DIR: "./cache/cache_data"
+  #
+  TUNER:
+  #
+  MODEL:
+    NAME: LatentDiffusionCogVideoX
+    PRETRAINED_MODEL:
+    PARAMETERIZATION: v
+    TIMESTEPS: 1000
+    MIN_SNR_GAMMA: 3.0
+    ZERO_TERMINAL_SNR: True
+    SCALE_FACTOR_SPATIAL: 8
+    SCALE_FACTOR_TEMPORAL: 4
+    SCALING_FACTOR_IMAGE: 0.7  # 5b diff
+    IGNORE_KEYS: [ ]
+    DEFAULT_N_PROMPT:
+    USE_EMA: False
+    EVAL_EMA: False
+    DIFFUSION:
+      NAME: BaseDiffusion
+      PREDICTION_TYPE: v
+      NOISE_SCHEDULER:
+        NAME: ScaledLinearScheduler
+        BETA_MIN: 0.00085
+        BETA_MAX: 0.012
+        SNR_SHIFT_SCALE: 1.0  # 5b diff
+        RESCALE_BETAS_ZERO_SNR: True
+      DIFFUSION_SAMPLERS:
+        NAME: DDIMSampler
+        DISCRETIZATION_TYPE: trailing
+        ETA: 0.0
+    #
+    DIFFUSION_MODEL:
+      NAME: CogVideoXTransformer3DModel
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: # 5b diff
+        - ms://AI-ModelScope/CogVideoX-5b@transformer/diffusion_pytorch_model-00001-of-00002.safetensors
+        - ms://AI-ModelScope/CogVideoX-5b@transformer/diffusion_pytorch_model-00002-of-00002.safetensors
+      NUM_ATTENTION_HEADS: 48  # 5b diff
+      ATTENTION_HEAD_DIM: 64
+      IN_CHANNELS: 16
+      OUT_CHANNELS: 16
+      FLIP_SIN_TO_COS: True
+      FREQ_SHIFT: 0
+      TIME_EMBED_DIM: 512
+      TEXT_EMBED_DIM: 4096
+      NUM_LAYERS: 42  # 5b diff
+      DROPOUT: 0.0
+      ATTENTION_BIAS: True
+      SAMPLE_WIDTH: 90
+      SAMPLE_HEIGHT: 60
+      SAMPLE_FRAMES: 49
+      PATCH_SIZE: 2
+      TEMPORAL_COMPRESSION_RATIO: 4
+      MAX_TEXT_SEQ_LENGTH: 226
+      ACTIVATION_FN: "gelu-approximate"
+      TIMESTEP_ACTIVATION_FN: "silu"
+      NORM_ELEMENTWISE_AFFINE: True
+      NORM_EPS: 1e-5
+      SPATIAL_INTERPOLATION_SCALE: 1.875
+      TEMPORAL_INTERPOLATION_SCALE: 1.0
+      USE_ROTARY_POSITIONAL_EMBEDDINGS: True # 5b diff
+      USE_LEARNED_POSITIONAL_EMBEDDINGS: False
+      GRADIENT_CHECKPOINTING: True
+    #
+    FIRST_STAGE_MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-5b@vae/diffusion_pytorch_model.safetensors # 5b diff
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+    #
+    COND_STAGE_MODEL:
+      NAME: T5EmbedderHF
+      PRETRAINED_MODEL: ms://AI-ModelScope/t5-v1_1-xxl
+      TOKENIZER_PATH: ms://AI-ModelScope/t5-v1_1-xxl
+      LENGTH: 226
+      CLEAN:
+      USE_GRAD: False
+    #
+    LOSS:
+      NAME: ReconstructLoss
+      LOSS_TYPE: l2
+  #
+  SAMPLE_ARGS:
+    SAMPLER: ddim
+    SAMPLE_STEPS: 50
+    SEED: 42
+    GUIDE_SCALE: 6.0
+    GUIDE_RESCALE: 0.0
+    NUM_FRAMES: 49
+  #
+  OPTIMIZER:
+    NAME: Adam
+    LEARNING_RATE: 1e-3
+    BETAS: [ 0.9, 0.95 ]
+    EPS: 1e-8
+    WEIGHT_DECAY: 0.0
+    AMSGRAD: False
+  #
+#  LR_SCHEDULER:
+#    NAME: StepAnnealingLR
+#    WARMUP_STEPS: 200
+#    TOTAL_STEPS: 2000
+#    DECAY_MODE: 'cosine'
+  #
+  TRAIN_DATA:
+    NAME: VideoGenDatasetOTF
+    MODE: train
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    NUM_WORKERS: 4
+    PROMPT_PREFIX: ''
+    DELIMITER: '#;#'
+    FIELDS: [ 'video_path', 'width', 'height', 'prompt' ]
+    PATH_PREFIX:
+    DATA_FILE:
+    SAMPLER:
+      NAME: LoopSampler
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'video', 'video_latent', "prompt" ]
+        META_KEYS: [ ]
+    MODEL:
+      NAME: AutoencoderKLCogVideoX
+      DTYPE: bfloat16
+      PRETRAINED_MODEL: ms://AI-ModelScope/CogVideoX-5b@vae/diffusion_pytorch_model.safetensors
+      SAMPLE_HEIGHT: 480
+      SAMPLE_WIDTH: 720
+      USE_QUANT_CONV: False
+      USE_POST_QUANT_CONV: False
+      USE_SLICING: True
+      USE_TILING: True
+      GRADIENT_CHECKPOINTING: True
+      ENCODER:
+        NAME: CogVideoXEncoder3D
+        IN_CHANNELS: 3
+        OUT_CHANNELS: 16
+        UP_BLOCK_TYPES: [ "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+      DECODER:
+        NAME: CogVideoXDecoder3D
+        IN_CHANNELS: 16
+        OUT_CHANNELS: 3
+        UP_BLOCK_TYPES: [ "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D" ]
+        BLOCK_OUT_CHANNELS: [ 128, 256, 256, 512 ]
+        LAYERS_PER_BLOCK: 3
+        ACT_FN: "silu"
+        NORM_EPS: 1e-6
+        NORM_NUM_GROUPS: 32
+        DROPOUT: 0.0
+        PAD_MODE: "first"
+        TEMPORAL_COMPRESSION_RATIO: 4
+        GRADIENT_CHECKPOINTING: True
+  #
+  EVAL_DATA:
+    NAME: Text2ImageDataset
+    MODE: eval
+    PROMPT_FILE:
+    PROMPT_DATA: [ "A girl riding a bike.", "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance." ]
+    IMAGE_SIZE: [ 480, 720 ]
+    FIELDS: [ "prompt" ]
+    DELIMITER: '#;#'
+    PROMPT_PREFIX: ''
+    PIN_MEMORY: True
+    BATCH_SIZE: 1
+    #    USE_NUM: 8
+    NUM_WORKERS: 4
+    TRANSFORMS:
+      - NAME: Select
+        KEYS: [ 'index', 'prompt' ]
+        META_KEYS: [ 'image_size' ]
+  #
+  TRAIN_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+    - NAME: BackwardHook
+      PRIORITY: 10
+    - NAME: LogHook
+      LOG_INTERVAL: 10
+      PRIORITY: 20
+    - NAME: CheckpointHook
+      INTERVAL: 1000
+      PRIORITY: 40
+      SAVE_LAST: True
+      SAVE_NAME_PREFIX: 'step'
+      DISABLE_SNAPSHOT: True
+  #
+  EVAL_HOOKS:
+    - NAME: ProbeDataHook
+      PROB_INTERVAL: 100
+      PRIORITY: 0
+      SAVE_LAST: True
+      SAVE_NAME_PREFIX: 'step'
+      SAVE_PROBE_PREFIX: 'image'

scepter/methods/studio/self_train/dit/flux1.0_dv_pro.yaml

@@ -3,7 +3,7 @@ ENV:
 META:
   VERSION: 'FLUX1.0_DEV'
   DESCRIPTION: "flux 1.0 dev"
-  IS_DEFAULT: False
+  IS_DEFAULT: True
   IS_SHARE: True
   INFERENCE_PARAS:
     INFERENCE_BATCH_SIZE: 1
@@ -50,43 +50,33 @@ META:
 #
 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'
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
   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'] #
+  SAVE_MODULES: [ 'model']
   TRAIN_MODULES: ['model']
-  #
+
   FILE_SYSTEM:
     NAME: "ModelscopeFs"
     TEMP_DIR: "./cache/cache_data"
-  #
+
   FREEZE:
-  #
+
   TUNER:
-  #
+
   MODEL:
     NAME: LatentDiffusionFlux
     PARAMETERIZATION: rf
@@ -99,65 +89,39 @@ SOLVER:
     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
@@ -167,7 +131,7 @@ SOLVER:
       USE_CONV: False
       SCALE_FACTOR: 0.3611
       SHIFT_FACTOR: 0.1159
-      #
+
       ENCODER:
         NAME: Encoder
         USE_CHECKPOINT: True
@@ -181,7 +145,7 @@ SOLVER:
         DOUBLE_Z: True
         DROPOUT: 0.0
         RESAMP_WITH_CONV: True
-      #
+
       DECODER:
         NAME: Decoder
         USE_CHECKPOINT: True
@@ -196,61 +160,40 @@ SOLVER:
         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
+    SAMPLER: flow_euler
     SEED: 2024
     IMAGE_SIZE: [ 1024, 1024 ]
     SHIFT: True
     GUIDE_SCALE: 3.5
-  #
+
   OPTIMIZER:
     NAME: AdamW
     LEARNING_RATE: 4e-4
@@ -258,7 +201,7 @@ SOLVER:
     EPS: 1e-8
     WEIGHT_DECAY: 1e-2
     AMSGRAD: False
-  #
+
   TRAIN_DATA:
     NAME: ImageTextPairMSDataset
     MODE: train
@@ -302,7 +245,7 @@ SOLVER:
       - NAME: Select
         KEYS: [ 'image', 'prompt' ]
         META_KEYS: [ 'data_key' ]
-  #
+
   EVAL_DATA:
     NAME: Text2ImageDataset
     MODE: eval
@@ -319,13 +262,12 @@ SOLVER:
       - 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

scepter/methods/studio/self_train/dit/flux1.0_schnell_pro.yaml

@@ -50,43 +50,33 @@ META:
 #
 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'
+  ENABLE_GRADSCALER: False
+  USE_SCALER: False
   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'] #
+  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
@@ -99,65 +89,39 @@ SOLVER:
     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
@@ -167,7 +131,7 @@ SOLVER:
       USE_CONV: False
       SCALE_FACTOR: 0.3611
       SHIFT_FACTOR: 0.1159
-      #
+
       ENCODER:
         NAME: Encoder
         USE_CHECKPOINT: True
@@ -181,7 +145,7 @@ SOLVER:
         DOUBLE_Z: True
         DROPOUT: 0.0
         RESAMP_WITH_CONV: True
-      #
+
       DECODER:
         NAME: Decoder
         USE_CHECKPOINT: True
@@ -196,60 +160,39 @@ SOLVER:
         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
+    SAMPLER: flow_euler
     SEED: 2024
     IMAGE_SIZE: [ 1024, 1024 ]
     GUIDE_SCALE: 3.5
-  #
+
   OPTIMIZER:
     NAME: AdamW
     LEARNING_RATE: 4e-4
@@ -257,7 +200,7 @@ SOLVER:
     EPS: 1e-8
     WEIGHT_DECAY: 1e-2
     AMSGRAD: False
-  #
+
   TRAIN_DATA:
     NAME: ImageTextPairMSDataset
     MODE: train
@@ -301,7 +244,7 @@ SOLVER:
       - NAME: Select
         KEYS: [ 'image', 'prompt' ]
         META_KEYS: [ 'data_key' ]
-  #
+
   EVAL_DATA:
     NAME: Text2ImageDataset
     MODE: eval
@@ -318,13 +261,12 @@ SOLVER:
       - 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
@@ -339,4 +281,4 @@ SOLVER:
       PROB_INTERVAL: 100
       SAVE_LAST: True
       SAVE_NAME_PREFIX: 'step'
-      SAVE_PROBE_PREFIX: 'image'
+      SAVE_PROBE_PREFIX: 'image'

scepter/methods/studio/self_train/dit/pixart_alpha_pro.yaml

@@ -35,7 +35,7 @@ META:
       SAVE_INTERVAL: 25
       EPSEC: 0.818
       LEARNING_RATE: 0.0001
-      IS_DEFAULT: False
+      IS_DEFAULT: True
       TUNER: LORA
   #
   TUNERS:

scepter/methods/studio/self_train/self_train.yaml

@@ -13,8 +13,10 @@ TRAIN_PARAS:
     VALUES: [[256, 256], [320, 180], [180, 320],
              [512, 512], [640, 360], [360, 640],
              [768, 768], [960, 540], [540, 960],
-             [1024, 1024], [1280, 720], [720, 1280]]
+             [1024, 1024], [1280, 720], [720, 1280],
+             [720, 480], [480, 720]]
     DEFAULT: [1024, 1024]
   EVAL_PROMPTS:
     - a boy wearing a jacket
     - a dog running on the lawn
+SAVE_FILE_LOCAL_PATH: "cache/scepter_ui/datasets/train_data_from_list"

scepter/modules/__init__.py

@@ -1,4 +1,23 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from scepter.modules import (data, inference, model, opt, solver, transform,
-                             utils)
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
+
+
+if TYPE_CHECKING:
+    from scepter.modules import (data, inference, model, opt, solver, transform,
+                                 utils)
+else:
+    _import_structure = {
+        'modules': ['data', 'inference', 'model', 'opt', 'solver',
+                    'transform', 'utils']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/annotator/__init__.py

@@ -1,23 +1,60 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-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
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+    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
+else:
+    _import_structure = {
+        'base_annotator': ['GeneralAnnotator'],
+        'canny': ['CannyAnnotator'],
+        'color': ['ColorAnnotator'],
+        'degradation': ['DegradationAnnotator'],
+        'doodle': ['DoodleAnnotator'],
+        'gray': ['GrayAnnotator'],
+        'hed': ['HedAnnotator'],
+        'identity': ['IdentityAnnotator'],
+        'informative_drawing': ['InfoDrawAnimeAnnotator',
+                                'InfoDrawContourAnnotator',
+                                'InfoDrawOpenSketchAnnotator'],
+        'inpainting': ['InpaintingAnnotator'],
+        'invert': ['InvertAnnotator'],
+        'midas_op': ['MidasDetector'],
+        'mlsd_op': ['MLSDdetector'],
+        'openpose': ['OpenposeAnnotator'],
+        'outpainting': ['OutpaintingAnnotator', 'OutpaintingResize'],
+        'pidinet': ['PiDiAnnotator'],
+        'segmentation': ['ESAMAnnotator'],
+        'sketch': ['SketchAnnotator'],
+        'lama': ['LamaAnnotator'],
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/annotator/doodle.py

@@ -1,20 +1,11 @@
 # -*- coding: utf-8 -*-
-
-import math
+# 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 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()

scepter/modules/annotator/hed.py

@@ -114,7 +114,7 @@ class HedAnnotator(BaseAnnotator, metaclass=ABCMeta):
         pretrained_model = cfg.get('PRETRAINED_MODEL', None)
         if pretrained_model:
             with FS.get_from(pretrained_model, wait_finish=True) as local_path:
-                self.netNetwork.load_state_dict(torch.load(local_path))
+                self.netNetwork.load_state_dict(torch.load(local_path, weights_only=True))
 
     @torch.no_grad()
     @torch.inference_mode()

scepter/modules/annotator/informative_drawing.py

@@ -2,20 +2,17 @@
 # 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 torchvision.transforms import InterpolationMode
+
 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
 
@@ -123,7 +120,7 @@ class InfoDrawContourAnnotator(BaseAnnotator, metaclass=ABCMeta):
         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.load_state_dict(torch.load(local_path, weights_only=True))
         self.model = self.model.eval().requires_grad_(False).to(we.device_id)
 
     @torch.no_grad()

scepter/modules/annotator/inpainting.py

@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-import math
 import random
 from abc import ABCMeta
 from enum import Enum
@@ -8,7 +7,8 @@ from enum import Enum
 import cv2
 import numpy as np
 import torch
-from PIL import Image, ImageDraw
+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
@@ -226,7 +226,12 @@ class InpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
         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):
+    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
@@ -247,18 +252,17 @@ class InpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
         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)
+            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)
+            image = np.where(mask[:, :, np.newaxis] == 255, colored_mask,
+                             image)
 
         if return_mask:
-            ret_data = {
-                'image': np.array(image),
-                'mask': np.array(mask)
-            }
+            ret_data = {'image': np.array(image), 'mask': np.array(mask)}
         else:
             ret_data = np.array(image)
         return ret_data

scepter/modules/annotator/lama.py

@@ -1,27 +1,31 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 from abc import ABCMeta
 
-import torch
-import cv2
 import numpy as np
+
+import cv2
+import torch
 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
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+
 
 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
-    )
+    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
@@ -32,26 +36,29 @@ class LamaAnnotator(BaseAnnotator, metaclass=ABCMeta):
         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")
+        @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'
+                    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()
+                    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)
@@ -69,9 +76,13 @@ class LamaAnnotator(BaseAnnotator, metaclass=ABCMeta):
                 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))
+        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)
@@ -93,6 +104,3 @@ class LamaAnnotator(BaseAnnotator, metaclass=ABCMeta):
                             __class__.__name__,
                             LamaAnnotator.para_dict,
                             set_name=True)
-
-
-

scepter/modules/annotator/midas/base_model.py

@@ -10,7 +10,7 @@ class BaseModel(torch.nn.Module):
         Args:
             path (str): file path
         """
-        parameters = torch.load(path, map_location=torch.device('cpu'))
+        parameters = torch.load(path, map_location=torch.device('cpu'), weights_only=True)
 
         if 'optimizer' in parameters:
             parameters = parameters['model']

scepter/modules/annotator/mlsd_op.py

@@ -29,7 +29,7 @@ class MLSDdetector(BaseAnnotator, metaclass=ABCMeta):
         pretrained_model = cfg.get('PRETRAINED_MODEL', None)
         if pretrained_model:
             with FS.get_from(pretrained_model, wait_finish=True) as local_path:
-                model.load_state_dict(torch.load(local_path), strict=True)
+                model.load_state_dict(torch.load(local_path, weights_only=True), strict=True)
         self.model = model.eval()
         self.thr_v = cfg.get('THR_V', 0.1)
         self.thr_d = cfg.get('THR_D', 0.1)

scepter/modules/annotator/openpose.py

@@ -423,7 +423,7 @@ class Hand(object):
         self.model = handpose_model()
         if torch.cuda.is_available():
             self.model = self.model.to(device)
-        model_dict = transfer(self.model, torch.load(model_path))
+        model_dict = transfer(self.model, torch.load(model_path, weights_only=True))
         self.model.load_state_dict(model_dict)
         self.model.eval()
         self.device = device
@@ -503,7 +503,7 @@ class Body(object):
         self.model = bodypose_model()
         if torch.cuda.is_available():
             self.model = self.model.to(device)
-        model_dict = transfer(self.model, torch.load(model_path))
+        model_dict = transfer(self.model, torch.load(model_path, weights_only=True))
         self.model.load_state_dict(model_dict)
         self.model.eval()
         self.device = device
@@ -789,7 +789,7 @@ class OpenposeAnnotator(BaseAnnotator, metaclass=ABCMeta):
             raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
         image = image[:, :, ::-1]
         candidate, subset = self.body_estimation(image)
-        canvas = np.zeros_like(image)
+        canvas = np.zeros_like(image, order='C')   # to check
         canvas = draw_bodypose(canvas, candidate, subset)
         if self.use_hand:
             hands_list = handDetect(candidate, subset, image)

scepter/modules/annotator/outpainting.py

@@ -4,10 +4,10 @@ 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
@@ -87,7 +87,8 @@ class OutpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
             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)
+                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))
@@ -108,20 +109,30 @@ class OutpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
                 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_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 = 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)}
+                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)}
+                ret_data = {
+                    'src_image': np.array(init_image),
+                    'image': np.array(img)
+                }
             else:
                 ret_data = np.array(img)
         return ret_data
@@ -133,6 +144,7 @@ class OutpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta):
                             OutpaintingAnnotator.para_dict,
                             set_name=True)
 
+
 @ANNOTATORS.register_class()
 class OutpaintingResize(BaseAnnotator, metaclass=ABCMeta):
     para_dict = {}
@@ -148,11 +160,7 @@ class OutpaintingResize(BaseAnnotator, metaclass=ABCMeta):
         top, bottom = np.min(locs[0]), np.max(locs[0])
         return [left, top, right, bottom]
 
-    def forward(self,
-                image,
-                target_image,
-                mask=None
-                ):
+    def forward(self, image, target_image, mask=None):
         if isinstance(image, Image.Image):
             image = image
         elif isinstance(image, torch.Tensor):
@@ -175,8 +183,10 @@ class OutpaintingResize(BaseAnnotator, metaclass=ABCMeta):
         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')
+            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

scepter/modules/annotator/pidinet.py

@@ -1,14 +1,13 @@
 # -*- coding: utf-8 -*-
-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 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
@@ -883,7 +882,7 @@ class PiDiAnnotator(BaseAnnotator, metaclass=ABCMeta):
         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']
+                                   map_location='cpu', weights_only=True)['state_dict']
                 if vanilla_cnn:
                     state = convert_pidinet(state, 'carv4')
                 state = {

scepter/modules/annotator/registry.py

@@ -15,7 +15,7 @@ def build_annotator(cfg, registry, logger=None, *args, **kwargs):
         raise TypeError(f'Config must be type dict, got {type(cfg)}')
     if cfg.have('PRETRAINED_MODEL'):
         pretrain_cfg = cfg.PRETRAINED_MODEL
-        if pretrain_cfg is not None and not isinstance(pretrain_cfg, (str)):
+        if pretrain_cfg is not None and not isinstance(pretrain_cfg, (str, list)):
             raise TypeError('Pretrain parameter must be a string')
     else:
         pretrain_cfg = None

scepter/modules/annotator/segmentation.py

@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-import math
 import random
 from abc import ABCMeta
 
@@ -9,15 +8,20 @@ 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 scipy import ndimage
+try:
+    from sklearn.cluster import KMeans
+except:
+    import warnings
+    warnings.warn("ignore sklearn import, please pip install scikit-learn.")
+from torchvision.ops.boxes import batched_nms
+
 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):
@@ -28,7 +32,7 @@ def find_dominant_color(image, k=1):
         kmeans = KMeans(n_clusters=k, n_init='auto')
         kmeans.fit(pixels)
         dominant_color = kmeans.cluster_centers_.astype(int)[0]
-    except:
+    except Exception:
         dominant_color = np.array([255, 255, 255])
     return dominant_color
 
@@ -62,16 +66,9 @@ class ESAMAnnotator(BaseAnnotator, metaclass=ABCMeta):
         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:
+        except Exception:
             raise NotImplementedError(
-                f'Please install efficient_sam and segment_anything modules.')
+                'Please install efficient_sam and segment_anything modules.')
 
         pretrained_model = cfg.get('PRETRAINED_MODEL', None)
         if pretrained_model:
@@ -294,8 +291,6 @@ class ESAMAnnotator(BaseAnnotator, metaclass=ABCMeta):
                             set_name=True)
 
 
-
-
 @ANNOTATORS.register_class()
 class SAMAnnotatorDraw(BaseAnnotator, metaclass=ABCMeta):
     para_dict = {}
@@ -312,10 +307,16 @@ class SAMAnnotatorDraw(BaseAnnotator, metaclass=ABCMeta):
 
         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)
+                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):
+    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):
@@ -337,21 +338,25 @@ class SAMAnnotatorDraw(BaseAnnotator, metaclass=ABCMeta):
             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))
+            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}
+            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 = ndimage.center_of_mass(scribble, labeled_array,
+                                             range(1, num_features + 1))
             centers = np.array(centers)
-            ### (x1, y1, x2, y2)
+            # (x1, y1, x2, y2)
             x_min = centers[:, 0].min()
             x_max = centers[:, 0].max()
             y_min = centers[:, 1].min()
@@ -365,12 +370,13 @@ class SAMAnnotatorDraw(BaseAnnotator, metaclass=ABCMeta):
             sample = {'box': input_box}
 
         self.sam_predictor.set_image(image)
-        masks, scores, logits = self.sam_predictor.predict(**sample, multimask_output=True)
+        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]
+            'mask': (masks[index] * 255).astype(np.uint8),
+            'score': scores[index]
         }
         return ret_data
 
@@ -379,4 +385,4 @@ class SAMAnnotatorDraw(BaseAnnotator, metaclass=ABCMeta):
         return dict_to_yaml('ANNOTATORS',
                             __class__.__name__,
                             SAMAnnotatorDraw.para_dict,
-                            set_name=True)
+                            set_name=True)

scepter/modules/annotator/sketch.py

@@ -1,14 +1,11 @@
 # -*- coding: utf-8 -*-
-
-import math
+# 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 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
@@ -89,7 +86,7 @@ class SketchAnnotator(BaseAnnotator, metaclass=ABCMeta):
                                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')
+                state = torch.load(local_path, map_location='cpu', weights_only=True)
                 self.model.load_state_dict(state)
 
     @torch.no_grad()

scepter/modules/data/__init__.py

@@ -1,4 +1,21 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-from scepter.modules.data import dataset, sampler
+
+if TYPE_CHECKING:
+    from scepter.modules.data import dataset, sampler
+else:
+    _import_structure = {
+        'data': ['dataset', 'sampler']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/data/dataset/__init__.py

@@ -1,11 +1,35 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-from scepter.modules.data.dataset.base_dataset import BaseDataset
-from scepter.modules.data.dataset.dataset import (Image2ImageDataset,
-                                                  ImageClassifyPublicDataset,
-                                                  ImageTextPairDataset,
-                                                  Text2ImageDataset)
-from scepter.modules.data.dataset.ms_dataset import (
-    ImageTextPairFolderDataset, ImageTextPairMSDataset)
-from scepter.modules.data.dataset.registry import DATASETS
+
+if TYPE_CHECKING:
+    from scepter.modules.data.dataset.base_dataset import BaseDataset
+    from scepter.modules.data.dataset.dataset import (Image2ImageDataset,
+                                                      ImageClassifyPublicDataset,
+                                                      ImageTextPairDataset,
+                                                      Text2ImageDataset)
+    from scepter.modules.data.dataset.ms_dataset import (
+        ImageTextPairFolderDataset, ImageTextPairMSDataset)
+    from scepter.modules.data.dataset.registry import DATASETS
+    from scepter.modules.data.dataset.video_gen_dataset import VideoGenDataset
+else:
+    _import_structure = {
+        'base_dataset': ['BaseDataset'],
+        'dataset': ['Image2ImageDataset', 'ImageClassifyPublicDataset',
+                    'ImageTextPairDataset', 'Text2ImageDataset'],
+        'ms_dataset': ['ImageTextPairFolderDataset',
+                       'ImageTextPairMSDataset'],
+        'registry': ['DATASETS'],
+        'video_gen_dataset': ['VideoGenDataset']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/data/dataset/base_dataset.py

@@ -82,7 +82,7 @@ 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.local_we["seed"] += (worker_id + self.local_we['rank'] * 1234)
         self.seed = self.local_we["seed"]
         we.set_env(self.local_we)
 

scepter/modules/data/dataset/dataset.py

@@ -4,6 +4,7 @@
 import numbers
 import os
 import sys
+import copy
 from collections.abc import Iterable
 
 import numpy as np
@@ -242,6 +243,8 @@ class Text2ImageDataset(BaseDataset):
         prompt_prefix = cfg.get('PROMPT_PREFIX', '')
         path_prefix = cfg.get('PATH_PREFIX', '')
         use_num = cfg.get('USE_NUM', -1)
+        meta_cfg = cfg.get('META_CFG', None)
+        meta_cfg = meta_cfg.get_lowercase_dict() if meta_cfg is not None else None
 
         image_size = cfg.get('IMAGE_SIZE', 1024)
         if isinstance(image_size, numbers.Number):
@@ -264,7 +267,12 @@ class Text2ImageDataset(BaseDataset):
 
         self.items = list()
         for i, row in enumerate(rows):
-            item = {'index': i, 'meta': {'image_size': image_size}}
+            if meta_cfg is not None:
+                meta_cfg_copy = copy.deepcopy(meta_cfg)
+                meta_cfg_copy['image_size'] = image_size
+                item = {'index': i, 'meta': meta_cfg_copy}
+            else:
+                item = {'index': i, 'meta': {'image_size': image_size}}
             for key, value in zip(fields, row):
                 if key in ['prompt', 'caption', 'text']:
                     item['ori_prompt'] = value

scepter/modules/data/dataset/ms_dataset.py

@@ -1,16 +1,28 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
+import io
+import math
 import numbers
 import os
 import sys
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
 
 from scepter.modules.data.dataset.base_dataset import BaseDataset
 from scepter.modules.data.dataset.registry import DATASETS
+from scepter.modules.transform.io import pillow_convert
 from scepter.modules.utils.config import dict_to_yaml
 from scepter.modules.utils.distribute import we
 from scepter.modules.utils.file_system import FS
 
+Image.MAX_IMAGE_PIXELS = None
+
 
 @DATASETS.register_class()
 class ImageTextPairMSDataset(BaseDataset):
@@ -102,7 +114,7 @@ class ImageTextPairMSDataset(BaseDataset):
                 self.output_size = [self.output_size, self.output_size]
         # Use modelscope dataset
         if not ms_dataset_name:
-            raise (
+            raise ValueError(
                 'Your must set MS_DATASET_NAME as modelscope dataset or your local dataset orignized '
                 'as modelscope dataset.')
         if FS.exists(ms_dataset_name):
@@ -125,7 +137,7 @@ class ImageTextPairMSDataset(BaseDataset):
                     split=ms_dataset_split,
                     download_mode=DownloadMode.FORCE_REDOWNLOAD)
             except Exception as sec_e:
-                raise f'Load Modelscope dataset failed {sec_e}.'
+                raise ValueError(f'Load Modelscope dataset failed {sec_e}.')
         if ms_remap_keys:
             self.data = self.data.remap_columns(ms_remap_keys.get_dict())
 
@@ -245,7 +257,7 @@ class ImageTextPairFolderDataset(BaseDataset):
                 self.output_size = [self.output_size, self.output_size]
         # Use modelscope dataset
         if not data_folder or not FS.exists(data_folder):
-            raise ('Your must set datafolder for local dataset.')
+            raise ValueError('Your must set datafolder for local dataset.')
         data_folder = FS.get_dir_to_local_dir(data_folder)
         all_lines = open(os.path.join(data_folder, 'train.csv'),
                          'r').read().split('\n')
@@ -311,3 +323,254 @@ class ImageTextPairFolderDataset(BaseDataset):
                             __class__.__name__,
                             ImageTextPairMSDataset.para_dict,
                             set_name=True)
+
+
+@DATASETS.register_class()
+class ImageTextPairMSDatasetForACE(BaseDataset):
+    para_dict = {
+        'MS_DATASET_NAME': {
+            'value': '',
+            'description': 'Modelscope dataset name.'
+        },
+        'MS_DATASET_NAMESPACE': {
+            'value': '',
+            'description': 'Modelscope dataset namespace.'
+        },
+        'MS_DATASET_SUBNAME': {
+            'value': '',
+            'description': 'Modelscope dataset subname.'
+        },
+        'MS_DATASET_SPLIT': {
+            'value': '',
+            'description':
+            'Modelscope dataset split set name, default is train.'
+        },
+        'MS_REMAP_KEYS': {
+            'value':
+            None,
+            'description':
+            'Modelscope dataset header of list file, the default is Target:FILE; '
+            'If your file is not this header, please set this field, which is a map dict.'
+            "For example, { 'Image:FILE': 'Target:FILE' } will replace the filed Image:FILE to Target:FILE"
+        },
+        'MS_REMAP_PATH': {
+            'value':
+            None,
+            'description':
+            'When modelscope dataset name is not None, that means you use the dataset from modelscope,'
+            ' default is None. But if you want to use the datalist from modelscope and the file from '
+            'local device, you can use this field to set the root path of your images. '
+        },
+        'TRIGGER_WORDS': {
+            'value':
+            '',
+            'description':
+            'The words used to describe the common features of your data, especially when you customize a '
+            'tuner. Use these words you can get what you want.'
+        },
+        'REPLACE_STYLE': {
+            'value':
+            False,
+            'description':
+            'Whether use the MS_DATASET_SUBNAME to replace the word in your description, default is False.'
+        },
+        'HIGHLIGHT_KEYWORDS': {
+            'value':
+            '',
+            'description':
+            'The keywords you want to highlight in prompt, which will be replace by <HIGHLIGHT_KEYWORDS>.'
+        },
+        'KEYWORDS_SIGN': {
+            'value':
+            '',
+            'description':
+            'The keywords sign you want to add, which is like <{HIGHLIGHT_KEYWORDS}{KEYWORDS_SIGN}>'
+        },
+        'ALIGN_SIZE': {
+            'value': False,
+            'description':
+                'Whether ensure the size align between the source image and target image.'
+        },
+        'OUTPUT_SIZE': {
+            'value':
+            None,
+            'description':
+            'If you use the FlexibleResize transforms, this filed will output the image_size as [h, w],'
+            'which will be used to set the output size of images used to train the model.'
+        },
+    }
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg=cfg, logger=logger)
+        from modelscope import MsDataset
+        from modelscope.utils.constant import DownloadMode
+        ms_dataset_name = cfg.get('MS_DATASET_NAME', None)
+        ms_dataset_namespace = cfg.get('MS_DATASET_NAMESPACE', None)
+        ms_dataset_subname = cfg.get('MS_DATASET_SUBNAME', None)
+        ms_dataset_split = cfg.get('MS_DATASET_SPLIT', 'train')
+        ms_remap_keys = cfg.get('MS_REMAP_KEYS', None)
+        ms_remap_path = cfg.get('MS_REMAP_PATH', None)
+
+        self.max_seq_len = cfg.get('MAX_SEQ_LEN', 1024)
+        self.max_aspect_ratio = cfg.get('MAX_ASPECT_RATIO', 4)
+        self.d = cfg.get('DOWNSAMPLE_RATIO', 16)
+        self.replace_style = cfg.get('REPLACE_STYLE', False)
+        self.trigger_words = cfg.get('TRIGGER_WORDS', '')
+        self.replace_keywords = cfg.get('HIGHLIGHT_KEYWORDS', '')
+        self.keywords_sign = cfg.get('KEYWORDS_SIGN', '')
+        self.add_indicator = cfg.get('ADD_INDICATOR', False)
+
+        self.align_size = cfg.get('ALIGN_SIZE', False)
+        # Use modelscope dataset
+        if not ms_dataset_name:
+            raise ValueError(
+                'Your must set MS_DATASET_NAME as modelscope dataset or your local dataset orignized '
+                'as modelscope dataset.')
+        if FS.exists(ms_dataset_name):
+            ms_dataset_name = FS.get_dir_to_local_dir(ms_dataset_name)
+            self.ms_dataset_name = ms_dataset_name
+            # ms_remap_path = ms_dataset_name
+        try:
+            self.data = MsDataset.load(str(ms_dataset_name),
+                                       namespace=ms_dataset_namespace,
+                                       subset_name=ms_dataset_subname,
+                                       split=ms_dataset_split)
+        except Exception:
+            self.logger.info(
+                "Load Modelscope dataset failed, retry with download_mode='force_redownload'."
+            )
+            try:
+                self.data = MsDataset.load(
+                    str(ms_dataset_name),
+                    namespace=ms_dataset_namespace,
+                    subset_name=ms_dataset_subname,
+                    split=ms_dataset_split,
+                    download_mode=DownloadMode.FORCE_REDOWNLOAD)
+            except Exception as sec_e:
+                raise ValueError(f'Load Modelscope dataset failed {sec_e}.')
+        if ms_remap_keys:
+            self.data = self.data.remap_columns(ms_remap_keys.get_dict())
+
+        if ms_remap_path:
+
+            def map_func(example):
+                return {
+                    k: os.path.join(ms_remap_path, v)
+                    if k.endswith(':FILE') else v
+                    for k, v in example.items()
+                }
+
+            self.data = self.data.ds_instance.map(map_func)
+
+        self.transforms = T.Compose([
+            T.ToTensor(),
+            T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+
+    def __len__(self):
+        if self.mode == 'train':
+            return sys.maxsize
+        else:
+            return len(self.data)
+
+    def _get(self, index: int):
+        current_data = self.data[index % len(self.data)]
+
+        tar_image_path = current_data.get('Target:FILE', '')
+        src_image_path = current_data.get('Source:FILE', '')
+
+        style = current_data.get('Style', '')
+        prompt = current_data.get('Prompt', current_data.get('prompt', ''))
+        if self.replace_style and not style == '':
+            prompt = prompt.replace(style, f'<{self.keywords_sign}>')
+
+        elif not self.replace_keywords.strip() == '':
+            prompt = prompt.replace(
+                self.replace_keywords,
+                '<' + self.replace_keywords + f'{self.keywords_sign}>')
+
+        if not self.trigger_words == '':
+            prompt = self.trigger_words.strip() + ' ' + prompt
+
+        src_image = self.load_image(self.ms_dataset_name,
+                                    src_image_path,
+                                    cvt_type='RGB')
+        tar_image = self.load_image(self.ms_dataset_name,
+                                    tar_image_path,
+                                    cvt_type='RGB')
+        src_image = self.image_preprocess(src_image)
+        tar_image = self.image_preprocess(tar_image, size = src_image.shape[:2]  if self.align_size else None)
+
+        tar_image = self.transforms(tar_image)
+        src_image = self.transforms(src_image)
+        src_mask = torch.ones_like(src_image[[0]])
+        tar_mask = torch.ones_like(tar_image[[0]])
+        if self.add_indicator:
+            if '{image}' not in prompt:
+                prompt = '{image}, ' + prompt
+        return {
+            'src_image_list': [src_image],
+            'src_mask_list': [src_mask],
+            'image': tar_image,
+            'image_mask': tar_mask,
+            'prompt': [prompt],
+            'edit_id': [0]
+        }
+
+    def load_image(self, prefix, img_path, cvt_type=None):
+        if img_path is None or img_path == '':
+            return None
+        img_path = os.path.join(prefix, img_path)
+        with FS.get_object(img_path) as image_bytes:
+            image = Image.open(io.BytesIO(image_bytes))
+            if cvt_type is not None:
+                image = pillow_convert(image, cvt_type)
+        return image
+
+    def image_preprocess(self,
+                         img,
+                         size=None,
+                         interpolation=InterpolationMode.BILINEAR):
+        H, W = img.height, img.width
+        if H / W > self.max_aspect_ratio:
+            img = T.CenterCrop((self.max_aspect_ratio * W, W))(img)
+        elif W / H > self.max_aspect_ratio:
+            img = T.CenterCrop((H, self.max_aspect_ratio * H))(img)
+
+        if size is None:
+            # resize image for max_seq_len, while keep the aspect ratio
+            H, W = img.height, img.width
+            scale = min(
+                1.0,
+                math.sqrt(self.max_seq_len / ((H / self.d) * (W / self.d))))
+            rH = int(
+                H * scale) // self.d * self.d  # ensure divisible by self.d
+            rW = int(W * scale) // self.d * self.d
+        else:
+            rH, rW = size
+        img = T.Resize((rH, rW), interpolation=interpolation,
+                       antialias=True)(img)
+        return np.array(img, dtype=np.uint8)
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('DATASet',
+                            __class__.__name__,
+                            ImageTextPairMSDatasetForACE.para_dict,
+                            set_name=True)
+
+    @staticmethod
+    def collate_fn(batch):
+        collect = defaultdict(list)
+        for sample in batch:
+            for k, v in sample.items():
+                collect[k].append(v)
+
+        new_batch = dict()
+        for k, v in collect.items():
+            if all([i is None for i in v]):
+                new_batch[k] = None
+            else:
+                new_batch[k] = v
+
+        return new_batch

scepter/modules/data/dataset/registry.py

@@ -337,8 +337,12 @@ def build_dataset_config(cfg, registry, logger=None, *args, **kwargs):
             f'registry must be type Registry, got {type(registry)}')
 
     cfg = deep_copy(cfg)
-
     req_type = cfg.get('NAME')
+
+    from scepter.modules.utils.import_utils import LazyImportModule
+    sig = (registry.name.upper(), req_type)
+    LazyImportModule.import_module(sig)
+
     if isinstance(req_type, str):
         req_type_entry = registry.get(req_type)
         if req_type_entry is None:

scepter/modules/data/dataset/video_gen_dataset.py

@@ -0,0 +1,202 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import io
+import os
+import random
+import sys
+import warnings
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from scepter.modules.data.dataset import DATASETS, BaseDataset
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+
+try:
+    import decord
+    decord.bridge.set_bridge('torch')
+except ImportError:
+    warnings.warn(
+        'The `decord` package is required for loading the video dataset. Install with `pip install decord`'
+    )
+
+
+@DATASETS.register_class()
+class VideoGenDataset(BaseDataset):
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        self.prompt_prefix = cfg.get('PROMPT_PREFIX', '')
+        self.path_prefix = cfg.get('PATH_PREFIX', '')
+        self.p_zero = cfg.get('P_ZERO', 0.0)
+        self.max_num_frames = cfg.get('NUM_FRAMES', 49)
+        self.fps = cfg.get('FPS', 8)
+        self.height = cfg.get('HEIGHT', 480)
+        self.width = cfg.get('WIDTH', 720)
+        self.skip_frames_start = cfg.get('SKIP_FRAMES_START', 0)
+        self.skip_frames_end = cfg.get('SKIP_FRAMES_END', 0)
+        self.data_type = cfg.get('DATA_TYPE', 't2v')
+
+    def worker_init_fn(self, worker_id, num_workers=1):
+        super().worker_init_fn(worker_id, num_workers=num_workers)
+        randseed = np.random.randint(0, 2**32 - num_workers - 1)
+        workerseed = randseed + worker_id
+        random.seed(workerseed)
+        np.random.seed(workerseed)
+
+    def _preprocess_video_data(self, video_path):
+
+        with FS.get_object(video_path) as video_data:
+            video_reader = decord.VideoReader(io.BytesIO(video_data),
+                                              width=self.width,
+                                              height=self.height)
+        video_num_frames = len(video_reader)
+
+        start_frame = min(self.skip_frames_start, video_num_frames)
+        end_frame = max(0, video_num_frames - self.skip_frames_end)
+        if end_frame <= start_frame:
+            frames = video_reader.get_batch([start_frame])
+        elif end_frame - start_frame <= self.max_num_frames:
+            frames = video_reader.get_batch(list(range(start_frame,
+                                                       end_frame)))
+        else:
+            indices = list(
+                range(start_frame, end_frame,
+                      (end_frame - start_frame) // self.max_num_frames))
+            frames = video_reader.get_batch(indices)
+
+        # Ensure that we don't go over the limit
+        frames = frames[:self.max_num_frames]
+        selected_num_frames = frames.shape[0]
+
+        # Choose first (4k + 1) frames as this is how many is required by the VAE
+        remainder = (3 + (selected_num_frames % 4)) % 4
+        if remainder != 0:
+            frames = frames[:-remainder]
+        selected_num_frames = frames.shape[0]
+
+        assert (selected_num_frames - 1) % 4 == 0
+
+        # Training transforms
+        frames = frames.float().div_(127.5).sub_(1.)
+        frames = frames.permute(3, 0, 1, 2).contiguous()  # [C, F, H, W]
+        return frames
+
+    def _parse_index(self, index):
+        meta = dict()
+        for key, value in zip(index[-1], index[:-1]):
+            if key in ['oss_key', 'path', 'video_path', 'target_video_path']:
+                meta['video_path'] = value
+            elif key in ['source_video_path', 'src_video_path']:
+                meta['src_video_path'] = value
+            elif key in ['prompt', 'caption', 'text']:
+                meta['prompt'] = value
+            elif key in ['width', 'height']:
+                meta[key] = int(value)
+            else:
+                meta[key] = value
+        return meta
+
+    def _get(self, index):
+        meta = self._parse_index(index)
+
+        video_path = os.path.join(self.path_prefix, meta.get('video_path', ''))
+        video = self._preprocess_video_data(video_path)
+
+        prompt = self.prompt_prefix + meta.get('prompt', '')
+        if self.mode == 'train' and np.random.uniform() < self.p_zero:
+            prompt = ''
+
+        item = {
+            'video': video,
+            'prompt': prompt,
+            'meta': meta,
+        }
+        if 'i2v' in self.data_type:
+            item['image'] = item['video'][:, :1, :, :]
+        if 'v2v' in self.data_type:
+            src_video_path = os.path.join(self.path_prefix,
+                                          meta.get('src_video_path', ''))
+            src_video = self._preprocess_video_data(src_video_path)
+            item['src_video'] = src_video
+        return item
+
+    def __len__(self):
+        return sys.maxsize
+
+    @staticmethod
+    def collate_fn(batch):
+        collect = {}
+        for sample in batch:
+            for k, v in sample.items():
+                if k not in collect:
+                    collect[k] = []
+                collect[k].append(v)
+        return collect
+
+
+@DATASETS.register_class()
+class VideoGenDatasetOTF(VideoGenDataset):
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger)
+        self.data_file = cfg.DATA_FILE
+        self.delimiter = cfg.get('DELIMITER', '#;#')
+        self.fields = cfg.get('FIELDS', ['video_path', 'prompt'])
+        self.use_num = cfg.get('USE_NUM', -1)
+
+        from scepter.modules.model.registry import MODELS
+        model_cfg = cfg.get('MODEL', None)
+        if model_cfg is not None:
+            self.model = MODELS.build(
+                cfg.MODEL,
+                logger=logger).eval().requires_grad_(False).to(we.device_id)
+        self.items = self.parse_data(self.data_file, self.delimiter,
+                                     self.fields)
+        if self.use_num and self.use_num > 0:
+            self.items = self.items[:self.use_num]
+        self.data = self.encode(self.items)
+        self.real_number = len(self.data)
+        if model_cfg is not None:
+            self.model.to('cpu')
+            del self.model
+            torch.cuda.empty_cache()
+
+    def parse_data(self, data_file, delimiter, fields):
+        items = list()
+        with FS.get_object(data_file) as local_data:
+            rows = [
+                i.split(delimiter,
+                        len(fields) - 1)
+                for i in local_data.decode('utf-8').strip().split('\n')
+            ]
+            for i, row in enumerate(rows):
+                item = {}
+                for key, value in zip(self.fields, row):
+                    if key in ['oss_key', 'path', 'video_path']:
+                        item['video_path'] = value
+                    elif key in ['prompt', 'caption', 'text']:
+                        item['prompt'] = value
+                    elif key in ['width', 'height']:
+                        item[key] = int(value)
+                    else:
+                        item[key] = value
+                items.append(item)
+        return items
+
+    def encode(self, items):
+        self.logger.info('Start to encode video data [{}]!'.format(len(items)))
+        for item in tqdm(items):
+            video_path = os.path.join(self.path_prefix,
+                                      item.get('video_path', ''))
+            video = self._preprocess_video_data(video_path)
+            latent = self.model.encode_first_stage(
+                video.unsqueeze(0).to(we.device_id)).squeeze(0)
+            item['video_latent'] = latent.detach().cpu()
+            item['video'] = video
+            if self.data_type == 'i2v':
+                item['image'] = item['video'][:, :1, :, :]
+        return items
+
+    def _get(self, index):
+        return self.data[index % self.real_number]

scepter/modules/data/sampler/__init__.py

@@ -1,9 +1,31 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-from scepter.modules.data.sampler.base_sampler import BaseSampler
-from scepter.modules.data.sampler.registry import SAMPLERS
-from scepter.modules.data.sampler.sampler import (
-    EvalDistributedSampler, LoopSampler, MixtureOfSamplers,
-    MultiFoldDistributedSampler, MultiLevelBatchSampler,
-    MultiLevelBatchSamplerMultiSource, ResolutionBatchSampler)
+
+if TYPE_CHECKING:
+    from scepter.modules.data.sampler.base_sampler import BaseSampler
+    from scepter.modules.data.sampler.registry import SAMPLERS
+    from scepter.modules.data.sampler.sampler import (
+        EvalDistributedSampler, LoopSampler, MixtureOfSamplers,
+        MultiFoldDistributedSampler, MultiLevelBatchSampler,
+        MultiLevelBatchSamplerMultiSource, ResolutionBatchSampler)
+else:
+    _import_structure = {
+        'base_sampler': ['BaseSampler'],
+        'registry': ['SAMPLERS'],
+        'sampler': ['EvalDistributedSampler', 'LoopSampler',
+                    'MixtureOfSamplers', 'MultiFoldDistributedSampler',
+                    'MultiLevelBatchSampler', 'MultiLevelBatchSamplerMultiSource',
+                    'ResolutionBatchSampler']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/data/sampler/registry.py

@@ -35,8 +35,12 @@ def build_sampler_config(cfg, registry, logger=None, **kwargs):
             f'registry must be type Registry, got {type(registry)}')
 
     cfg = deep_copy(cfg)
-
     req_type = cfg.get('NAME')
+
+    from scepter.modules.utils.import_utils import LazyImportModule
+    sig = (registry.name.upper(), req_type)
+    LazyImportModule.import_module(sig)
+
     if isinstance(req_type, str):
         req_type_entry = registry.get(req_type)
         if req_type_entry is None:

scepter/modules/data/utils/__init__.py

@@ -1,4 +1,21 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-from scepter.modules.data.utils.data_bucket import BucketManager
+
+if TYPE_CHECKING:
+    from scepter.modules.data.utils.data_bucket import BucketManager
+else:
+    _import_structure = {
+        'data_bucket': ['BucketManager']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/inference/__init__.py

@@ -1,3 +1,39 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from scepter.modules.inference.diffusion_inference import DiffusionInference
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
+
+
+if TYPE_CHECKING:
+    from scepter.modules.inference.diffusion_inference import DiffusionInference
+    from scepter.modules.inference.ace_inference import ACEInference
+    from scepter.modules.inference.cogvideox_inference import CogVideoXInference
+    from scepter.modules.inference.control_inference import ControlInference
+    from scepter.modules.inference.flux_inference import FluxInference
+    from scepter.modules.inference.largen_inference import LargenInference
+    from scepter.modules.inference.pixart_inference import PixArtInference
+    from scepter.modules.inference.sd3_inference import SD3Inference
+    from scepter.modules.inference.stylebooth_inference import StyleboothInference
+    from scepter.modules.inference.tuner_inference import TunerInference
+else:
+    _import_structure = {
+        'diffusion_inference': ['DiffusionInference'],
+        'ace_inference': ['ACEInference'],
+        'cogvideox_inference': ['CogVideoXInference'],
+        'control_inference': ['ControlInference'],
+        'flux_inference': ['FluxInference'],
+        'largen_inference': ['LargenInference'],
+        'pixart_inference': ['PixArtInference'],
+        'sd3_inference': ['SD3Inference'],
+        'stylebooth_inference': ['StyleboothInference'],
+        'tuner_inference': ['TunerInference']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/inference/ace_inference.py

@@ -0,0 +1,551 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import copy
+import math
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from PIL import Image
+import torchvision.transforms as T
+from scepter.modules.model.registry import DIFFUSIONS
+from scepter.modules.model.utils.basic_utils import check_list_of_list
+from scepter.modules.model.utils.basic_utils import \
+    pack_imagelist_into_tensor_v2 as pack_imagelist_into_tensor
+from scepter.modules.model.utils.basic_utils import (
+    to_device, unpack_tensor_into_imagelist)
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.logger import get_logger
+
+from .diffusion_inference import DiffusionInference, get_model
+
+
+def process_edit_image(images,
+                       masks,
+                       tasks,
+                       max_seq_len=1024,
+                       max_aspect_ratio=4,
+                       d=16,
+                       **kwargs):
+
+    if not isinstance(images, list):
+        images = [images]
+    if not isinstance(masks, list):
+        masks = [masks]
+    if not isinstance(tasks, list):
+        tasks = [tasks]
+
+    img_tensors = []
+    mask_tensors = []
+    for img, mask, task in zip(images, masks, tasks):
+        if mask is None or mask == '':
+            mask = Image.new('L', img.size, 0)
+        W, H = img.size
+        if H / W > max_aspect_ratio:
+            img = TF.center_crop(img, [int(max_aspect_ratio * W), W])
+            mask = TF.center_crop(mask, [int(max_aspect_ratio * W), W])
+        elif W / H > max_aspect_ratio:
+            img = TF.center_crop(img, [H, int(max_aspect_ratio * H)])
+            mask = TF.center_crop(mask, [H, int(max_aspect_ratio * H)])
+
+        H, W = img.height, img.width
+        scale = min(1.0, math.sqrt(max_seq_len / ((H / d) * (W / d))))
+        rH = int(H * scale) // d * d  # ensure divisible by self.d
+        rW = int(W * scale) // d * d
+
+        img = TF.resize(img, (rH, rW),
+                        interpolation=TF.InterpolationMode.BICUBIC)
+        mask = TF.resize(mask, (rH, rW),
+                         interpolation=TF.InterpolationMode.NEAREST_EXACT)
+
+        mask = np.asarray(mask)
+        mask = np.where(mask > 128, 1, 0)
+        mask = mask.astype(
+            np.float32) if np.any(mask) else np.ones_like(mask).astype(
+                np.float32)
+
+        img_tensor = TF.to_tensor(img).to(we.device_id)
+        img_tensor = TF.normalize(img_tensor,
+                                  mean=[0.5, 0.5, 0.5],
+                                  std=[0.5, 0.5, 0.5])
+        mask_tensor = TF.to_tensor(mask).to(we.device_id)
+        if task in ['inpainting', 'Try On', 'Inpainting']:
+            mask_indicator = mask_tensor.repeat(3, 1, 1)
+            img_tensor[mask_indicator == 1] = -1.0
+        img_tensors.append(img_tensor)
+        mask_tensors.append(mask_tensor)
+    return img_tensors, mask_tensors
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, embedding_shape):
+        super().__init__()
+        self.pos = nn.Parameter(data=torch.zeros(embedding_shape))
+
+class RefinerInference(DiffusionInference):
+    def init_from_cfg(self, cfg):
+        self.use_dynamic_model = cfg.get('USE_DYNAMIC_MODEL', True)
+        super().init_from_cfg(cfg)
+        self.diffusion = DIFFUSIONS.build(cfg.MODEL.DIFFUSION, logger=self.logger) \
+            if cfg.MODEL.have('DIFFUSION') else None
+        self.max_seq_length = cfg.MODEL.get("MAX_SEQ_LENGTH", 4096)
+        assert self.diffusion is not None
+        if not self.use_dynamic_model:
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+    @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)):
+            def run_one_image(u):
+                zu = get_model(self.first_stage_model).encode(u)
+                if isinstance(zu, (tuple, list)):
+                    zu = zu[0]
+                return zu
+            z = [run_one_image(u.unsqueeze(0) if u.dim == 3 else u) for u in x]
+            return z
+    def upscale_resize(self, image, interpolation=T.InterpolationMode.BILINEAR):
+        c, H, W = image.shape
+        scale = max(1.0, math.sqrt(self.max_seq_length / ((H / 16) * (W / 16))))
+        rH = int(H * scale) // 16 * 16  # ensure divisible by self.d
+        rW = int(W * scale) // 16 * 16
+        image = T.Resize((rH, rW), interpolation=interpolation, antialias=True)(image)
+        return image
+    @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(zu) for zu in z]
+
+    def noise_sample(self, num_samples, h, w, seed, device = None, dtype = torch.bfloat16):
+        noise = torch.randn(
+            num_samples,
+            16,
+            # allow for packing
+            2 * math.ceil(h / 16),
+            2 * math.ceil(w / 16),
+            device=device,
+            dtype=dtype,
+            generator=torch.Generator(device=device).manual_seed(seed),
+        )
+        return noise
+    def refine(self,
+               x_samples=None,
+               prompt=None,
+               reverse_scale=-1.,
+               seed = 2024,
+               **kwargs
+               ):
+        print(prompt)
+        value_input = copy.deepcopy(self.input)
+        x_samples = [self.upscale_resize(x) for x in x_samples]
+
+        noise = []
+        for i, x in enumerate(x_samples):
+            noise_ = self.noise_sample(1, x.shape[1],
+                                       x.shape[2], seed,
+                                       device = x.device)
+            noise.append(noise_)
+        noise, x_shapes = pack_imagelist_into_tensor(noise)
+        if reverse_scale > 0:
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            x_samples = [x.unsqueeze(0) for x in x_samples]
+            x_start = self.encode_first_stage(x_samples, **kwargs)
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=not self.use_dynamic_model)
+            x_start, _ = pack_imagelist_into_tensor(x_start)
+        else:
+            x_start = None
+        # 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)(prompt)
+            ctx["x_shapes"] = x_shapes
+        self.dynamic_unload(self.cond_stage_model,
+                            'cond_stage_model',
+                            skip_loaded=not self.use_dynamic_model)
+
+
+        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_euler')
+            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,
+                reverse_scale=reverse_scale,
+                x=x_start,
+                **kwargs).float()
+        latent = unpack_tensor_into_imagelist(latent, x_shapes)
+        self.dynamic_unload(self.diffusion_model,
+                            'diffusion_model',
+                            skip_loaded=not self.use_dynamic_model)
+        self.dynamic_load(self.first_stage_model, 'first_stage_model')
+        x_samples = self.decode_first_stage(latent)
+        self.dynamic_unload(self.first_stage_model,
+                            'first_stage_model',
+                            skip_loaded=not self.use_dynamic_model)
+        return x_samples
+
+
+class ACEInference(DiffusionInference):
+    def __init__(self, logger=None):
+        if logger is None:
+            logger = get_logger(name='scepter')
+        self.logger = logger
+        self.loaded_model = {}
+        self.loaded_model_name = [
+            'diffusion_model', 'first_stage_model', 'cond_stage_model'
+        ]
+
+    def init_from_cfg(self, cfg):
+        self.name = cfg.NAME
+        self.is_default = cfg.get('IS_DEFAULT', False)
+        self.use_dynamic_model = cfg.get('USE_DYNAMIC_MODEL', True)
+        module_paras = self.load_default(cfg.get('DEFAULT_PARAS', None))
+        assert cfg.have('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_model_cfg = cfg.get('REFINER_MODEL', None)
+        # self.refiner_scale = cfg.get('REFINER_SCALE', 0.)
+        # self.refiner_prompt = cfg.get('REFINER_PROMPT', "")
+        self.ace_prompt = cfg.get("ACE_PROMPT", [])
+        if self.refiner_model_cfg:
+            self.refiner_model_cfg.USE_DYNAMIC_MODEL = self.use_dynamic_model
+            self.refiner_module = RefinerInference(self.logger)
+            self.refiner_module.init_from_cfg(self.refiner_model_cfg)
+        else:
+            self.refiner_module = None
+
+        self.diffusion = DIFFUSIONS.build(cfg.MODEL.DIFFUSION,
+                                          logger=self.logger)
+
+
+        self.interpolate_func = lambda x: (F.interpolate(
+            x.unsqueeze(0),
+            scale_factor=1 / self.size_factor,
+            mode='nearest-exact') if x is not None else None)
+        self.text_indentifers = cfg.MODEL.get('TEXT_IDENTIFIER', [])
+        self.use_text_pos_embeddings = cfg.MODEL.get('USE_TEXT_POS_EMBEDDINGS',
+                                                     False)
+        if self.use_text_pos_embeddings:
+            self.text_position_embeddings = TextEmbedding(
+                (10, 4096)).eval().requires_grad_(False).to(we.device_id)
+        else:
+            self.text_position_embeddings = None
+
+        self.max_seq_len = cfg.MODEL.DIFFUSION_MODEL.MAX_SEQ_LEN
+        self.scale_factor = cfg.get('SCALE_FACTOR', 0.18215)
+        self.size_factor = cfg.get('SIZE_FACTOR', 8)
+        self.decoder_bias = cfg.get('DECODER_BIAS', 0)
+        self.default_n_prompt = cfg.get('DEFAULT_N_PROMPT', '')
+        if not self.use_dynamic_model:
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+
+    @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 != 'float32'),
+                            dtype=getattr(torch, dtype)):
+            z = [
+                self.scale_factor * get_model(self.first_stage_model)._encode(
+                    i.unsqueeze(0).to(getattr(torch, dtype))) for i in x
+            ]
+        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 != 'float32'),
+                            dtype=getattr(torch, dtype)):
+            x = [
+                get_model(self.first_stage_model)._decode(
+                    1. / self.scale_factor * i.to(getattr(torch, dtype)))
+                for i in z
+            ]
+        return x
+
+
+
+    @torch.no_grad()
+    def __call__(self,
+                 image=None,
+                 mask=None,
+                 prompt='',
+                 task=None,
+                 negative_prompt='',
+                 output_height=512,
+                 output_width=512,
+                 sampler='ddim',
+                 sample_steps=20,
+                 guide_scale=4.5,
+                 guide_rescale=0.5,
+                 seed=-1,
+                 history_io=None,
+                 tar_index=0,
+                 **kwargs):
+        input_image, input_mask = image, mask
+        g = torch.Generator(device=we.device_id)
+        seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
+        g.manual_seed(int(seed))
+        if input_image is not None:
+            # assert isinstance(input_image, list) and isinstance(input_mask, list)
+            if task is None:
+                task = [''] * len(input_image)
+            if not isinstance(prompt, list):
+                prompt = [prompt] * len(input_image)
+            if history_io is not None and len(history_io) > 0:
+                his_image, his_maks, his_prompt, his_task = history_io[
+                    'image'], history_io['mask'], history_io[
+                        'prompt'], history_io['task']
+                assert len(his_image) == len(his_maks) == len(
+                    his_prompt) == len(his_task)
+                input_image = his_image + input_image
+                input_mask = his_maks + input_mask
+                task = his_task + task
+                prompt = his_prompt + [prompt[-1]]
+                prompt = [
+                    pp.replace('{image}', f'{{image{i}}}') if i > 0 else pp
+                    for i, pp in enumerate(prompt)
+                ]
+
+            edit_image, edit_image_mask = process_edit_image(
+                input_image, input_mask, task, max_seq_len=self.max_seq_len)
+
+            image, image_mask = edit_image[tar_index], edit_image_mask[
+                tar_index]
+            edit_image, edit_image_mask = [edit_image], [edit_image_mask]
+
+        else:
+            edit_image = edit_image_mask = [[]]
+            image = torch.zeros(
+                size=[3, int(output_height),
+                      int(output_width)])
+            image_mask = torch.ones(
+                size=[1, int(output_height),
+                      int(output_width)])
+            if not isinstance(prompt, list):
+                prompt = [prompt]
+
+        image, image_mask, prompt = [image], [image_mask], [prompt]
+        assert check_list_of_list(prompt) and check_list_of_list(
+            edit_image) and check_list_of_list(edit_image_mask)
+        # Assign Negative Prompt
+        if isinstance(negative_prompt, list):
+            negative_prompt = negative_prompt[0]
+        assert isinstance(negative_prompt, str)
+
+        n_prompt = copy.deepcopy(prompt)
+        for nn_p_id, nn_p in enumerate(n_prompt):
+            assert isinstance(nn_p, list)
+            n_prompt[nn_p_id][-1] = negative_prompt
+
+        is_txt_image = sum([len(e_i) for e_i in edit_image]) < 1
+        image = to_device(image)
+
+        refiner_scale = kwargs.pop("refiner_scale", 0.0)
+        refiner_prompt = kwargs.pop("refiner_prompt", "")
+        use_ace = kwargs.pop("use_ace", True)
+        # <= 0 use ace as the txt2img generator.
+        if use_ace and (not is_txt_image or refiner_scale <= 0):
+            ctx, null_ctx = {}, {}
+            # Get Noise Shape
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            x = self.encode_first_stage(image)
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=not self.use_dynamic_model)
+            noise = [
+                torch.empty(*i.shape, device=we.device_id).normal_(generator=g)
+                for i in x
+            ]
+            noise, x_shapes = pack_imagelist_into_tensor(noise)
+            ctx['x_shapes'] = null_ctx['x_shapes'] = x_shapes
+
+            image_mask = to_device(image_mask, strict=False)
+            cond_mask = [self.interpolate_func(i) for i in image_mask
+                         ] if image_mask is not None else [None] * len(image)
+            ctx['x_mask'] = null_ctx['x_mask'] = cond_mask
+
+            # Encode Prompt
+            self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
+            function_name, dtype = self.get_function_info(self.cond_stage_model)
+            cont, cont_mask = getattr(get_model(self.cond_stage_model),
+                                      function_name)(prompt)
+            cont, cont_mask = self.cond_stage_embeddings(prompt, edit_image, cont,
+                                                         cont_mask)
+            null_cont, null_cont_mask = getattr(get_model(self.cond_stage_model),
+                                                function_name)(n_prompt)
+            null_cont, null_cont_mask = self.cond_stage_embeddings(
+                prompt, edit_image, null_cont, null_cont_mask)
+            self.dynamic_unload(self.cond_stage_model,
+                                'cond_stage_model',
+                                skip_loaded=not self.use_dynamic_model)
+            ctx['crossattn'] = cont
+            null_ctx['crossattn'] = null_cont
+
+            # Encode Edit Images
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            edit_image = [to_device(i, strict=False) for i in edit_image]
+            edit_image_mask = [to_device(i, strict=False) for i in edit_image_mask]
+            e_img, e_mask = [], []
+            for u, m in zip(edit_image, edit_image_mask):
+                if u is None:
+                    continue
+                if m is None:
+                    m = [None] * len(u)
+                e_img.append(self.encode_first_stage(u, **kwargs))
+                e_mask.append([self.interpolate_func(i) for i in m])
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=not self.use_dynamic_model)
+            null_ctx['edit'] = ctx['edit'] = e_img
+            null_ctx['edit_mask'] = ctx['edit_mask'] = e_mask
+
+            # Diffusion Process
+            self.dynamic_load(self.diffusion_model, 'diffusion_model')
+            function_name, dtype = self.get_function_info(self.diffusion_model)
+            with torch.autocast('cuda',
+                                enabled=dtype in ('float16', 'bfloat16'),
+                                dtype=getattr(torch, dtype)):
+                latent = self.diffusion.sample(
+                    noise=noise,
+                    sampler=sampler,
+                    model=get_model(self.diffusion_model),
+                    model_kwargs=[{
+                        'cond':
+                        ctx,
+                        'mask':
+                        cont_mask,
+                        'text_position_embeddings':
+                        self.text_position_embeddings.pos if hasattr(
+                            self.text_position_embeddings, 'pos') else None
+                    }, {
+                        'cond':
+                        null_ctx,
+                        'mask':
+                        null_cont_mask,
+                        'text_position_embeddings':
+                        self.text_position_embeddings.pos if hasattr(
+                            self.text_position_embeddings, 'pos') else None
+                    }] if guide_scale is not None and guide_scale > 1 else {
+                        'cond':
+                        null_ctx,
+                        'mask':
+                        cont_mask,
+                        'text_position_embeddings':
+                        self.text_position_embeddings.pos if hasattr(
+                            self.text_position_embeddings, 'pos') else None
+                    },
+                    steps=sample_steps,
+                    show_progress=True,
+                    seed=seed,
+                    guide_scale=guide_scale,
+                    guide_rescale=guide_rescale,
+                    return_intermediate=None,
+                    **kwargs)
+            self.dynamic_unload(self.diffusion_model,
+                                'diffusion_model',
+                                skip_loaded=not self.use_dynamic_model)
+
+            # Decode to Pixel Space
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            samples = unpack_tensor_into_imagelist(latent, x_shapes)
+            x_samples = self.decode_first_stage(samples)
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=not self.use_dynamic_model)
+            x_samples = [x.squeeze(0) for x in x_samples]
+        else:
+            x_samples = image
+        if self.refiner_module and refiner_scale > 0:
+            if is_txt_image:
+                random.shuffle(self.ace_prompt)
+                input_refine_prompt = [self.ace_prompt[0] + refiner_prompt if p[0] == "" else p[0] for p in prompt]
+                input_refine_scale = -1.
+            else:
+                input_refine_prompt = [p[0].replace("{image}", "") + " " + refiner_prompt for p in prompt]
+                input_refine_scale = refiner_scale
+                print(input_refine_prompt)
+
+            x_samples = self.refiner_module.refine(x_samples,
+                                                   reverse_scale = input_refine_scale,
+                                                   prompt= input_refine_prompt,
+                                                   seed=seed,
+                                                   use_dynamic_model=self.use_dynamic_model)
+
+        imgs = [
+            torch.clamp((x_i.float() + 1.0) / 2.0 + self.decoder_bias / 255,
+                        min=0.0,
+                        max=1.0).squeeze(0).permute(1, 2, 0).cpu().numpy()
+            for x_i in x_samples
+        ]
+        imgs = [Image.fromarray((img * 255).astype(np.uint8)) for img in imgs]
+        return imgs
+
+    def cond_stage_embeddings(self, prompt, edit_image, cont, cont_mask):
+        if self.use_text_pos_embeddings and not torch.sum(
+                self.text_position_embeddings.pos) > 0:
+            identifier_cont, _ = getattr(get_model(self.cond_stage_model),
+                                         'encode')(self.text_indentifers,
+                                                   return_mask=True)
+            self.text_position_embeddings.load_state_dict(
+                {'pos': identifier_cont[:, 0, :]})
+
+        cont_, cont_mask_ = [], []
+        for pp, edit, c, cm in zip(prompt, edit_image, cont, cont_mask):
+            if isinstance(pp, list):
+                cont_.append([c[-1], *c] if len(edit) > 0 else [c[-1]])
+                cont_mask_.append([cm[-1], *cm] if len(edit) > 0 else [cm[-1]])
+            else:
+                raise NotImplementedError
+
+        return cont_, cont_mask_

scepter/modules/inference/cogvideox_inference.py

@@ -0,0 +1,183 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import copy
+import numpy as np
+from typing import Tuple
+import random
+
+import torch
+
+from scepter.modules.utils.file_system import FS
+from scepter.modules.utils.distribute import we
+from scepter.modules.model.backbone.cogvideox.utils import get_3d_rotary_pos_embed, get_resize_crop_region_for_grid
+from .diffusion_inference import DiffusionInference, get_model
+from .tuner_inference import TunerInference
+
+class CogVideoXInference(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)
+
+    @torch.no_grad()
+    def decode_first_stage(self, latents):
+        _, dtype = self.get_function_info(self.first_stage_model, 'decode')
+        with torch.autocast('cuda',
+                            enabled=dtype in ('bfloat16'),
+                            dtype=getattr(torch, dtype)):
+            latents = latents.permute(0, 2, 1, 3, 4)
+            latents = 1 / self.first_stage_model['paras']['scaling_factor_image'] * latents
+            frames = get_model(self.first_stage_model).decode(latents)
+        return frames
+
+    def _prepare_rotary_positional_embeddings(
+        self,
+        height: int,
+        width: int,
+        num_frames: int,
+        device: torch.device,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        grid_height = height // (self.diffusion_model['paras']['scale_factor_spatial'] * self.diffusion_model['paras']['patch_size'])
+        grid_width = width // (self.diffusion_model['paras']['scale_factor_spatial'] * self.diffusion_model['paras']['patch_size'])
+        base_size_width = self.diffusion_model['paras']['sample_width'] // (self.diffusion_model['paras']['scale_factor_spatial'] * self.diffusion_model['paras']['patch_size'])
+        base_size_height = self.diffusion_model['paras']['sample_height'] // (self.diffusion_model['paras']['scale_factor_spatial'] * self.diffusion_model['paras']['patch_size'])
+
+        grid_crops_coords = get_resize_crop_region_for_grid(
+            (grid_height, grid_width), base_size_width, base_size_height
+        )
+        freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
+            embed_dim=self.diffusion_model['paras']['attention_head_dim'],
+            crops_coords=grid_crops_coords,
+            grid_size=(grid_height, grid_width),
+            temporal_size=num_frames,
+        )
+
+        freqs_cos = freqs_cos.to(device=device)
+        freqs_sin = freqs_sin.to(device=device)
+        return freqs_cos, freqs_sin
+
+    @torch.no_grad()
+    def __call__(self,
+                 input,
+                 num_samples=1,
+                 cat_uc=True,
+                 tuner_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.tuner_infer.register_tuner(tuner_model, self.diffusion_model,
+                                            cond_stage_model=None)
+            self.dynamic_unload(self.diffusion_model,
+                                'diffusion_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(device_type='cuda', enabled=True, dtype=torch.bfloat16):
+            cont = getattr(get_model(self.cond_stage_model),
+                                  function_name)(value_input['prompt'], return_mask=False, use_mask=False)
+            null_cont = getattr(get_model(self.cond_stage_model),
+                                        function_name)(value_input['negative_prompt'] * num_samples, return_mask=False, use_mask=False)
+        self.dynamic_unload(self.cond_stage_model,
+                            'cond_stage_model',
+                            skip_loaded=True)
+
+        # get noise
+        seed = kwargs.pop('seed', -1)
+        seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
+        generator = torch.Generator().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_shape = (1,
+                         (value_input['num_frames'] - 1) // self.diffusion_model['paras']['scale_factor_temporal'] + 1,
+                         self.diffusion_model['paras']['latent_channels'],
+                         height // self.diffusion_model['paras']['scale_factor_spatial'],
+                         width // self.diffusion_model['paras']['scale_factor_spatial']
+                         )
+                noise = torch.randn(noise_shape, generator=generator, dtype=getattr(torch, dtype), device='cpu').to(we.device_id)
+
+                self.dynamic_load(self.diffusion_model, 'diffusion_model')
+
+                image_rotary_emb = (
+                    self._prepare_rotary_positional_embeddings(height, width, noise.size(1), we.device_id)
+                    if self.diffusion_model['paras']['use_rotary_positional_embeddings']
+                    else None
+                )
+                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', '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)
+
+                    latent = self.diffusion.sample(noise=noise,
+                                        sampler=solver_sample,
+                                        model=get_model(self.diffusion_model),
+                                        model_kwargs=[{
+                                            'cond': cont,
+                                            'image_latent': None,
+                                            'image_rotary_emb': image_rotary_emb,
+                                        }, {
+                                            'cond': null_cont,
+                                            'image_latent': None,
+                                            'image_rotary_emb': image_rotary_emb,
+                                        }],
+                                        steps=sample_steps,
+                                        show_progress=True,
+                                        guide_scale=guide_scale,
+                                        guide_rescale=guide_rescale,
+                                        return_intermediate=None,
+                                        **kwargs).float()
+                self.dynamic_unload(self.diffusion_model,
+                                    'diffusion_model',
+                                    skip_loaded=True)
+            self.dynamic_load(self.first_stage_model, 'first_stage_model')
+            x_samples = self.decode_first_stage(latent).float()  # [B, C, F, H, W]
+            self.dynamic_unload(self.first_stage_model,
+                                'first_stage_model',
+                                skip_loaded=True)
+
+            x_frames = torch.clamp(x_samples / 2 + 0.5, min=0.0, max=1.0)
+            if 'videos' in value_output:
+                if value_output['videos'] is None or (
+                        isinstance(value_output['videos'], list)
+                        and len(value_output['videos']) < 1):
+                    value_output['videos'] = []
+                value_output['videos'].append(x_frames)
+
+        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,
+                                              cond_stage_model=None)
+        return value_output

scepter/modules/inference/diffusion_inference.py

@@ -11,9 +11,10 @@ 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)
+                                            TOKENIZERS, DIFFUSIONS)
 from scepter.modules.utils.distribute import we
 from scepter.modules.utils.file_system import FS
+from scepter.modules.utils.config import Config
 from scepter.studio.utils.env import get_available_memory
 
 from .control_inference import ControlInference
@@ -49,7 +50,10 @@ class DiffusionInference():
         assert cfg.have('MODEL')
         if self.is_redefine_paras:
             cfg.MODEL = self.redefine_paras(cfg.MODEL)
-        self.diffusion = self.load_schedule(cfg.MODEL.SCHEDULE)
+        if 'DIFFUSION' in cfg.MODEL:
+            self.diffusion = DIFFUSIONS.build(cfg.MODEL.DIFFUSION, logger=self.logger)
+        else:
+            self.diffusion = self.load_schedule(cfg.MODEL.SCHEDULE)
         self.diffusion_model = self.infer_model(
             cfg.MODEL.DIFFUSION_MODEL, module_paras.get(
                 'DIFFUSION_MODEL',
@@ -93,7 +97,7 @@ class DiffusionInference():
                     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')
+                        sd = torch.load(local_path, map_location='cpu', weights_only=True)
                 first_stage_model_path = os.path.join(
                     os.path.dirname(local_path), 'first_stage_model.pth')
                 cond_stage_model_path = os.path.join(
@@ -199,7 +203,7 @@ class DiffusionInference():
             from safetensors.torch import load_file as load_safetensors
             sd = load_safetensors(path)
         else:
-            sd = torch.load(path, map_location='cpu')
+            sd = torch.load(path, map_location='cpu', weights_only=True)
 
         new_sd = OrderedDict()
         for k, v in sd.items():
@@ -226,16 +230,22 @@ class DiffusionInference():
 
     def load(self, module):
         if module['device'] == 'offline':
-            if module['cfg'].NAME in MODELS.class_map:
+            from scepter.modules.utils.import_utils import LazyImportModule
+            if (LazyImportModule.get_module_type(('MODELS', module['cfg'].NAME)) or
+                    module['cfg'].NAME in MODELS.class_map):
                 model = MODELS.build(module['cfg'], logger=self.logger).eval()
-            elif module['cfg'].NAME in BACKBONES.class_map:
+            elif (LazyImportModule.get_module_type(('BACKBONES', module['cfg'].NAME)) or
+                    module['cfg'].NAME in BACKBONES.class_map):
                 model = BACKBONES.build(module['cfg'],
                                         logger=self.logger).eval()
-            elif module['cfg'].NAME in EMBEDDERS.class_map:
+            elif (LazyImportModule.get_module_type(('EMBEDDERS', module['cfg'].NAME)) or
+                    module['cfg'].NAME in EMBEDDERS.class_map):
                 model = EMBEDDERS.build(module['cfg'],
                                         logger=self.logger).eval()
             else:
                 raise NotImplementedError
+            if 'DTYPE' in module['cfg'] and module['cfg']['DTYPE'] is not None:
+                model = model.to(getattr(torch, module['cfg'].DTYPE))
             if module['cfg'].get('RELOAD_MODEL', None):
                 self.init_from_ckpt(module['cfg'].RELOAD_MODEL, model)
             module['model'] = model
@@ -264,8 +274,9 @@ class DiffusionInference():
                 module['device'] = 'cpu'
             else:
                 module['device'] = 'offline'
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
         return module
 
     def dynamic_load(self, module=None, name=''):
@@ -313,7 +324,8 @@ class DiffusionInference():
         module_paras = {}
         if cfg is not None:
             self.paras = cfg.PARAS
-            self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict)) else v for k, v in cfg.INPUT.items()}
+            self.input_cfg = {k.lower(): v for k, v in cfg.INPUT.items()}
+            self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict, Config)) 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

scepter/modules/inference/flux_inference.py

@@ -151,7 +151,7 @@ class FluxInference(DiffusionInference):
                 with torch.autocast('cuda',
                                     enabled= dtype in ('float16', 'bfloat16'),
                                     dtype=getattr(torch, dtype)):
-                    solver_sample = value_input.get('sample', 'flow_eluer')
+                    solver_sample = value_input.get('sample', 'flow_euler')
                     sample_steps = value_input.get('sample_steps', 20)
                     guide_scale = value_input.get('guide_scale', 3.5)
                     if guide_scale is not None:

scepter/modules/inference/largen_inference.py

@@ -43,7 +43,7 @@ class LargenInference(DiffusionInference):
                     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')
+                        sd = torch.load(local_path, map_location='cpu', weights_only=True)
                 if 'model' in sd:
                     sd = sd['model']
 

scepter/modules/inference/pixart_inference.py

@@ -1,20 +1,12 @@
 # -*- 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

scepter/modules/inference/sd3_inference.py

@@ -3,10 +3,8 @@
 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

scepter/modules/inference/tuner_inference.py

@@ -29,11 +29,11 @@ class TunerInference():
             warnings.warn(f'Import swift error, please deal with this problem: {e}')
 
         self.logger.info('Unloading tuner model')
-        if isinstance(diffusion_model['model'], SwiftModel):
+        if diffusion_model is not None and isinstance(diffusion_model['model'], SwiftModel):
             for adapter_name in diffusion_model['model'].adapters:
                 diffusion_model['model'].deactivate_adapter(adapter_name,
                                                             offload='cpu')
-        if isinstance(cond_stage_model['model'], SwiftModel):
+        if cond_stage_model is not None and isinstance(cond_stage_model['model'], SwiftModel):
             for adapter_name in cond_stage_model['model'].adapters:
                 cond_stage_model['model'].deactivate_adapter(adapter_name,
                                                              offload='cpu')
@@ -144,9 +144,9 @@ class TunerInference():
                         is_bin_file = True
                         if os.path.isfile(bin_file):
                             if 'weights_only' in torch.load.__code__.co_varnames:
-                                state_dict = torch.load(bin_file, weights_only=True)
+                                state_dict = torch.load(bin_file, weights_only=True, map_location="cpu")
                             else:
-                                state_dict = torch.load(bin_file)
+                                state_dict = torch.load(bin_file, map_location="cpu")
                         elif os.path.isfile(safe_file):
                             is_bin_file = False
                             from safetensors.torch import \

scepter/modules/model/__init__.py

@@ -1,5 +1,23 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-from scepter.modules.model import (backbone, embedder, head, loss, metric,
-                                   neck, network, tokenizer, tuner, diffusion)
+
+if TYPE_CHECKING:
+    from scepter.modules.model import (backbone, embedder, head, loss, metric,
+                                       neck, network, tokenizer, tuner, diffusion)
+else:
+    _import_structure = {
+        'model': ['backbone', 'embedder', 'head', 'loss', 'metric',
+                  'neck', 'network', 'tokenizer', 'tuner', 'diffusion']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/model/backbone/__init__.py

@@ -1,4 +1,23 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from scepter.modules.model.backbone import (autoencoder, image, mmdit, pixart,
-                                            unet, utils, video, flux)
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
+
+
+if TYPE_CHECKING:
+    from scepter.modules.model.backbone import (ace, autoencoder, flux, image, cogvideox,
+                                                mmdit, pixart, unet, utils, video)
+else:
+    _import_structure = {
+        'backbone': ['ace', 'autoencoder', 'flux', 'image', 'cogvideox',
+                     'mmdit', 'pixart', 'unet', 'utils', 'video']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/model/backbone/ace/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .ace import ACE

scepter/modules/model/backbone/ace/ace.py

@@ -0,0 +1,372 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import re
+from collections import OrderedDict
+from functools import partial
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.checkpoint import checkpoint_sequential
+
+from scepter.modules.model.backbone.transformer.layers import (Mlp,
+                                                               T2IFinalLayer,
+                                                               TimestepEmbedder
+                                                               )
+from scepter.modules.model.backbone.transformer.patchify import PatchEmbed
+from scepter.modules.model.backbone.transformer.pos_embed import rope_params
+from scepter.modules.model.base_model import BaseModel
+from scepter.modules.model.registry import BACKBONES
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.file_system import FS
+
+from .layers import ACEBlock
+
+
+@BACKBONES.register_class()
+class ACE(BaseModel):
+
+    para_dict = {
+        'PATCH_SIZE': {
+            'value': 2,
+            'description': ''
+        },
+        'IN_CHANNELS': {
+            'value': 4,
+            'description': ''
+        },
+        'HIDDEN_SIZE': {
+            'value': 1152,
+            'description': ''
+        },
+        'DEPTH': {
+            'value': 28,
+            'description': ''
+        },
+        'NUM_HEADS': {
+            'value': 16,
+            'description': ''
+        },
+        'MLP_RATIO': {
+            'value': 4.0,
+            'description': ''
+        },
+        'PRED_SIGMA': {
+            'value': True,
+            'description': ''
+        },
+        'DROP_PATH': {
+            'value': 0.,
+            'description': ''
+        },
+        'WINDOW_SIZE': {
+            'value': 0,
+            'description': ''
+        },
+        'WINDOW_BLOCK_INDEXES': {
+            'value': None,
+            'description': ''
+        },
+        'Y_CHANNELS': {
+            'value': 4096,
+            'description': ''
+        },
+        'ATTENTION_BACKEND': {
+            'value': None,
+            'description': ''
+        },
+        'QK_NORM': {
+            'value': True,
+            'description': 'Whether to use RMSNorm for query and key.',
+        },
+    }
+    para_dict.update(BaseModel.para_dict)
+
+    def __init__(self, cfg, logger):
+        super().__init__(cfg, logger=logger)
+        self.window_block_indexes = cfg.get('WINDOW_BLOCK_INDEXES', None)
+        if self.window_block_indexes is None:
+            self.window_block_indexes = []
+        self.pred_sigma = cfg.get('PRED_SIGMA', True)
+        self.in_channels = cfg.get('IN_CHANNELS', 4)
+        self.out_channels = self.in_channels * 2 if self.pred_sigma else self.in_channels
+        self.patch_size = cfg.get('PATCH_SIZE', 2)
+        self.num_heads = cfg.get('NUM_HEADS', 16)
+        self.hidden_size = cfg.get('HIDDEN_SIZE', 1152)
+        self.y_channels = cfg.get('Y_CHANNELS', 4096)
+        self.drop_path = cfg.get('DROP_PATH', 0.)
+        self.depth = cfg.get('DEPTH', 28)
+        self.mlp_ratio = cfg.get('MLP_RATIO', 4.0)
+        self.use_grad_checkpoint = cfg.get('USE_GRAD_CHECKPOINT', False)
+        self.attention_backend = cfg.get('ATTENTION_BACKEND', None)
+        self.max_seq_len = cfg.get('MAX_SEQ_LEN', 1024)
+        self.qk_norm = cfg.get('QK_NORM', False)
+        self.ignore_keys = cfg.get('IGNORE_KEYS', [])
+        assert (self.hidden_size % self.num_heads
+                ) == 0 and (self.hidden_size // self.num_heads) % 2 == 0
+        d = self.hidden_size // self.num_heads
+        self.freqs = torch.cat(
+            [
+                rope_params(self.max_seq_len, d - 4 * (d // 6)),  # T (~1/3)
+                rope_params(self.max_seq_len, 2 * (d // 6)),  # H (~1/3)
+                rope_params(self.max_seq_len, 2 * (d // 6))  # W (~1/3)
+            ],
+            dim=1)
+
+        # init embedder
+        self.x_embedder = PatchEmbed(self.patch_size,
+                                     self.in_channels + 1,
+                                     self.hidden_size,
+                                     bias=True,
+                                     flatten=False)
+        self.t_embedder = TimestepEmbedder(self.hidden_size)
+        self.y_embedder = Mlp(in_features=self.y_channels,
+                              hidden_features=self.hidden_size,
+                              out_features=self.hidden_size,
+                              act_layer=lambda: nn.GELU(approximate='tanh'),
+                              drop=0)
+        self.t_block = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(self.hidden_size, 6 * self.hidden_size, bias=True))
+        # init blocks
+        drop_path = [
+            x.item() for x in torch.linspace(0, self.drop_path, self.depth)
+        ]
+        self.blocks = nn.ModuleList([
+            ACEBlock(self.hidden_size,
+                     self.num_heads,
+                     mlp_ratio=self.mlp_ratio,
+                     drop_path=drop_path[i],
+                     window_size=self.window_size
+                     if i in self.window_block_indexes else 0,
+                     backend=self.attention_backend,
+                     use_condition=True,
+                     qk_norm=self.qk_norm) for i in range(self.depth)
+        ])
+        self.final_layer = T2IFinalLayer(self.hidden_size, self.patch_size,
+                                         self.out_channels)
+        self.initialize_weights()
+
+    def load_pretrained_model(self, pretrained_model):
+        if pretrained_model:
+            with FS.get_from(pretrained_model, wait_finish=True) as local_path:
+                model = torch.load(local_path, map_location='cpu', weights_only=True)
+                if 'state_dict' in model:
+                    model = model['state_dict']
+                new_ckpt = OrderedDict()
+                for k, v in model.items():
+                    if self.ignore_keys is not None:
+                        if (isinstance(self.ignore_keys, str) and re.match(self.ignore_keys, k)) or \
+                                (isinstance(self.ignore_keys, list) and k in self.ignore_keys):
+                            continue
+                    k = k.replace('.cross_attn.q_linear.', '.cross_attn.q.')
+                    k = k.replace('.cross_attn.proj.',
+                                  '.cross_attn.o.').replace(
+                                      '.attn.proj.', '.attn.o.')
+                    if '.cross_attn.kv_linear.' in k:
+                        k_p, v_p = torch.split(v, v.shape[0] // 2)
+                        new_ckpt[k.replace('.cross_attn.kv_linear.',
+                                           '.cross_attn.k.')] = k_p
+                        new_ckpt[k.replace('.cross_attn.kv_linear.',
+                                           '.cross_attn.v.')] = v_p
+                    elif '.attn.qkv.' in k:
+                        q_p, k_p, v_p = torch.split(v, v.shape[0] // 3)
+                        new_ckpt[k.replace('.attn.qkv.', '.attn.q.')] = q_p
+                        new_ckpt[k.replace('.attn.qkv.', '.attn.k.')] = k_p
+                        new_ckpt[k.replace('.attn.qkv.', '.attn.v.')] = v_p
+                    elif 'y_embedder.y_proj.' in k:
+                        new_ckpt[k.replace('y_embedder.y_proj.',
+                                           'y_embedder.')] = v
+                    elif k in ('x_embedder.proj.weight'):
+                        model_p = self.state_dict()[k]
+                        if v.shape != model_p.shape:
+                            model_p.zero_()
+                            model_p[:, :4, :, :].copy_(v)
+                            new_ckpt[k] = torch.nn.parameter.Parameter(model_p)
+                        else:
+                            new_ckpt[k] = v
+                    elif k in ('x_embedder.proj.bias'):
+                        new_ckpt[k] = v
+                    else:
+                        new_ckpt[k] = v
+                missing, unexpected = self.load_state_dict(new_ckpt,
+                                                           strict=False)
+                print(
+                    f'Restored from {pretrained_model} with {len(missing)} missing and {len(unexpected)} unexpected keys'
+                )
+                if len(missing) > 0:
+                    print(f'Missing Keys:\n {missing}')
+                if len(unexpected) > 0:
+                    print(f'\nUnexpected Keys:\n {unexpected}')
+
+    def forward(self,
+                x,
+                t=None,
+                cond=dict(),
+                mask=None,
+                text_position_embeddings=None,
+                gc_seg=-1,
+                **kwargs):
+        if self.freqs.device != x.device:
+            self.freqs = self.freqs.to(x.device)
+        if isinstance(cond, dict):
+            context = cond.get('crossattn', None)
+        else:
+            context = cond
+        if text_position_embeddings is not None:
+            # default use the text_position_embeddings in state_dict
+            # if state_dict doesn't including this key, use the arg: text_position_embeddings
+            proj_position_embeddings = self.y_embedder(
+                text_position_embeddings)
+        else:
+            proj_position_embeddings = None
+
+        ctx_batch, txt_lens = [], []
+        if mask is not None and isinstance(mask, list):
+            for ctx, ctx_mask in zip(context, mask):
+                for frame_id, one_ctx in enumerate(zip(ctx, ctx_mask)):
+                    u, m = one_ctx
+                    t_len = m.flatten().sum()  # l
+                    u = u[:t_len]
+                    u = self.y_embedder(u)
+                    if frame_id == 0:
+                        u = u + proj_position_embeddings[
+                            len(ctx) -
+                            1] if proj_position_embeddings is not None else u
+                    else:
+                        u = u + proj_position_embeddings[
+                            frame_id -
+                            1] if proj_position_embeddings is not None else u
+                    ctx_batch.append(u)
+                    txt_lens.append(t_len)
+        else:
+            raise TypeError
+        y = torch.cat(ctx_batch, dim=0)
+        txt_lens = torch.LongTensor(txt_lens).to(x.device, non_blocking=True)
+
+        batch_frames = []
+        for u, shape, m in zip(x, cond['x_shapes'], cond['x_mask']):
+            u = u[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
+            m = torch.ones_like(u[[0], :, :]) if m is None else m.squeeze(0)
+            batch_frames.append([torch.cat([u, m], dim=0).unsqueeze(0)])
+        if 'edit' in cond:
+            for i, (edit, edit_mask) in enumerate(
+                    zip(cond['edit'], cond['edit_mask'])):
+                if edit is None:
+                    continue
+                for u, m in zip(edit, edit_mask):
+                    u = u.squeeze(0)
+                    m = torch.ones_like(
+                        u[[0], :, :]) if m is None else m.squeeze(0)
+                    batch_frames[i].append(
+                        torch.cat([u, m], dim=0).unsqueeze(0))
+
+        patch_batch, shape_batch, self_x_len, cross_x_len = [], [], [], []
+        for frames in batch_frames:
+            patches, patch_shapes = [], []
+            self_x_len.append(0)
+            for frame_id, u in enumerate(frames):
+                u = self.x_embedder(u)
+                h, w = u.size(2), u.size(3)
+                u = rearrange(u, '1 c h w -> (h w) c')
+                if frame_id == 0:
+                    u = u + proj_position_embeddings[
+                        len(frames) -
+                        1] if proj_position_embeddings is not None else u
+                else:
+                    u = u + proj_position_embeddings[
+                        frame_id -
+                        1] if proj_position_embeddings is not None else u
+                patches.append(u)
+                patch_shapes.append([h, w])
+                cross_x_len.append(h * w)  # b*s, 1
+                self_x_len[-1] += h * w  # b, 1
+            # u = torch.cat(patches, dim=0)
+            patch_batch.extend(patches)
+            shape_batch.append(
+                torch.LongTensor(patch_shapes).to(x.device, non_blocking=True))
+        # repeat t to align with x
+        t = torch.cat([t[i].repeat(l) for i, l in enumerate(self_x_len)])
+        self_x_len, cross_x_len = (torch.LongTensor(self_x_len).to(
+            x.device, non_blocking=True), torch.LongTensor(cross_x_len).to(
+                x.device, non_blocking=True))
+        # x = pad_sequence(tuple(patch_batch), batch_first=True)  # b, s*max(cl), c
+        x = torch.cat(patch_batch, dim=0)
+        x_shapes = pad_sequence(tuple(shape_batch),
+                                batch_first=True)  # b, max(len(frames)), 2
+        t = self.t_embedder(t)  # (N, D)
+        t0 = self.t_block(t)
+        # y = self.y_embedder(context)
+
+        kwargs = dict(y=y,
+                      t=t0,
+                      x_shapes=x_shapes,
+                      self_x_len=self_x_len,
+                      cross_x_len=cross_x_len,
+                      freqs=self.freqs,
+                      txt_lens=txt_lens)
+        if self.use_grad_checkpoint and gc_seg >= 0:
+            x = checkpoint_sequential(
+                functions=[partial(block, **kwargs) for block in self.blocks],
+                segments=gc_seg if gc_seg > 0 else len(self.blocks),
+                input=x,
+                use_reentrant=False)
+        else:
+            for block in self.blocks:
+                x = block(x, **kwargs)
+        x = self.final_layer(x, t)  # b*s*n, d
+        outs, cur_length = [], 0
+        p = self.patch_size
+        for seq_length, shape in zip(self_x_len, shape_batch):
+            x_i = x[cur_length:cur_length + seq_length]
+            h, w = shape[0].tolist()
+            u = x_i[:h * w].view(h, w, p, p, -1)
+            u = rearrange(u, 'h w p q c -> (h p w q) c'
+                          )  # dump into sequence for following tensor ops
+            cur_length = cur_length + seq_length
+            outs.append(u)
+        x = pad_sequence(tuple(outs), batch_first=True).permute(0, 2, 1)
+        if self.pred_sigma:
+            return x.chunk(2, dim=1)[0]
+        else:
+            return x
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        nn.init.normal_(self.t_block[1].weight, std=0.02)
+        # Initialize caption embedding MLP:
+        if hasattr(self, 'y_embedder'):
+            nn.init.normal_(self.y_embedder.fc1.weight, std=0.02)
+            nn.init.normal_(self.y_embedder.fc2.weight, std=0.02)
+        # Zero-out adaLN modulation layers
+        for block in self.blocks:
+            nn.init.constant_(block.cross_attn.o.weight, 0)
+            nn.init.constant_(block.cross_attn.o.bias, 0)
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('BACKBONE',
+                            __class__.__name__,
+                            ACE.para_dict,
+                            set_name=True)

scepter/modules/model/backbone/ace/layers.py

@@ -0,0 +1,205 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import math
+import warnings
+
+import torch
+import torch.nn as nn
+
+from scepter.modules.model.backbone.transformer.attention import RMSNorm
+from scepter.modules.model.backbone.transformer.layers import (DropPath, Mlp,
+                                                               modulate)
+from scepter.modules.model.backbone.transformer.pos_embed import \
+    rope_apply_multires as rope_apply
+
+try:
+    from flash_attn import (flash_attn_varlen_func)
+    FLASHATTN_IS_AVAILABLE = True
+except ImportError as e:
+    FLASHATTN_IS_AVAILABLE = False
+    flash_attn_varlen_func = None
+    warnings.warn(f'{e}')
+
+
+class ACEBlock(nn.Module):
+    def __init__(self,
+                 hidden_size,
+                 num_heads,
+                 mlp_ratio=4.0,
+                 drop_path=0.,
+                 window_size=0,
+                 backend=None,
+                 use_condition=True,
+                 qk_norm=False,
+                 **block_kwargs):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.use_condition = use_condition
+        self.norm1 = nn.LayerNorm(hidden_size,
+                                  elementwise_affine=False,
+                                  eps=1e-6)
+        self.attn = MultiHeadAttention(hidden_size,
+                                       num_heads=num_heads,
+                                       qkv_bias=True,
+                                       backend=backend,
+                                       qk_norm=qk_norm,
+                                       **block_kwargs)
+        if self.use_condition:
+            self.cross_attn = MultiHeadAttention(hidden_size,
+                                                 context_dim=hidden_size,
+                                                 num_heads=num_heads,
+                                                 qkv_bias=True,
+                                                 backend=backend,
+                                                 qk_norm=qk_norm,
+                                                 **block_kwargs)
+        self.norm2 = nn.LayerNorm(hidden_size,
+                                  elementwise_affine=False,
+                                  eps=1e-6)
+        # to be compatible with lower version pytorch
+        approx_gelu = lambda: nn.GELU(approximate='tanh')
+        self.mlp = Mlp(in_features=hidden_size,
+                       hidden_features=int(hidden_size * mlp_ratio),
+                       act_layer=approx_gelu,
+                       drop=0)
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+        self.window_size = window_size
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(6, hidden_size) / hidden_size**0.5)
+
+    def forward(self, x, y, t, **kwargs):
+        B = x.size(0)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            shift_msa.squeeze(1), scale_msa.squeeze(1), gate_msa.squeeze(1),
+            shift_mlp.squeeze(1), scale_mlp.squeeze(1), gate_mlp.squeeze(1))
+        x = x + self.drop_path(gate_msa * self.attn(
+            modulate(self.norm1(x), shift_msa, scale_msa, unsqueeze=False), **
+            kwargs))
+        if self.use_condition:
+            x = x + self.cross_attn(x, context=y, **kwargs)
+
+        x = x + self.drop_path(gate_mlp * self.mlp(
+            modulate(self.norm2(x), shift_mlp, scale_mlp, unsqueeze=False)))
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self,
+                 dim,
+                 context_dim=None,
+                 num_heads=None,
+                 head_dim=None,
+                 attn_drop=0.0,
+                 qkv_bias=False,
+                 dropout=0.0,
+                 backend=None,
+                 qk_norm=False,
+                 eps=1e-6,
+                 **block_kwargs):
+        super().__init__()
+        # consider head_dim first, then num_heads
+        num_heads = dim // head_dim if head_dim else num_heads
+        head_dim = dim // num_heads
+        assert num_heads * head_dim == dim
+        context_dim = context_dim or dim
+        self.dim = dim
+        self.context_dim = context_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.scale = math.pow(head_dim, -0.25)
+        # layers
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k = nn.Linear(context_dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(context_dim, dim, bias=qkv_bias)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+        self.dropout = nn.Dropout(dropout)
+        self.attention_op = None
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.backend = backend
+        assert self.backend in ('flash_attn', 'xformer_attn', 'pytorch_attn',
+                                None)
+        if FLASHATTN_IS_AVAILABLE and self.backend in ('flash_attn', None):
+            self.backend = 'flash_attn'
+            self.softmax_scale = block_kwargs.get('softmax_scale', None)
+            self.causal = block_kwargs.get('causal', False)
+            self.window_size = block_kwargs.get('window_size', (-1, -1))
+            self.deterministic = block_kwargs.get('deterministic', False)
+        else:
+            raise NotImplementedError
+
+    def flash_attn(self, x, context=None, **kwargs):
+        '''
+         The implementation will be very slow when mask is not None,
+         because we need rearange the x/context features according to mask.
+        Args:
+            x:
+            context:
+            mask:
+            **kwargs:
+        Returns: x
+        '''
+        dtype = kwargs.get('dtype', torch.float16)
+
+        def half(x):
+            return x if x.dtype in [torch.float16, torch.bfloat16
+                                    ] else x.to(dtype)
+
+        x_shapes = kwargs['x_shapes']
+        freqs = kwargs['freqs']
+        self_x_len = kwargs['self_x_len']
+        cross_x_len = kwargs['cross_x_len']
+        txt_lens = kwargs['txt_lens']
+        n, d = self.num_heads, self.head_dim
+
+        if context is None:
+            # self-attn
+            q = self.norm_q(self.q(x)).view(-1, n, d)
+            k = self.norm_q(self.k(x)).view(-1, n, d)
+            v = self.v(x).view(-1, n, d)
+            q = rope_apply(q, self_x_len, x_shapes, freqs, pad=False)
+            k = rope_apply(k, self_x_len, x_shapes, freqs, pad=False)
+            q_lens = k_lens = self_x_len
+        else:
+            # cross-attn
+            q = self.norm_q(self.q(x)).view(-1, n, d)
+            k = self.norm_q(self.k(context)).view(-1, n, d)
+            v = self.v(context).view(-1, n, d)
+            q_lens = cross_x_len
+            k_lens = txt_lens
+
+        cu_seqlens_q = torch.cat([q_lens.new_zeros([1]),
+                                  q_lens]).cumsum(0, dtype=torch.int32)
+        cu_seqlens_k = torch.cat([k_lens.new_zeros([1]),
+                                  k_lens]).cumsum(0, dtype=torch.int32)
+        max_seqlen_q = q_lens.max()
+        max_seqlen_k = k_lens.max()
+
+        out_dtype = q.dtype
+        q, k, v = half(q), half(k), half(v)
+        x = flash_attn_varlen_func(q,
+                                   k,
+                                   v,
+                                   cu_seqlens_q=cu_seqlens_q,
+                                   cu_seqlens_k=cu_seqlens_k,
+                                   max_seqlen_q=max_seqlen_q,
+                                   max_seqlen_k=max_seqlen_k,
+                                   dropout_p=self.attn_drop.p,
+                                   softmax_scale=self.softmax_scale,
+                                   causal=self.causal,
+                                   window_size=self.window_size,
+                                   deterministic=self.deterministic)
+
+        x = x.type(out_dtype)
+        x = x.reshape(-1, n * d)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+    def forward(self, x, context=None, **kwargs):
+        x = getattr(self, self.backend)(x, context=context, **kwargs)
+        return x

scepter/modules/model/backbone/ace_plus/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.

scepter/modules/model/backbone/ace_plus/ace_plus.py

@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from einops import rearrange
+from scepter.modules.model.backbone.flux import FluxMR
+from scepter.modules.model.registry import BACKBONES
+from scepter.modules.utils.config import dict_to_yaml
+
+
+@BACKBONES.register_class()
+class FluxMRACEPlus(FluxMR):
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger)
+
+    def prepare_input(self, x, cond):
+        context, y = cond['context'], cond['y']
+        batch_frames, batch_frames_ids = [], []
+        for ix, shape, imask, ie, ie_mask in zip(x, cond['x_shapes'],
+                                                 cond['x_mask'], cond['edit'],
+                                                 cond['edit_mask']):
+            # unpack image from sequence
+            ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
+            imask = torch.ones_like(
+                ix[[0], :, :]) if imask is None else imask.squeeze(0)
+            if len(ie) > 0:
+                ie = [iie.squeeze(0) for iie in ie]
+                ie_mask = [
+                    torch.ones(
+                        (ix.shape[0] * 4, ix.shape[1],
+                         ix.shape[2])) if iime is None else iime.squeeze(0)
+                    for iime in ie_mask
+                ]
+                ie = torch.cat(ie, dim=-1)
+                ie_mask = torch.cat(ie_mask, dim=-1)
+            else:
+                ie, ie_mask = torch.zeros_like(ix).to(x), torch.ones_like(
+                    imask).to(x)
+            ix = torch.cat([ix, ie, ie_mask], dim=0)
+            c, h, w = ix.shape
+            ix = rearrange(ix,
+                           'c (h ph) (w pw) -> (h w) (c ph pw)',
+                           ph=2,
+                           pw=2)
+            ix_id = torch.zeros(h // 2, w // 2, 3)
+            ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None]
+            ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :]
+            ix_id = rearrange(ix_id, 'h w c -> (h w) c')
+            batch_frames.append([ix])
+            batch_frames_ids.append([ix_id])
+        x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], []
+        for frames, frame_ids in zip(batch_frames, batch_frames_ids):
+            proj_frames = []
+            for idx, one_frame in enumerate(frames):
+                one_frame = self.img_in(one_frame)
+                proj_frames.append(one_frame)
+            ix = torch.cat(proj_frames, dim=0)
+            if_id = torch.cat(frame_ids, dim=0)
+            x_list.append(ix)
+            x_id_list.append(if_id)
+            mask_x_list.append(
+                torch.ones(ix.shape[0]).to(ix.device,
+                                           non_blocking=True).bool())
+            x_seq_length.append(ix.shape[0])
+        # if len(x_list) < 1: import pdb;pdb.set_trace()
+        x = pad_sequence(tuple(x_list), batch_first=True)
+        x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(
+            x)  # [b,pad_seq,2] pad (0.,0.) at dim2
+        mask_x = pad_sequence(tuple(mask_x_list), batch_first=True)
+        # import pdb;pdb.set_trace()
+        if isinstance(context, list):
+            txt_list, mask_txt_list, y_list = [], [], []
+            for sample_id, (ctx, yy) in enumerate(zip(context, y)):
+                txt_list.append(self.txt_in(ctx.to(x)))
+                mask_txt_list.append(
+                    torch.ones(txt_list[-1].shape[0]).to(
+                        ctx.device, non_blocking=True).bool())
+                y_list.append(yy.to(x))
+            txt = pad_sequence(tuple(txt_list), batch_first=True)
+            txt_ids = torch.zeros(txt.shape[0], txt.shape[1], 3).to(x)
+            mask_txt = pad_sequence(tuple(mask_txt_list), batch_first=True)
+            y = torch.cat(y_list, dim=0)
+            assert y.ndim == 2 and txt.ndim == 3
+        else:
+            txt = self.txt_in(context)
+            txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x)
+            mask_txt = torch.ones(context.shape[0], context.shape[1]).to(
+                x.device, non_blocking=True).bool()
+        return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('MODEL',
+                            __class__.__name__,
+                            FluxMRACEPlus.para_dict,
+                            set_name=True)

scepter/modules/model/backbone/cogvideox/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from scepter.modules.model.backbone.cogvideox.cogvideox import CogVideoXTransformer3DModel

scepter/modules/model/backbone/cogvideox/cogvideox.py

@@ -0,0 +1,357 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import OrderedDict
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from scepter.modules.model.base_model import BaseModel
+from scepter.modules.model.registry import BACKBONES
+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 .layers import CogVideoXBlock, CogVideoXPatchEmbed, TimestepEmbedding, Timesteps, AdaLayerNorm
+
+
+@BACKBONES.register_class()
+class CogVideoXTransformer3DModel(BaseModel):
+    """
+    A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo).
+
+    Parameters:
+        num_attention_heads (`int`, defaults to `30`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, defaults to `64`):
+            The number of channels in each head.
+        in_channels (`int`, defaults to `16`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `16`):
+            The number of channels in the output.
+        flip_sin_to_cos (`bool`, defaults to `True`):
+            Whether to flip the sin to cos in the time embedding.
+        time_embed_dim (`int`, defaults to `512`):
+            Output dimension of timestep embeddings.
+        ofs_embed_dim (`int`, defaults to `512`):
+            Output dimension of "ofs" embeddings used in CogVideoX-5b-I2B in version 1.5
+        text_embed_dim (`int`, defaults to `4096`):
+            Input dimension of text embeddings from the text encoder.
+        num_layers (`int`, defaults to `30`):
+            The number of layers of Transformer blocks to use.
+        dropout (`float`, defaults to `0.0`):
+            The dropout probability to use.
+        attention_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in the attention projection layers.
+        sample_width (`int`, defaults to `90`):
+            The width of the input latents.
+        sample_height (`int`, defaults to `60`):
+            The height of the input latents.
+        sample_frames (`int`, defaults to `49`):
+            The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49
+            instead of 13 because CogVideoX processed 13 latent frames at once in its default and recommended settings,
+            but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with
+            K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1).
+        patch_size (`int`, defaults to `2`):
+            The size of the patches to use in the patch embedding layer.
+        temporal_compression_ratio (`int`, defaults to `4`):
+            The compression ratio across the temporal dimension. See documentation for `sample_frames`.
+        max_text_seq_length (`int`, defaults to `226`):
+            The maximum sequence length of the input text embeddings.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to use in feed-forward.
+        timestep_activation_fn (`str`, defaults to `"silu"`):
+            Activation function to use when generating the timestep embeddings.
+        norm_elementwise_affine (`bool`, defaults to `True`):
+            Whether or not to use elementwise affine in normalization layers.
+        norm_eps (`float`, defaults to `1e-5`):
+            The epsilon value to use in normalization layers.
+        spatial_interpolation_scale (`float`, defaults to `1.875`):
+            Scaling factor to apply in 3D positional embeddings across spatial dimensions.
+        temporal_interpolation_scale (`float`, defaults to `1.0`):
+            Scaling factor to apply in 3D positional embeddings across temporal dimensions.
+    """
+
+    def __init__(
+        self,
+        cfg,
+        logger=None
+    ):
+        super().__init__(cfg, logger=logger)
+        num_attention_heads = cfg.get("NUM_ATTENTION_HEADS", 30)
+        attention_head_dim = cfg.get("ATTENTION_HEAD_DIM", 64)
+        in_channels = cfg.get("IN_CHANNELS", 16)
+        out_channels = cfg.get("OUT_CHANNELS", 16)
+        flip_sin_to_cos = cfg.get("FLIP_SIN_TO_COS", True)
+        freq_shift = cfg.get("FREQ_SHIFT", 0)
+        time_embed_dim = cfg.get("TIME_EMBED_DIM", 512)
+        ofs_embed_dim = cfg.get("OFS_EMBED_DIM", None)  # 1.5
+        text_embed_dim = cfg.get("TEXT_EMBED_DIM", 4096)
+        num_layers = cfg.get("NUM_LAYERS", 30)
+        dropout = cfg.get("DROPOUT", 0.0)
+        attention_bias = cfg.get("ATTENTION_BIAS", True)
+        sample_width = cfg.get("SAMPLE_WIDTH", 90)
+        sample_height = cfg.get("SAMPLE_HEIGHT", 60)
+        sample_frames = cfg.get("SAMPLE_FRAMES", 49)
+        patch_size = cfg.get("PATCH_SIZE", 2)
+        patch_size_t = cfg.get("PATCH_SIZE_T", None)
+        patch_bias = cfg.get("PATCH_BIAS", True)
+        temporal_compression_ratio = cfg.get("TEMPORAL_COMPRESSION_RATIO", 4)
+        max_text_seq_length = cfg.get("MAX_TEXT_SEQ_LENGTH", 226)
+        activation_fn = cfg.get("ACTIVATION_FN", "gelu-approximate")
+        timestep_activation_fn = cfg.get("TIMESTEP_ACTIVATION_FN", "silu")
+        norm_elementwise_affine = cfg.get("NORM_ELEMENTWISE_AFFINE", True)
+        norm_eps = cfg.get("NORM_EPS", 1e-5)
+        spatial_interpolation_scale = cfg.get("SPATIAL_INTERPOLATION_SCALE", 1.875)
+        temporal_interpolation_scale = cfg.get("TEMPORAL_INTERPOLATION_SCALE", 1.0)
+        use_rotary_positional_embeddings = cfg.get("USE_ROTARY_POSITIONAL_EMBEDDINGS", False)
+        use_learned_positional_embeddings = cfg.get("USE_LEARNED_POSITIONAL_EMBEDDINGS", False)
+        self.gradient_checkpointing = cfg.get("GRADIENT_CHECKPOINTING", False)
+        inner_dim = num_attention_heads * attention_head_dim
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.use_rotary_positional_embeddings = use_rotary_positional_embeddings
+
+        if not use_rotary_positional_embeddings and use_learned_positional_embeddings:
+            raise ValueError(
+                "There are no CogVideoX checkpoints available with disable rotary embeddings and learned positional "
+                "embeddings. If you're using a custom model and/or believe this should be supported, please open an "
+                "issue at https://github.com/huggingface/diffusers/issues."
+            )
+
+        # 1. Patch embedding
+        self.patch_embed = CogVideoXPatchEmbed(
+            patch_size=patch_size,
+            patch_size_t=patch_size_t,
+            in_channels=in_channels,
+            embed_dim=inner_dim,
+            text_embed_dim=text_embed_dim,
+            bias=patch_bias,
+            sample_width=sample_width,
+            sample_height=sample_height,
+            sample_frames=sample_frames,
+            temporal_compression_ratio=temporal_compression_ratio,
+            max_text_seq_length=max_text_seq_length,
+            spatial_interpolation_scale=spatial_interpolation_scale,
+            temporal_interpolation_scale=temporal_interpolation_scale,
+            use_positional_embeddings=not use_rotary_positional_embeddings,
+            use_learned_positional_embeddings=use_learned_positional_embeddings,
+        )
+        self.embedding_dropout = nn.Dropout(dropout)
+
+        # 2. Time embeddings and ofs embedding(Only CogVideoX1.5-5B I2V have)
+        self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift)
+        self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn)
+
+        self.ofs_proj = None
+        self.ofs_embedding = None
+        if ofs_embed_dim:
+            self.ofs_proj = Timesteps(ofs_embed_dim, flip_sin_to_cos, freq_shift)
+            self.ofs_embedding = TimestepEmbedding(
+                ofs_embed_dim, ofs_embed_dim, timestep_activation_fn
+            )  # same as time embeddings, for ofs
+
+        # 3. Define spatio-temporal transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                CogVideoXBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    time_embed_dim=time_embed_dim,
+                    dropout=dropout,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine)
+
+        # 4. Output blocks
+        self.norm_out = AdaLayerNorm(
+            embedding_dim=time_embed_dim,
+            output_dim=2 * inner_dim,
+            norm_elementwise_affine=norm_elementwise_affine,
+            norm_eps=norm_eps,
+            chunk_dim=1,
+        )
+
+        if patch_size_t is None:
+            # For CogVideox 1.0
+            output_dim = patch_size * patch_size * out_channels
+        else:
+            # For CogVideoX 1.5
+            output_dim = patch_size * patch_size * patch_size_t * out_channels
+
+        self.proj_out = nn.Linear(inner_dim, output_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor = None,
+        t: Union[int, float, torch.LongTensor] = None,
+        cond: torch.Tensor = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        ofs: Optional[Union[int, float, torch.LongTensor]] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs
+    ):
+        if 'image_latent' in kwargs and kwargs['image_latent'] is not None:
+            hidden_states = torch.cat([x, kwargs['image_latent']], dim=2)
+        else:
+            hidden_states = x
+        timestep = t
+        encoder_hidden_states = cond
+
+        batch_size, num_frames, channels, height, width = hidden_states.shape
+
+        # 1. Time embedding
+        timesteps = timestep
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=encoder_hidden_states.dtype)
+        emb = self.time_embedding(t_emb, timestep_cond)
+
+        if self.ofs_embedding is not None:
+            ofs_emb = self.ofs_proj(ofs)
+            ofs_emb = ofs_emb.to(dtype=hidden_states.dtype)
+            ofs_emb = self.ofs_embedding(ofs_emb)
+            emb = emb + ofs_emb
+
+        # 2. Patch embedding
+        hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
+        hidden_states = self.embedding_dropout(hidden_states)
+
+        text_seq_length = encoder_hidden_states.shape[1]
+        encoder_hidden_states = hidden_states[:, :text_seq_length]
+        hidden_states = hidden_states[:, text_seq_length:]
+
+        # 3. Transformer blocks
+        for i, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False}
+                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    emb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+            else:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=emb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+        if not self.use_rotary_positional_embeddings:
+            # CogVideoX-2B
+            hidden_states = self.norm_final(hidden_states)
+        else:
+            # CogVideoX-5B
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+            hidden_states = self.norm_final(hidden_states)
+            hidden_states = hidden_states[:, text_seq_length:]
+
+        # 4. Final block
+        hidden_states = self.norm_out(hidden_states, temb=emb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 5. Unpatchify
+        # Note: we use `-1` instead of `channels`:
+        #   - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels)
+        #   - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels)
+        p = self.patch_size
+        p_t = self.patch_size_t
+
+        if p_t is None:
+            output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
+            output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
+        else:
+            output = hidden_states.reshape(
+                batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p
+            )
+            output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
+
+        return output
+
+    def load_pretrained_model(self, pretrained_model):
+        if pretrained_model is not None:
+            pretrained_model_list = [pretrained_model] if isinstance(pretrained_model, str) else pretrained_model
+            ckpt_all = OrderedDict()
+            for pretrained_model in pretrained_model_list:
+                with FS.get_from(pretrained_model,
+                                 wait_finish=True) as local_model:
+                    if local_model.endswith('safetensors'):
+                        from safetensors.torch import load_file as load_safetensors
+                        ckpt = load_safetensors(local_model)
+                    else:
+                        ckpt = torch.load(local_model, map_location='cpu', weights_only=True)
+                    ckpt_all.update(ckpt)
+            missing, unexpected = self.load_state_dict(ckpt_all, strict=False)
+            if we.rank == 0:
+                self.logger.info(
+                    f'Restored from {pretrained_model_list} with {len(missing)} missing and {len(unexpected)} unexpected keys'
+                )
+                if len(missing) > 0:
+                    self.logger.info(f'Missing Keys:\n {missing}')
+                if len(unexpected) > 0:
+                    self.logger.info(f'\nUnexpected Keys:\n {unexpected}')
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('MODEL',
+                            __class__.__name__,
+                            CogVideoXTransformer3DModel.para_dict,
+                            set_name=True)
+
+
+if __name__ == "__main__":
+    import argparse
+    from scepter.modules.utils.file_system import FS
+    from scepter.modules.utils.config import Config
+    from scepter.modules.utils.logger import get_logger
+
+    parser = argparse.ArgumentParser()
+    cfg = Config(parser_ins=parser)
+    for file_sys in cfg.FILE_SYSTEM:
+        FS.init_fs_client(file_sys)
+    model = BACKBONES.build(cfg.DIFFUSION_MODEL, logger=get_logger()).eval().requires_grad_(False).to('cuda').to(torch.bfloat16)
+
+    hidden_states = torch.load(FS.get_from(cfg.HIDDEN_STATES), weights_only=True)
+    encoder_hidden_states = torch.load(FS.get_from(cfg.ENCODER_HIDDEN_STATES), weights_only=True)
+    timestep = torch.load(FS.get_from(cfg.TIMESTEP), weights_only=True)
+    timestep_cond = None
+    image_rotary_emb = None
+    attention_kwargs = None
+    output = model(hidden_states, encoder_hidden_states, timestep, timestep_cond, image_rotary_emb, attention_kwargs)
+    print(output, torch.sum(output))

scepter/modules/model/backbone/cogvideox/layers.py

@@ -0,0 +1,574 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .utils import get_activation, get_timestep_embedding, get_3d_sincos_pos_embed, apply_rotary_emb
+from .utils import GELU, GEGLU, ApproximateGELU, SwiGLU
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+        post_act_fn: Optional[str] = None,
+        cond_proj_dim=None,
+        sample_proj_bias=True,
+    ):
+        super().__init__()
+
+        self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
+
+        if cond_proj_dim is not None:
+            self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
+        else:
+            self.cond_proj = None
+
+        self.act = get_activation(act_fn)
+
+        if out_dim is not None:
+            time_embed_dim_out = out_dim
+        else:
+            time_embed_dim_out = time_embed_dim
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
+
+        if post_act_fn is None:
+            self.post_act = None
+        else:
+            self.post_act = get_activation(post_act_fn)
+
+    def forward(self, sample, condition=None):
+        if condition is not None:
+            sample = sample + self.cond_proj(condition)
+        sample = self.linear_1(sample)
+
+        if self.act is not None:
+            sample = self.act(sample)
+
+        sample = self.linear_2(sample)
+
+        if self.post_act is not None:
+            sample = self.post_act(sample)
+        return sample
+
+
+
+class Timesteps(nn.Module):
+    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+
+    def forward(self, timesteps):
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+        )
+        return t_emb
+
+
+class CogVideoXLayerNormZero(nn.Module):
+    def __init__(
+        self,
+        conditioning_dim: int,
+        embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(conditioning_dim, 6 * embedding_dim, bias=bias)
+        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
+
+    def forward(
+        self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1)
+        hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+        encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale)[:, None, :] + enc_shift[:, None, :]
+        return hidden_states, encoder_hidden_states, gate[:, None, :], enc_gate[:, None, :]
+
+
+class AdaLayerNorm(nn.Module):
+    r"""
+    Norm layer modified to incorporate timestep embeddings.
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+        num_embeddings (`int`, *optional*): The size of the embeddings dictionary.
+        output_dim (`int`, *optional*):
+        norm_elementwise_affine (`bool`, defaults to `False):
+        norm_eps (`bool`, defaults to `False`):
+        chunk_dim (`int`, defaults to `0`):
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_embeddings: Optional[int] = None,
+        output_dim: Optional[int] = None,
+        norm_elementwise_affine: bool = False,
+        norm_eps: float = 1e-5,
+        chunk_dim: int = 0,
+    ):
+        super().__init__()
+
+        self.chunk_dim = chunk_dim
+        output_dim = output_dim or embedding_dim * 2
+
+        if num_embeddings is not None:
+            self.emb = nn.Embedding(num_embeddings, embedding_dim)
+        else:
+            self.emb = None
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, output_dim)
+        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
+
+    def forward(
+        self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        if self.emb is not None:
+            temb = self.emb(timestep)
+
+        temb = self.linear(self.silu(temb))
+
+        if self.chunk_dim == 1:
+            # This is a bit weird why we have the order of "shift, scale" here and "scale, shift" in the
+            # other if-branch. This branch is specific to CogVideoX for now.
+            shift, scale = temb.chunk(2, dim=1)
+            shift = shift[:, None, :]
+            scale = scale[:, None, :]
+        else:
+            scale, shift = temb.chunk(2, dim=0)
+
+        x = self.norm(x) * (1 + scale) + shift
+        return x
+
+
+class CogVideoXPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 2,
+        patch_size_t: Optional[int] = None,
+        in_channels: int = 16,
+        embed_dim: int = 1920,
+        text_embed_dim: int = 4096,
+        bias: bool = True,
+        sample_width: int = 90,
+        sample_height: int = 60,
+        sample_frames: int = 49,
+        temporal_compression_ratio: int = 4,
+        max_text_seq_length: int = 226,
+        spatial_interpolation_scale: float = 1.875,
+        temporal_interpolation_scale: float = 1.0,
+        use_positional_embeddings: bool = True,
+        use_learned_positional_embeddings: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.embed_dim = embed_dim
+        self.sample_height = sample_height
+        self.sample_width = sample_width
+        self.sample_frames = sample_frames
+        self.temporal_compression_ratio = temporal_compression_ratio
+        self.max_text_seq_length = max_text_seq_length
+        self.spatial_interpolation_scale = spatial_interpolation_scale
+        self.temporal_interpolation_scale = temporal_interpolation_scale
+        self.use_positional_embeddings = use_positional_embeddings
+        self.use_learned_positional_embeddings = use_learned_positional_embeddings
+
+        if patch_size_t is None:
+            # CogVideoX 1.0 checkpoints
+            self.proj = nn.Conv2d(
+                in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias
+            )
+        else:
+            # CogVideoX 1.5 checkpoints
+            self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim)
+
+        self.text_proj = nn.Linear(text_embed_dim, embed_dim)
+
+        if use_positional_embeddings or use_learned_positional_embeddings:
+            persistent = use_learned_positional_embeddings
+            pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames)
+            self.register_buffer("pos_embedding", pos_embedding, persistent=persistent)
+
+    def _get_positional_embeddings(self, sample_height: int, sample_width: int, sample_frames: int) -> torch.Tensor:
+        post_patch_height = sample_height // self.patch_size
+        post_patch_width = sample_width // self.patch_size
+        post_time_compression_frames = (sample_frames - 1) // self.temporal_compression_ratio + 1
+        num_patches = post_patch_height * post_patch_width * post_time_compression_frames
+
+        pos_embedding = get_3d_sincos_pos_embed(
+            self.embed_dim,
+            (post_patch_width, post_patch_height),
+            post_time_compression_frames,
+            self.spatial_interpolation_scale,
+            self.temporal_interpolation_scale,
+        )
+        pos_embedding = torch.from_numpy(pos_embedding).flatten(0, 1)
+        joint_pos_embedding = torch.zeros(
+            1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False
+        )
+        joint_pos_embedding.data[:, self.max_text_seq_length :].copy_(pos_embedding)
+
+        return joint_pos_embedding
+
+    def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):
+        r"""
+        Args:
+            text_embeds (`torch.Tensor`):
+                Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).
+            image_embeds (`torch.Tensor`):
+                Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).
+        """
+        text_embeds = self.text_proj(text_embeds)
+
+        batch_size, num_frames, channels, height, width = image_embeds.shape
+
+        if self.patch_size_t is None:
+            image_embeds = image_embeds.reshape(-1, channels, height, width)
+            image_embeds = self.proj(image_embeds)
+            image_embeds = image_embeds.view(batch_size, num_frames, *image_embeds.shape[1:])
+            image_embeds = image_embeds.flatten(3).transpose(2, 3)  # [batch, num_frames, height x width, channels]
+            image_embeds = image_embeds.flatten(1, 2)  # [batch, num_frames x height x width, channels]
+        else:
+            p = self.patch_size
+            p_t = self.patch_size_t
+
+            image_embeds = image_embeds.permute(0, 1, 3, 4, 2)
+            image_embeds = image_embeds.reshape(
+                batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels
+            )
+            image_embeds = image_embeds.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(4, 7).flatten(1, 3)
+            image_embeds = self.proj(image_embeds)
+
+        embeds = torch.cat(
+            [text_embeds, image_embeds], dim=1
+        ).contiguous()  # [batch, seq_length + num_frames x height x width, channels]
+
+        if self.use_positional_embeddings or self.use_learned_positional_embeddings:
+            if self.use_learned_positional_embeddings and (self.sample_width != width or self.sample_height != height):
+                raise ValueError(
+                    "It is currently not possible to generate videos at a different resolution that the defaults. This should only be the case with 'THUDM/CogVideoX-5b-I2V'."
+                    "If you think this is incorrect, please open an issue at https://github.com/huggingface/diffusers/issues."
+                )
+
+            pre_time_compression_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+
+            if (
+                self.sample_height != height
+                or self.sample_width != width
+                or self.sample_frames != pre_time_compression_frames
+            ):
+                pos_embedding = self._get_positional_embeddings(height, width, pre_time_compression_frames)
+                pos_embedding = pos_embedding.to(embeds.device, dtype=embeds.dtype)
+            else:
+                pos_embedding = self.pos_embedding
+
+            embeds = embeds + pos_embedding
+
+        return embeds
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+        inner_dim=None,
+        bias: bool = True,
+    ):
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim, bias=bias)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim, bias=bias)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim, bias=bias)
+        elif activation_fn == "swiglu":
+            act_fn = SwiGLU(dim, inner_dim, bias=bias)
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # project dropout
+        self.net.append(nn.Dropout(dropout))
+        # project out
+        self.net.append(nn.Linear(inner_dim, dim_out, bias=bias))
+        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
+        if final_dropout:
+            self.net.append(nn.Dropout(dropout))
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            print(deprecation_message)
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        dim_head: int = 64,
+        heads: int = 8,
+        kv_heads: Optional[int] = None,
+        qk_norm: Optional[str] = None,
+        eps: float = 1e-5,
+        bias: bool = False,
+        out_bias: bool = True,
+        dropout: float = 0.0,
+        out_dim: int = None,
+        cross_attention_dim: Optional[int] = None,
+    ):
+        super().__init__()
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        self.is_cross_attention = cross_attention_dim is not None
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "layer_norm":
+            self.norm_q = nn.LayerNorm(dim_head, eps=eps)
+            self.norm_k = nn.LayerNorm(dim_head, eps=eps)
+
+        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(nn.Dropout(dropout))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // self.heads
+
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+        if self.norm_q is not None:
+            query = self.norm_q(query)
+        if self.norm_k is not None:
+            key = self.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+            if not self.is_cross_attention:
+                key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
+
+        # linear proj
+        hidden_states = self.to_out[0](hidden_states)
+        # dropout
+        hidden_states = self.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class CogVideoXBlock(nn.Module):
+    r"""
+    Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model.
+
+    Parameters:
+        dim (`int`):
+            The number of channels in the input and output.
+        num_attention_heads (`int`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`):
+            The number of channels in each head.
+        time_embed_dim (`int`):
+            The number of channels in timestep embedding.
+        dropout (`float`, defaults to `0.0`):
+            The dropout probability to use.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to be used in feed-forward.
+        attention_bias (`bool`, defaults to `False`):
+            Whether or not to use bias in attention projection layers.
+        qk_norm (`bool`, defaults to `True`):
+            Whether or not to use normalization after query and key projections in Attention.
+        norm_elementwise_affine (`bool`, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_eps (`float`, defaults to `1e-5`):
+            Epsilon value for normalization layers.
+        final_dropout (`bool` defaults to `False`):
+            Whether to apply a final dropout after the last feed-forward layer.
+        ff_inner_dim (`int`, *optional*, defaults to `None`):
+            Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used.
+        ff_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in Feed-forward layer.
+        attention_out_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in Attention output projection layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        time_embed_dim: int,
+        dropout: float = 0.0,
+        activation_fn: str = "gelu-approximate",
+        attention_bias: bool = False,
+        qk_norm: bool = True,
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        final_dropout: bool = True,
+        ff_inner_dim: Optional[int] = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+    ):
+        super().__init__()
+
+        # 1. Self Attention
+        self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            qk_norm="layer_norm" if qk_norm else None,
+            eps=1e-6,
+            bias=attention_bias,
+            out_bias=attention_out_bias
+        )
+
+        # 2. Feed Forward
+        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
+
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        # norm & modulate
+        norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
+            hidden_states, encoder_hidden_states, temb
+        )
+
+        # attention
+        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        hidden_states = hidden_states + gate_msa * attn_hidden_states
+        encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states
+
+        # norm & modulate
+        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
+            hidden_states, encoder_hidden_states, temb
+        )
+
+        # feed-forward
+        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
+        ff_output = self.ff(norm_hidden_states)
+
+        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
+        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]
+
+        return hidden_states, encoder_hidden_states

scepter/modules/model/backbone/cogvideox/utils.py

@@ -0,0 +1,570 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Optional, Tuple, Union, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+ACTIVATION_FUNCTIONS = {
+    "swish": nn.SiLU(),
+    "silu": nn.SiLU(),
+    "mish": nn.Mish(),
+    "gelu": nn.GELU(),
+    "relu": nn.ReLU(),
+}
+
+
+def get_activation(act_fn: str) -> nn.Module:
+    """Helper function to get activation function from string.
+
+    Args:
+        act_fn (str): Name of activation function.
+
+    Returns:
+        nn.Module: Activation function.
+    """
+
+    act_fn = act_fn.lower()
+    if act_fn in ACTIVATION_FUNCTIONS:
+        return ACTIVATION_FUNCTIONS[act_fn]
+    else:
+        raise ValueError(f"Unsupported activation function: {act_fn}")
+
+
+class FP32SiLU(nn.Module):
+    r"""
+    SiLU activation function with input upcasted to torch.float32.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return F.silu(inputs.float(), inplace=False).to(inputs.dtype)
+
+
+class GELU(nn.Module):
+    r"""
+    GELU activation function with tanh approximation support with `approximate="tanh"`.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out, bias=bias)
+        self.approximate = approximate
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate, approximate=self.approximate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype)
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_states = self.gelu(hidden_states)
+        return hidden_states
+
+
+class GEGLU(nn.Module):
+    r"""
+    A [variant](https://arxiv.org/abs/2002.05202) of the gated linear unit activation function.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
+
+    def forward(self, hidden_states, *args, **kwargs):
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            print("scale", "1.0.0", deprecation_message)
+        hidden_states = self.proj(hidden_states)
+        hidden_states, gate = hidden_states.chunk(2, dim=-1)
+        return hidden_states * self.gelu(gate)
+
+
+class SwiGLU(nn.Module):
+    r"""
+    A [variant](https://arxiv.org/abs/2002.05202) of the gated linear unit activation function. It's similar to `GEGLU`
+    but uses SiLU / Swish instead of GeLU.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)
+        self.activation = nn.SiLU()
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_states, gate = hidden_states.chunk(2, dim=-1)
+        return hidden_states * self.activation(gate)
+
+
+class ApproximateGELU(nn.Module):
+    r"""
+    The approximate form of the Gaussian Error Linear Unit (GELU). For more details, see section 2 of this
+    [paper](https://arxiv.org/abs/1606.08415).
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return x * torch.sigmoid(1.702 * x)
+
+
+def randn_tensor(
+    shape: Union[Tuple, List],
+    generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None,
+    device: Optional["torch.device"] = None,
+    dtype: Optional["torch.dtype"] = None,
+    layout: Optional["torch.layout"] = None,
+):
+    """A helper function to create random tensors on the desired `device` with the desired `dtype`. When
+    passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor
+    is always created on the CPU.
+    """
+    # device on which tensor is created defaults to device
+    rand_device = device
+    batch_size = shape[0]
+
+    layout = layout or torch.strided
+    device = device or torch.device("cpu")
+
+    if generator is not None:
+        gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type
+        if gen_device_type != device.type and gen_device_type == "cpu":
+            rand_device = "cpu"
+            if device != "mps":
+                print(
+                    f"The passed generator was created on 'cpu' even though a tensor on {device} was expected."
+                    f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"
+                    f" slighly speed up this function by passing a generator that was created on the {device} device."
+                )
+        elif gen_device_type != device.type and gen_device_type == "cuda":
+            raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.")
+
+    # make sure generator list of length 1 is treated like a non-list
+    if isinstance(generator, list) and len(generator) == 1:
+        generator = generator[0]
+
+    if isinstance(generator, list):
+        shape = (1,) + shape[1:]
+        latents = [
+            torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout)
+            for i in range(batch_size)
+        ]
+        latents = torch.cat(latents, dim=0).to(device)
+    else:
+        latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device)
+
+    return latents
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+    Args
+        timesteps (torch.Tensor):
+            a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        embedding_dim (int):
+            the dimension of the output.
+        flip_sin_to_cos (bool):
+            Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+        downscale_freq_shift (float):
+            Controls the delta between frequencies between dimensions
+        scale (float):
+            Scaling factor applied to the embeddings.
+        max_period (int):
+            Controls the maximum frequency of the embeddings
+    Returns
+        torch.Tensor: an [N x dim] Tensor of positional embeddings.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+def get_3d_sincos_pos_embed(
+    embed_dim: int,
+    spatial_size: Union[int, Tuple[int, int]],
+    temporal_size: int,
+    spatial_interpolation_scale: float = 1.0,
+    temporal_interpolation_scale: float = 1.0,
+) -> np.ndarray:
+    r"""
+    Args:
+        embed_dim (`int`):
+        spatial_size (`int` or `Tuple[int, int]`):
+        temporal_size (`int`):
+        spatial_interpolation_scale (`float`, defaults to 1.0):
+        temporal_interpolation_scale (`float`, defaults to 1.0):
+    """
+    if embed_dim % 4 != 0:
+        raise ValueError("`embed_dim` must be divisible by 4")
+    if isinstance(spatial_size, int):
+        spatial_size = (spatial_size, spatial_size)
+
+    embed_dim_spatial = 3 * embed_dim // 4
+    embed_dim_temporal = embed_dim // 4
+
+    # 1. Spatial
+    grid_h = np.arange(spatial_size[1], dtype=np.float32) / spatial_interpolation_scale
+    grid_w = np.arange(spatial_size[0], dtype=np.float32) / spatial_interpolation_scale
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, spatial_size[1], spatial_size[0]])
+    pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid)
+
+    # 2. Temporal
+    grid_t = np.arange(temporal_size, dtype=np.float32) / temporal_interpolation_scale
+    pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t)
+
+    # 3. Concat
+    pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]
+    pos_embed_spatial = np.repeat(pos_embed_spatial, temporal_size, axis=0)  # [T, H*W, D // 4 * 3]
+
+    pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]
+    pos_embed_temporal = np.repeat(pos_embed_temporal, spatial_size[0] * spatial_size[1], axis=1)  # [T, H*W, D // 4]
+
+    pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1)  # [T, H*W, D]
+    return pos_embed
+
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    use_real: bool = True,
+    use_real_unbind_dim: int = -1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+    tensors contain rotary embeddings and are returned as real tensors.
+
+    Args:
+        x (`torch.Tensor`):
+            Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+    else:
+        # used for lumina
+        x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.unsqueeze(2)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+        return x_out.type_as(x)
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[np.ndarray, int],
+    theta: float = 10000.0,
+    use_real=False,
+    linear_factor=1.0,
+    ntk_factor=1.0,
+    repeat_interleave_real=True,
+    freqs_dtype=torch.float32,  #  torch.float32, torch.float64 (flux)
+):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+    index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+    data type.
+
+    Args:
+        dim (`int`): Dimension of the frequency tensor.
+        pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+        theta (`float`, *optional*, defaults to 10000.0):
+            Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (`bool`, *optional*):
+            If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+        linear_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor for the context extrapolation. Defaults to 1.0.
+        ntk_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.
+        repeat_interleave_real (`bool`, *optional*, defaults to `True`):
+            If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.
+            Otherwise, they are concateanted with themselves.
+        freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):
+            the dtype of the frequency tensor.
+    Returns:
+        `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+    """
+    assert dim % 2 == 0
+
+    if isinstance(pos, int):
+        pos = torch.arange(pos)
+    if isinstance(pos, np.ndarray):
+        pos = torch.from_numpy(pos)  # type: ignore  # [S]
+
+    theta = theta * ntk_factor
+    freqs = (
+        1.0
+        / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[: (dim // 2)] / dim))
+        / linear_factor
+    )  # [D/2]
+    freqs = torch.outer(pos, freqs)  # type: ignore   # [S, D/2]
+    if use_real and repeat_interleave_real:
+        # flux, hunyuan-dit, cogvideox
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+    elif use_real:
+        # stable audio
+        freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float()  # [S, D]
+        freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        # lumina
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64     # [S, D/2]
+        return freqs_cis
+
+
+def get_3d_rotary_pos_embed(
+    embed_dim,
+    crops_coords,
+    grid_size,
+    temporal_size,
+    theta: int = 10000,
+    use_real: bool = True,
+    grid_type: str = "linspace",
+    max_size: Optional[Tuple[int, int]] = None,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    RoPE for video tokens with 3D structure.
+
+    Args:
+    embed_dim: (`int`):
+        The embedding dimension size, corresponding to hidden_size_head.
+    crops_coords (`Tuple[int]`):
+        The top-left and bottom-right coordinates of the crop.
+    grid_size (`Tuple[int]`):
+        The grid size of the spatial positional embedding (height, width).
+    temporal_size (`int`):
+        The size of the temporal dimension.
+    theta (`float`):
+        Scaling factor for frequency computation.
+    grid_type (`str`):
+        Whether to use "linspace" or "slice" to compute grids.
+
+    Returns:
+        `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
+    """
+    if use_real is not True:
+        raise ValueError(" `use_real = False` is not currently supported for get_3d_rotary_pos_embed")
+
+    if grid_type == "linspace":
+        start, stop = crops_coords
+        grid_size_h, grid_size_w = grid_size
+        grid_h = np.linspace(start[0], stop[0], grid_size_h, endpoint=False, dtype=np.float32)
+        grid_w = np.linspace(start[1], stop[1], grid_size_w, endpoint=False, dtype=np.float32)
+        grid_t = np.arange(temporal_size, dtype=np.float32)
+        grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)
+    elif grid_type == "slice":
+        max_h, max_w = max_size
+        grid_size_h, grid_size_w = grid_size
+        grid_h = np.arange(max_h, dtype=np.float32)
+        grid_w = np.arange(max_w, dtype=np.float32)
+        grid_t = np.arange(temporal_size, dtype=np.float32)
+    else:
+        raise ValueError("Invalid value passed for `grid_type`.")
+
+    # Compute dimensions for each axis
+    dim_t = embed_dim // 4
+    dim_h = embed_dim // 8 * 3
+    dim_w = embed_dim // 8 * 3
+
+    # Temporal frequencies
+    freqs_t = get_1d_rotary_pos_embed(dim_t, grid_t, use_real=True)
+    # Spatial frequencies for height and width
+    freqs_h = get_1d_rotary_pos_embed(dim_h, grid_h, use_real=True)
+    freqs_w = get_1d_rotary_pos_embed(dim_w, grid_w, use_real=True)
+
+    # BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor
+    def combine_time_height_width(freqs_t, freqs_h, freqs_w):
+        freqs_t = freqs_t[:, None, None, :].expand(
+            -1, grid_size_h, grid_size_w, -1
+        )  # temporal_size, grid_size_h, grid_size_w, dim_t
+        freqs_h = freqs_h[None, :, None, :].expand(
+            temporal_size, -1, grid_size_w, -1
+        )  # temporal_size, grid_size_h, grid_size_2, dim_h
+        freqs_w = freqs_w[None, None, :, :].expand(
+            temporal_size, grid_size_h, -1, -1
+        )  # temporal_size, grid_size_h, grid_size_2, dim_w
+
+        freqs = torch.cat(
+            [freqs_t, freqs_h, freqs_w], dim=-1
+        )  # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w)
+        freqs = freqs.view(
+            temporal_size * grid_size_h * grid_size_w, -1
+        )  # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w)
+        return freqs
+
+    t_cos, t_sin = freqs_t  # both t_cos and t_sin has shape: temporal_size, dim_t
+    h_cos, h_sin = freqs_h  # both h_cos and h_sin has shape: grid_size_h, dim_h
+    w_cos, w_sin = freqs_w  # both w_cos and w_sin has shape: grid_size_w, dim_w
+
+    if grid_type == "slice":
+        t_cos, t_sin = t_cos[:temporal_size], t_sin[:temporal_size]
+        h_cos, h_sin = h_cos[:grid_size_h], h_sin[:grid_size_h]
+        w_cos, w_sin = w_cos[:grid_size_w], w_sin[:grid_size_w]
+
+    cos = combine_time_height_width(t_cos, h_cos, w_cos)
+    sin = combine_time_height_width(t_sin, h_sin, w_sin)
+    return cos, sin
+
+
+def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
+    tw = tgt_width
+    th = tgt_height
+    h, w = src
+    r = h / w
+    if r > (th / tw):
+        resize_height = th
+        resize_width = int(round(th / h * w))
+    else:
+        resize_width = tw
+        resize_height = int(round(tw / w * h))
+
+    crop_top = int(round((th - resize_height) / 2.0))
+    crop_left = int(round((tw - resize_width) / 2.0))
+
+    return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)

scepter/modules/model/backbone/flux/__init__.py

@@ -1 +1,3 @@
-from .flux import Flux
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .flux import Flux, FluxMR, FluxMRFill, FluxMRRedux, FluxMRControl

scepter/modules/model/backbone/flux/flux.py

@@ -1,4 +1,9 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# This file contains code that is adapted from
+# https://github.com/black-forest-labs/flux.git
 import math
+from collections import OrderedDict
 from functools import partial
 
 import torch
@@ -10,84 +15,81 @@ from scepter.modules.utils.distribute import we
 from scepter.modules.utils.file_system import FS
 from torch import Tensor, nn
 from torch.utils.checkpoint import checkpoint_sequential
-
-from .layers import (DoubleStreamBlock, EmbedND, LastLayer,
-                                 MLPEmbedder, SingleStreamBlock,
-                                 timestep_embedding)
-
+from torch.nn.utils.rnn import pad_sequence
+from .layers import (DoubleStreamBlock, EmbedND, LastLayer, MLPEmbedder,
+                     SingleStreamBlock, timestep_embedding)
 @BACKBONES.register_class()
 class Flux(BaseModel):
     """
     Transformer backbone Diffusion model with RoPE.
     """
     para_dict = {
-        "IN_CHANNELS": {
-            "value": 64,
-            "description": "model's input channels."
+        'IN_CHANNELS': {
+            'value': 64,
+            'description': "model's input channels."
         },
-        "OUT_CHANNELS": {
-            "value": 64,
-            "description": "model's output channels."
+        'OUT_CHANNELS': {
+            'value': 64,
+            'description': "model's output channels."
         },
-        "HIDDEN_SIZE": {
-            "value": 1024,
-            "description": "model's hidden size."
+        'HIDDEN_SIZE': {
+            'value': 1024,
+            'description': "model's hidden size."
         },
-        "NUM_HEADS": {
-            "value": 16,
-            "description": "number of heads in the transformer."
+        'NUM_HEADS': {
+            'value': 16,
+            'description': 'number of heads in the transformer.'
         },
-        "AXES_DIM": {
-            "value": [16, 56, 56],
-            "description": "dimensions of the axes of the positional encoding."
+        'AXES_DIM': {
+            'value': [16, 56, 56],
+            'description': 'dimensions of the axes of the positional encoding.'
         },
-        "THETA": {
-            "value": 10_000,
-            "description": "theta for positional encoding."
+        'THETA': {
+            'value': 10_000,
+            'description': 'theta for positional encoding.'
         },
-        "VEC_IN_DIM": {
-            "value": 768,
-            "description": "dimension of the vector input."
+        'VEC_IN_DIM': {
+            'value': 768,
+            'description': 'dimension of the vector input.'
         },
-        "GUIDANCE_EMBED": {
-            "value": False,
-            "description": "whether to use guidance embedding."
+        'GUIDANCE_EMBED': {
+            'value': False,
+            'description': 'whether to use guidance embedding.'
         },
-        "CONTEXT_IN_DIM": {
-            "value": 4096,
-            "description": "dimension of the context input."
+        'CONTEXT_IN_DIM': {
+            'value': 4096,
+            'description': 'dimension of the context input.'
         },
-        "MLP_RATIO": {
-            "value": 4.0,
-            "description": "ratio of mlp hidden size to hidden size."
+        'MLP_RATIO': {
+            'value': 4.0,
+            'description': 'ratio of mlp hidden size to hidden size.'
         },
-        "QKV_BIAS": {
-            "value": True,
-            "description": "whether to use bias in qkv projection."
+        'QKV_BIAS': {
+            'value': True,
+            'description': 'whether to use bias in qkv projection.'
         },
-        "DEPTH": {
-            "value": 19,
-            "description": "number of transformer blocks."
+        'DEPTH': {
+            'value': 19,
+            'description': 'number of transformer blocks.'
         },
-        "DEPTH_SINGLE_BLOCKS": {
-            "value": 38,
-            "description": "number of transformer blocks in the single stream block."
+        'DEPTH_SINGLE_BLOCKS': {
+            'value':
+            38,
+            'description':
+            'number of transformer blocks in the single stream block.'
         },
-        "USE_GRAD_CHECKPOINT": {
-            "value": False,
-            "description": "whether to use gradient checkpointing."
+        'USE_GRAD_CHECKPOINT': {
+            'value': False,
+            'description': 'whether to use gradient checkpointing.'
         }
     }
-    def __init__(
-            self,
-            cfg,
-            logger = None
-    ):
+
+    def __init__(self, cfg, logger=None):
         super().__init__(cfg, logger=logger)
         self.in_channels = cfg.IN_CHANNELS
-        self.out_channels = cfg.get("OUT_CHANNELS", self.in_channels)
-        hidden_size = cfg.get("HIDDEN_SIZE", 1024)
-        num_heads = cfg.get("NUM_HEADS", 16)
+        self.out_channels = cfg.get('OUT_CHANNELS', self.in_channels)
+        hidden_size = cfg.get('HIDDEN_SIZE', 1024)
+        num_heads = cfg.get('NUM_HEADS', 16)
         axes_dim = cfg.AXES_DIM
         theta = cfg.THETA
         vec_in_dim = cfg.VEC_IN_DIM
@@ -98,6 +100,13 @@ class Flux(BaseModel):
         depth = cfg.DEPTH
         depth_single_blocks = cfg.DEPTH_SINGLE_BLOCKS
         self.use_grad_checkpoint = cfg.get("USE_GRAD_CHECKPOINT", False)
+        self.attn_backend = cfg.get("ATTN_BACKEND", "pytorch")
+        self.cache_pretrain_model = cfg.get("CACHE_PRETRAIN_MODEL", False)
+        self.lora_model = cfg.get("DIFFUSERS_LORA_MODEL", None)
+        self.comfyui_lora_model = cfg.get("COMFYUI_LORA_MODEL", None)
+        self.swift_lora_model = cfg.get("SWIFT_LORA_MODEL", None)
+        self.blackforest_lora_model = cfg.get("BLACKFOREST_LORA_MODEL", None)
+        self.pretrain_adapter = cfg.get("PRETRAIN_ADAPTER", None)
 
         if hidden_size % num_heads != 0:
             raise ValueError(
@@ -105,16 +114,17 @@ class Flux(BaseModel):
             )
         pe_dim = hidden_size // num_heads
         if sum(axes_dim) != pe_dim:
-            raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
+            raise ValueError(
+                f"Got {axes_dim} but expected positional dim {pe_dim}")
         self.hidden_size = hidden_size
         self.num_heads = num_heads
-        self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim= axes_dim)
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
         self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
         self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
         self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)
-        self.guidance_in = (
-            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if self.guidance_embed else nn.Identity()
-        )
+        self.guidance_in = (MLPEmbedder(in_dim=256,
+                                        hidden_dim=self.hidden_size)
+                            if self.guidance_embed else nn.Identity())
         self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
 
         self.double_blocks = nn.ModuleList(
@@ -124,6 +134,7 @@ class Flux(BaseModel):
                     self.num_heads,
                     mlp_ratio=mlp_ratio,
                     qkv_bias=qkv_bias,
+                    backend=self.attn_backend
                 )
                 for _ in range(depth)
             ]
@@ -131,7 +142,7 @@ class Flux(BaseModel):
 
         self.single_blocks = nn.ModuleList(
             [
-                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio)
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, backend=self.attn_backend)
                 for _ in range(depth_single_blocks)
             ]
         )
@@ -159,19 +170,284 @@ class Flux(BaseModel):
             pw=2,
         )
 
-    def load_pretrained_model(self, pretrained_model):
+    def merge_diffuser_lora(self, ori_sd, lora_sd, scale=1.0):
+        key_map = {
+            "single_blocks.{}.linear1.weight": {"key_list": [
+                ["transformer.single_transformer_blocks.{}.attn.to_q.lora_A.weight",
+                 "transformer.single_transformer_blocks.{}.attn.to_q.lora_B.weight", [0, 3072]],
+                ["transformer.single_transformer_blocks.{}.attn.to_k.lora_A.weight",
+                 "transformer.single_transformer_blocks.{}.attn.to_k.lora_B.weight", [3072, 6144]],
+                ["transformer.single_transformer_blocks.{}.attn.to_v.lora_A.weight",
+                 "transformer.single_transformer_blocks.{}.attn.to_v.lora_B.weight", [6144, 9216]],
+                ["transformer.single_transformer_blocks.{}.proj_mlp.lora_A.weight",
+                 "transformer.single_transformer_blocks.{}.proj_mlp.lora_B.weight", [9216, 21504]]
+            ], "num": 38},
+            "single_blocks.{}.modulation.lin.weight": {"key_list": [
+                ["transformer.single_transformer_blocks.{}.norm.linear.lora_A.weight",
+                 "transformer.single_transformer_blocks.{}.norm.linear.lora_B.weight", [0, 9216]],
+            ], "num": 38},
+            "single_blocks.{}.linear2.weight": {"key_list": [
+                ["transformer.single_transformer_blocks.{}.proj_out.lora_A.weight",
+                 "transformer.single_transformer_blocks.{}.proj_out.lora_B.weight", [0, 3072]],
+            ], "num": 38},
+            "double_blocks.{}.txt_attn.qkv.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.attn.add_q_proj.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.add_q_proj.lora_B.weight", [0, 3072]],
+                ["transformer.transformer_blocks.{}.attn.add_k_proj.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.add_k_proj.lora_B.weight", [3072, 6144]],
+                ["transformer.transformer_blocks.{}.attn.add_v_proj.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.add_v_proj.lora_B.weight", [6144, 9216]],
+            ], "num": 19},
+            "double_blocks.{}.img_attn.qkv.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.attn.to_q.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.to_q.lora_B.weight", [0, 3072]],
+                ["transformer.transformer_blocks.{}.attn.to_k.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.to_k.lora_B.weight", [3072, 6144]],
+                ["transformer.transformer_blocks.{}.attn.to_v.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.to_v.lora_B.weight", [6144, 9216]],
+            ], "num": 19},
+            "double_blocks.{}.img_attn.proj.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.attn.to_out.0.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.to_out.0.lora_B.weight", [0, 3072]]
+            ], "num": 19},
+            "double_blocks.{}.txt_attn.proj.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.attn.to_add_out.lora_A.weight",
+                 "transformer.transformer_blocks.{}.attn.to_add_out.lora_B.weight", [0, 3072]]
+            ], "num": 19},
+            "double_blocks.{}.img_mlp.0.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.ff.net.0.proj.lora_A.weight",
+                 "transformer.transformer_blocks.{}.ff.net.0.proj.lora_B.weight", [0, 12288]]
+            ], "num": 19},
+            "double_blocks.{}.img_mlp.2.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.ff.net.2.lora_A.weight",
+                 "transformer.transformer_blocks.{}.ff.net.2.lora_B.weight", [0, 3072]]
+            ], "num": 19},
+            "double_blocks.{}.txt_mlp.0.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.ff_context.net.0.proj.lora_A.weight",
+                 "transformer.transformer_blocks.{}.ff_context.net.0.proj.lora_B.weight", [0, 12288]]
+            ], "num": 19},
+            "double_blocks.{}.txt_mlp.2.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.ff_context.net.2.lora_A.weight",
+                 "transformer.transformer_blocks.{}.ff_context.net.2.lora_B.weight", [0, 3072]]
+            ], "num": 19},
+            "double_blocks.{}.img_mod.lin.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.norm1.linear.lora_A.weight",
+                 "transformer.transformer_blocks.{}.norm1.linear.lora_B.weight", [0, 18432]]
+            ], "num": 19},
+            "double_blocks.{}.txt_mod.lin.weight": {"key_list": [
+                ["transformer.transformer_blocks.{}.norm1_context.linear.lora_A.weight",
+                 "transformer.transformer_blocks.{}.norm1_context.linear.lora_B.weight", [0, 18432]]
+            ], "num": 19}
+        }
+        cover_lora_keys = set()
+        cover_ori_keys = set()
+        for k, v in key_map.items():
+            key_list = v["key_list"]
+            block_num = v["num"]
+            for block_id in range(block_num):
+                for k_list in key_list:
+                    if k_list[0].format(block_id) in lora_sd and k_list[1].format(block_id) in lora_sd:
+                        cover_lora_keys.add(k_list[0].format(block_id))
+                        cover_lora_keys.add(k_list[1].format(block_id))
+                        current_weight = torch.matmul(lora_sd[k_list[0].format(block_id)].permute(1, 0),
+                                                      lora_sd[k_list[1].format(block_id)].permute(1, 0)).permute(1, 0)
+                        ori_sd[k.format(block_id)][k_list[2][0]:k_list[2][1], ...] += scale * current_weight
+                        cover_ori_keys.add(k.format(block_id))
+                        # lora_sd.pop(k_list[0].format(block_id))
+                        # lora_sd.pop(k_list[1].format(block_id))
+        self.logger.info(f"merge_blackforest_lora loads lora'parameters lora-paras: \n"
+                         f"cover-{len(cover_lora_keys)} vs total {len(lora_sd)} \n"
+                         f"cover ori-{len(cover_ori_keys)} vs total {len(ori_sd)}")
+        return ori_sd
+
+    def merge_swift_lora(self, ori_sd, lora_sd, scale = 1.0):
+        have_lora_keys = {}
+        for k, v in lora_sd.items():
+            k = k[len("model."):] if k.startswith("model.") else k
+            ori_key = k.split("lora")[0] + "weight"
+            if ori_key not in ori_sd:
+                raise f"{ori_key} should in the original statedict"
+            if ori_key not in have_lora_keys:
+                have_lora_keys[ori_key] = {}
+            if "lora_A" in k:
+                have_lora_keys[ori_key]["lora_A"] = v
+            elif "lora_B" in k:
+                have_lora_keys[ori_key]["lora_B"] = v
+            else:
+                raise NotImplementedError
+        self.logger.info(f"merge_swift_lora loads lora'parameters {len(have_lora_keys)}")
+        for key, v in have_lora_keys.items():
+            current_weight = torch.matmul(v["lora_A"].permute(1, 0), v["lora_B"].permute(1, 0)).permute(1, 0)
+            ori_sd[key] += scale * current_weight
+        return ori_sd
+
+
+    def merge_blackforest_lora(self, ori_sd, lora_sd, scale = 1.0):
+        have_lora_keys = {}
+        cover_lora_keys = set()
+        cover_ori_keys = set()
+        for k, v in lora_sd.items():
+            if "lora" in k:
+                ori_key = k.split("lora")[0] + "weight"
+                if ori_key not in ori_sd:
+                    raise f"{ori_key} should in the original statedict"
+                if ori_key not in have_lora_keys:
+                    have_lora_keys[ori_key] = {}
+                if "lora_A" in k:
+                    have_lora_keys[ori_key]["lora_A"] = v
+                    cover_lora_keys.add(k)
+                    cover_ori_keys.add(ori_key)
+                elif "lora_B" in k:
+                    have_lora_keys[ori_key]["lora_B"] = v
+                    cover_lora_keys.add(k)
+                    cover_ori_keys.add(ori_key)
+            else:
+                if k in ori_sd:
+                    ori_sd[k] = v
+                    cover_lora_keys.add(k)
+                    cover_ori_keys.add(k)
+                else:
+                    print("unsurpport keys: ", k)
+        self.logger.info(f"merge_blackforest_lora loads lora'parameters lora-paras: \n"
+                         f"cover-{len(cover_lora_keys)} vs total {len(lora_sd)} \n"
+                         f"cover ori-{len(cover_ori_keys)} vs total {len(ori_sd)}")
+
+        for key, v in have_lora_keys.items():
+            current_weight = torch.matmul(v["lora_A"].permute(1, 0), v["lora_B"].permute(1, 0)).permute(1, 0)
+            # print(key, ori_sd[key].shape, current_weight.shape)
+            ori_sd[key] += scale * current_weight
+        return ori_sd
+
+    def merge_comfyui_lora(self, ori_sd, lora_sd, scale = 1.0):
+        ori_key_map = {key.replace("_", ".") : key for key in ori_sd.keys()}
+        parse_ckpt = OrderedDict()
+        for k, v in lora_sd.items():
+            if "alpha" in k:
+                continue
+            k = k.replace("lora_unet_", "").replace("_", ".")
+            map_k = ori_key_map[k.split(".lora")[0] + ".weight"]
+            if map_k not in parse_ckpt:
+                parse_ckpt[map_k] = {}
+            if "lora.up" in k:
+                parse_ckpt[map_k]["lora_up"] = v
+            elif "lora.down" in k:
+                parse_ckpt[map_k]["lora_down"] = v
+        if self.cache_pretrain_model:
+            self.lora_dict[self.comfyui_lora_model] = {}
+
+        for key, v in parse_ckpt.items():
+            current_weight = torch.matmul(v["lora_down"].permute(1, 0), v["lora_up"].permute(1, 0)).permute(1, 0)
+            self.lora_dict[self.comfyui_lora_model] = current_weight
+            ori_sd[key] += scale * current_weight
+        return ori_sd
+
+    def easy_lora_merge(self, ori_sd, lora_sd, scale = 1.0):
+        for key, v in lora_sd.items():
+            ori_sd[key] += scale * v
+        return ori_sd
+
+    def load_pretrained_model(self, pretrained_model, lora_scale = 1.0):
         if next(self.parameters()).device.type == 'meta':
-            map_location = we.device_id
+            map_location = torch.device(we.device_id)
+            safe_device = we.device_id
         else:
             map_location = "cpu"
+            safe_device = "cpu"
+
         if pretrained_model is not None:
-            with FS.get_from(pretrained_model, wait_finish=True) as local_model:
-                if local_model.endswith('safetensors'):
-                    from safetensors.torch import load_file as load_safetensors
-                    sd = load_safetensors(local_model, device=map_location)
+            if not hasattr(self, "ckpt"):
+                with FS.get_from(pretrained_model, wait_finish=True) as local_model:
+                    if local_model.endswith('safetensors'):
+                        from safetensors.torch import load_file as load_safetensors
+                        ckpt = load_safetensors(local_model, device=safe_device)
+                    else:
+                        ckpt = torch.load(local_model, map_location=map_location, weights_only=True)
+                if "state_dict" in ckpt:
+                    ckpt = ckpt["state_dict"]
+                if "model" in ckpt:
+                    ckpt = ckpt["model"]["model"]
+                if self.cache_pretrain_model:
+                    self.ckpt = ckpt
+                    self.lora_dict = {}
+            else:
+                ckpt = self.ckpt
+
+            new_ckpt = OrderedDict()
+            for k, v in ckpt.items():
+                if k in ("img_in.weight"):
+                    model_p = self.state_dict()[k]
+                    if v.shape != model_p.shape:
+                        expanded_state_dict_weight = torch.zeros_like(model_p, device=v.device)
+                        slices = tuple(slice(0, dim) for dim in v.shape)
+                        expanded_state_dict_weight[slices] = v
+                        new_ckpt[k] = expanded_state_dict_weight
+                    else:
+                        new_ckpt[k] = v
                 else:
-                    sd = torch.load(local_model, map_location=map_location)
-            missing, unexpected = self.load_state_dict(sd, strict=False, assign=True)
+                    new_ckpt[k] = v
+
+
+            if self.lora_model is not None:
+                with FS.get_from(self.lora_model, wait_finish=True) as local_model:
+                    if local_model.endswith('safetensors'):
+                        from safetensors.torch import load_file as load_safetensors
+                        lora_sd = load_safetensors(local_model, device=safe_device)
+                    else:
+                        lora_sd = torch.load(local_model, map_location=map_location, weights_only=True)
+                new_ckpt = self.merge_diffuser_lora(new_ckpt, lora_sd, scale=lora_scale)
+            if self.swift_lora_model is not None:
+                if not isinstance(self.swift_lora_model, list):
+                    self.swift_lora_model = [(self.swift_lora_model, 1.0)]
+                for lora_model in self.swift_lora_model:
+                    if isinstance(lora_model, str):
+                        lora_model = (lora_model, 1.0/len(self.swift_lora_model))
+                    print(lora_model)
+                    self.logger.info(f"load swift lora model: {lora_model}")
+                    with FS.get_from(lora_model[0], wait_finish=True) as local_model:
+                        if local_model.endswith('safetensors'):
+                            from safetensors.torch import load_file as load_safetensors
+                            lora_sd = load_safetensors(local_model, device=safe_device)
+                        else:
+                            lora_sd = torch.load(local_model, map_location=map_location, weights_only=True)
+                    new_ckpt = self.merge_swift_lora(new_ckpt, lora_sd, scale=lora_model[1])
+
+            if self.blackforest_lora_model is not None:
+                with FS.get_from(self.blackforest_lora_model, wait_finish=True) as local_model:
+                    if local_model.endswith('safetensors'):
+                        from safetensors.torch import load_file as load_safetensors
+                        lora_sd = load_safetensors(local_model, device=safe_device)
+                    else:
+                        lora_sd = torch.load(local_model, map_location=map_location, weights_only=True)
+                new_ckpt = self.merge_blackforest_lora(new_ckpt, lora_sd, scale=lora_scale)
+
+            if self.comfyui_lora_model is not None:
+                if hasattr(self, "current_lora") and self.current_lora == self.comfyui_lora_model:
+                    return
+                if hasattr(self, "lora_dict") and self.comfyui_lora_model in self.lora_dict:
+                    new_ckpt = self.easy_lora_merge(new_ckpt, self.lora_dict[self.comfyui_lora_model], scale=lora_scale)
+                else:
+                    with FS.get_from(self.comfyui_lora_model, wait_finish=True) as local_model:
+                        if local_model.endswith('safetensors'):
+                            from safetensors.torch import load_file as load_safetensors
+                            lora_sd = load_safetensors(local_model, device=safe_device)
+                        else:
+                            lora_sd = torch.load(local_model, map_location=map_location, weights_only=True)
+                    new_ckpt = self.merge_comfyui_lora(new_ckpt, lora_sd, scale=lora_scale)
+                if self.comfyui_lora_model:
+                    self.current_lora = self.comfyui_lora_model
+
+
+            adapter_ckpt = {}
+            if self.pretrain_adapter is not None:
+                with FS.get_from(self.pretrain_adapter, wait_finish=True) as local_adapter:
+                    if local_adapter.endswith('safetensors'):
+                        from safetensors.torch import load_file as load_safetensors
+                        adapter_ckpt = load_safetensors(local_adapter, device=safe_device)
+                    else:
+                        adapter_ckpt = torch.load(local_adapter, map_location=map_location, weights_only=True)
+            new_ckpt.update(adapter_ckpt)
+
+            missing, unexpected = self.load_state_dict(new_ckpt, strict=False, assign=True)
             self.logger.info(
                 f'Restored from {pretrained_model} with {len(missing)} missing and {len(unexpected)} unexpected keys'
             )
@@ -239,7 +515,298 @@ class Flux(BaseModel):
 
     @staticmethod
     def get_config_template():
-        return dict_to_yaml('MODEL',
+        return dict_to_yaml('BACKBONE',
                             __class__.__name__,
                             Flux.para_dict,
                             set_name=True)
+@BACKBONES.register_class()
+class FluxMR(Flux):
+    def prepare_input(self, x, cond):
+        if isinstance(cond['context'], list):
+            context, y = torch.cat(cond["context"], dim=0).to(x), torch.cat(cond["y"], dim=0).to(x)
+        else:
+            context, y = cond['context'].to(x), cond['y'].to(x)
+        batch_frames, batch_frames_ids = [], []
+        for ix, shape in zip(x, cond["x_shapes"]):
+            # unpack image from sequence
+            ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
+            c, h, w = ix.shape
+            ix = rearrange(ix, "c (h ph) (w pw) -> (h w) (c ph pw)", ph=2, pw=2)
+            ix_id = torch.zeros(h // 2, w // 2, 3)
+            ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None]
+            ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :]
+            ix_id = rearrange(ix_id, "h w c -> (h w) c")
+            batch_frames.append([ix])
+            batch_frames_ids.append([ix_id])
+
+        x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], []
+        for frames, frame_ids in zip(batch_frames, batch_frames_ids):
+            proj_frames = []
+            for idx, one_frame in enumerate(frames):
+                one_frame = self.img_in(one_frame)
+                proj_frames.append(one_frame)
+            ix = torch.cat(proj_frames, dim=0)
+            if_id = torch.cat(frame_ids, dim=0)
+            x_list.append(ix)
+            x_id_list.append(if_id)
+            mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool())
+            x_seq_length.append(ix.shape[0])
+        x = pad_sequence(tuple(x_list), batch_first=True)
+        x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x)  # [b,pad_seq,2] pad (0.,0.) at dim2
+        mask_x = pad_sequence(tuple(mask_x_list), batch_first=True)
+
+        txt = self.txt_in(context)
+        txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x)
+        mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool()
+
+        return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length
+
+    def unpack(self, x: Tensor, cond: dict = None, x_seq_length: list = None) -> Tensor:
+        x_list = []
+        image_shapes = cond["x_shapes"]
+        for u, shape, seq_length in zip(x, image_shapes, x_seq_length):
+            height, width = shape
+            h, w = math.ceil(height / 2), math.ceil(width / 2)
+            u = rearrange(
+                u[seq_length-h*w:seq_length, ...],
+                "(h w) (c ph pw) -> (h ph w pw) c",
+                h=h,
+                w=w,
+                ph=2,
+                pw=2,
+            )
+            x_list.append(u)
+        x = pad_sequence(tuple(x_list), batch_first=True).permute(0, 2, 1)
+        return x
+
+    def forward(
+            self,
+            x: Tensor,
+            t: Tensor,
+            cond: dict = {},
+            guidance: Tensor | None = None,
+            gc_seg: int = 0,
+            **kwargs
+    ) -> Tensor:
+        x, x_ids, txt, txt_ids, y, mask_x, mask_txt, seq_length_list = self.prepare_input(x, cond)
+        # running on sequences img
+        vec = self.time_in(timestep_embedding(t, 256))
+        if self.guidance_embed and guidance[-1] >= 0:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        ids = torch.cat((txt_ids, x_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        mask_aside = torch.cat((mask_txt, mask_x), dim=1)
+        mask = mask_aside[:, None, :] * mask_aside[:, :, None]
+
+        kwargs = dict(
+            vec=vec,
+            pe=pe,
+            mask=mask,
+            txt_length = txt.shape[1],
+        )
+        x = torch.cat((txt, x), 1)
+        if self.use_grad_checkpoint and gc_seg >= 0:
+            x = checkpoint_sequential(
+                functions=[partial(block, **kwargs) for block in self.double_blocks],
+                segments=gc_seg if gc_seg > 0 else len(self.double_blocks),
+                input=x,
+                use_reentrant=False
+            )
+        else:
+            for block in self.double_blocks:
+                x = block(x, **kwargs)
+
+        kwargs = dict(
+            vec=vec,
+            pe=pe,
+            mask=mask,
+        )
+
+        if self.use_grad_checkpoint and gc_seg >= 0:
+            x = checkpoint_sequential(
+                functions=[partial(block, **kwargs) for block in self.single_blocks],
+                segments=gc_seg if gc_seg > 0 else len(self.single_blocks),
+                input=x,
+                use_reentrant=False
+            )
+        else:
+            for block in self.single_blocks:
+                x = block(x, **kwargs)
+        x = x[:, txt.shape[1]:, ...]
+        x = self.final_layer(x, vec)  # (N, T, patch_size ** 2 * out_channels) 6 64 64
+        x = self.unpack(x, cond, seq_length_list)
+        return x
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('MODEL',
+                            __class__.__name__,
+                            FluxMR.para_dict,
+                            set_name=True)
+@BACKBONES.register_class()
+class FluxMRFill(FluxMR):
+    def __init__(self, cfg, logger = None):
+        super().__init__(cfg, logger)
+    def prepare_input(self, x, cond):
+        context, y = cond["context"], cond["y"]
+        batch_frames, batch_frames_ids = [], []
+        for ix, shape, imask, ie, ie_mask in zip(x, cond["x_shapes"], cond["x_mask"],
+                                                 cond["edit"], cond["edit_mask"]):
+            # unpack image from sequence
+            ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
+            imask = torch.ones_like(ix[[0], :, :]) if imask is None else imask.squeeze(0)
+            if len(ie) > 0:
+                ie = ie[0].squeeze(0)
+                ie_mask = torch.ones((ix.shape[0] * 4, ix.shape[1], ix.shape[2])) if ie_mask is None else ie_mask[0].squeeze(0)
+            else:
+                ie, ie_mask = torch.zeros_like(ix).to(x), torch.ones_like(imask).to(x)
+            ix = torch.cat([ix, ie, ie_mask], dim=0)
+            c, h, w = ix.shape
+            ix = rearrange(ix, "c (h ph) (w pw) -> (h w) (c ph pw)", ph=2, pw=2)
+            ix_id = torch.zeros(h // 2, w // 2, 3)
+            ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None]
+            ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :]
+            ix_id = rearrange(ix_id, "h w c -> (h w) c")
+            batch_frames.append([ix])
+            batch_frames_ids.append([ix_id])
+        x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], []
+        for frames, frame_ids in zip(batch_frames, batch_frames_ids):
+            proj_frames = []
+            for idx, one_frame in enumerate(frames):
+                one_frame = self.img_in(one_frame)
+                proj_frames.append(one_frame)
+            ix = torch.cat(proj_frames, dim=0)
+            if_id = torch.cat(frame_ids, dim=0)
+            x_list.append(ix)
+            x_id_list.append(if_id)
+            mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool())
+            x_seq_length.append(ix.shape[0])
+        # if len(x_list) < 1: import pdb;pdb.set_trace()
+        x = pad_sequence(tuple(x_list), batch_first=True)
+        x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x)  # [b,pad_seq,2] pad (0.,0.) at dim2
+        mask_x = pad_sequence(tuple(mask_x_list), batch_first=True)
+        # import pdb;pdb.set_trace()
+        if isinstance(context, list):
+            txt_list, mask_txt_list, y_list = [], [], []
+            for sample_id, (ctx, yy) in enumerate(zip(context, y)):
+                txt_list.append(self.txt_in(ctx.to(x)))
+                mask_txt_list.append(torch.ones(txt_list[-1].shape[0]).to(ctx.device, non_blocking=True).bool())
+                y_list.append(yy.to(x))
+            txt = pad_sequence(tuple(txt_list), batch_first=True)
+            txt_ids = torch.zeros(txt.shape[0], txt.shape[1], 3).to(x)
+            mask_txt = pad_sequence(tuple(mask_txt_list), batch_first=True)
+            y = torch.cat(y_list, dim=0)
+            assert y.ndim == 2 and txt.ndim == 3
+        else:
+            txt = self.txt_in(context)
+            txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x)
+            mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool()
+        return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('MODEL',
+                            __class__.__name__,
+                            FluxMRFill.para_dict,
+                            set_name=True)
+@BACKBONES.register_class()
+class FluxMRRedux(FluxMR):
+    '''
+        ref_image_siglip + projector
+    '''
+    def __init__(self, cfg, logger = None):
+        super().__init__(cfg, logger)
+        self.redux_dim = cfg.get("REDUX_DIM", 1152)
+        self.context_in_dim = cfg.CONTEXT_IN_DIM
+        self.redux_up = nn.Linear(self.redux_dim, self.context_in_dim * 3)
+        self.redux_down = nn.Linear(self.context_in_dim * 3, self.context_in_dim)
+
+
+    def prepare_input(self, x, cond):
+        ref_x = cond.get("ref_x", None)
+        context, y = torch.cat(cond["context"], dim=0).to(x), torch.cat(cond["y"], dim=0).to(x)
+        if ref_x is not None:
+            ref_x = [torch.cat(ref_ix, dim=0).mean(dim=0, keepdim=True) for ref_ix in ref_x]
+            ref_x = self.redux_down(nn.functional.silu(self.redux_up(torch.cat(ref_x, dim=0))))
+            context = torch.cat((context, ref_x), dim=-2)
+
+        batch_frames, batch_frames_ids = [], []
+        for ix, shape in zip(x, cond["x_shapes"]):
+            # unpack image from sequence
+            ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
+            c, h, w = ix.shape
+            ix = rearrange(ix, "c (h ph) (w pw) -> (h w) (c ph pw)", ph=2, pw=2)
+            ix_id = torch.zeros(h // 2, w // 2, 3)
+            ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None]
+            ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :]
+            ix_id = rearrange(ix_id, "h w c -> (h w) c")
+            batch_frames.append([ix])
+            batch_frames_ids.append([ix_id])
+
+        x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], []
+        for frames, frame_ids in zip(batch_frames, batch_frames_ids):
+            proj_frames = []
+            for idx, one_frame in enumerate(frames):
+                one_frame = self.img_in(one_frame)
+                proj_frames.append(one_frame)
+            ix = torch.cat(proj_frames, dim=0)
+            if_id = torch.cat(frame_ids, dim=0)
+            x_list.append(ix)
+            x_id_list.append(if_id)
+            mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool())
+            x_seq_length.append(ix.shape[0])
+        x = pad_sequence(tuple(x_list), batch_first=True)
+        x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x)  # [b,pad_seq,2] pad (0.,0.) at dim2
+        mask_x = pad_sequence(tuple(mask_x_list), batch_first=True)
+
+        txt = self.txt_in(context)
+        txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x)
+        mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool()
+
+        return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('MODEL',
+                            __class__.__name__,
+                            FluxMRRedux.para_dict,
+                            set_name=True)
+@BACKBONES.register_class()
+class FluxMRControl(FluxMR):
+    '''
+        cat([x, ie]) ensure the same size bettwn the x and ie
+    '''
+    def prepare_input(self, x, cond, *args, **kwargs ):
+        context, y = torch.cat(cond["context"], dim=0).to(x), torch.cat(cond["y"], dim=0).to(x)
+        x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], []
+        for ix, shape, ie in zip(x, cond["x_shapes"], cond["edit"]):
+            ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
+            ix = torch.cat([ix, ie], dim=0)
+            c, h, w = ix.shape
+            ix = rearrange(ix, "c (h ph) (w pw) -> (h w) (c ph pw)", ph=2, pw=2)
+            ix_id = torch.zeros(h // 2, w // 2, 3)
+            ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None]
+            ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :]
+            ix_id = rearrange(ix_id, "h w c -> (h w) c")
+            x_list.append(self.img_in(ix))
+            x_id_list.append(ix_id)
+            mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool())
+            x_seq_length.append(ix.shape[0])
+        # if len(x_list) < 1: import pdb;pdb.set_trace()
+        x = pad_sequence(tuple(x_list), batch_first=True)
+        x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x)  # [b,pad_seq,2] pad (0.,0.) at dim2
+        mask_x = pad_sequence(tuple(mask_x_list), batch_first=True)
+        txt = self.txt_in(context)
+        txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x)
+        mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool()
+        return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length
+
+    @staticmethod
+    def get_config_template():
+        return dict_to_yaml('MODEL',
+                            __class__.__name__,
+                            FluxMRControl.para_dict,
+                            set_name=True)

scepter/modules/model/backbone/flux/layers.py

@@ -1,3 +1,7 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# This file contains code that is adapted from
+# https://github.com/black-forest-labs/flux.git
 from __future__ import annotations
 
 import math
@@ -6,32 +10,89 @@ from torch import Tensor, nn
 import torch
 from einops import rearrange, repeat
 from torch import Tensor
+from torch.nn.utils.rnn import pad_sequence
 
+try:
+    from flash_attn import (
+        flash_attn_varlen_func
+    )
+    FLASHATTN_IS_AVAILABLE = True
+except ImportError:
+    FLASHATTN_IS_AVAILABLE = False
+    flash_attn_varlen_func = None
 
-def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor | None = None) -> Tensor:
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor | None = None, backend = 'pytorch') -> Tensor:
     q, k = apply_rope(q, k, pe)
-    x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-    x = torch.nan_to_num(x, nan=0.0, posinf=1e10, neginf=-1e10)
-    x = rearrange(x, "B H L D -> B L (H D)")
+    if backend == 'pytorch':
+        if mask is not None and mask.dtype == torch.bool:
+            mask = torch.zeros_like(mask).to(q).masked_fill_(mask.logical_not(), -1e20)
+        x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        # x = torch.nan_to_num(x, nan=0.0, posinf=1e10, neginf=-1e10)
+        x = rearrange(x, "B H L D -> B L (H D)")
+    elif backend == 'flash_attn':
+        # q: (B, H, L, D)
+        # k: (B, H, S, D) now L = S
+        # v: (B, H, S, D)
+        b, h, lq, d = q.shape
+        _, _, lk, _ = k.shape
+        q = rearrange(q, "B H L D -> B L H D")
+        k = rearrange(k, "B H S D -> B S H D")
+        v = rearrange(v, "B H S D -> B S H D")
+        if mask is None:
+            q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(q.device, non_blocking=True)
+            k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(k.device, non_blocking=True)
+        else:
+            q_lens = torch.sum(mask[:, 0, :, 0], dim=1).int()
+            k_lens = torch.sum(mask[:, 0, 0, :], dim=1).int()
+        q = torch.cat([q_v[:q_l] for q_v, q_l in zip(q, q_lens)])
+        k = torch.cat([k_v[:k_l] for k_v, k_l in zip(k, k_lens)])
+        v = torch.cat([v_v[:v_l] for v_v, v_l in zip(v, k_lens)])
+        cu_seqlens_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32)
+        cu_seqlens_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32)
+        max_seqlen_q = q_lens.max()
+        max_seqlen_k = k_lens.max()
+
+        x = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k
+        )
+        x_list = [x[cu_seqlens_q[i]:cu_seqlens_q[i+1]] for i in range(b)]
+        x = pad_sequence(tuple(x_list), batch_first=True)
+        x = rearrange(x, "B L H D -> B L (H D)")
+    else:
+        raise NotImplementedError
     return x
 
 
 def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
     assert dim % 2 == 0
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    scale = torch.arange(0, dim, 2, dtype=torch.float64,
+                         device=pos.device) / dim
     omega = 1.0 / (theta**scale)
-    out = torch.einsum("...n,d->...nd", pos, omega)
-    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
-    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    out = torch.einsum('...n,d->...nd', pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out),
+         torch.sin(out),
+         torch.cos(out)],
+        dim=-1)
+    out = rearrange(out, 'b n d (i j) -> b n d i j', i=2, j=2)
     return out.float()
 
 
-def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+def apply_rope(xq: Tensor, xk: Tensor,
+               freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
     xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
     xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
     xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
     xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
-    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(
+        *xk.shape).type_as(xk)
+
 
 class EmbedND(nn.Module):
     def __init__(self, dim: int, theta: int, axes_dim: list[int]):
@@ -43,14 +104,20 @@ class EmbedND(nn.Module):
     def forward(self, ids: Tensor) -> Tensor:
         n_axes = ids.shape[-1]
         emb = torch.cat(
-            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            [
+                rope(ids[..., i], self.axes_dim[i], self.theta)
+                for i in range(n_axes)
+            ],
             dim=-3,
         )
 
         return emb.unsqueeze(1)
 
 
-def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+def timestep_embedding(t: Tensor,
+                       dim,
+                       max_period=10000,
+                       time_factor: float = 1000.0):
     """
     Create sinusoidal timestep embeddings.
     :param t: a 1-D Tensor of N indices, one per batch element.
@@ -61,14 +128,15 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
     """
     t = time_factor * t
     half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-        t.device
-    )
+    freqs = torch.exp(-math.log(max_period) *
+                      torch.arange(start=0, end=half, dtype=torch.float32) /
+                      half).to(t.device)
 
     args = t[:, None].float() * freqs[None]
     embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
     if dim % 2:
-        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        embedding = torch.cat(
+            [embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
     if torch.is_floating_point(t):
         embedding = embedding.to(t)
     return embedding
@@ -103,7 +171,8 @@ class QKNorm(torch.nn.Module):
         self.query_norm = RMSNorm(dim)
         self.key_norm = RMSNorm(dim)
 
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+    def forward(self, q: Tensor, k: Tensor,
+                v: Tensor) -> tuple[Tensor, Tensor]:
         q = self.query_norm(q)
         k = self.key_norm(k)
         return q.to(v), k.to(v)
@@ -119,9 +188,15 @@ class SelfAttention(nn.Module):
         self.norm = QKNorm(head_dim)
         self.proj = nn.Linear(dim, dim)
 
-    def forward(self, x: Tensor, pe: Tensor, mask: Tensor | None = None) -> Tensor:
+    def forward(self,
+                x: Tensor,
+                pe: Tensor,
+                mask: Tensor | None = None) -> Tensor:
         qkv = self.qkv(x)
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k, v = rearrange(qkv,
+                            'B L (K H D) -> K B H L D',
+                            K=3,
+                            H=self.num_heads)
         q, k = self.norm(q, k, v)
         x = attention(q, k, v, pe=pe, mask=mask)
         x = self.proj(x)
@@ -152,7 +227,7 @@ class Modulation(nn.Module):
 
 
 class DoubleStreamBlock(nn.Module):
-    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, backend = 'pytorch'):
         super().__init__()
 
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
@@ -169,6 +244,8 @@ class DoubleStreamBlock(nn.Module):
             nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
         )
 
+        self.backend = backend
+
         self.txt_mod = Modulation(hidden_size, double=True)
         self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
         self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
@@ -205,7 +282,7 @@ class DoubleStreamBlock(nn.Module):
         v = torch.cat((txt_v, img_v), dim=2)
         if mask is not None:
             mask = repeat(mask, 'B L S->  B H L S', H=self.num_heads)
-        attn = attention(q, k, v, pe=pe, mask = mask)
+        attn = attention(q, k, v, pe=pe, mask = mask, backend = self.backend)
         txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
 
         # calculate the img bloks
@@ -224,13 +301,13 @@ class SingleStreamBlock(nn.Module):
     A DiT block with parallel linear layers as described in
     https://arxiv.org/abs/2302.05442 and adapted modulation interface.
     """
-
     def __init__(
         self,
         hidden_size: int,
         num_heads: int,
         mlp_ratio: float = 4.0,
         qk_scale: float | None = None,
+        backend='pytorch'
     ):
         super().__init__()
         self.hidden_dim = hidden_size
@@ -240,29 +317,43 @@ class SingleStreamBlock(nn.Module):
 
         self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
         # qkv and mlp_in
-        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        self.linear1 = nn.Linear(hidden_size,
+                                 hidden_size * 3 + self.mlp_hidden_dim)
         # proj and mlp_out
-        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim,
+                                 hidden_size)
 
         self.norm = QKNorm(head_dim)
 
         self.hidden_size = hidden_size
-        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.pre_norm = nn.LayerNorm(hidden_size,
+                                     elementwise_affine=False,
+                                     eps=1e-6)
 
-        self.mlp_act = nn.GELU(approximate="tanh")
+        self.mlp_act = nn.GELU(approximate='tanh')
         self.modulation = Modulation(hidden_size, double=False)
+        self.backend = backend
 
-    def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None) -> Tensor:
+    def forward(self,
+                x: Tensor,
+                vec: Tensor,
+                pe: Tensor,
+                mask: Tensor = None) -> Tensor:
         mod, _ = self.modulation(vec)
         x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
-        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+        qkv, mlp = torch.split(self.linear1(x_mod),
+                               [3 * self.hidden_size, self.mlp_hidden_dim],
+                               dim=-1)
 
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k, v = rearrange(qkv,
+                            'B L (K H D) -> K B H L D',
+                            K=3,
+                            H=self.num_heads)
         q, k = self.norm(q, k, v)
         if mask is not None:
             mask = repeat(mask, 'B L S->  B H L S', H=self.num_heads)
         # compute attention
-        attn = attention(q, k, v, pe=pe, mask = mask)
+        attn = attention(q, k, v, pe=pe, mask=mask, backend=self.backend)
         # compute activation in mlp stream, cat again and run second linear layer
         output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
         return x + mod.gate * output
@@ -271,9 +362,14 @@ class SingleStreamBlock(nn.Module):
 class LastLayer(nn.Module):
     def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
         super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+        self.norm_final = nn.LayerNorm(hidden_size,
+                                       elementwise_affine=False,
+                                       eps=1e-6)
+        self.linear = nn.Linear(hidden_size,
+                                patch_size * patch_size * out_channels,
+                                bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
 
     def forward(self, x: Tensor, vec: Tensor) -> Tensor:
         shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)

scepter/modules/model/backbone/image/vit_modify.py

@@ -74,7 +74,7 @@ class VisualTransformer(BaseModel):
             with FS.get_from(self.pretrain_path,
                              wait_finish=True) as local_file:
                 logger.info(f'Loading checkpoint from {self.pretrain_path}')
-                visual_pre = torch.load(local_file, map_location='cpu')
+                visual_pre = torch.load(local_file, map_location='cpu', weights_only=True)
                 if not use_proj:
                     visual_pre.pop('proj')
                 if visual_pre['conv1.weight'].dtype == torch.float16:
@@ -145,7 +145,7 @@ class SomeFTVisualTransformer(BaseModel):
             with FS.get_from(self.pretrain_path,
                              wait_finish=True) as local_file:
                 logger.info(f'Loading checkpoint from {self.pretrain_path}')
-                visual_pre = torch.load(local_file, map_location='cpu')
+                visual_pre = torch.load(local_file, map_location='cpu', weights_only=True)
             state_dict_update = self.reformat_state_dict(visual_pre)
             self.visual.load_state_dict(state_dict_update, strict=True)
 

scepter/modules/model/backbone/mmdit/__init__.py

@@ -1,2 +1,3 @@
 # -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 from .sd3 import MMDiT

scepter/modules/model/backbone/mmdit/sd3.py

@@ -5,12 +5,11 @@
 # diffusers: https://github.com/huggingface/diffusers
 # ComfyUI: https://github.com/comfyanonymous/ComfyUI
 
-import logging
 import math
 import re
 from collections import OrderedDict
 from functools import partial
-from typing import Dict, Optional
+from typing import Optional
 
 import numpy as np
 import torch
@@ -26,7 +25,7 @@ try:
     import xformers
     import xformers.ops
     XFORMERS_IS_AVAILBLE = True
-except:
+except Exception:
     XFORMERS_IS_AVAILBLE = False
 
 BROKEN_XFORMERS = False
@@ -35,7 +34,7 @@ try:
     # XFormers bug confirmed on all versions from 0.0.21 to 0.0.26 (q with bs bigger than 65535 gives CUDA error)
     BROKEN_XFORMERS = x_vers.startswith(
         '0.0.2') and not x_vers.startswith('0.0.20')
-except:
+except Exception:
     pass
 
 
@@ -1137,7 +1136,7 @@ class MMDiT(BaseModel):
                     from safetensors.torch import load_file as load_safetensors
                     model = load_safetensors(local_path)
                 else:
-                    model = torch.load(local_path, map_location='cpu')
+                    model = torch.load(local_path, map_location='cpu', weights_only=True)
                 if 'state_dict' in model:
                     model = model['state_dict']
                 new_ckpt = OrderedDict()
@@ -1145,7 +1144,7 @@ class MMDiT(BaseModel):
                 for k, v in model.items():
                     if self.ignore_keys is not None:
                         if (isinstance(self.ignore_keys, str) and re.match(self.ignore_keys, k)) or \
-                            (isinstance(self.ignore_keys, list) and k in self.ignore_keys):
+                           (isinstance(self.ignore_keys, list) and k in self.ignore_keys):
                             ignore_ckpt[k] = v
                             continue
                     k = k.replace('model.diffusion_model.', '')
@@ -1185,11 +1184,6 @@ class MMDiT(BaseModel):
         spatial_pos_embed = spatial_pos_embed[:, top:top + h, left:left + w, :]
         spatial_pos_embed = rearrange(spatial_pos_embed,
                                       '1 h w c -> 1 (h w) c')
-        # print(spatial_pos_embed, top, left, h, w)
-        # # t = get_2d_sincos_pos_embed_torch(self.hidden_size, w, h, 7.875, 7.875, device=device) #matches exactly for 1024 res
-        # t = get_2d_sincos_pos_embed_torch(self.hidden_size, w, h, 7.5, 7.5, device=device) #scales better
-        # # print(t)
-        # return t
         return spatial_pos_embed
 
     def unpatchify(self, x, hw=None):

scepter/modules/model/backbone/pixart/__init__.py

@@ -1,2 +1,3 @@
 # -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 from .pixart_alpha import PixArt

scepter/modules/model/backbone/pixart/pixart_alpha.py

@@ -354,7 +354,7 @@ class PixArt(BaseModel):
     def load_pretrained_model(self, pretrained_model):
         if pretrained_model:
             with FS.get_from(pretrained_model, wait_finish=True) as local_path:
-                model = torch.load(local_path, map_location='cpu')
+                model = torch.load(local_path, map_location='cpu', weights_only=True)
                 if 'state_dict' in model:
                     model = model['state_dict']
                 new_ckpt = OrderedDict()

scepter/modules/model/backbone/transformer/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.

scepter/modules/model/backbone/transformer/attention.py

@@ -10,7 +10,7 @@ import warnings
 
 import torch
 import torch.nn as nn
-from torch.cuda import amp
+from torch import amp
 from torch.nn import functional as F
 from torch.nn.utils.rnn import pad_sequence
 from tqdm import tqdm
@@ -440,7 +440,7 @@ def multi_head_varlen_attention(q_img,
         k = k.type(flash_dtype)
         v = v.type(flash_dtype)
 
-    with amp.autocast():
+    with amp.autocast("cuda"):
         x = flash_attn_varlen_func(q=q,
                                    k=k,
                                    v=v,
@@ -537,7 +537,7 @@ class FullAttention(nn.Module):
         k_img, k_txt = self.k_img_norm(k_img).view(
             b, -1, n * d), self.k_txt_norm(k_txt).view(b, -1, n * d)
 
-        ### add position
+        # add position
         q_img, k_img = apply_2d_rope(q_img, k_img, padded_pos_index, n, d)
 
         # support varying length

scepter/modules/model/backbone/transformer/layers.py

@@ -9,6 +9,7 @@ import math
 import torch
 import torch.nn as nn
 from einops import rearrange
+
 from scepter.modules.model.backbone.transformer.attention import drop_path
 
 
@@ -299,3 +300,28 @@ class Mlp(nn.Module):
         x = self.fc2(x)
         x = self.drop(x)
         return x
+
+
+class T2IFinalLayer(nn.Module):
+    """
+    The final layer of PixArt.
+    """
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size,
+                                       elementwise_affine=False,
+                                       eps=1e-6)
+        self.linear = nn.Linear(hidden_size,
+                                patch_size * patch_size * out_channels,
+                                bias=True)
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, hidden_size) / hidden_size**0.5)
+        self.out_channels = out_channels
+
+    def forward(self, x, t):
+        shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2,
+                                                                         dim=1)
+        shift, scale = shift.squeeze(1), scale.squeeze(1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x

scepter/modules/model/backbone/transformer/pos_embed.py

@@ -8,8 +8,13 @@ from itertools import repeat as iter_repeat
 from typing import Iterable
 
 import numpy as np
-
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+from torch import amp
+from torch.nn.utils.rnn import pad_sequence
 
 
 def _ntuple(n):
@@ -127,3 +132,218 @@ def apply_2d_rope(xq,
         2)  # point_wise mul, then flatten eg[[1,2],[3,4],[5,6]]->[1,2,3,4,5,6]
     xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2)
     return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(
+        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.float()
+
+
+def frame_pad(x, seq_len, shapes):
+    max_h, max_w = np.max(shapes, 0)
+    frames = []
+    cur_len = 0
+    for h, w in shapes:
+        frame_len = h * w
+        frames.append(
+            F.pad(
+                x[cur_len:cur_len + frame_len].view(h, w, -1),
+                (0, 0, 0, max_w - w, 0, max_h - h))  # .view(max_h * max_w, -1)
+        )
+        cur_len += frame_len
+        if cur_len >= seq_len:
+            break
+    return torch.stack(frames)
+
+
+def frame_unpad(x, shapes):
+    max_h, max_w = np.max(shapes, 0)
+    x = rearrange(x, '(b h w) n c -> b h w n c', h=max_h, w=max_w)
+    frames = []
+    for i, (h, w) in enumerate(shapes):
+        if i >= len(x):
+            break
+        frames.append(rearrange(x[i, :h, :w], 'h w n c -> (h w) n c'))
+    return torch.concat(frames)
+
+
+@amp.autocast("cuda", enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    """
+    Precompute the frequency tensor for complex exponentials.
+    """
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta,
+                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+@amp.autocast("cuda", enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    """
+    x:          [B, L, N, C].
+    grid_sizes: [B, 3].
+    freqs:      [M, C // 2].
+    """
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
+            seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2).type_as(x)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output)
+
+
+@amp.autocast("cuda", enabled=False)
+def rope_apply_multires_pad(x, x_lens, x_shapes, freqs, pad=True):
+    """
+    x:          [B, L, N, C].
+    x_lens:     [B].
+    x_shapes:   [B, F, 2].
+    freqs:      [M, C // 2].
+    """
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (seq_len,
+            shapes) in enumerate(zip(x_lens.tolist(), x_shapes.tolist())):
+        x_i = frame_pad(x[i], seq_len, shapes)  # f, h, w, c
+        f, h, w = x_i.shape[:3]
+        pad_seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(
+            x_i.to(torch.float64).reshape(pad_seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(pad_seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2).type_as(x)
+        x_i = frame_unpad(x_i, shapes)
+        if pad:
+            x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output) if pad else torch.concat(output)
+
+
+@amp.autocast("cuda", enabled=False)
+def rope_apply_multires(x, x_lens, x_shapes, freqs, pad=True):
+    """
+    x:          [B*L, N, C].
+    x_lens:     [B].
+    x_shapes:   [B, F, 2].
+    freqs:      [M, C // 2].
+    """
+    n, c = x.size(1), x.size(2) // 2
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+    # loop over samples
+    output = []
+    st = 0
+    for i, (seq_len,
+            shapes) in enumerate(zip(x_lens.tolist(), x_shapes.tolist())):
+        x_i = frame_pad(x[st:st + seq_len], seq_len, shapes)  # f, h, w, c
+        f, h, w = x_i.shape[:3]
+        pad_seq_len = f * h * w
+        # precompute multipliers
+        x_i = torch.view_as_complex(
+            x_i.to(torch.float64).reshape(pad_seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(pad_seq_len, 1, -1)
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2).type_as(x)
+        x_i = frame_unpad(x_i, shapes)
+        # append to collection
+        output.append(x_i)
+        st += seq_len
+    return pad_sequence(output) if pad else torch.concat(output)
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64,
+                         device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum('...n,d->...nd', pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out),
+         torch.sin(out),
+         torch.cos(out)],
+        dim=-1)
+    out = rearrange(out, 'b n d (i j) -> b n d i j', i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor,
+               freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(
+        *xk.shape).type_as(xk)
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [
+                rope(ids[..., i], self.axes_dim[i], self.theta)
+                for i in range(n_axes)
+            ],
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)

scepter/modules/model/backbone/unet/unet_module.py

@@ -459,7 +459,7 @@ class DiffusionUNet(BaseModel):
             from safetensors.torch import load_file as load_safetensors
             sd = load_safetensors(path)
         else:
-            sd = torch.load(path, map_location='cpu')
+            sd = torch.load(path, map_location='cpu', weights_only=True)
 
         new_sd = OrderedDict()
         for k, v in sd.items():
@@ -1231,7 +1231,7 @@ class LargenUNetXL(DiffusionUNetXL):
             from safetensors.torch import load_file as load_safetensors
             sd = load_safetensors(path)
         else:
-            sd = torch.load(path, map_location='cpu')
+            sd = torch.load(path, map_location='cpu', weights_only=True)
 
         new_sd = OrderedDict()
         for k, v in sd.items():

scepter/modules/model/diffusion/__init__.py

@@ -1,3 +1,27 @@
-from .samplers import BaseDiffusionSampler, FlowEluerSampler, DDIMSampler
-from .schedules import BaseNoiseScheduler, ScaledLinearScheduler, FlowMatchShiftScheduler
-from .diffusions import BaseDiffusion, DiffusionFluxRF
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
+
+
+if TYPE_CHECKING:
+    from .diffusions import BaseDiffusion, DiffusionFluxRF
+    from .samplers import BaseDiffusionSampler, DDIMSampler, FlowEluerSampler
+    from .schedules import (BaseNoiseScheduler, FlowMatchShiftScheduler,
+                            ScaledLinearScheduler)
+else:
+    _import_structure = {
+        'diffusions': ['BaseDiffusion', 'DiffusionFluxRF'],
+        'samplers': ['BaseDiffusionSampler', 'DDIMSampler', 'FlowEluerSampler'],
+        'schedules': ['BaseNoiseScheduler', 'FlowMatchShiftScheduler',
+                      'ScaledLinearScheduler']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/model/diffusion/diffusions.py

@@ -1,28 +1,31 @@
-import os
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 import math
-import torch
+import os
 from collections import OrderedDict
 
-from scepter.modules.utils.config import dict_to_yaml, Config
+import torch
+from tqdm import trange
+
+from scepter.modules.model.registry import (DIFFUSION_SAMPLERS, DIFFUSIONS,
+                                            NOISE_SCHEDULERS)
+from scepter.modules.utils.config import Config, dict_to_yaml
 from scepter.modules.utils.distribute import we
 from scepter.modules.utils.file_system import FS
-from scepter.modules.model.registry import DIFFUSIONS, NOISE_SCHEDULERS, DIFFUSION_SAMPLERS
-from tqdm import trange
+
 
 @DIFFUSIONS.register_class()
 class BaseDiffusion(object):
     para_dict = {
-        "NOISE_SCHEDULER": {},
-        "SAMPLER_SCHEDULER": {},
-        "MIN_SNR_GAMMA": {
-            "value": None,
-            "description": "The minimum SNR gamma value for the loss function."
-        },
-        "PREDICTION_TYPE": {
-            "value": "eps",
-            "description": "The type of prediction to use for the loss function."
+        'NOISE_SCHEDULER': {},
+        'SAMPLER_SCHEDULER': {},
+        'PREDICTION_TYPE': {
+            'value': 'eps',
+            'description':
+            'The type of prediction to use for the loss function.'
         }
     }
+
     def __init__(self, cfg, logger=None):
         super(BaseDiffusion, self).__init__()
         self.logger = logger
@@ -30,39 +33,56 @@ class BaseDiffusion(object):
         self.init_params()
 
     def init_params(self):
-        self.min_snr_gamma = self.cfg.get("MIN_SNR_GAMMA", None)
-        self.prediction_type = self.cfg.get("PREDICTION_TYPE", "eps")
-        self.noise_scheduler = NOISE_SCHEDULERS.build(self.cfg.NOISE_SCHEDULER, logger=self.logger)
-        self.sampler_scheduler = NOISE_SCHEDULERS.build(self.cfg.get("SAMPLER_SCHEDULER", self.cfg.NOISE_SCHEDULER),
+        self.prediction_type = self.cfg.get('PREDICTION_TYPE', 'eps')
+        self.use_dynamic_cfg = self.cfg.get('USE_DYNAMIC_CFG', False)
+        self.noise_scheduler = NOISE_SCHEDULERS.build(self.cfg.NOISE_SCHEDULER,
+                                                      logger=self.logger)
+        self.sampler_scheduler = NOISE_SCHEDULERS.build(self.cfg.get(
+            'SAMPLER_SCHEDULER', self.cfg.NOISE_SCHEDULER),
                                                         logger=self.logger)
         self.num_timesteps = self.noise_scheduler.num_timesteps
-        if self.cfg.have("WORK_DIR") and we.rank == 0:
-            schedule_visualization = os.path.join(self.cfg.WORK_DIR, "noise_schedule.png")
+        if self.cfg.have('WORK_DIR') and we.rank == 0:
+            schedule_visualization = os.path.join(self.cfg.WORK_DIR,
+                                                  'noise_schedule.png')
             with FS.put_to(schedule_visualization) as local_path:
                 self.noise_scheduler.plot_noise_sampling_map(local_path)
-            schedule_visualization = os.path.join(self.cfg.WORK_DIR, "sampler_schedule.png")
+            schedule_visualization = os.path.join(self.cfg.WORK_DIR,
+                                                  'sampler_schedule.png')
             with FS.put_to(schedule_visualization) as local_path:
                 self.sampler_scheduler.plot_noise_sampling_map(local_path)
 
-
-    def sample(self, noise, model, model_kwargs={}, steps=20, sampler=None, use_dynamic_cfg=False, guide_scale=None, guide_rescale=None,
-               show_progress=False, return_intermediate=None, intermediate_callback=None, **kwargs):
+    def sample(self,
+               noise,
+               model,
+               model_kwargs={},
+               steps=20,
+               sampler=None,
+               guide_scale=None,
+               guide_rescale=None,
+               show_progress=False,
+               return_intermediate=None,
+               intermediate_callback=None,
+               reverse_scale = -1.,
+               x = None,
+               **kwargs):
         assert isinstance(steps, (int, torch.LongTensor))
         assert return_intermediate in (None, 'x0', 'xt')
         assert isinstance(sampler, (str, dict, Config))
         intermediates = []
 
-        def callback_fn(x_t, t, sigma=None, alpha=None):
+        def callback_fn(x_t, t, sigma=None, alpha_bar=None):
             timestamp = t
             t = t.repeat(len(x_t)).round().long().to(x_t.device)
             sigma = sigma.repeat(len(x_t), *([1] * (len(sigma.shape) - 1)))
-            alpha = alpha.repeat(len(x_t), *([1] * (len(alpha.shape) - 1)))
+            alpha_bar = alpha_bar.repeat(len(x_t), *([1] * (len(alpha_bar.shape) - 1)))
 
             if guide_scale is None or guide_scale == 1.0:
                 out = model(x=x_t, t=t, **model_kwargs)
             else:
-                if use_dynamic_cfg:
-                    guidance_scale = 1 + guide_scale * ((1 - math.cos(math.pi * ((steps - timestamp.item()) / steps) ** 5.0)) / 2)
+                if self.use_dynamic_cfg:
+                    guidance_scale = 1 + guide_scale * (
+                        (1 - math.cos(math.pi * (
+                            (steps - timestamp.item()) / steps)**5.0)) / 2)
                 else:
                     guidance_scale = guide_scale
                 y_out = model(x=x_t, t=t, **model_kwargs[0])
@@ -78,15 +98,12 @@ class BaseDiffusion(object):
             if self.prediction_type == 'x0':
                 x0 = out
             elif self.prediction_type == 'eps':
-                x0 = (x_t - sigma * out) / alpha
+                x0 = (x_t - sigma * out) / alpha_bar
             elif self.prediction_type == 'v':
-                x0 = alpha * x_t - sigma * out
+                x0 = alpha_bar * x_t - sigma * out
             else:
                 raise NotImplementedError(
                     f'prediction_type {self.prediction_type} not implemented')
-
-            # print("torch.sum(y_out):", torch.sum(y_out), "torch.sum(u_out):", torch.sum(u_out), "torch.sum(out):",
-            #       torch.sum(out), "torch.sum(x0):", torch.sum(x0), "sigmas", sigma, "alphas", alpha)
             return x0
 
         sampler_ins = self.get_sampler(sampler)
@@ -94,13 +111,14 @@ class BaseDiffusion(object):
         # this is ignored for schnell
         sampler_output = sampler_ins.preprare_sampler(
             noise,
+            x = x,
             steps=steps,
+            reverse_scale= reverse_scale,
             prediction_type=self.prediction_type,
             scheduler_ins=self.sampler_scheduler,
-            callback_fn=callback_fn
-        )
+            callback_fn=callback_fn)
 
-        for _ in trange(steps, disable=not show_progress):
+        for _ in trange(sampler_output.steps, disable=not show_progress):
             trange.desc = sampler_output.msg
             sampler_output = sampler_ins.step(sampler_output)
             if return_intermediate == 'x_0':
@@ -109,53 +127,58 @@ class BaseDiffusion(object):
                 intermediates.append(sampler_output.x_t)
             if intermediate_callback is not None:
                 intermediate_callback(intermediates[-1])
-        return (sampler_output.x_0, intermediates) if return_intermediate is not None else sampler_output.x_0
+        return (sampler_output.x_0, intermediates
+                ) if return_intermediate is not None else sampler_output.x_0
 
-
-    def loss(self, x_0, model, model_kwargs={}, reduction='mean', noise=None, **kwargs):
+    def loss(self,
+             x_0,
+             model,
+             model_kwargs={},
+             reduction='mean',
+             noise=None,
+             **kwargs):
         # use noise scheduler to add noise
         if noise is None:
             noise = torch.randn_like(x_0)
-        schedule_output = self.noise_scheduler.add_noise(x_0, noise)
-        x_t, t, sigma, alpha = schedule_output.x_t, schedule_output.t, schedule_output.sigma, schedule_output.alpha
+        schedule_output = self.noise_scheduler.add_noise(x_0, noise, **kwargs)
+        x_t, t, sigma, alpha_bar = schedule_output.x_t, schedule_output.t, schedule_output.sigma, schedule_output.alpha_bar
         out = model(x=x_t, t=t, **model_kwargs)
 
         # mse loss
         target = {
             'eps': noise,
             'x0': x_0,
-            'v': alpha * noise - sigma * x_0
+            'v': alpha_bar * noise - sigma * x_0
         }[self.prediction_type]
 
         loss = (out - target).pow(2)
         if reduction == 'mean':
             loss = loss.flatten(1).mean(dim=1)
-
-        if self.min_snr_gamma is not None:
-            alphas = self.noise_scheduler.alphas.to(x_0.device)[t]
-            sigmas = self.noise_scheduler.sigmas.pow(2).to(x_0.device)[t]
-            snrs = (alphas / sigmas).clamp(min=1e-20)
-            min_snrs = snrs.clamp(max=self.min_snr_gamma)
-            weights = min_snrs / snrs
-        else:
-            weights = 1
-
-        loss = loss * weights
         return loss
 
     def get_sampler(self, sampler):
         if isinstance(sampler, str):
-            if sampler not in DIFFUSION_SAMPLERS.class_map:
+            from scepter.modules.utils.import_utils import LazyImportModule
+            if (not LazyImportModule.get_module_type(('DIFFUSION_SAMPLERS', sampler))) and (
+                    sampler not in DIFFUSION_SAMPLERS.class_map):
                 if self.logger is not None:
-                    self.logger.info(f"{sampler} not in the defined samplers list {DIFFUSION_SAMPLERS.class_map.keys()}")
+                    self.logger.info(
+                        f'{sampler} not in the defined samplers list.'
+                    )
                 else:
-                    print(f"{sampler} not in the defined samplers list {DIFFUSION_SAMPLERS.class_map.keys()}")
+                    print(
+                        f'{sampler} not in the defined samplers list.'
+                    )
                 return None
-            sampler_cfg = Config(cfg_dict={"NAME": sampler}, load=False)
-            sampler_ins = DIFFUSION_SAMPLERS.build(sampler_cfg, logger=self.logger)
+            sampler_cfg = Config(cfg_dict={'NAME': sampler}, load=False)
+            sampler_ins = DIFFUSION_SAMPLERS.build(sampler_cfg,
+                                                   logger=self.logger)
         elif isinstance(sampler, (Config, dict, OrderedDict)):
             if isinstance(sampler, (dict, OrderedDict)):
-                sampler = Config(cfg_dict={k.upper():v for k, v in dict(sampler).items()}, load=False)
+                sampler = Config(
+                    cfg_dict={k.upper(): v
+                              for k, v in dict(sampler).items()},
+                    load=False)
             sampler_ins = DIFFUSION_SAMPLERS.build(sampler, logger=self.logger)
         else:
             raise NotImplementedError
@@ -171,49 +194,47 @@ class BaseDiffusion(object):
                             BaseDiffusion.para_dict,
                             set_name=True)
 
+
 @DIFFUSIONS.register_class()
 class DiffusionFluxRF(BaseDiffusion):
     para_dict = {
-        "PREDICTION_TYPE": {
-            "value": "raw",
-            "description": "The type of prediction to use for the loss function."
+        'PREDICTION_TYPE': {
+            'value': 'raw',
+            'description':
+            'The type of prediction to use for the loss function.'
         }
     }
     para_dict.update(BaseDiffusion.para_dict)
+
     def __init__(self, cfg, logger=None):
         super(DiffusionFluxRF, self).__init__(cfg, logger=logger)
-        self.prediction_type = self.cfg.get("PREDICTION_TYPE", "raw")
+        self.prediction_type = self.cfg.get('PREDICTION_TYPE', 'raw')
 
-    def loss(self, x_0, model, model_kwargs={}, reduction='mean', noise=None, **kwargs):
+    def loss(self,
+             x_0,
+             model,
+             model_kwargs={},
+             reduction='mean',
+             noise=None,
+             **kwargs):
         if noise is None:
             noise = torch.randn_like(x_0)
-        schedule_output = self.noise_scheduler.add_noise(x_0, noise)
+        schedule_output = self.noise_scheduler.add_noise(x_0, noise, **kwargs)
         x_t, t, sigma = schedule_output.x_t, schedule_output.t, schedule_output.sigma
         out = model(x=x_t, t=sigma, **model_kwargs)
         # raw
-        if self.prediction_type == "raw":
+        if self.prediction_type == 'raw':
             target = noise - x_0
             out = out
-        elif self.prediction_type == "sigma_scaled":
+        elif self.prediction_type == 'sigma_scaled':
             target = x_0
             out = out * (-sigma) + x_t
         else:
             raise NotImplementedError
 
-        loss = (target - out) ** 2
+        loss = (target - out)**2
         if reduction == 'mean':
             loss = loss.flatten(1).mean(dim=1)
-
-        if self.min_snr_gamma is not None:
-            alphas = self.noise_scheduler.alphas.to(x_0.device)[t]
-            sigmas = self.noise_scheduler.sigmas.pow(2).to(x_0.device)[t]
-            snrs = (alphas / sigmas).clamp(min=1e-20)
-            min_snrs = snrs.clamp(max=self.min_snr_gamma)
-            weights = min_snrs / snrs
-        else:
-            weights = 1
-
-        loss = loss * weights
         return loss
 
     @torch.no_grad()
@@ -222,10 +243,12 @@ class DiffusionFluxRF(BaseDiffusion):
                model,
                model_kwargs={},
                steps=20,
-               sampler = None,
+               sampler=None,
                show_progress=False,
                return_intermediate=None,
                intermediate_callback=None,
+               reverse_scale=-1.,
+               x=None,
                **kwargs):
         # sanity check
         assert isinstance(steps, (int, torch.LongTensor))
@@ -233,9 +256,12 @@ class DiffusionFluxRF(BaseDiffusion):
         assert isinstance(sampler, (str, dict, Config))
         intermediates = []
 
-        def callback_fn(x_t, t, sigma):
-            sigma = torch.full((x_t.shape[0],), sigma, dtype=x_t.dtype, device=x_t.device)
-            x_0 = model(x = x_t, t=sigma, **model_kwargs)
+        def callback_fn(x_t, t, sigma=None, alpha_bar=None):
+            sigma = torch.full((x_t.shape[0], ),
+                               sigma,
+                               dtype=x_t.dtype,
+                               device=x_t.device)
+            x_0 = model(x=x_t, t=sigma, **model_kwargs)
             return x_0
 
         sampler_ins = self.get_sampler(sampler)
@@ -243,13 +269,14 @@ class DiffusionFluxRF(BaseDiffusion):
         # this is ignored for schnell
         sampler_output = sampler_ins.preprare_sampler(
             noise,
-            steps = steps,
+            x=x,
+            steps=steps,
+            reverse_scale=reverse_scale,
             prediction_type=self.prediction_type,
-            scheduler_ins = self.sampler_scheduler,
-            callback_fn=callback_fn
-            )
+            scheduler_ins=self.sampler_scheduler,
+            callback_fn=callback_fn)
 
-        for _ in trange(steps, disable=not show_progress):
+        for _ in trange(sampler_output.steps, disable=not show_progress):
             trange.desc = sampler_output.msg
             sampler_output = sampler_ins.step(sampler_output)
             if return_intermediate == 'x_0':
@@ -258,11 +285,12 @@ class DiffusionFluxRF(BaseDiffusion):
                 intermediates.append(sampler_output.x_t)
             if intermediate_callback is not None:
                 intermediate_callback(intermediates[-1])
-        return (sampler_output.x_0, intermediates) if return_intermediate is not None else sampler_output.x_t
+        return (sampler_output.x_0, intermediates
+                ) if return_intermediate is not None else sampler_output.x_t
 
     @staticmethod
     def get_config_template():
         return dict_to_yaml('DIFFUSIONS',
                             __class__.__name__,
                             DiffusionFluxRF.para_dict,
-                            set_name=True)
+                            set_name=True)

scepter/modules/model/diffusion/samplers.py

@@ -1,32 +1,42 @@
-from dataclasses import dataclass, field
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from dataclasses import dataclass
+
 import torch
-from scepter.modules.utils.config import dict_to_yaml
+
 from scepter.modules.model.registry import DIFFUSION_SAMPLERS
+from scepter.modules.utils.config import dict_to_yaml
+
 from .util import _i
 
+
 @dataclass
 class SamplerOutput(object):
     callback_fn: callable
     prediction_type: str
     alphas: torch.Tensor
+    alphas_bar: torch.Tensor
     betas: torch.Tensor
     sigmas: torch.Tensor
     alphas_init: torch.Tensor
+    alphas_bar_init: torch.Tensor
     betas_init: torch.Tensor
     sigmas_init: torch.Tensor
     ts: torch.Tensor
     x_t: torch.Tensor
     x_0: torch.Tensor
     step: int
+    steps: int
     msg: str
 
     def add_custom_field(self, key: str, value) -> None:
         self.__setattr__(key, value)
 
 
-@DIFFUSION_SAMPLERS.register_class("base")
+@DIFFUSION_SAMPLERS.register_class('base')
 class BaseDiffusionSampler(object):
     para_dict = {}
+
     def __init__(self, cfg, logger=None):
         super(BaseDiffusionSampler, self).__init__()
         self.logger = logger
@@ -34,13 +44,15 @@ class BaseDiffusionSampler(object):
         self.init_params()
 
     def init_params(self):
-        self.discretization_type = self.cfg.get("DISCRETIZATION_TYPE", "linspace")
-        self.discard_penultimate_step = self.cfg.get("DISCARD_PENULTIMATE_STEP", False)
-        self.free_steps = self.cfg.get("FREE_STEPS", None)
-        self.t_max = self.cfg.get("T_MAX", None)
-        self.t_min = self.cfg.get("T_MIN", None)
+        self.discretization_type = self.cfg.get('DISCRETIZATION_TYPE',
+                                                'linspace')
+        self.discard_penultimate_step = self.cfg.get(
+            'DISCARD_PENULTIMATE_STEP', False)
+        self.free_steps = self.cfg.get('FREE_STEPS', None)
+        self.t_max = self.cfg.get('T_MAX', None)
+        self.t_min = self.cfg.get('T_MIN', None)
 
-    def discretization(self, steps=20, num_timesteps=1000, **kwargs):
+    def discretization(self, steps=20, num_timesteps=1000, reverse_scale = -1., **kwargs):
         # get timesteps
         if isinstance(steps, int):
             steps += 1 if self.discard_penultimate_step else 0
@@ -60,15 +72,29 @@ class BaseDiffusionSampler(object):
                 steps = torch.tensor(self.free_steps)
             else:
                 raise NotImplementedError(
-                    f'{self.discretization_type} discretization not implemented')
+                    f'{self.discretization_type} discretization not implemented'
+                )
             steps = steps.clamp_(t_min, t_max)
         elif isinstance(steps, list):
             steps = torch.tensor(steps)
-        timesteps = torch.as_tensor(steps, dtype=torch.float32)
-        return timesteps
+        if reverse_scale >=0:
+            img2img_step = int((1 - reverse_scale) * len(steps))
+            timesteps = torch.as_tensor(steps[img2img_step:], dtype=torch.float32)
+            return timesteps
+        return torch.as_tensor(steps, dtype=torch.float32)
 
-    def preprare_sampler(self, noise, steps=20, scheduler_ins=None, prediction_type="",
-                         sigmas=None, betas=None, alphas=None, callback_fn = None,
+    def preprare_sampler(self,
+                         noise,
+                         x=None,
+                         steps=20,
+                         reverse_scale=-1.,
+                         scheduler_ins=None,
+                         prediction_type='',
+                         sigmas=None,
+                         betas=None,
+                         alphas=None,
+                         alphas_bar=None,
+                         callback_fn=None,
                          **kwargs):
         '''
             1. Control the model's inputs and outputs externally in the solver by callback_fn,
@@ -79,34 +105,57 @@ class BaseDiffusionSampler(object):
             4. To ensure the safety of threading, use the instance of SamplerOutput as the manager,
             which manage all necessary information.
         '''
+        if reverse_scale >= 0:
+            assert x is not None
         num_timesteps = scheduler_ins.num_timesteps if scheduler_ins is not None else 1000
-        timestamps = self.discretization(steps, num_timesteps=num_timesteps, **kwargs)
-        alphas = scheduler_ins.t_to_alpha(timestamps, **kwargs) if scheduler_ins is not None else alphas
-        betas = scheduler_ins.t_to_beta(timestamps, **kwargs) if scheduler_ins is not None else betas
-        sigmas = scheduler_ins.t_to_sigma(timestamps, **kwargs) if scheduler_ins is not None else sigmas
-        alphas_init = scheduler_ins.t_to_alpha_init(timestamps, **kwargs) if scheduler_ins is not None else alphas
-        betas_init = scheduler_ins.t_to_beta_init(timestamps, **kwargs) if scheduler_ins is not None else betas
-        sigmas_init = scheduler_ins.t_to_sigma_init(timestamps, **kwargs) if scheduler_ins is not None else sigmas
+        timestamps = self.discretization(steps,
+                                         num_timesteps=num_timesteps,
+                                         reverse_scale=reverse_scale,
+                                         **kwargs)
+        alphas = scheduler_ins.t_to_alpha(
+            timestamps, **kwargs) if scheduler_ins is not None else alphas
+        alphas_bar = scheduler_ins.t_to_alpha_bar(
+            timestamps, **kwargs) if scheduler_ins is not None else alphas_bar
+        betas = scheduler_ins.t_to_beta(
+            timestamps, **kwargs) if scheduler_ins is not None else betas
+        sigmas = scheduler_ins.t_to_sigma(
+            timestamps, **kwargs) if scheduler_ins is not None else sigmas
+        alphas_init = scheduler_ins.t_to_alpha_init(
+            timestamps, **kwargs) if scheduler_ins is not None else alphas
 
-        output = SamplerOutput(
-            callback_fn=callback_fn,
-            prediction_type=prediction_type,
-            alphas=alphas,
-            betas=betas,
-            sigmas=sigmas,
-            alphas_init=alphas_init,
-            betas_init=betas_init,
-            sigmas_init=sigmas_init,
-            ts=timestamps,
-            x_t=noise,
-            x_0=noise,
-            step=0,
-            msg=f"step 0"
-        )
+        alphas_bar_init = scheduler_ins.t_to_alpha_bar_init(
+            timestamps, **kwargs) if scheduler_ins is not None else alphas_bar
+
+        betas_init = scheduler_ins.t_to_beta_init(
+            timestamps, **kwargs) if scheduler_ins is not None else betas
+        sigmas_init = scheduler_ins.t_to_sigma_init(
+            timestamps, **kwargs) if scheduler_ins is not None else sigmas
+        if reverse_scale >= 0:
+            x_t = x_0 = scheduler_ins.add_noise(x, noise=noise, t=timestamps[0].repeat(x.size(0)).to(x.device)).x_t if len(timestamps) > 0 else x
+        else:
+            x_t = x_0 = noise
+        # Consider the sigma's list is from sigma_ to zero. the steps equal to len(timestamps)
+        output = SamplerOutput(callback_fn=callback_fn,
+                               prediction_type=prediction_type,
+                               alphas=alphas,
+                               alphas_bar=alphas_bar,
+                               betas=betas,
+                               sigmas=sigmas,
+                               alphas_init=alphas_init,
+                               alphas_bar_init=alphas_bar_init,
+                               betas_init=betas_init,
+                               sigmas_init=sigmas_init,
+                               ts=timestamps,
+                               x_t=x_t,
+                               x_0=x_0,
+                               step=0,
+                               msg='step 0',
+                               steps=len(timestamps) - 1)
         return output
 
     def step(self, sampler_ouput):
-        raise NotImplementedError(f'DiffusionSampler step function not implemented')
+        raise NotImplementedError(
+            'DiffusionSampler step function not implemented')
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}' + ' ' + super().__repr__()
@@ -119,29 +168,51 @@ class BaseDiffusionSampler(object):
                             set_name=True)
 
 
-@DIFFUSION_SAMPLERS.register_class("eluer")
+@DIFFUSION_SAMPLERS.register_class('eluer')
 class EulerSampler(BaseDiffusionSampler):
     def step(self, sampler_ouput):
         pass
 
 
-@DIFFUSION_SAMPLERS.register_class("ddim")
+@DIFFUSION_SAMPLERS.register_class('ddim')
 class DDIMSampler(BaseDiffusionSampler):
-
     def init_params(self):
         super().init_params()
         self.eta = self.cfg.get('ETA', 0.)
-        self.discretization_type = self.cfg.get("DISCRETIZATION_TYPE", "trailing")
+        self.discretization_type = self.cfg.get('DISCRETIZATION_TYPE',
+                                                'trailing')
 
-    def preprare_sampler(self, noise, steps=20, scheduler_ins=None, prediction_type="",
-                         sigmas=None, betas=None, alphas=None, callback_fn = None,
+    def preprare_sampler(self,
+                         noise,
+                         x=None,
+                         steps=20,
+                         reverse_scale = -1.,
+                         scheduler_ins=None,
+                         prediction_type='',
+                         sigmas=None,
+                         betas=None,
+                         alphas=None,
+                         alphas_bar=None,
+                         callback_fn=None,
                          **kwargs):
-        output = super().preprare_sampler(noise, steps, scheduler_ins, prediction_type, sigmas, betas, alphas, callback_fn, **kwargs)
+        output = super().preprare_sampler(noise,
+                                          x = x,
+                                          steps = steps,
+                                          reverse_scale = reverse_scale,
+                                          scheduler_ins = scheduler_ins,
+                                          prediction_type = prediction_type,
+                                          sigmas = sigmas,
+                                          betas = betas,
+                                          alphas = alphas,
+                                          alphas_bar = alphas_bar,
+                                          callback_fn = callback_fn,
+                                          **kwargs)
         sigmas = output.sigmas
         sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
         sigmas_vp = (sigmas**2 / (1 + sigmas**2))**0.5
         sigmas_vp[sigmas == float('inf')] = 1.
         output.add_custom_field('sigmas_vp', sigmas_vp)
+        output.steps += 1
         return output
 
     def step(self, sampler_output):
@@ -149,46 +220,68 @@ class DDIMSampler(BaseDiffusionSampler):
         step = sampler_output.step
         t = sampler_output.ts[step]
         sigmas_vp = sampler_output.sigmas_vp.to(x_t.device)
-        alpha_init = _i(sampler_output.alphas_init, step, x_t[:1])
+        alpha_bar_init = _i(sampler_output.alphas_bar_init, step, x_t[:1])
         sigma_init = _i(sampler_output.sigmas_init, step, x_t[:1])
 
-        x = sampler_output.callback_fn(x_t, t, sigma_init, alpha_init)
-        noise_factor = self.eta * (sigmas_vp[step + 1] ** 2 / sigmas_vp[step] ** 2 *
-                              (1 - (1 - sigmas_vp[step] ** 2) /
-                               (1 - sigmas_vp[step + 1] ** 2)))
-        d = (x_t - (1 - sigmas_vp[step] ** 2) ** 0.5 * x) / sigmas_vp[step]
+        x = sampler_output.callback_fn(x_t, t, sigma_init, alpha_bar_init)
+        noise_factor = self.eta * (sigmas_vp[step + 1]**2 /
+                                   sigmas_vp[step]**2 *
+                                   (1 - (1 - sigmas_vp[step]**2) /
+                                    (1 - sigmas_vp[step + 1]**2)))
+        d = (x_t - (1 - sigmas_vp[step]**2)**0.5 * x) / sigmas_vp[step]
         x = (1 - sigmas_vp[step + 1] ** 2) ** 0.5 * x + \
             (sigmas_vp[step + 1] ** 2 - noise_factor ** 2) ** 0.5 * d
         sampler_output.x_0 = x
         if sigmas_vp[step + 1] > 0:
             x += noise_factor * torch.randn_like(x)
-        # print("i:", step, "sigma_init:", sigma_init, "alpha_init", alpha_init, "sigmas_vp[i]", sigmas_vp[step], "torch.sum(x_0):", torch.sum(x_0), "torch.sum(x):", torch.sum(x))
         sampler_output.x_t = x
         sampler_output.step += 1
         sampler_output.msg = f'step {step}'
         return sampler_output
 
 
-@DIFFUSION_SAMPLERS.register_class("flow_eluer")
+@DIFFUSION_SAMPLERS.register_class('flow_euler')
 class FlowEluerSampler(BaseDiffusionSampler):
-    def preprare_sampler(self, noise, steps=20, scheduler_ins=None, prediction_type="",
-                         sigmas=None, betas=None, alphas=None, callback_fn = None,
+    def preprare_sampler(self,
+                         noise,
+                         x=None,
+                         steps=20,
+                         reverse_scale = -1.,
+                         scheduler_ins=None,
+                         prediction_type='',
+                         sigmas=None,
+                         betas=None,
+                         alphas=None,
+                         alphas_bar=None,
+                         callback_fn=None,
                          **kwargs):
         if noise.ndim == 3:
             seq_len = noise.shape[2] // 4
         else:
             n, _, h, w = noise.shape
             seq_len = (h // 2 * w // 2)
-        kwargs["seq_len"] = seq_len
-        output = super().preprare_sampler(noise, steps, scheduler_ins, prediction_type, sigmas, betas, alphas, callback_fn, **kwargs)
+        kwargs['seq_len'] = seq_len
+        output = super().preprare_sampler(noise,
+                                          x = x,
+                                          steps = steps,
+                                          reverse_scale = reverse_scale,
+                                          scheduler_ins = scheduler_ins,
+                                          prediction_type = prediction_type,
+                                          sigmas = sigmas,
+                                          betas = betas,
+                                          alphas = alphas,
+                                          alphas_bar = alphas_bar,
+                                          callback_fn = callback_fn,
+                                          **kwargs)
         return output
 
     def step(self, sampler_output):
         step = sampler_output.step
         x_t = sampler_output.x_t
-        sigma_curr, sigma_prev = sampler_output.sigmas[step], sampler_output.sigmas[step + 1]
+        sigma_curr, sigma_prev = sampler_output.sigmas[
+            step], sampler_output.sigmas[step + 1]
         prediction_type = sampler_output.prediction_type
-        assert prediction_type in ("raw", "sigma_scaled")
+        assert prediction_type in ('raw', 'sigma_scaled')
         t = sampler_output.ts[step]
         x_0 = sampler_output.callback_fn(x_t, t, sigma_curr)
         x_t = x_t + (sigma_prev - sigma_curr) * x_0
@@ -198,9 +291,13 @@ class FlowEluerSampler(BaseDiffusionSampler):
         sampler_output.msg = f'step {step}, sigma_curr: {sigma_curr}, sigma_prev: {sigma_prev}'
         return sampler_output
 
-    def discretization(self, steps=20, num_timesteps = 1000, **kwargs):
+    def discretization(self, steps=20, num_timesteps=1000, reverse_scale=-1., **kwargs):
         # extra step for zero
         timesteps = torch.linspace(num_timesteps, 0, steps + 1)
+        if reverse_scale >= 0:
+            img2img_step = int((1 - reverse_scale) * len(timesteps))
+            timesteps = timesteps[img2img_step:]
+            return timesteps
         return timesteps
 
     @staticmethod
@@ -208,4 +305,4 @@ class FlowEluerSampler(BaseDiffusionSampler):
         return dict_to_yaml('DIFFUSION_SAMPLERS',
                             __class__.__name__,
                             FlowEluerSampler.para_dict,
-                            set_name=True)
+                            set_name=True)

scepter/modules/model/diffusion/schedules.py

@@ -1,24 +1,28 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 import math
+import random
 from dataclasses import dataclass, field
 from typing import Callable
 
-import torch
 import numpy as np
-
-from scepter.modules.utils.config import dict_to_yaml
-from scepter.modules.utils.math_plot import plot_multi_curves
-from scepter.modules.model.registry import NOISE_SCHEDULERS
+import torch
 from torch import Tensor
 
+from scepter.modules.model.registry import NOISE_SCHEDULERS
+from scepter.modules.utils.config import dict_to_yaml
+from scepter.modules.utils.math_plot import plot_multi_curves
+
 from .util import _i
 
+
 @dataclass
 class ScheduleOutput(object):
     x_t: torch.Tensor
     x_0: torch.Tensor
     t: torch.Tensor
     sigma: torch.Tensor
-    alpha: torch.Tensor
+    alpha_bar: torch.Tensor
     custom_fields: dict = field(default_factory=dict)
 
     def add_custom_field(self, key: str, value) -> None:
@@ -27,12 +31,13 @@ class ScheduleOutput(object):
 
 @NOISE_SCHEDULERS.register_class()
 class BaseNoiseScheduler(object):
-    '''
+    r'''
     In the diffusion model, the parameters related to the noise schedule are alpha, beta,
     and sigma. The following are the definitions of the above three parameters, which should
     be the basic property for the instance of noise scheduler.
         \alpha_{t} = \sqrt{1 - \beta_{t}^2} \alpha is the strength of signal and \beta is the strength of noise
         \sigma_{t} = \sqrt{1 - \overline\alpha} = \sqrt{1 - \prod_{i=1}^{t}\alpha^2_{i}} (P(x_{t}|x_{0}) ~ N(\overline\alpha x_{0}, \sigma^2))
+        \alpha_bar_{t} = \sqrt{\overline\alpha} = \sqrt{\prod_{i=1}^{t}\alpha^2_{i}} (P(x_{t}|x_{0}) ~ N(\overline\alpha x_{0}, \sigma^2))
 
     where sigma_{t} is the var of p(x_{t-1}|x_{t}, x_{0}).
 
@@ -42,54 +47,68 @@ class BaseNoiseScheduler(object):
 
     '''
     para_dict = {
-        "NUM_TIMESTEPS": {
-            "value": 1000,
-            "description": "The number of timesteps for sampling."
+        'NUM_TIMESTEPS': {
+            'value': 1000,
+            'description': 'The number of timesteps for sampling.'
         },
     }
+
     def __init__(self, cfg, logger=None):
         super(BaseNoiseScheduler, self).__init__()
         self.logger = logger
         self.cfg = cfg
         self.init_params()
         self.get_schedule()
-        # self.check_function()
 
     def init_params(self):
-        self.num_timesteps = self.cfg.get("NUM_TIMESTEPS", 1000)
-        self._sample_steps = torch.arange(self.num_timesteps, dtype=torch.float32)
-        self._sigmas, self._betas, self._alphas, self._timesteps = None, None, None, None
+        self.num_timesteps = self.cfg.get('NUM_TIMESTEPS', 1000)
+        self._sample_steps = torch.arange(self.num_timesteps,
+                                          dtype=torch.float32)
+        self._sigmas, self._betas, self._alphas, self._alphas_bar, self._timesteps = None, None, None, None, None
+
     def check_function(self):
-        # for the same t, we should gurantee t_to_sigma and sigma_to_t is aligned
         try:
             predict_timestamps = self.sigma_to_t(self.sigmas)
             predict_sigmas = self.t_to_sigma(self._timesteps)
             diff_sigmas = torch.sum(torch.abs(predict_sigmas - self.sigmas))
-            diff_timestamps = torch.sum(torch.abs(predict_timestamps - self._timesteps))
+            diff_timestamps = torch.sum(
+                torch.abs(predict_timestamps - self._timesteps))
             if diff_sigmas > 1e-3 or diff_timestamps > 1:
-                self.logger.info(f"The noise scheduler {self.__class__.__name__} is not correct, "
-                                 f"please check the function sigma_to_t or t_to_sigma."
-                                 f"Info: diff sigmas {diff_sigmas}, diff timestamps {diff_timestamps}")
-                raise "The noise scheduler checked failed."
+                self.logger.info(
+                    f'The noise scheduler {self.__class__.__name__} is not correct, '
+                    f'please check the function sigma_to_t or t_to_sigma.'
+                    f'Info: diff sigmas {diff_sigmas}, diff timestamps {diff_timestamps}'
+                )
+                raise 'The noise scheduler checked failed.'
             else:
-                self.logger.info(f"The noise scheduler {self.__class__.__name__} is checked and passed.")
+                self.logger.info(
+                    f'The noise scheduler {self.__class__.__name__} is checked and passed.'
+                )
         except Exception as e:
             if isinstance(e, NotImplementedError):
-                self.logger.info("Not implemented function sigma_to_t or t_to_sigma, skip check.")
+                self.logger.info(
+                    'Not implemented function sigma_to_t or t_to_sigma, skip check.'
+                )
             else:
-                self.logger.info(f"The noise scheduler {self.__class__.__name__} is not correct, "
-                                   f"please check the function sigma_to_t or t_to_sigma. Error: {e}")
+                self.logger.info(
+                    f'The noise scheduler {self.__class__.__name__} is not correct, '
+                    f'please check the function sigma_to_t or t_to_sigma. Error: {e}'
+                )
                 raise e
 
     def get_schedule(self):
-        raise NotImplementedError(f'NoiseScheduler get_schedule function not implemented')
+        raise NotImplementedError(
+            'NoiseScheduler get_schedule function not implemented')
 
     def square_betas_to_sigmas(self, square_betas):
         return torch.sqrt(1 - torch.cumprod(1 - square_betas, dim=0))
+
     def sigmas_to_square_betas(self, sigmas):
-        square_alphas = 1 - sigmas ** 2
-        betas = 1 - torch.cat([square_alphas[:1], square_alphas[1:] / square_alphas[:-1]])
+        square_alphas = 1 - sigmas**2
+        betas = 1 - torch.cat(
+            [square_alphas[:1], square_alphas[1:] / square_alphas[:-1]])
         return betas
+
     def sigma_to_t(self, sigma, **kwargs):
         if sigma == float('inf'):
             t = torch.full_like(sigma, len(self._sigmas) - 1)
@@ -109,6 +128,7 @@ class BaseNoiseScheduler(object):
         if t.ndim == 0:
             t = t.unsqueeze(0)
         return t
+
     def t_to_sigma(self, t, **kwargs):
         t = t.float()
         low_idx, high_idx, w = t.floor().long(), t.ceil().long(), t.frac()
@@ -124,33 +144,51 @@ class BaseNoiseScheduler(object):
         square_beta = self.sigmas_to_square_betas(sigma)
         return torch.sqrt(1 - square_beta)
 
+    def t_to_alpha_bar(self, t, **kwargs):
+        sigma = self.t_to_sigma(t)
+        return torch.sqrt(1 - sigma**2)
+
     def t_to_beta(self, t, **kwargs):
         sigma = self.t_to_sigma(t)
         square_beta = self.sigmas_to_square_betas(sigma)
         return torch.sqrt(square_beta)
 
-    def add_noise(self, x_0, noise = None, t = None):
+    def add_noise(self, x_0, noise=None, t=None, **kwargs):
         if t is None:
-            t = torch.randint(0, self.num_timesteps, (x_0.shape[0],), device=x_0.device).long()
-        alpha = _i(self.alphas, t, x_0)
+            t = torch.randint(0,
+                              self.num_timesteps, (x_0.shape[0], ),
+                              device=x_0.device).long()
+        alpha = _i(self.alphas_bar, t, x_0)
         sigma = _i(self.sigmas, t, x_0)
         x_t = alpha * x_0 + sigma * noise
 
-        return ScheduleOutput(x_0 = x_0, x_t = x_t, t = t, alpha=alpha, sigma=sigma)
+        return ScheduleOutput(x_0=x_0, x_t=x_t, t=t, alpha_bar=alpha, sigma=sigma)
 
     def t_to_alpha_init(self, t, **kwargs):
         indices = t.long()
         indices[indices >= self.num_timesteps] = self.num_timesteps - 1
         timesteps = self.timesteps.to(t)[indices]
-        step_indices = [(self.timesteps.to(t) == t).nonzero().item() for t in timesteps]
+        step_indices = [(self.timesteps.to(t) == t).nonzero().item()
+                        for t in timesteps]
         alpha = self.alphas[step_indices].flatten().to(t)
         return alpha
 
+    def t_to_alpha_bar_init(self, t, **kwargs):
+        indices = t.long()
+        indices[indices >= self.num_timesteps] = self.num_timesteps - 1
+        timesteps = self.timesteps.to(t)[indices]
+        step_indices = [(self.timesteps.to(t) == t).nonzero().item()
+                        for t in timesteps]
+        alpha_bar = self.alphas_bar[step_indices].flatten().to(t)
+        return alpha_bar
+
+
     def t_to_beta_init(self, t, **kwargs):
         indices = t.long()
         indices[indices >= self.num_timesteps] = self.num_timesteps - 1
         timesteps = self.timesteps.to(t)[indices]
-        step_indices = [(self.timesteps.to(t) == t).nonzero().item() for t in timesteps]
+        step_indices = [(self.timesteps.to(t) == t).nonzero().item()
+                        for t in timesteps]
         beta = self.betas[step_indices].flatten().to(t)
         return beta
 
@@ -158,7 +196,8 @@ class BaseNoiseScheduler(object):
         indices = t.long()
         indices[indices >= self.num_timesteps] = self.num_timesteps - 1
         timesteps = self.timesteps.to(t)[indices]
-        step_indices = [(self.timesteps.to(t) == t).nonzero().item() for t in timesteps]
+        step_indices = [(self.timesteps.to(t) == t).nonzero().item()
+                        for t in timesteps]
         sigma = self.sigmas[step_indices].flatten().to(t)
         return sigma
 
@@ -178,9 +217,10 @@ class BaseNoiseScheduler(object):
         # Shift so the last timestep is zero.
         alphas_bar_sqrt -= alphas_bar_sqrt_T
         # Scale so the first timestep is back to the old value.
-        alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+        alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 -
+                                                alphas_bar_sqrt_T)
         # Convert alphas_bar_sqrt to betas
-        alphas_bar = alphas_bar_sqrt ** 2  # Revert sqrt
+        alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
         return alphas_bar
 
     @property
@@ -195,26 +235,40 @@ class BaseNoiseScheduler(object):
     def alphas(self):
         return self._alphas
 
+    @property
+    def alphas_bar(self):
+        return self._alphas_bar
+
     @property
     def timesteps(self):
         return self._timesteps
 
     # plot the noise sampling map
     def plot_noise_sampling_map(self, save_path):
-        y = [
-            {"data": self._sigmas.cpu().numpy(), "label": "sigmas"},
-            {"data": self._betas.cpu().numpy(), "label": "betas"},
-            {"data": self._alphas.cpu().numpy(), "label": "alphas"},
-            {"data": self._timesteps.cpu().numpy()/self.num_timesteps, "label": "timesteps"}
-            ]
+        y = [{
+            'data': self._sigmas.cpu().numpy(),
+            'label': 'sigmas'
+        }, {
+            'data': self._betas.cpu().numpy(),
+            'label': 'betas'
+        }, {
+            'data': self._alphas.cpu().numpy(),
+            'label': 'alphas'
+        }, {
+            'data': self._alphas_bar.cpu().numpy(),
+            'label': 'alphas_bar'
+        },
+            {
+            'data': self._timesteps.cpu().numpy() / self.num_timesteps,
+            'label': 'timesteps'
+        }]
         plot_multi_curves(
             x=self._sample_steps.cpu().numpy(),
             y=y,
             x_label='timesteps',
             y_label=None,
             title=f"{self.__class__.__name__}'s noise sampling map",
-            save_path=save_path
-        )
+            save_path=save_path)
         return save_path
 
     def __repr__(self) -> str:
@@ -231,18 +285,24 @@ class BaseNoiseScheduler(object):
 @NOISE_SCHEDULERS.register_class()
 class ScaledLinearScheduler(BaseNoiseScheduler):
     para_dict = {}
+
     def init_params(self):
         super().init_params()
         self.beta_min = self.cfg.get('BETA_MIN', 0.00085)
         self.beta_max = self.cfg.get('BETA_MAX', 0.012)
         self.snr_shift_scale = self.cfg.get('SNR_SHIFT_SCALE', None)
-        self.rescale_betas_zero_snr = self.cfg.get('RESCALE_BETAS_ZERO_SNR', False)
+        self.rescale_betas_zero_snr = self.cfg.get('RESCALE_BETAS_ZERO_SNR',
+                                                   False)
 
-    def square_betas_to_sigmas(self, square_betas, snr_shift_scale=None, rescale_betas_zero_snr=False):
+    def square_betas_to_sigmas(self,
+                               square_betas,
+                               snr_shift_scale=None,
+                               rescale_betas_zero_snr=False):
         if snr_shift_scale is not None or rescale_betas_zero_snr:
             alphas_cumprod = torch.cumprod(1 - square_betas, dim=0)
             if snr_shift_scale is not None and snr_shift_scale > 0:
-                alphas_cumprod = alphas_cumprod / (snr_shift_scale + (1 - snr_shift_scale) * alphas_cumprod)
+                alphas_cumprod = alphas_cumprod / (
+                    snr_shift_scale + (1 - snr_shift_scale) * alphas_cumprod)
             if rescale_betas_zero_snr:
                 alphas_cumprod = self.rescale_zero_terminal_snr(alphas_cumprod)
             return torch.sqrt(1 - alphas_cumprod)
@@ -250,36 +310,76 @@ class ScaledLinearScheduler(BaseNoiseScheduler):
             return torch.sqrt(1 - torch.cumprod(1 - square_betas, dim=0))
 
     def get_schedule(self):
-        square_betas = torch.linspace(self.beta_min**0.5, self.beta_max**0.5, self.num_timesteps, dtype=torch.float32) ** 2
-        self._sigmas = self.square_betas_to_sigmas(square_betas, self.snr_shift_scale, self.rescale_betas_zero_snr)
+        square_betas = torch.linspace(self.beta_min**0.5,
+                                      self.beta_max**0.5,
+                                      self.num_timesteps,
+                                      dtype=torch.float32)**2
+        self._sigmas = self.square_betas_to_sigmas(square_betas,
+                                                   self.snr_shift_scale,
+                                                   self.rescale_betas_zero_snr)
         self._betas = torch.sqrt(square_betas)
-        self._alphas = torch.sqrt(1 - self._sigmas ** 2)
+        self._alphas = torch.sqrt(1 - square_betas)
+        self._alphas_bar = torch.sqrt(1 - self._sigmas**2)
         self._timesteps = torch.arange(len(self._sigmas), dtype=torch.float32)
 
 
+@NOISE_SCHEDULERS.register_class()
+class LinearScheduler(BaseNoiseScheduler):
+    para_dict = {}
+
+    def init_params(self):
+        super().init_params()
+        self.beta_min = self.cfg.get('BETA_MIN', 0.00085)
+        self.beta_max = self.cfg.get('BETA_MAX', 0.012)
+
+    def betas_to_sigmas(self, betas):
+        return torch.sqrt(1 - torch.cumprod(1 - betas, dim=0))
+
+    def get_schedule(self):
+        betas = torch.linspace(self.beta_min,
+                               self.beta_max,
+                               self.num_timesteps,
+                               dtype=torch.float32)
+        sigmas = self.betas_to_sigmas(betas)
+        self._sigmas = sigmas
+        self._betas = betas
+        self._alphas = torch.sqrt(1 - betas**2)
+        self._alphas_bar = torch.sqrt(1 - sigmas**2)
+        self._timesteps = torch.arange(len(sigmas), dtype=torch.float32)
+
+
 @NOISE_SCHEDULERS.register_class()
 class FlowMatchUniformScheduler(BaseNoiseScheduler):
     def get_schedule(self):
-        timesteps = np.linspace(1, self.num_timesteps, self.num_timesteps, dtype=np.float32).copy()
+        timesteps = np.linspace(1,
+                                self.num_timesteps,
+                                self.num_timesteps,
+                                dtype=np.float32).copy()
         timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
         self._timesteps = timesteps
         self._sigmas = self.t_to_sigma(timesteps)
         self._betas = torch.sqrt(self.sigmas_to_square_betas(self._sigmas))
-        self._alphas = torch.sqrt(1 - self.betas ** 2)
+        self._alphas = torch.sqrt(1 - self._betas**2)
+        self._alphas_bar = torch.sqrt(1 - self._sigmas ** 2)
 
-    def add_noise(self, x_0, noise = None, t = None):
+    def add_noise(self, x_0, noise=None, t=None, **kwargs):
         if t is None:
-            t = torch.rand((x_0.shape[0],), device=x_0.device)
+            t = torch.rand(
+                (x_0.shape[0], ), device=x_0.device) * self.num_timesteps
         sigma = self.t_to_sigma(t)
-        shape = (x_0.size(0),) + (1,) * (x_0.ndim - 1)
+        shape = (x_0.size(0), ) + (1, ) * (x_0.ndim - 1)
         x_t = (1 - sigma.view(shape)) * x_0 + sigma.view(shape) * noise
-        return ScheduleOutput(x_0 = x_0, x_t = x_t, t = t, sigma=sigma, alpha=self.t_to_alpha(t))
+        return ScheduleOutput(x_0=x_0,
+                              x_t=x_t,
+                              t=t,
+                              sigma=sigma,
+                              alpha_bar=self.t_to_alpha_bar(t))
 
     def sigma_to_t(self, sigma, **kwargs):
         return sigma * self.num_timesteps
 
     def t_to_sigma(self, t, **kwargs):
-        return t/self.num_timesteps
+        return t / self.num_timesteps
 
     @staticmethod
     def get_config_template():
@@ -292,21 +392,22 @@ class FlowMatchUniformScheduler(BaseNoiseScheduler):
 @NOISE_SCHEDULERS.register_class()
 class FlowMatchSigmoidScheduler(FlowMatchUniformScheduler):
     para_dict = {
-        "SIGMOID_SCALE": {
-            "value": 1,
-            "description": "The scale for the sigmoid function."
+        'SIGMOID_SCALE': {
+            'value': 1,
+            'description': 'The scale for the sigmoid function.'
         }
     }
+
     def init_params(self):
         super().init_params()
-        self.sigmoid_scale = self.cfg.get("SIGMOID_SCALE", 1)
+        self.sigmoid_scale = self.cfg.get('SIGMOID_SCALE', 1)
 
     def sigma_to_t(self, sigma, **kwargs):
-        t = - torch.log(1/sigma - 1)/self.sigmoid_scale
+        t = -torch.log(1 / sigma - 1) / self.sigmoid_scale
         return t * self.num_timesteps
 
     def t_to_sigma(self, t, **kwargs):
-        return torch.sigmoid(self.sigmoid_scale * t/self.num_timesteps)
+        return torch.sigmoid(self.sigmoid_scale * t / self.num_timesteps)
 
     @staticmethod
     def get_config_template():
@@ -315,41 +416,46 @@ class FlowMatchSigmoidScheduler(FlowMatchUniformScheduler):
                             FlowMatchSigmoidScheduler.para_dict,
                             set_name=True)
 
+
 @NOISE_SCHEDULERS.register_class()
 class FlowMatchShiftScheduler(FlowMatchUniformScheduler):
     para_dict = {
-        "SHIFT": {
-            "value": 3,
-            "description": "The shift factor for the timestamp."
+        'SHIFT': {
+            'value': 3,
+            'description': 'The shift factor for the timestamp.'
         },
-        "SIGMOID_SCALE": {
-            "value": 1,
-            "description": "The scale for the sigmoid function."
+        'SIGMOID_SCALE': {
+            'value': 1,
+            'description': 'The scale for the sigmoid function.'
         }
     }
 
     def init_params(self):
         super().init_params()
-        self.shift = self.cfg.get("SHIFT", 3)
-        self.sigmoid_scale = self.cfg.get("SIGMOID_SCALE", 1)
+        self.shift = self.cfg.get('SHIFT', 3)
+        self.sigmoid_scale = self.cfg.get('SIGMOID_SCALE', 1)
 
-    def add_noise(self, x_0, noise = None, t = None):
+    def add_noise(self, x_0, noise=None, t=None, **kwargs):
         if t is None:
             logits_norm = torch.randn(x_0.shape[0], device=x_0.device)
             logits_norm = logits_norm * self.sigmoid_scale  # larger scale for more uniform sampling
             t = logits_norm.sigmoid() * self.num_timesteps
         sigma = self.t_to_sigma(t)
-        shape = (x_0.size(0),) + (1,) * (x_0.ndim - 1)
+        shape = (x_0.size(0), ) + (1, ) * (x_0.ndim - 1)
         x_t = (1 - sigma.view(shape)) * x_0 + sigma.view(shape) * noise
-        return ScheduleOutput(x_0 = x_0, x_t = x_t, t = t, sigma=sigma, alpha=self.t_to_alpha(t))
+        return ScheduleOutput(x_0=x_0,
+                              x_t=x_t,
+                              t=t,
+                              sigma=sigma,
+                              alpha_bar=self.t_to_alpha_bar(t))
 
     def sigma_to_t(self, sigma, **kwargs):
-        t = sigma/(sigma - self.shift * sigma + self.shift)
+        t = sigma / (sigma - self.shift * sigma + self.shift)
         return t * self.num_timesteps
 
     def t_to_sigma(self, t, **kwargs):
         t = t / self.num_timesteps
-        return  (t * self.shift) / (1 + (self.shift - 1) * t)
+        return (t * self.shift) / (1 + (self.shift - 1) * t)
 
     @staticmethod
     def get_config_template():
@@ -358,74 +464,122 @@ class FlowMatchShiftScheduler(FlowMatchUniformScheduler):
                             FlowMatchShiftScheduler.para_dict,
                             set_name=True)
 
+
 @NOISE_SCHEDULERS.register_class()
 class FlowMatchFluxShiftScheduler(FlowMatchUniformScheduler):
     para_dict = {
-        "SHIFT": {
-            "value": True,
-            "description": "Use timestamp shift or not, default is True."
+        'SHIFT': {
+            'value': True,
+            'description': 'Use timestamp shift or not, default is True.'
         },
-        "SIGMOID_SCALE": {
-            "value": 1,
-            "description": "The scale of sigmoid function for sampling timesteps."
+        'SIGMOID_SCALE': {
+            'value': 1,
+            'description':
+            'The scale of sigmoid function for sampling timesteps.'
         },
-        "BASE_SHIFT": {
-            "value": 0.5,
-            "description": "The base shift factor for the timestamp."
+        'BASE_SHIFT': {
+            'value': 0.5,
+            'description': 'The base shift factor for the timestamp.'
         },
-        "MAX_SHIFT": {
-            "value": 1.15,
-            "description": "The max shift factor for the timestamp."
+        'MAX_SHIFT': {
+            'value': 1.15,
+            'description': 'The max shift factor for the timestamp.'
+        },
+        'PRE_T_SAMPLE': {
+            'value': False,
+            'description': 'Use pre-sampled timesteps or not, default is False.'
+        },
+        'PRE_T_SAMPLE_FOLD': {
+            'value': 1,
+            'description': 'The folds of pre-sampled timesteps.'
         }
     }
 
     def init_params(self):
         super().init_params()
-        self.shift = self.cfg.get("SHIFT", True)
-        self.sigmoid_scale = self.cfg.get("SIGMOID_SCALE", 1)
+        self.shift = self.cfg.get('SHIFT', True)
+        self.sigmoid_scale = self.cfg.get('SIGMOID_SCALE', 1)
         self.base_shift = self.cfg.get('BASE_SHIFT', 0.5)
         self.max_shift = self.cfg.get('MAX_SHIFT', 1.15)
+        self.pre_t_sample = self.cfg.get('PRE_T_SAMPLE', False)
+        self.pre_t_sample_fold = self.cfg.get('PRE_T_SAMPLE_FOLD', 1)
+        if self.pre_t_sample:
+            t = torch.sigmoid(torch.randn((self.num_timesteps * self.pre_t_sample_fold,)))
+            # Scale and reverse the values to go from 1000 to 0
+            timesteps = ((1 - t) * 1000)
+            # Sort the timesteps in descending order
+            self.pre_sample_timesteps, _ = torch.sort(timesteps, descending=True)
+        else:
+            self.pre_sample_timesteps = None
+
+    @property
+    def pre_timesteps(self):
+        fold_id = random.randint(0, self.pre_t_sample_fold - 1)
+        # print("fold_id", fold_id)
+        return self.pre_sample_timesteps[fold_id::self.pre_t_sample_fold]
+
 
     def time_shift(self, mu: float, sigma_scale: float, t: Tensor):
-        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma_scale)
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma_scale)
+
     def sigma_shift(self, mu: float, sigma_scale: float, sigma: Tensor):
-        return 1/(torch.pow((1-sigma) * math.exp(mu)/sigma, sigma_scale) + 1)
+        return 1 / (torch.pow(
+            (1 - sigma) * math.exp(mu) / sigma, sigma_scale) + 1)
 
     def get_lin_function(self,
-            x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
-    ) -> Callable[[float], float]:
+                         x1: float = 256,
+                         y1: float = 0.5,
+                         x2: float = 4096,
+                         y2: float = 1.15) -> Callable[[float], float]:
         m = (y2 - y1) / (x2 - x1)
         b = y1 - m * x1
         return lambda x: m * x + b
 
-    def add_noise(self, x_0, noise = None, t = None):
+    def add_noise(self, x_0, noise=None, t=None, **kwargs):
         if x_0.ndim == 3:
             seq_len = x_0.shape[2] // 4
         else:
             n, _, h, w = x_0.shape
             seq_len = (h // 2 * w // 2)
         if t is None:
-            logits_norm = torch.randn(x_0.shape[0], device=x_0.device)
-            logits_norm = logits_norm * self.sigmoid_scale  # larger scale for more uniform sampling
-            t = logits_norm.sigmoid() * self.num_timesteps
+            if self.pre_t_sample:
+                timestep_indices = torch.randint(
+                     1,
+                    self.num_timesteps - 1,
+                    (x_0.shape[0],)
+                )
+                timestep_indices = timestep_indices.long()
+                t = [self.pre_timesteps[x.item()].to(x_0.device) for x in timestep_indices]
+                t = torch.stack(t, dim=0)
+            else:
+                logits_norm = torch.randn(x_0.shape[0], device=x_0.device)
+                logits_norm = logits_norm * self.sigmoid_scale  # larger scale for more uniform sampling
+                t = logits_norm.sigmoid() * self.num_timesteps
         sigma = self.t_to_sigma(t, seq_len=seq_len)
-        shape = (x_0.size(0),) + (1,) * (x_0.ndim - 1)
+        shape = (x_0.size(0), ) + (1, ) * (x_0.ndim - 1)
+        # print(sigma)
         x_t = (1 - sigma.view(shape)) * x_0 + sigma.view(shape) * noise
-        return ScheduleOutput(x_0 = x_0, x_t = x_t, t = t, sigma=sigma, alpha=self.t_to_alpha(t))
+        return ScheduleOutput(x_0=x_0,
+                              x_t=x_t,
+                              t=t,
+                              sigma=sigma,
+                              alpha_bar=self.t_to_alpha_bar(t))
 
     def sigma_to_t(self, sigma, **kwargs):
         seq_len = kwargs.get('seq_len', 256)
         if self.shift:
-            mu = self.get_lin_function(y1=self.base_shift, y2=self.max_shift)(seq_len)
+            mu = self.get_lin_function(y1=self.base_shift,
+                                       y2=self.max_shift)(seq_len)
             sigma = self.sigma_shift(mu, 1.0, sigma)
         t = torch.as_tensor(sigma, dtype=torch.float32)
         return t * self.num_timesteps
 
     def t_to_sigma(self, t, **kwargs):
         seq_len = kwargs.get('seq_len', 256)
-        t = t/self.num_timesteps
+        t = t / self.num_timesteps
         if self.shift:
-            mu = self.get_lin_function(y1=self.base_shift, y2=self.max_shift)(seq_len)
+            mu = self.get_lin_function(y1=self.base_shift,
+                                       y2=self.max_shift)(seq_len)
             t = self.time_shift(mu, 1.0, t)
         sigma = torch.as_tensor(t, dtype=torch.float32)
         return sigma
@@ -437,71 +591,95 @@ class FlowMatchFluxShiftScheduler(FlowMatchUniformScheduler):
                             FlowMatchFluxShiftScheduler.para_dict,
                             set_name=True)
 
+
 @NOISE_SCHEDULERS.register_class()
 class FlowMatchSigmaScheduler(FlowMatchUniformScheduler):
     para_dict = {
-       "WEIGHTING_SCHEME" : {
-            "value": "logit_normal",
-            "description": "The weighting scheme for sampling timesteps, "
-                           "choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']."
+        'WEIGHTING_SCHEME': {
+            'value':
+            'logit_normal',
+            'description':
+            'The weighting scheme for sampling timesteps, '
+            "choose from ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none']."
         },
-        "SHIFT": {
-            "value": 3.0,
-            "description": "The shift factor for the timestamp."
+        'SHIFT': {
+            'value': 3.0,
+            'description': 'The shift factor for the timestamp.'
         },
-        "LOGIT_MEAN" : {
-            "value": 0.0,
-            "description": "The mean of the logit distribution for sampling timesteps."
+        'LOGIT_MEAN': {
+            'value':
+            0.0,
+            'description':
+            'The mean of the logit distribution for sampling timesteps.'
         },
-        "LOGIT_STD" : {
-            "value": 1.0,
-            "description": "The standard deviation of the logit distribution for sampling timesteps."
+        'LOGIT_STD': {
+            'value':
+            1.0,
+            'description':
+            'The standard deviation of the logit distribution for sampling timesteps.'
         },
-        "MODE_SCALE" : {
-            "value": 1.29,
-            "description": "The scale factor for the mode of the logit distribution for sampling timesteps."
+        'MODE_SCALE': {
+            'value':
+            1.29,
+            'description':
+            'The scale factor for the mode of the logit distribution for sampling timesteps.'
         }
     }
 
     def init_params(self):
         super().init_params()
-        self.weighting_scheme = self.cfg.get("WEIGHTING_SCHEME", "logit_normal")
-        self.logit_mean = self.cfg.get("LOGIT_MEAN", 0.0)
-        self.logit_std = self.cfg.get("LOGIT_STD", 1.0)
-        self.mode_scale = self.cfg.get("MODE_SCALE", 1.29)
-        self.shift = self.cfg.get("SHIFT", 1.0)
+        self.weighting_scheme = self.cfg.get('WEIGHTING_SCHEME',
+                                             'logit_normal')
+        self.logit_mean = self.cfg.get('LOGIT_MEAN', 0.0)
+        self.logit_std = self.cfg.get('LOGIT_STD', 1.0)
+        self.mode_scale = self.cfg.get('MODE_SCALE', 1.29)
+        self.shift = self.cfg.get('SHIFT', 1.0)
 
     def get_schedule(self):
-        timesteps = np.linspace(1, self.num_timesteps, self.num_timesteps, dtype=np.float32).copy()
+        timesteps = np.linspace(1,
+                                self.num_timesteps,
+                                self.num_timesteps,
+                                dtype=np.float32).copy()
         timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
         self._timesteps = timesteps
         timesteps = timesteps / self.num_timesteps
-        self._sigmas = self.shift * timesteps / (1 + (self.shift - 1) * timesteps)
+        self._sigmas = self.shift * timesteps / (1 +
+                                                 (self.shift - 1) * timesteps)
         self._betas = torch.sqrt(self.sigmas_to_square_betas(self._sigmas))
-        self._alphas = torch.sqrt(1 - self.betas ** 2)
+        self._alphas = torch.sqrt(1 - self.betas**2)
+        self._alphas_bar = torch.sqrt(1 - self._sigmas ** 2)
 
-    def add_noise(self, x_0, noise=None, t=None):
+    def add_noise(self, x_0, noise=None, t=None, **kwargs):
         if t is None:
-            if self.weighting_scheme == "logit_normal":
-                t = torch.normal(mean=self.logit_mean, std=self.logit_std, size=(x_0.shape[0],), device=x_0.device)
+            if self.weighting_scheme == 'logit_normal':
+                t = torch.normal(mean=self.logit_mean,
+                                 std=self.logit_std,
+                                 size=(x_0.shape[0], ),
+                                 device=x_0.device)
             else:
                 t = torch.rand(x_0.shape[0], device=x_0.device)
-            t = self.compute_density_for_timestep_sampling(t) * self.num_timesteps
+            t = self.compute_density_for_timestep_sampling(
+                t) * self.num_timesteps
         sigma = self.t_to_sigma(t)
-        shape = (x_0.size(0),) + (1,) * (x_0.ndim - 1)
+        shape = (x_0.size(0), ) + (1, ) * (x_0.ndim - 1)
         x_t = (1 - sigma.view(shape)) * x_0 + sigma.view(shape) * noise
-        return ScheduleOutput(x_0=x_0, x_t=x_t, t=t, sigma=sigma, alpha=self.t_to_alpha(t))
+        return ScheduleOutput(x_0=x_0,
+                              x_t=x_t,
+                              t=t,
+                              sigma=sigma,
+                              alpha_bar=self.t_to_alpha_bar(t))
 
     def compute_density_for_timestep_sampling(self, t):
         """Compute the density for sampling the timesteps when doing SD3 training.
         Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
         SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
         """
-        if self.weighting_scheme == "logit_normal":
+        if self.weighting_scheme == 'logit_normal':
             # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
             t = torch.nn.functional.sigmoid(t)
-        elif self.weighting_scheme == "mode":
-            t = 1 - t - self.mode_scale * (torch.cos(math.pi * t / 2) ** 2 - 1 + t)
+        elif self.weighting_scheme == 'mode':
+            t = 1 - t - self.mode_scale * (torch.cos(math.pi * t / 2)**2 - 1 +
+                                           t)
         return t
 
     def sigma_to_t(self, sigma, **kwargs):
@@ -511,7 +689,8 @@ class FlowMatchSigmaScheduler(FlowMatchUniformScheduler):
         indices = t.long()
         indices[indices >= self.num_timesteps] = self.num_timesteps - 1
         timesteps = self.timesteps.to(t)[indices]
-        step_indices = [(self.timesteps.to(t) == t).nonzero().item() for t in timesteps]
+        step_indices = [(self.timesteps.to(t) == t).nonzero().item()
+                        for t in timesteps]
         sigma = self.sigmas[step_indices].flatten().to(t)
         return sigma
 
@@ -526,8 +705,9 @@ class FlowMatchSigmaScheduler(FlowMatchUniformScheduler):
 if __name__ == '__main__':
     from scepter.modules.utils.config import Config
     cfg = Config(cfg_dict={
-        "NAME": "FlowMatchShiftScheduler",
-        "SHIFT": 1.15
-    }, load=False)
+        'NAME': 'FlowMatchShiftScheduler',
+        'SHIFT': 1.15
+    },
+                 load=False)
 
     scheduler = NOISE_SCHEDULERS.build(cfg)

scepter/modules/model/diffusion/util.py

@@ -1,10 +1,13 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
 import torch
 
+
 def _i(tensor, t, x):
     """
     Index tensor using t and format the output according to x.
     """
-    shape = (x.size(0),) + (1,) * (x.ndim - 1)
+    shape = (x.size(0), ) + (1, ) * (x.ndim - 1)
     if isinstance(t, torch.Tensor):
         t = t.to(tensor.device)
-    return tensor[t].view(shape).to(x.device)
+    return tensor[t].view(shape).to(x.device)

scepter/modules/model/embedder/__init__.py

@@ -1,8 +1,30 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+from scepter.modules.utils.import_utils import LazyImportModule
 
-from scepter.modules.model.embedder.embedder import (
-    ConcatTimestepEmbedderND, FrozenCLIPEmbedder, FrozenCLIPEmbedder2,
-    FrozenOpenCLIPEmbedder, FrozenOpenCLIPEmbedder2, GeneralConditioner,
-    IPAdapterPlusEmbedder, RefCrossEmbedder, SD3TextEmbedder, T5EmbedderHF)
-from scepter.modules.model.embedder.flux_embedder import HFEmbedder
+
+if TYPE_CHECKING:
+    from scepter.modules.model.embedder.embedder import (
+        ConcatTimestepEmbedderND, FrozenCLIPEmbedder, FrozenCLIPEmbedder2,
+        FrozenOpenCLIPEmbedder, FrozenOpenCLIPEmbedder2, GeneralConditioner,
+        IPAdapterPlusEmbedder, RefCrossEmbedder, SD3TextEmbedder, T5EmbedderHF)
+    from scepter.modules.model.embedder.flux_embedder import HFEmbedder
+else:
+    _import_structure = {
+        'embedder': ['ConcatTimestepEmbedderND', 'FrozenCLIPEmbedder',
+                     'FrozenCLIPEmbedder2', 'FrozenOpenCLIPEmbedder',
+                     'FrozenOpenCLIPEmbedder2', 'GeneralConditioner',
+                     'IPAdapterPlusEmbedder', 'RefCrossEmbedder',
+                     'SD3TextEmbedder', 'T5EmbedderHF'],
+        'flux_embedder': ['HFEmbedder']
+    }
+
+    import sys
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

scepter/modules/model/embedder/embedder.py

@@ -9,8 +9,11 @@ import numpy as np
 import open_clip
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import torch.utils.dlpack
 from einops import rearrange
+from torch.utils.checkpoint import checkpoint
+
 from scepter.modules.model.backbone.unet.unet_utils import Timestep
 from scepter.modules.model.embedder.base_embedder import BaseEmbedder
 from scepter.modules.model.embedder.resampler import Resampler
@@ -21,7 +24,6 @@ from scepter.modules.model.utils.basic_utils import expand_dims_like
 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 torch.utils.checkpoint import checkpoint
 
 try:
     from transformers import (CLIPTextModel, CLIPTokenizer,
@@ -34,7 +36,8 @@ except Exception as e:
 
 def autocast(f, enabled=True):
     def do_autocast(*args, **kwargs):
-        with torch.cuda.amp.autocast(
+        with torch.amp.autocast(
+                "cuda",
                 enabled=enabled,
                 dtype=torch.get_autocast_gpu_dtype(),
                 cache_enabled=torch.is_autocast_cache_enabled(),
@@ -237,7 +240,7 @@ class FrozenOpenCLIPEmbedder(BaseEmbedder):
         if cfg.PRETRAINED_MODEL is not None:
             with FS.get_from(cfg.PRETRAINED_MODEL,
                              wait_finish=True) as local_path:
-                model.load_state_dict(torch.load(local_path), strict=False)
+                model.load_state_dict(torch.load(local_path, weights_only=True), strict=False)
         self.model = model
 
         self.use_grad = cfg.get('USE_GRAD', False)
@@ -536,7 +539,7 @@ class IPAdapterPlusEmbedder(BaseEmbedder):
         )
 
         with FS.get_from(cfg.PRETRAINED_MODEL, wait_finish=True) as local_path:
-            ckpt = torch.load(local_path, map_location='cpu')
+            ckpt = torch.load(local_path, map_location='cpu', weights_only=True)
             self.image_proj_model.load_state_dict(ckpt['image_proj'],
                                                   strict=True)
 
@@ -643,7 +646,7 @@ class GeneralConditioner(BaseEmbedder):
             from safetensors.torch import load_file as load_safetensors
             sd = load_safetensors(path)
         else:
-            sd = torch.load(path, map_location='cpu')
+            sd = torch.load(path, map_location='cpu', weights_only=True)
         new_sd = OrderedDict()
         for k, v in sd.items():
             ignored = False
@@ -830,23 +833,37 @@ class T5EmbedderHF(BaseEmbedder):
 
     def __init__(self, cfg, logger=None):
         super().__init__(cfg, logger=logger)
-        pretrained_path = cfg.get('PRETRAINED_MODEL', None)
-        t5_dtype = cfg.get('T5_DTYPE', None)
-        assert pretrained_path
-        with FS.get_dir_to_local_dir(pretrained_path,
-                                     wait_finish=True) as local_path:
-            if t5_dtype is not None:
-                self.model = T5EncoderModel.from_pretrained(
-                    local_path, torch_dtype=getattr(torch, t5_dtype))
-            else:
-                self.model = T5EncoderModel.from_pretrained(local_path)
-        tokenizer_path = cfg.get('TOKENIZER_PATH', None)
         self.length = cfg.get('LENGTH', 77)
+        self.t5_dtype = cfg.get('T5_DTYPE', 'bfloat16')
+        self.use_grad = cfg.get('USE_GRAD', False)
+        self.clean = cfg.get('CLEAN', 'whitespace')
+        self.added_identifier = cfg.get('ADDED_IDENTIFIER', None)
+        tokenizer_path = cfg.get('TOKENIZER_PATH', None)
+        pretrained_path = cfg.get('PRETRAINED_MODEL', None)
+
+        if pretrained_path:
+            with FS.get_dir_to_local_dir(pretrained_path,
+                                         wait_finish=True) as local_path:
+                if self.t5_dtype is not None:
+                    self.model = T5EncoderModel.from_pretrained(
+                        local_path,
+                        torch_dtype=getattr(
+                            torch,
+                            'float' if self.t5_dtype == 'float32' else self.t5_dtype))
+                else:
+                    self.model = T5EncoderModel.from_pretrained(local_path)
+        else:
+            self.model = None
+
         if tokenizer_path:
             self.tokenize_kargs = {'return_tensors': 'pt'}
             with FS.get_dir_to_local_dir(tokenizer_path,
                                          wait_finish=True) as local_path:
-                self.tokenizer = AutoTokenizer.from_pretrained(local_path)
+                if self.added_identifier is not None and isinstance(
+                        self.added_identifier, list):
+                    self.tokenizer = AutoTokenizer.from_pretrained(local_path)
+                else:
+                    self.tokenizer = AutoTokenizer.from_pretrained(local_path)
             if self.length is not None:
                 self.tokenize_kargs.update({
                     'padding': 'max_length',
@@ -859,26 +876,22 @@ class T5EmbedderHF(BaseEmbedder):
             self.tokenizer = None
             self.tokenize_kargs = {}
 
-        self.use_grad = cfg.get('USE_GRAD', False)
-        self.clean = cfg.get('CLEAN', 'whitespace')
-
     def freeze(self):
         self.model = self.model.eval()
         for param in self.parameters():
             param.requires_grad = False
 
     # encode && encode_text
-    def forward(self, tokens, return_mask=False):
+    def forward(self, tokens, return_mask=False, use_mask=True):
         # tokenization
         embedding_context = nullcontext if self.use_grad else torch.no_grad
         with embedding_context():
-            x = self.model(tokens.input_ids.to(we.device_id),
-                           tokens.attention_mask.to(we.device_id))
+            if use_mask:
+                x = self.model(tokens.input_ids.to(we.device_id),
+                               tokens.attention_mask.to(we.device_id))
+            else:
+                x = self.model(tokens.input_ids.to(we.device_id))
             x = x.last_hidden_state
-            # if not self.return_pooled:
-            #     return x.detach()
-            # else:
-            #     return x.detach(), self.pool(x, tokens.input_ids)
             if return_mask:
                 return x.detach() + 0.0, tokens.attention_mask.to(we.device_id)
             else:
@@ -897,7 +910,7 @@ class T5EmbedderHF(BaseEmbedder):
         elif self.clean == 'canonicalize':
             text = canonicalize(basic_clean(text))
         elif self.clean == 'heavy':
-            text = heavy_clean(heavy_clean(text))
+            text = heavy_clean(basic_clean(text))
         return text
 
     def encode_text(self,
@@ -907,14 +920,48 @@ class T5EmbedderHF(BaseEmbedder):
                     return_mask=False):
         return self(tokens, return_mask=return_mask)
 
-    def encode(self, text, return_mask=False):
+    def encode(self, text, return_mask=False, use_mask=True):
         if isinstance(text, str):
             text = [text]
         if self.clean:
             text = [self._clean(u) for u in text]
         assert self.tokenizer is not None
         tokens = self.tokenizer(text, **self.tokenize_kargs)
-        return self(tokens, return_mask=return_mask)
+        return self(tokens, return_mask=return_mask, use_mask=use_mask)
+
+    def encode_list(self, text, return_mask=False, use_mask=True):
+        if isinstance(text, str):
+            text = [text]
+        if self.clean:
+            text = [self._clean(u) for u in text]
+        assert self.tokenizer is not None
+        cont, mask = [], []
+        with torch.autocast(device_type='cuda',
+                            enabled=self.t5_dtype in ('float16', 'bfloat16'),
+                            dtype=getattr(torch, self.t5_dtype)):
+            for tt in text:
+                tokens = self.tokenizer([tt], **self.tokenize_kargs)
+                one_cont, one_mask = self(tokens,
+                                          return_mask=return_mask,
+                                          use_mask=use_mask)
+                cont.append(one_cont)
+                mask.append(one_mask)
+        if return_mask:
+            return torch.cat(cont, dim=0), torch.cat(mask, dim=0)
+        else:
+            return torch.cat(cont, dim=0)
+
+    def encode_list_of_list(self, text_list, return_mask=True, use_mask=True):
+        cont_list = []
+        mask_list = []
+        for pp in text_list:
+            cont, cont_mask = self.encode_list(pp, return_mask=return_mask, use_mask=use_mask)
+            cont_list.append(cont)
+            mask_list.append(cont_mask)
+        if return_mask:
+            return cont_list, mask_list
+        else:
+            return cont_list
 
     @staticmethod
     def get_config_template():