jiangzeyinzi
@@ -3,14 +3,12 @@
|
||||
*.pt
|
||||
*.pkl
|
||||
*.ckpt
|
||||
*.png
|
||||
*.DS_Store
|
||||
*__pycache__*
|
||||
*.cache*
|
||||
*.bin
|
||||
*.idea
|
||||
*.csv
|
||||
#*.txt
|
||||
build
|
||||
dist
|
||||
dev
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
repos:
|
||||
- repo: https://github.com/pycqa/flake8.git
|
||||
rev: 4.0.0
|
||||
hooks:
|
||||
- id: flake8
|
||||
args: ['--max-line-length=120', '--per-file-ignores=__init__.py:F401']
|
||||
- repo: https://github.com/PyCQA/isort.git
|
||||
rev: 4.3.21
|
||||
hooks:
|
||||
- id: isort
|
||||
- repo: https://github.com/pre-commit/mirrors-yapf.git
|
||||
rev: v0.30.0
|
||||
hooks:
|
||||
- id: yapf
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks.git
|
||||
rev: v3.1.0
|
||||
hooks:
|
||||
- id: trailing-whitespace
|
||||
- id: check-yaml
|
||||
- id: end-of-file-fixer
|
||||
- id: requirements-txt-fixer
|
||||
- id: double-quote-string-fixer
|
||||
- id: check-merge-conflict
|
||||
- id: fix-encoding-pragma
|
||||
- id: mixed-line-ending
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 881 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 74 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 222 KiB |
@@ -356,7 +356,7 @@ Used to instantiate a standard logging instance for printing information.
|
||||
```python
|
||||
from scepter.utils.logger import get_logger, init_logger
|
||||
|
||||
std_logger = get_logger(name="std_torch")
|
||||
std_logger = get_logger(name="scepter")
|
||||
init_logger(std_logger, log_file="", dist_launcher="pytorch")
|
||||
```
|
||||
<hr/>
|
||||
|
||||
@@ -358,7 +358,7 @@ input_type 一一对应。
|
||||
```python
|
||||
from scepter.utils.logger import get_logger, init_logger
|
||||
|
||||
std_logger = get_logger(name="std_torch")
|
||||
std_logger = get_logger(name="scepter")
|
||||
init_logger(std_logger, log_file="", dist_launcher="pytorch")
|
||||
```
|
||||
<hr/>
|
||||
|
||||
@@ -3,14 +3,18 @@
|
||||
<p align="center">
|
||||
<img src="https://img.shields.io/badge/python-%E2%89%A53.8-5be.svg">
|
||||
<img src="https://img.shields.io/badge/pytorch-%E2%89%A51.12%20%7C%20%E2%89%A52.0-orange.svg">
|
||||
<a href="https://pypi.org/project/scepter/"><img src="https://badge.fury.io/py/scepter.svg"></a>
|
||||
<a href="https://github.com/modelscope/scepter/blob/main/LICENSE"><img src="https://img.shields.io/github/license/modelscope/scepter"></a>
|
||||
<a href="https://github.com/modelscope/scepter/"><img src="https://img.shields.io/badge/scepter-Build from source-6FEBB9.svg"></a>
|
||||
</p>
|
||||
|
||||
## 📖 Table of Contents
|
||||
- [Introduction](#-introduction)
|
||||
- [News](#-news)
|
||||
- [Installation](#-installation)
|
||||
- [Installation](#-Installation)
|
||||
- [Getting Started](#-getting-started)
|
||||
- [SCEPTER Studio](#-scepter-studio)
|
||||
- [Features](#-features)
|
||||
- [Learn More](#-learn-more)
|
||||
- [License](#license)
|
||||
|
||||
@@ -20,14 +24,17 @@ SCEPTER is an open-source code repository dedicated to generative training, fine
|
||||
|
||||
Main Feature:
|
||||
|
||||
- Training:
|
||||
- distribute: DDP / FSDP / FairScale
|
||||
- Inference
|
||||
- text-to-image generation
|
||||
- controllable image synthesis (TODO)
|
||||
- Deploy-Gradio (TODO)
|
||||
- fine-tuning
|
||||
- inference
|
||||
- Task:
|
||||
- Text-to-image generation
|
||||
- Controllable image synthesis
|
||||
- Image editing (TODO)
|
||||
- Training / Inference:
|
||||
- Distribute: DDP / FSDP / FairScale / Xformers
|
||||
- File system: Local / Http / OSS / Modelscope
|
||||
- Deploy:
|
||||
- Data management
|
||||
- Training
|
||||
- Inference
|
||||
|
||||
Currently supported approches (and counting):
|
||||
|
||||
@@ -36,6 +43,8 @@ Currently supported approches (and counting):
|
||||
3. Res-Tuning(TODO): [Res-Tuning: A Flexible and Efficient Tuning Paradigm via Unbinding Tuner from Backbone](https://arxiv.org/abs/2310.19859) [](https://arxiv.org/abs/2310.19859) [](https://res-tuning.github.io/)
|
||||
|
||||
## 🎉 News
|
||||
- [2024.01]: We release **SCEPTER Studio**, an integrated toolkit for data management, model training and inference based on [Gradio](https://www.gradio.app/).
|
||||
- [2024.01]: [SCEdit](https://arxiv.org/abs/2312.11392) support controllable image synthesis for training and inference.
|
||||
- [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.
|
||||
|
||||
@@ -51,14 +60,15 @@ conda activate scepter
|
||||
- Install SCEPTER by the `pip` command:
|
||||
|
||||
```shell
|
||||
pip install -e .
|
||||
pip install scepter
|
||||
```
|
||||
- PS: We recommend installing PyTorch follwing [official documentation](https://pytorch.org/get-started/locally/)
|
||||
|
||||
## 🚀 Getting Started
|
||||
|
||||
### Dataset
|
||||
|
||||
#### Text-to-Image generation
|
||||
#### Modelscope Format
|
||||
|
||||
We use a [custom-stylized dataset](https://modelscope.cn/datasets/damo/style_custom_dataset/summary), which included classes 3D, anime, flat illustration, oil painting, sketch, and watercolor, each with 30 image-text pairs.
|
||||
|
||||
@@ -69,54 +79,148 @@ ms_train_dataset = MsDataset.load('style_custom_dataset', namespace='damo', subs
|
||||
print(next(iter(ms_train_dataset)))
|
||||
```
|
||||
|
||||
#### CSV Format
|
||||
|
||||
For the data format used by SCEPTER Studio, please refer to [3D_example_csv.zip](https://modelscope.cn/api/v1/models/damo/scepter/repo?Revision=master&FilePath=datasets/3D_example_csv.zip).
|
||||
|
||||
#### TXT Format
|
||||
|
||||
To facilitate starting training in command-line mode, you can use a dataset in text format, please refer to [3D_example_txt.zip](https://modelscope.cn/api/v1/models/damo/scepter/repo?Revision=master&FilePath=datasets/3D_example_txt.zip)
|
||||
|
||||
```shell
|
||||
mkdir -p cache/dataset/ && wget 'https://modelscope.cn/api/v1/models/damo/scepter_scedit/repo?Revision=master&FilePath=dataset/3D_example_txt.zip' -O cache/dataset/3D_example_txt.zip && unzip cache/dataset/3D_example_txt.zip -d cache/dataset/ && rm cache/dataset/3D_example_txt.zip
|
||||
```
|
||||
|
||||
### Training
|
||||
|
||||
#### Text-to-Image generation
|
||||
We provide a framework for training and inference, so the script below is just for illustration purposes. To achieve better results, you can modify the corresponding parameters as needed.
|
||||
|
||||
#### Text-to-Image Generation
|
||||
|
||||
- SCEdit
|
||||
```python
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/t2i/sd15_512_sce_t2i.yaml # SD v1.5
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/t2i/sd21_768_sce_t2i.yaml # SD v2.1
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i.yaml # SD XL
|
||||
```
|
||||
|
||||
- Existing Tuning Strategies
|
||||
```python
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml # fully-tuning on SD v1.5
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/examples/generation/stable_diffusion_2.1_768_lora.yaml # lora-tuning on SD v2.1
|
||||
```
|
||||
|
||||
- Data Text Format
|
||||
```python
|
||||
# Download the 3D_example_txt.zip as previously mentioned
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i_datatxt.yaml
|
||||
```
|
||||
|
||||
#### Controllable Image Synthesis
|
||||
|
||||
- SCEdit
|
||||
|
||||
The YAML configuration can be modified to combine different base models and conditions. The following is provided as an example.
|
||||
```python
|
||||
# SD v1.5
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/SCEdit/t2i_sd15_512_sce.yaml
|
||||
# SD v2.1
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/SCEdit/t2i_sd21_768_sce.yaml
|
||||
# SD XL
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/SCEdit/t2i_sdxl_1024_sce.yaml
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/ctr/sd15_512_sce_ctr_hed.yaml # SD v1.5 + hed
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_canny.yaml # SD v2.1 + canny
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_pose.yaml # SD v2.1 + pose
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/ctr/sdxl_1024_sce_ctr_depth.yaml # SD XL + depth
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/ctr/sdxl_1024_sce_ctr_color.yaml # SD XL + color
|
||||
```
|
||||
|
||||
- Existing strategies
|
||||
- Data Text Format
|
||||
```python
|
||||
# fully-tuning on SD v1.5
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml
|
||||
# lora-tuning on SD v2.1
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/examples/generation/stable_diffusion_2.1_768_lora.yaml
|
||||
# Download the 3D_example_txt.zip as previously mentioned
|
||||
python scepter/tools/run_train.py --cfg scepter/methods/scedit/ctr/sdxl_1024_sce_ctr_color_datatxt.yaml
|
||||
```
|
||||
#### Controllable Image Synthesis
|
||||
|
||||
TODO
|
||||
|
||||
### Inference
|
||||
|
||||
#### Base Model Inference
|
||||
|
||||
```python
|
||||
# generation on SD v1.5
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml --prompt 'a cute dog' --save_folder 'inference'
|
||||
# generation on SD v2.1
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml --prompt 'a cute dog' --save_folder 'inference'
|
||||
# generation on SD XL
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml --prompt 'a cute dog' --save_folder 'inference'
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml --prompt 'a cute dog' --save_folder 'inference' # generation on SD v1.5
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml --prompt 'a cute dog' --save_folder 'inference' # generation on SD v2.1
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml --prompt 'a cute dog' --save_folder 'inference' # generation on SD XL
|
||||
```
|
||||
|
||||
#### Fine-tuned Model Inference
|
||||
|
||||
```python
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/scedit/t2i/sd15_512_sce_t2i_swift.yaml --pretrained_model 'cache/save_data/sd15_512_sce_t2i_swift/checkpoints/ldm_step-100.pth' --prompt 'A close up of a small rabbit wearing a hat and scarf' --save_folder 'trained_test_prompt_rabbit'
|
||||
```
|
||||
|
||||
#### Controllable Image Synthesis Inference
|
||||
|
||||
- SCEdit
|
||||
```python
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_canny.yaml --num_samples 1 --prompt 'a single flower is shown in front of a tree' --save_folder 'test_flower_canny' --image_size 768 --task control --image 'asset/images/flower.jpg' --control_mode canny --pretrained_model ms://damo/scepter_scedit@controllable_model/SD2.1/canny_control/0_SwiftSCETuning/pytorch_model.bin # canny
|
||||
python scepter/tools/run_inference.py --cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_pose.yaml --num_samples 1 --prompt 'super mario' --save_folder 'test_mario_pose' --image_size 768 --task control --image 'asset/images/pose_source.png' --control_mode source --pretrained_model ms://damo/scepter_scedit@controllable_model/SD2.1/pose_control/0_SwiftSCETuning/pytorch_model.bin # pose
|
||||
```
|
||||
|
||||
|
||||
## 🖥️ SCEPTER Studio
|
||||
|
||||
### Launch
|
||||
|
||||
To fully experience **SCEPTER Studio**, you can launch the following command line:
|
||||
|
||||
```shell
|
||||
PYTHONPATH=. python scepter/tools/webui.py --cfg scepter/methods/studio/scepter_ui.yaml
|
||||
```
|
||||
|
||||
### Modelscope Studio
|
||||
|
||||
We deploy a work studio on Modelscope that includes only the inference tab, please refer to [ms_scepter_studio](https://www.modelscope.cn/studios/damo/scepter_studio/summary)
|
||||
|
||||
## ✨ Features
|
||||
|
||||
### Text-to-Image Generation
|
||||
|
||||
| **Model** | **SCEdit** | **Full** | **LoRA** |
|
||||
|:---------:|:----------:|:--------:|:--------:|
|
||||
| SD 1.5 | 🪄 | ✅ | ✅ |
|
||||
| SD 2.1 | 🪄 | ✅ | ✅ |
|
||||
| SD XL | 🪄 | ✅ | ✅ |
|
||||
|
||||
### Controllable Image Synthesis
|
||||
- SCEdit
|
||||
|
||||
| **Model** | **Canny** | **HED** | **Depth** | **Pose** | **Color** |
|
||||
|:---------:|:---------:|:-------:|:---------:|:--------:|:---------:|
|
||||
| SD 1.5 | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| SD 2.1 | 🪄 | ✅ | ✅ | 🪄 | 🪄 |
|
||||
| SD XL | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
|
||||
### Model URL
|
||||
|
||||
- ✅ indicates support for both training and inference.
|
||||
- 🪄 denotes that the model has been published.
|
||||
- More models will be released in the future.
|
||||
|
||||
| Model | URL |
|
||||
|--------|-------------------------------------------------------------------------------------|
|
||||
| SCEdit | [ModelCard](https://modelscope.cn/models/damo/scepter_scedit/summary) |
|
||||
|
||||
PS: Scripts running within the SCEPTER framework will automatically fetch and load models based on the required dependency files, eliminating the need for manual downloads.
|
||||
|
||||
|
||||
## 🔍 Learn More
|
||||
|
||||
- [ModelScope library](https://github.com/modelscope/modelscope/)
|
||||
|
||||
ModelScope Library is the model library of ModelScope project, which contains a large number of popular models.
|
||||
|
||||
- [Alibaba TongYi Vision Intelligence Lab](https://github.com/damo-vilab)
|
||||
|
||||
Discover more about open-source projects on image generation, video generation, and editing tasks.
|
||||
|
||||
- [ModelScope library](https://github.com/modelscope/modelscope/)
|
||||
|
||||
ModelScope Library is the model library of ModelScope project, which contains a large number of popular models.
|
||||
|
||||
- [SWIFT library](https://github.com/modelscope/swift/)
|
||||
|
||||
SWIFT (Scalable lightWeight Infrastructure for Fine-Tuning) is an extensible framwork designed to faciliate lightweight model fine-tuning and inference.
|
||||
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under the [Apache License (Version 2.0)](https://github.com/modelscope/modelscope/blob/master/LICENSE).
|
||||
|
||||
@@ -1 +1,2 @@
|
||||
-r requirements/framework.txt
|
||||
-r requirements/scepter_studio.txt
|
||||
|
||||
@@ -1,12 +1,11 @@
|
||||
einops
|
||||
modelscope
|
||||
ms_swift==1.5.2
|
||||
ms-swift>=1.5.2
|
||||
numpy
|
||||
open_clip_torch
|
||||
opencv-python
|
||||
opencv_transforms>=0.0.6
|
||||
oss2>=2.15.0
|
||||
pyyaml>=5.3.1
|
||||
torchvision==0.15.2
|
||||
transformers
|
||||
xformers>=0.0.21
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
gradio==3.50.2
|
||||
imagehash
|
||||
@@ -15,7 +15,7 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/sd15_512_full
|
||||
WORK_DIR: ./cache/save_data/sd15_512_full
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
|
||||
@@ -15,13 +15,13 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/sd15_512_lora
|
||||
WORK_DIR: ./cache/save_data/sd15_512_lora
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
TUNER:
|
||||
TUNER:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 64
|
||||
|
||||
@@ -15,7 +15,7 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/sd21_768_full
|
||||
WORK_DIR: ./cache/save_data/sd21_768_full
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
|
||||
@@ -15,14 +15,14 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/sd21_768_lora
|
||||
WORK_DIR: ./cache/save_data/sd21_768_lora
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TUNER:
|
||||
TUNER:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 64
|
||||
|
||||
@@ -15,7 +15,7 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/sdxl_1024_full
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_full
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
@@ -120,7 +120,7 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
@@ -133,21 +133,21 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
@@ -200,21 +200,21 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
@@ -287,8 +287,8 @@ SOLVER:
|
||||
KEYS: [ 'img', 'prompt', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
META_KEYS: [ 'data_key', 'img_path' ]
|
||||
- NAME: Rename
|
||||
IN_KEYS: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUT_KEYS: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
INPUT_KEY: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUTPUT_KEY: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
|
||||
@@ -15,14 +15,14 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/sdxl_1024_lora
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_lora
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TUNER:
|
||||
#
|
||||
TUNER:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 64
|
||||
@@ -129,7 +129,7 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
@@ -142,21 +142,21 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
@@ -209,21 +209,21 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
@@ -296,8 +296,8 @@ SOLVER:
|
||||
KEYS: [ 'img', 'prompt', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
META_KEYS: [ 'data_key', 'img_path' ]
|
||||
- NAME: Rename
|
||||
IN_KEYS: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUT_KEYS: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
INPUT_KEY: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUTPUT_KEY: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
@@ -334,4 +334,4 @@ SOLVER:
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
- NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
PROB_INTERVAL: 100
|
||||
|
||||
@@ -0,0 +1,265 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 200
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sd15_512_sce_ctr_hed
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "control_blocks" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionSCEControl
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-v1-5@v1-5-pruned-emaonly.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS: 8
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 768
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: False
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: ClipTokenizer
|
||||
PRETRAINED_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
LENGTH: 77
|
||||
CLEAN: True
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenCLIPEmbedder
|
||||
FREEZE: True
|
||||
LAYER: last
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
CONTROL_MODEL:
|
||||
NAME: CSCTuners
|
||||
PRE_HINT_IN_CHANNELS: 3
|
||||
PRE_HINT_OUT_CHANNELS: 256
|
||||
DENSE_HINT_KERNAL: 3
|
||||
SCALE: 1.0
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
CONTROL_ANNO:
|
||||
NAME: HedAnnotator
|
||||
PRETRAINED_MODEL: ms://damo/scepter_scedit@annotator/ckpts/ControlNetHED.pth
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [512, 512]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 10
|
||||
NUM_WORKERS: 4
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,264 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 200
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sd21_768_sce_ctr_canny
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "control_blocks" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionSCEControl
|
||||
PARAMETERIZATION: v
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: True
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-1@v2-1_768-ema-pruned.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 1024
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: OpenClipTokenizer
|
||||
LENGTH: 77
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenOpenCLIPEmbedder
|
||||
ARCH: ViT-H-14
|
||||
PRETRAINED_MODEL:
|
||||
LAYER: penultimate
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
CONTROL_MODEL:
|
||||
NAME: CSCTuners
|
||||
PRE_HINT_IN_CHANNELS: 3
|
||||
PRE_HINT_OUT_CHANNELS: 256
|
||||
DENSE_HINT_KERNAL: 3
|
||||
SCALE: 1.0
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
CONTROL_ANNO:
|
||||
NAME: CannyAnnotator
|
||||
LOW_THRESHOLD: 100
|
||||
HIGH_THRESHOLD: 200
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [768, 768]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 10
|
||||
NUM_WORKERS: 4
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,264 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 200
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sd21_768_sce_ctr_pose
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "control_blocks" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionSCEControl
|
||||
PARAMETERIZATION: v
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: True
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-1@v2-1_768-ema-pruned.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 1024
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: OpenClipTokenizer
|
||||
LENGTH: 77
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenOpenCLIPEmbedder
|
||||
ARCH: ViT-H-14
|
||||
PRETRAINED_MODEL:
|
||||
LAYER: penultimate
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
CONTROL_MODEL:
|
||||
NAME: CSCTuners
|
||||
PRE_HINT_IN_CHANNELS: 3
|
||||
PRE_HINT_OUT_CHANNELS: 256
|
||||
DENSE_HINT_KERNAL: 3
|
||||
SCALE: 1.0
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
CONTROL_ANNO:
|
||||
NAME: OpenposeAnnotator
|
||||
BODY_MODEL_PATH: ms://damo/scepter_scedit@annotator/ckpts/body_pose_model.pth
|
||||
HAND_MODEL_PATH: ms://damo/scepter_scedit@annotator/ckpts/hand_pose_model.pth
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [768, 768]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 10
|
||||
NUM_WORKERS: 4
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,378 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 200
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_color
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "control_blocks" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionXLSCEControl
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.13025
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.0120
|
||||
USE_EMA: False
|
||||
LOAD_REFINER: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 320
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: [ 1, 2, 10 ]
|
||||
CONTEXT_DIM: 2048
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2816
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenCLIPEmbedder
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
LAYER: hidden
|
||||
LAYER_IDX: 11
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "target_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
REFINER_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 384
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 4
|
||||
CONTEXT_DIM: [ 1280, 1280, 1280, 1280 ]
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2560
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
REFINER_COND_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "aesthetic_score" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
CONTROL_MODEL:
|
||||
NAME: CSCTuners
|
||||
PRE_HINT_IN_CHANNELS: 3
|
||||
PRE_HINT_OUT_CHANNELS: 256
|
||||
DENSE_HINT_KERNAL: 3
|
||||
SCALE: 1.0
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
CONTROL_ANNO:
|
||||
NAME: ColorAnnotator
|
||||
RATIO: 64
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleResize
|
||||
SIZE: 1024
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleCropXL
|
||||
SIZE: 1024
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 10
|
||||
NUM_WORKERS: 4
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 1024
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 1024
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,404 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 200
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_color_datatxt
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "control_blocks" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionXLSCEControl
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.13025
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.0120
|
||||
USE_EMA: False
|
||||
LOAD_REFINER: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 320
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: [ 1, 2, 10 ]
|
||||
CONTEXT_DIM: 2048
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2816
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenCLIPEmbedder
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
LAYER: hidden
|
||||
LAYER_IDX: 11
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "target_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
REFINER_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 384
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 4
|
||||
CONTEXT_DIM: [ 1280, 1280, 1280, 1280 ]
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2560
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
REFINER_COND_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "aesthetic_score" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
CONTROL_MODEL:
|
||||
NAME: CSCTuners
|
||||
PRE_HINT_IN_CHANNELS: 3
|
||||
PRE_HINT_OUT_CHANNELS: 256
|
||||
DENSE_HINT_KERNAL: 3
|
||||
SCALE: 1.0
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
CONTROL_ANNO:
|
||||
NAME: ColorAnnotator
|
||||
RATIO: 64
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairDataset
|
||||
MODE: train
|
||||
P_ZERO: 0.1
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: LocalFs
|
||||
AUTO_CLEAN: False
|
||||
SAMPLER:
|
||||
NAME: MixtureOfSamplers
|
||||
SUB_SAMPLERS:
|
||||
-
|
||||
NAME: MultiLevelBatchSampler
|
||||
PROB: 1.0
|
||||
IMAGE_SIZE: [ 1024, 1024 ]
|
||||
FIELDS: ["img_path", "width", "height", "prompt"]
|
||||
DELIMITER: '#;#'
|
||||
PATH_PREFIX: cache/datasets/3D_example_txt
|
||||
PROMPT_PREFIX: ''
|
||||
INDEX_FILE: cache/datasets/3D_example_txt/train.txt
|
||||
#
|
||||
TRANSFORMS:
|
||||
-
|
||||
NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: FlexibleResize
|
||||
INTERPOLATION: bilinear
|
||||
SIZE: 1024
|
||||
INPUT_KEY: ['img']
|
||||
OUTPUT_KEY: ['img']
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: FlexibleCropXL
|
||||
INPUT_KEY: ['img']
|
||||
OUTPUT_KEY: ['img']
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
-
|
||||
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: ['img']
|
||||
BACKEND: torchvision
|
||||
-
|
||||
NAME: Select
|
||||
KEYS: ['img', 'prompt', 'image_preprocess', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left']
|
||||
META_KEYS: [ 'img_path' ]
|
||||
-
|
||||
NAME: Rename
|
||||
INPUT_KEY: ['img', 'image_preprocess', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left']
|
||||
OUTPUT_KEY: ['image', 'image_preprocess', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left']
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: Text2ImageDataset
|
||||
MODE: eval
|
||||
PROMPT_FILE: cache/datasets/3D_example_txt/train.txt
|
||||
PATH_PREFIX: cache/datasets/3D_example_txt
|
||||
IMAGE_SIZE: [ 1024, 1024 ]
|
||||
FIELDS: ["img_path", "width", "height", "prompt"]
|
||||
DELIMITER: '#;#'
|
||||
PROMPT_PREFIX: ''
|
||||
USE_NUM: 8
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: LocalFs
|
||||
AUTO_CLEAN: False
|
||||
#
|
||||
TRANSFORMS:
|
||||
-
|
||||
NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: FlexibleResize
|
||||
INTERPOLATION: bilinear
|
||||
SIZE: 1024
|
||||
INPUT_KEY: ['img']
|
||||
OUTPUT_KEY: ['img']
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: FlexibleCenterCrop
|
||||
INPUT_KEY: ['img']
|
||||
OUTPUT_KEY: ['img']
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
-
|
||||
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: ['img']
|
||||
BACKEND: torchvision
|
||||
-
|
||||
NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
-
|
||||
NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,378 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 200
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_sce_ctr_depth
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "control_blocks" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionXLSCEControl
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.13025
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.0120
|
||||
USE_EMA: False
|
||||
LOAD_REFINER: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 320
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: [ 1, 2, 10 ]
|
||||
CONTEXT_DIM: 2048
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2816
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenCLIPEmbedder
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
LAYER: hidden
|
||||
LAYER_IDX: 11
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "target_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
REFINER_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 384
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 4
|
||||
CONTEXT_DIM: [ 1280, 1280, 1280, 1280 ]
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2560
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
REFINER_COND_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "aesthetic_score" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
CONTROL_MODEL:
|
||||
NAME: CSCTuners
|
||||
PRE_HINT_IN_CHANNELS: 3
|
||||
PRE_HINT_OUT_CHANNELS: 256
|
||||
DENSE_HINT_KERNAL: 3
|
||||
SCALE: 1.0
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
CONTROL_ANNO:
|
||||
NAME: MidasDetector
|
||||
PRETRAINED_MODEL: ms://damo/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleResize
|
||||
SIZE: 1024
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleCropXL
|
||||
SIZE: 1024
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 10
|
||||
NUM_WORKERS: 4
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 1024
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 1024
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ToNumpy
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image_preprocess' ]
|
||||
- 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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'image_preprocess' ]
|
||||
OUTPUT_KEY: [ 'image', 'image_preprocess' ]
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt', 'image_preprocess' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,228 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 2000
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sd15_512_sce_t2i
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "lsc_identity" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionSCETuning
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-v1-5@v1-5-pruned-emaonly.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
# DEFAULT_N_PROMPT: 'lowres, error, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS: 8
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 768
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: False
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: ClipTokenizer
|
||||
PRETRAINED_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
LENGTH: 77
|
||||
CLEAN: True
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenCLIPEmbedder
|
||||
FREEZE: True
|
||||
LAYER: last
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
TUNER_MODEL:
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [512, 512]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 512
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 512
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ImageToTensor
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: Normalize
|
||||
MEAN: [ 0.5, 0.5, 0.5 ]
|
||||
STD: [ 0.5, 0.5, 0.5 ]
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_REMAP_KEYS: { 'Image': 'Target:FILE' }
|
||||
MS_DATASET_SPLIT: test_short
|
||||
OUTPUT_SIZE: [512, 512]
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TRANSFORMS:
|
||||
-
|
||||
NAME: Select
|
||||
KEYS: ['prompt']
|
||||
META_KEYS: ['image_size']
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 1000
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
+3
-3
@@ -15,14 +15,14 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/t2i_sd15_512_sce
|
||||
WORK_DIR: ./cache/save_data/sd15_512_sce_t2i_swift
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TUNER:
|
||||
TUNER:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320]
|
||||
@@ -125,7 +125,7 @@ SOLVER:
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE:
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [512, 512]
|
||||
RUN_TRAIN_N: False
|
||||
@@ -0,0 +1,224 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 2000
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sd21_768_sce_t2i
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "lsc_identity" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionSCETuning
|
||||
PARAMETERIZATION: v
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-1@v2-1_768-ema-pruned.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 1024
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
PRETRAINED_MODEL:
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: OpenClipTokenizer
|
||||
LENGTH: 77
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenOpenCLIPEmbedder
|
||||
ARCH: ViT-H-14
|
||||
PRETRAINED_MODEL:
|
||||
LAYER: penultimate
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
TUNER_MODEL:
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [768, 768]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ImageToTensor
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: Normalize
|
||||
MEAN: [ 0.5, 0.5, 0.5 ]
|
||||
STD: [ 0.5, 0.5, 0.5 ]
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_REMAP_KEYS: { 'Image': 'Target:FILE' }
|
||||
MS_DATASET_SPLIT: test_short
|
||||
OUTPUT_SIZE: [768, 768]
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TRANSFORMS:
|
||||
-
|
||||
NAME: Select
|
||||
KEYS: ['prompt']
|
||||
META_KEYS: ['image_size']
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 1000
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
+3
-3
@@ -15,14 +15,14 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/t2i_sd21_768_sce
|
||||
WORK_DIR: ./cache/save_data/sd21_768_sce_t2i_swift
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TUNER:
|
||||
TUNER:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320]
|
||||
@@ -121,7 +121,7 @@ SOLVER:
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE:
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [768, 768]
|
||||
RUN_TRAIN_N: False
|
||||
@@ -0,0 +1,342 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 2000
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_sce_t2i
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "lsc_identity" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionXLSCETuning
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.13025
|
||||
SIZE_FACTOR: 8
|
||||
# DEFAULT_N_PROMPT: 'lowres, error, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.0120
|
||||
USE_EMA: False
|
||||
LOAD_REFINER: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 320
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: [ 1, 2, 10 ]
|
||||
CONTEXT_DIM: 2048
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2816
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenCLIPEmbedder
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
LAYER: hidden
|
||||
LAYER_IDX: 11
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "target_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
REFINER_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 384
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 4
|
||||
CONTEXT_DIM: [ 1280, 1280, 1280, 1280 ]
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2560
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
REFINER_COND_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "aesthetic_score" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
TUNER_MODEL:
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train_short
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleResize
|
||||
INTERPOLATION: bicubic
|
||||
SIZE: 1024
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleCropXL
|
||||
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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Select
|
||||
KEYS: [ 'img', 'prompt', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
META_KEYS: [ 'data_key', 'img_path' ]
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUTPUT_KEY: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: eval
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_REMAP_KEYS: { 'Image': 'Target:FILE' }
|
||||
MS_DATASET_SPLIT: test_short
|
||||
OUTPUT_SIZE: [ 1024, 1024 ]
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TRANSFORMS:
|
||||
- NAME: Select
|
||||
KEYS: [ 'prompt' ]
|
||||
META_KEYS: [ 'image_size' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 1000
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,355 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 2000
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_sce_t2i_datatxt
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
FREEZE_PART: [ "first_stage_model", "cond_stage_model", "model" ]
|
||||
TRAIN_PART: [ "lsc_identity" ]
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusionXLSCETuning
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.13025
|
||||
SIZE_FACTOR: 8
|
||||
# DEFAULT_N_PROMPT: 'lowres, error, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.0120
|
||||
USE_EMA: False
|
||||
LOAD_REFINER: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 320
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: [ 1, 2, 10 ]
|
||||
CONTEXT_DIM: 2048
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2816
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenCLIPEmbedder
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
LAYER: hidden
|
||||
LAYER_IDX: 11
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "target_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
REFINER_MODEL:
|
||||
NAME: DiffusionUNetXL
|
||||
PRETRAINED_MODEL:
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
NUM_RES_BLOCKS: 2
|
||||
MODEL_CHANNELS: 384
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
DROPOUT: 0
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
NUM_CLASSES: sequential
|
||||
USE_CHECKPOINT: False
|
||||
NUM_HEADS: -1
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 4
|
||||
CONTEXT_DIM: [ 1280, 1280, 1280, 1280 ]
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
ADM_IN_CHANNELS: 2560
|
||||
USE_SENTENCE_EMB: False
|
||||
USE_WORD_MAPPING: False
|
||||
REFINER_COND_MODEL:
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
EMBEDDERS:
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
MAX_LENGTH: 77
|
||||
FREEZE: True
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
LEGACY: False
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "aesthetic_score" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
TUNER_MODEL:
|
||||
SC_TUNER_CFG:
|
||||
NAME: SCTuner
|
||||
TUNER_NAME: SCEAdapter
|
||||
DOWN_RATIO: 1.0
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairDataset
|
||||
MODE: train
|
||||
P_ZERO: 0.0
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: LocalFs
|
||||
AUTO_CLEAN: False
|
||||
SAMPLER:
|
||||
NAME: MixtureOfSamplers
|
||||
SUB_SAMPLERS:
|
||||
-
|
||||
NAME: MultiLevelBatchSampler
|
||||
PROB: 1.0
|
||||
IMAGE_SIZE: [ 1024, 1024 ]
|
||||
FIELDS: ["img_path", "width", "height", "prompt"]
|
||||
DELIMITER: '#;#'
|
||||
PATH_PREFIX: cache/datasets/3D_example_txt
|
||||
PROMPT_PREFIX: '<sce> '
|
||||
INDEX_FILE: cache/datasets/3D_example_txt/train.txt
|
||||
#
|
||||
TRANSFORMS:
|
||||
-
|
||||
NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: FlexibleResize
|
||||
INTERPOLATION: bilinear
|
||||
SIZE: 1024
|
||||
INPUT_KEY: ['img']
|
||||
OUTPUT_KEY: ['img']
|
||||
BACKEND: pillow
|
||||
-
|
||||
NAME: FlexibleCropXL
|
||||
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: ['img']
|
||||
BACKEND: torchvision
|
||||
-
|
||||
NAME: Select
|
||||
KEYS: ['img', 'prompt', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left']
|
||||
META_KEYS: [ 'img_path' ]
|
||||
-
|
||||
NAME: Rename
|
||||
INPUT_KEY: ['img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left']
|
||||
OUTPUT_KEY: ['image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left']
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: Text2ImageDataset
|
||||
MODE: eval
|
||||
PROMPT_FILE: cache/datasets/3D_example_txt/test.txt
|
||||
IMAGE_SIZE: [ 1024, 1024 ]
|
||||
FIELDS: ["prompt"]
|
||||
DELIMITER: '#;#'
|
||||
PROMPT_PREFIX: '<sce> '
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 4
|
||||
NUM_WORKERS: 4
|
||||
FILE_SYSTEM:
|
||||
NAME: LocalFs
|
||||
AUTO_CLEAN: False
|
||||
#
|
||||
TRANSFORMS:
|
||||
-
|
||||
NAME: Select
|
||||
KEYS: ['index', 'prompt']
|
||||
META_KEYS: ['image_size']
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 1000
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
+25
-20
@@ -15,14 +15,14 @@ SOLVER:
|
||||
ACCU_STEP: 1
|
||||
EVAL_INTERVAL: 100
|
||||
#
|
||||
WORK_DIR: ./cache/t2i_sdxl_1024_sce
|
||||
WORK_DIR: ./cache/save_data/sdxl_1024_sce_t2i_swift
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
TUNER:
|
||||
#
|
||||
TUNER:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [1280, 1280, 640, 640, 640, 320, 320, 320, 320]
|
||||
@@ -128,7 +128,7 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
ARCH: ViT-bigG-14
|
||||
PRETRAINED_MODEL:
|
||||
@@ -141,21 +141,21 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
@@ -207,21 +207,21 @@ SOLVER:
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "prompt" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "original_size_as_tuple" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: [ "crop_coords_top_left" ]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
-
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
@@ -237,8 +237,8 @@ SOLVER:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 5.0
|
||||
GUIDE_RESCALE:
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
@@ -294,8 +294,8 @@ SOLVER:
|
||||
KEYS: [ 'img', 'prompt', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
META_KEYS: [ 'data_key', 'img_path' ]
|
||||
- NAME: Rename
|
||||
IN_KEYS: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUT_KEYS: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
INPUT_KEY: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
OUTPUT_KEY: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
#
|
||||
EVAL_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
@@ -321,15 +321,20 @@ SOLVER:
|
||||
META_KEYS: [ 'image_size' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
- NAME: BackwardHook
|
||||
-
|
||||
NAME: BackwardHook
|
||||
PRIORITY: 0
|
||||
- NAME: LogHook
|
||||
-
|
||||
NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
- NAME: CheckpointHook
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
INTERVAL: 1000
|
||||
- NAME: ProbeDataHook
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
#
|
||||
EVAL_HOOKS:
|
||||
- NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
-
|
||||
NAME: ProbeDataHook
|
||||
PROB_INTERVAL: 100
|
||||
@@ -0,0 +1,19 @@
|
||||
CONTROLLERS:
|
||||
- NAME: canny
|
||||
NAME_ZH:
|
||||
DESCRIPTION:
|
||||
BASE_MODEL: SD2.1
|
||||
TYPE: Canny
|
||||
MODEL_PATH: ms://damo/scepter_scedit@controllable_model/SD2.1/canny_control/0_SwiftSCETuning
|
||||
- NAME: openpose
|
||||
NAME_ZH:
|
||||
DESCRIPTION:
|
||||
BASE_MODEL: SD2.1
|
||||
TYPE: Openpose
|
||||
MODEL_PATH: ms://damo/scepter_scedit@controllable_model/SD2.1/pose_control/0_SwiftSCETuning
|
||||
- NAME: color
|
||||
NAME_ZH:
|
||||
DESCRIPTION:
|
||||
BASE_MODEL: SD2.1
|
||||
TYPE: Color
|
||||
MODEL_PATH: ms://damo/scepter_scedit@controllable_model/SD2.1/color_control/0_SwiftSCETuning
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,721 @@
|
||||
TUNERS:
|
||||
-
|
||||
NAME: Caricature
|
||||
NAME_ZH: 夸张漫画
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/894f40ed44b37c3372e6a22b8ae577a4.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/Caricature
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Caricature
|
||||
NAME_ZH: 夸张漫画
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/894f40ed44b37c3372e6a22b8ae577a4.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/Caricature
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Caricature
|
||||
NAME_ZH: 夸张漫画
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/894f40ed44b37c3372e6a22b8ae577a4.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/Caricature
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Color Field Painting
|
||||
NAME_ZH: 色域绘画
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/80e5b4075c572c04cbb4e48c37b8366b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/ColorFieldPainting
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Color Field Painting
|
||||
NAME_ZH: 色域绘画
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/80e5b4075c572c04cbb4e48c37b8366b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/ColorFieldPainting
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Color Field Painting
|
||||
NAME_ZH: 色域绘画
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/80e5b4075c572c04cbb4e48c37b8366b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/ColorFieldPainting
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Colored Pencil Art
|
||||
NAME_ZH: 彩色铅笔艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/9ae235d7f1a7c2a4edab52a5e9f9cbae.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/ColoredPencilArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Colored Pencil Art
|
||||
NAME_ZH: 彩色铅笔艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/9ae235d7f1a7c2a4edab52a5e9f9cbae.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/ColoredPencilArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Colored Pencil Art
|
||||
NAME_ZH: 彩色铅笔艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/9ae235d7f1a7c2a4edab52a5e9f9cbae.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/ColoredPencilArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Dark Moody Atmosphere
|
||||
NAME_ZH: 暗色忧郁氛围
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/3da915da2f5cedaf243e57e08163f35b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/DarkMoodyAtmosphere
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Dark Moody Atmosphere
|
||||
NAME_ZH: 暗色忧郁氛围
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/3da915da2f5cedaf243e57e08163f35b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/DarkMoodyAtmosphere
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Dark Moody Atmosphere
|
||||
NAME_ZH: 暗色忧郁氛围
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/3da915da2f5cedaf243e57e08163f35b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/DarkMoodyAtmosphere
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Dripping Paint Splatter Art
|
||||
NAME_ZH: 滴漆溅画艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/69fd81f5983107acc3d334af62915851.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/DrippingPaintSplatterArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Dripping Paint Splatter Art
|
||||
NAME_ZH: 滴漆溅画艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/69fd81f5983107acc3d334af62915851.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/DrippingPaintSplatterArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Dripping Paint Splatter Art
|
||||
NAME_ZH: 滴漆溅画艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/69fd81f5983107acc3d334af62915851.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/DrippingPaintSplatterArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Faded Polaroid Photo
|
||||
NAME_ZH: 褪色的宝丽来照片
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/f152edb4b3ca6248758b48115258ddfa.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/FadedPolaroidPhoto
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Faded Polaroid Photo
|
||||
NAME_ZH: 褪色的宝丽来照片
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/f152edb4b3ca6248758b48115258ddfa.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/FadedPolaroidPhoto
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Faded Polaroid Photo
|
||||
NAME_ZH: 褪色的宝丽来照片
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/f152edb4b3ca6248758b48115258ddfa.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/FadedPolaroidPhoto
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Flat 2D Art
|
||||
NAME_ZH: 扁平2D艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/940cfd34155634cf051e1b2942cca426.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/Flat2DArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Flat 2D Art
|
||||
NAME_ZH: 扁平2D艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/940cfd34155634cf051e1b2942cca426.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/Flat2DArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Flat 2D Art
|
||||
NAME_ZH: 扁平2D艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/940cfd34155634cf051e1b2942cca426.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/Flat2DArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Graffiti Art
|
||||
NAME_ZH: 涂鸦艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/57b751b11564cb22cd49ef21f2004a5f.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/GraffitiArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Graffiti Art
|
||||
NAME_ZH: 涂鸦艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/57b751b11564cb22cd49ef21f2004a5f.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/GraffitiArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Graffiti Art
|
||||
NAME_ZH: 涂鸦艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/57b751b11564cb22cd49ef21f2004a5f.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/GraffitiArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Impressionism
|
||||
NAME_ZH: 印象主义
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/0312b673dc6858a9864d7f45f0c5c1fc.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/Impressionism
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Impressionism
|
||||
NAME_ZH: 印象主义
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/0312b673dc6858a9864d7f45f0c5c1fc.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/Impressionism
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Impressionism
|
||||
NAME_ZH: 印象主义
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/0312b673dc6858a9864d7f45f0c5c1fc.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/Impressionism
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Logo Design
|
||||
NAME_ZH: 标志设计
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/9aa040b0c60d289da9610c91ad9b7c7e.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/LogoDesign
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Logo Design
|
||||
NAME_ZH: 标志设计
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/9aa040b0c60d289da9610c91ad9b7c7e.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/LogoDesign
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Logo Design
|
||||
NAME_ZH: 标志设计
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/9aa040b0c60d289da9610c91ad9b7c7e.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/LogoDesign
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Pencil Sketch Drawing
|
||||
NAME_ZH: 铅笔素描
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/a9056e1eac85e5e4fe96a93917d4cce4.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/PencilSketchDrawing
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Pencil Sketch Drawing
|
||||
NAME_ZH: 铅笔素描
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/a9056e1eac85e5e4fe96a93917d4cce4.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/PencilSketchDrawing
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Pencil Sketch Drawing
|
||||
NAME_ZH: 铅笔素描
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/a9056e1eac85e5e4fe96a93917d4cce4.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/PencilSketchDrawing
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Silhouette Art
|
||||
NAME_ZH: 剪影艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/568777f447fc02510b618152726d5002.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/SilhouetteArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Silhouette Art
|
||||
NAME_ZH: 剪影艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/568777f447fc02510b618152726d5002.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/SilhouetteArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Silhouette Art
|
||||
NAME_ZH: 剪影艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/568777f447fc02510b618152726d5002.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/SilhouetteArt
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Steampunk 2
|
||||
NAME_ZH: 蒸汽朋克
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/07d7b27cd73f2d43684003563511c15b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/Steampunk2
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Steampunk 2
|
||||
NAME_ZH: 蒸汽朋克
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/07d7b27cd73f2d43684003563511c15b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/Steampunk2
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Steampunk 2
|
||||
NAME_ZH: 蒸汽朋克
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/07d7b27cd73f2d43684003563511c15b.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/Steampunk2
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Sticker Designs
|
||||
NAME_ZH: 贴纸设计
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/2d1e9867058db2c57f2fe47530de3243.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/StickerDesigns
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Sticker Designs
|
||||
NAME_ZH: 贴纸设计
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/2d1e9867058db2c57f2fe47530de3243.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/StickerDesigns
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Sticker Designs
|
||||
NAME_ZH: 贴纸设计
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/2d1e9867058db2c57f2fe47530de3243.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/StickerDesigns
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Watercolor 2
|
||||
NAME_ZH: 水彩
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/8859d532ae5901cc8457d6118fb9b7da.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/Watercolor2
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Watercolor 2
|
||||
NAME_ZH: 水彩
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/8859d532ae5901cc8457d6118fb9b7da.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/Watercolor2
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: Watercolor 2
|
||||
NAME_ZH: 水彩
|
||||
DESCRIPTION:
|
||||
SOURCE: diva
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/8859d532ae5901cc8457d6118fb9b7da.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/Watercolor2
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-elemental-art
|
||||
NAME_ZH: MRE元素艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/5895d78cf58c1ca05178991f37cc48ff.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/mre-elemental-art
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-elemental-art
|
||||
NAME_ZH: MRE元素艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/5895d78cf58c1ca05178991f37cc48ff.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/mre-elemental-art
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-elemental-art
|
||||
NAME_ZH: MRE元素艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/5895d78cf58c1ca05178991f37cc48ff.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/mre-elemental-art
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-anime
|
||||
NAME_ZH: MRE动漫
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/a08149bc8e50f6bc65c0010d4cd416f8.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/mre-anime
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-anime
|
||||
NAME_ZH: MRE动漫
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/a08149bc8e50f6bc65c0010d4cd416f8.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/mre-anime
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-anime
|
||||
NAME_ZH: MRE动漫
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/a08149bc8e50f6bc65c0010d4cd416f8.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/mre-anime
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-comic
|
||||
NAME_ZH: MRE漫画书
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/48c65cebf1fa4284d7b8feb619412e65.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/mre-comic
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-comic
|
||||
NAME_ZH: MRE漫画书
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/48c65cebf1fa4284d7b8feb619412e65.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/mre-comic
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: mre-comic
|
||||
NAME_ZH: MRE漫画书
|
||||
DESCRIPTION:
|
||||
SOURCE: mre
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/48c65cebf1fa4284d7b8feb619412e65.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/mre-comic
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-craft clay
|
||||
NAME_ZH: SAI手工粘土
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/8fc51113f725f27326c4398a7457cd6d.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/sai-craftclay
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-craft clay
|
||||
NAME_ZH: SAI手工粘土
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/8fc51113f725f27326c4398a7457cd6d.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/sai-craftclay
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-craft clay
|
||||
NAME_ZH: SAI手工粘土
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/8fc51113f725f27326c4398a7457cd6d.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/sai-craftclay
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-fantasy art
|
||||
NAME_ZH: SAI幻想艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/a6f8d92afcd5803dfb2ebecbc92091b6.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/sai-fantasyart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-fantasy art
|
||||
NAME_ZH: SAI幻想艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/a6f8d92afcd5803dfb2ebecbc92091b6.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/sai-fantasyart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-fantasy art
|
||||
NAME_ZH: SAI幻想艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/a6f8d92afcd5803dfb2ebecbc92091b6.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/sai-fantasyart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-line art
|
||||
NAME_ZH: SAI线条艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/034a51b0dd34b018be8859bf45b4f7ed.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/sai-lineart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-line art
|
||||
NAME_ZH: SAI线条艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/034a51b0dd34b018be8859bf45b4f7ed.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/sai-lineart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-line art
|
||||
NAME_ZH: SAI线条艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/034a51b0dd34b018be8859bf45b4f7ed.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/sai-lineart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-neonpunk
|
||||
NAME_ZH: SAI霓虹朋克
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/7e9ed25bb34008beb5f417df63c4b2fe.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/sai-neonpunk
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-neonpunk
|
||||
NAME_ZH: SAI霓虹朋克
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/7e9ed25bb34008beb5f417df63c4b2fe.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/sai-neonpunk
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-neonpunk
|
||||
NAME_ZH: SAI霓虹朋克
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/7e9ed25bb34008beb5f417df63c4b2fe.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/sai-neonpunk
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-origami
|
||||
NAME_ZH: SAI折纸
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/924f46a8f276011a0953d7988e90ee25.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/sai-origami
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-origami
|
||||
NAME_ZH: SAI折纸
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/924f46a8f276011a0953d7988e90ee25.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/sai-origami
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-origami
|
||||
NAME_ZH: SAI折纸
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/924f46a8f276011a0953d7988e90ee25.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/sai-origami
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-pixel art
|
||||
NAME_ZH: SAI像素艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD_XL1.0
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD_XL1.0/a5ab89c0960be8c1216e65c98d92ae4a.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD_XL1.0/sai-pixelart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-pixel art
|
||||
NAME_ZH: SAI像素艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD2.1
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD2.1/a5ab89c0960be8c1216e65c98d92ae4a.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD2.1/sai-pixelart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
-
|
||||
NAME: sai-pixel art
|
||||
NAME_ZH: SAI像素艺术
|
||||
DESCRIPTION:
|
||||
SOURCE: sai
|
||||
BASE_MODEL: SD1.5
|
||||
IMAGE_PATH: ms://damo/scepter@mantra_images/SD1.5/a5ab89c0960be8c1216e65c98d92ae4a.png
|
||||
MODEL_PATH: ms://damo/scepter_scedit@tuners_model/SD1.5/sai-pixelart
|
||||
TUNER_TYPE: SwiftSCE
|
||||
PROMPT_EXAMPLE: a boy wearing green jacket
|
||||
@@ -0,0 +1,95 @@
|
||||
WORK_DIR: home
|
||||
FILE_SYSTEM:
|
||||
-
|
||||
NAME: LocalFs
|
||||
AUTO_CLEAN: False
|
||||
DESC_INFO:
|
||||
ZH_INFO: |
|
||||
<h2><center>基本介绍</center><h2>
|
||||
<p align="center">
|
||||
<table align="center">
|
||||
<tr>
|
||||
<td>
|
||||
<img src="https://modelscope.cn/api/v1/models/damo/scepter/repo?Revision=master&FilePath=assets/scepter_studio/scepter_studio_banner.jpg">
|
||||
<h3><center>SCEPTER Studio是基于开源基模型和自研微调编辑算法构建的生成定制和编辑工具箱，提供围绕生成、微调、编辑、数据处理等一系列的工具和插件。</center><h3>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</p>
|
||||
EN_INFO: |
|
||||
<h2><center>Introduction</center><h2>
|
||||
<p align="center">
|
||||
<table align="center">
|
||||
<tr>
|
||||
<td>
|
||||
<img src="https://modelscope.cn/api/v1/models/damo/scepter/repo?Revision=master&FilePath=assets/scepter_studio/scepter_studio_banner.jpg">
|
||||
<h3><center>SCEPTER Studio is a customized generation and editing toolkit built on the open-source base models and proprietary fine-tuning editing algorithms, offering a range of tools and plugins centered around generation, fine-tuning, editing, and data processing.</center><h3>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</p>
|
||||
GUIDE_INFO:
|
||||
ZH_INFO: |
|
||||
<h2><center>用户指南</center><h2>
|
||||
<style>
|
||||
.video-container {
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
width: 100%;
|
||||
}
|
||||
.video-wrapper {
|
||||
width: 75%;
|
||||
}
|
||||
video {
|
||||
width: 100%;
|
||||
display: block;
|
||||
}
|
||||
.description {
|
||||
text-align: center;
|
||||
margin-top: 10px;
|
||||
font-size: 0.8em;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="video-container">
|
||||
<div class="video-wrapper">
|
||||
<video controls>
|
||||
<source src="https://modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=assets/scepter_studio/scepter_studio_train_inference.webm" type="video/webm">
|
||||
</video>
|
||||
<div class="description">训练与推理演示</div>
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
EN_INFO: |
|
||||
<h2><center>User Guide</center><h2>
|
||||
<style>
|
||||
.video-container {
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
width: 100%;
|
||||
}
|
||||
.video-wrapper {
|
||||
width: 75%;
|
||||
}
|
||||
video {
|
||||
width: 100%;
|
||||
display: block;
|
||||
}
|
||||
.description {
|
||||
text-align: center;
|
||||
margin-top: 10px;
|
||||
font-size: 0.8em;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="video-container">
|
||||
<div class="video-wrapper">
|
||||
<video controls>
|
||||
<source src="https://modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=assets/scepter_studio/scepter_studio_train_inference.webm" type="video/webm">
|
||||
</video>
|
||||
<div class="description">Train & Inference Video</div>
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
@@ -0,0 +1,117 @@
|
||||
WORK_DIR: "inference"
|
||||
DIFFUSION_PARAS:
|
||||
SAMPLE:
|
||||
VALUES: ['ddim', 'euler', 'euler_ancestral', 'heun', 'dpm2',
|
||||
'dpm2_ancestral', 'dpmpp_2m', 'dpmpp_sde', 'dpmpp_2m_sde', 'dpmpp_2s_ancestral',
|
||||
'dpm2_karras', 'dpm2_ancestral_karras', 'dpmpp_2s_ancestral_karras', 'dpmpp_2m_karras',
|
||||
'dpmpp_sde_karras', 'dpmpp_2m_sde_karras']
|
||||
DEFAULT: 'dpmpp_2s_ancestral'
|
||||
NEGATIVE_PROMPT:
|
||||
DEFAULT:
|
||||
PROMPT_PREFIX:
|
||||
DEFAULT:
|
||||
SAMPLES:
|
||||
MIN: 1
|
||||
MAX: 4
|
||||
DEFAULT: 1
|
||||
SAMPLE_STEPS:
|
||||
MIN: 1
|
||||
MAX: 100
|
||||
DEFAULT: 30
|
||||
GUIDE_SCALE:
|
||||
MIN: 0
|
||||
MAX: 10
|
||||
DEFAULT: 5.0
|
||||
GUIDE_RESCALE:
|
||||
MIN: 0
|
||||
MAX: 1.0
|
||||
DEFAULT: 0.5
|
||||
DISCRETIZATION:
|
||||
VALUES: ["trailing", "leading", "linspace"]
|
||||
DEFAULT: "linspace"
|
||||
REFINE_SAMPLERS:
|
||||
VALUES: [ 'ddim', 'euler', 'euler_ancestral', 'heun', 'dpm2',
|
||||
'dpm2_ancestral', 'dpmpp_2m', 'dpmpp_sde', 'dpmpp_2m_sde', 'dpmpp_2s_ancestral',
|
||||
'dpm2_karras', 'dpm2_ancestral_karras', 'dpmpp_2s_ancestral_karras', 'dpmpp_2m_karras',
|
||||
'dpmpp_sde_karras', 'dpmpp_2m_sde_karras' ]
|
||||
DEFAULT: 'dpmpp_2s_ancestral'
|
||||
REFINE_SAMPLE_STEPS:
|
||||
MIN: 0
|
||||
MAX: 100
|
||||
DEFAULT: 30
|
||||
REFINE_GUIDE_SCALE:
|
||||
MIN: 0
|
||||
MAX: 10
|
||||
DEFAULT: 5.0
|
||||
REFINE_GUIDE_RESCALE:
|
||||
MIN: 0
|
||||
MAX: 1.0
|
||||
DEFAULT: 0.5
|
||||
REFINE_DISCRETIZATION:
|
||||
VALUES: [ "trailing", "leading", "linspace" ]
|
||||
DEFAULT: "linspace"
|
||||
AESTHETIC_SCORE:
|
||||
MIN: 0.0
|
||||
MAX: 10.0
|
||||
DEFAULT: 6.0
|
||||
NEGATIVE_AESTHETIC_SCORE:
|
||||
MIN: 0.0
|
||||
MAX: 10.0
|
||||
DEFAULT: 2.5
|
||||
REFINE_STRENGTH:
|
||||
MIN: 0
|
||||
MAX: 1.0
|
||||
DEFAULT: 0.15
|
||||
RESOLUTIONS:
|
||||
VALUES: [[704, 1408], [704, 1344], [768, 1344],
|
||||
[720, 1280],
|
||||
[768, 1280], [832, 1216], [832, 1152],
|
||||
[896, 1152], [896, 1088], [960, 1088],
|
||||
[960, 1024], [1024, 1024], [1024, 960],
|
||||
[1088, 960], [1088, 896], [1152, 896],
|
||||
[1152, 832], [1216, 832], [1280, 720],
|
||||
[1280, 768],
|
||||
[1344, 768], [1344, 704], [1408, 704],
|
||||
[1472, 704], [1536, 640], [1600, 640],
|
||||
[1664, 576], [1728, 576]]
|
||||
DEFAULT: [1024, 1024]
|
||||
EXTENSION_PARAS:
|
||||
MANTRA_BOOK: scepter/methods/studio/extensions/mantra_book/mantra_book.yaml
|
||||
OFFICIAL_TUNERS: scepter/methods/studio/extensions/tuners/official_tuners.yaml
|
||||
OFFICIAL_CONTROLLERS: scepter/methods/studio/extensions/controllers/official_controllers.yaml
|
||||
CONTROLABLE_ANNOTATORS:
|
||||
-
|
||||
NAME: "CannyAnnotator"
|
||||
LOW_THRESHOLD: 100
|
||||
HIGH_THRESHOLD: 200
|
||||
TYPE: Canny
|
||||
IS_DEFAULT: True
|
||||
-
|
||||
NAME: "HedAnnotator"
|
||||
PRETRAINED_MODEL: "ms://damo/scepter_scedit@annotator/ckpts/ControlNetHED.pth"
|
||||
TYPE: Hed
|
||||
IS_DEFAULT: False
|
||||
-
|
||||
NAME: "OpenposeAnnotator"
|
||||
BODY_MODEL_PATH: "ms://damo/scepter_scedit@annotator/ckpts/body_pose_model.pth"
|
||||
HAND_MODEL_PATH: "ms://damo/scepter_scedit@annotator/ckpts/hand_pose_model.pth"
|
||||
TYPE: Openpose
|
||||
IS_DEFAULT: False
|
||||
-
|
||||
NAME: "MidasDetector"
|
||||
PRETRAINED_MODEL: "ms://damo/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
|
||||
A: 6.2
|
||||
BG_TH: 0.1
|
||||
TYPE: Midas
|
||||
IS_DEFAULT: False
|
||||
-
|
||||
NAME: "ColorAnnotator"
|
||||
TYPE: Color
|
||||
IS_DEFAULT: False
|
||||
-
|
||||
NAME: "MLSDdetector"
|
||||
PRETRAINED_MODEL: "ms://damo/scepter_scedit@annotator/ckpts/mlsd_large_512_fp32.pth"
|
||||
THR_V: 0.1
|
||||
THR_D: 0.1
|
||||
TYPE: MLSD
|
||||
IS_DEFAULT: False
|
||||
@@ -0,0 +1,242 @@
|
||||
NAME: SD_XL1.0
|
||||
IS_DEFAULT: True
|
||||
DEFAULT_PARAS:
|
||||
PARAS:
|
||||
RESOLUTIONS: [[1024, 1024]]
|
||||
INPUT:
|
||||
IMAGE:
|
||||
ORIGINAL_SIZE_AS_TUPLE: [1024, 1024]
|
||||
TARGET_SIZE_AS_TUPLE: [1024, 1024]
|
||||
AESTHETIC_SCORE: 6.0
|
||||
NEGATIVE_AESTHETIC_SCORE: 2.5
|
||||
PROMPT: ""
|
||||
NEGATIVE_PROMPT: ""
|
||||
PROMPT_PREFIX: ""
|
||||
CROP_COORDS_TOP_LEFT: [0, 0]
|
||||
SAMPLE: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
REFINE_SAMPLE: ddim
|
||||
REFINE_GUIDE_SCALE: 7.5
|
||||
REFINE_GUIDE_RESCALE: 0.5
|
||||
REFINE_DISCRETIZATION: trailing
|
||||
OUTPUT:
|
||||
LATENT:
|
||||
BEFORE_REFINE_IMAGES:
|
||||
IMAGES:
|
||||
SEED:
|
||||
MODULES_PARAS:
|
||||
FIRST_STAGE_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode
|
||||
DTYPE: float32
|
||||
INPUT: ["IMAGE"]
|
||||
-
|
||||
NAME: decode
|
||||
DTYPE: float32
|
||||
INPUT: ["LATENT"]
|
||||
PARAS:
|
||||
# SCALE_FACTOR DESCRIPTION: The vae embeding scale. TYPE: float default: 0.18215
|
||||
SCALE_FACTOR: 0.13025
|
||||
SIZE_FACTOR: 8
|
||||
DIFFUSION_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: forward
|
||||
DTYPE: float16
|
||||
INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE", "GUIDE_RESCALE", "DISCRETIZATION"]
|
||||
COND_STAGE_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode
|
||||
DTYPE: float16
|
||||
INPUT: ["ORIGINAL_SIZE_AS_TUPLE", "CROP_COORDS_TOP_LEFT", "PROMPT", "NEGATIVE_PROMPT"]
|
||||
REFINER_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: forward
|
||||
DTYPE: float16
|
||||
INPUT: ["SAMPLE_STEPS", "REFINE_SAMPLE", "REFINE_GUIDE_SCALE", "REFINE_GUIDE_RESCALE", "REFINE_DISCRETIZATION"]
|
||||
REFINER_COND_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode
|
||||
DTYPE: float16
|
||||
INPUT: ["ORIGINAL_SIZE_AS_TUPLE", "AESTHETIC_SCORE", "NEGATIVE_AESTHETIC_SCORE", "CROP_COORDS_TOP_LEFT", "PROMPT", "NEGATIVE_PROMPT"]
|
||||
|
||||
MODEL:
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
# SCHEDULE_ARGS DESCRIPTION: TYPE: default: ''
|
||||
SCHEDULE:
|
||||
PARAMETERIZATION: "eps"
|
||||
TIMESTEPS: 1000
|
||||
ZERO_TERMINAL_SNR: False
|
||||
SCHEDULE_ARGS:
|
||||
# NAME DESCRIPTION: TYPE: default: ''
|
||||
NAME: "scaled_linear"
|
||||
BETA_MIN: 0.00085
|
||||
BETA_MAX: 0.0120
|
||||
# DIFFUSION_MODEL DESCRIPTION: TYPE: default: ''
|
||||
DIFFUSION_MODEL:
|
||||
# NAME DESCRIPTION: TYPE: default: 'DiffusionUNetXL'
|
||||
NAME: DiffusionUNetXL
|
||||
# PRETRAINED_MODEL DESCRIPTION: Whole model's pretrained model path. TYPE: NoneType default: None
|
||||
PRETRAINED_MODEL:
|
||||
# IN_CHANNELS DESCRIPTION: Unet channels for input, considering the input image's channels. TYPE: int default: 4
|
||||
IN_CHANNELS: 4
|
||||
# OUT_CHANNELS DESCRIPTION: Unet channels for output, considering the input image's channels. TYPE: int default: 4
|
||||
OUT_CHANNELS: 4
|
||||
# NUM_RES_BLOCKS DESCRIPTION: The blocks's number of res. TYPE: int default: 2
|
||||
NUM_RES_BLOCKS: 2
|
||||
# MODEL_CHANNELS DESCRIPTION: base channel count for the model. TYPE: int default: 320
|
||||
MODEL_CHANNELS: 320
|
||||
# ATTENTION_RESOLUTIONS DESCRIPTION: A collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attentio will be used. TYPE: list default: [4, 2]
|
||||
ATTENTION_RESOLUTIONS: [4, 2]
|
||||
# DROPOUT DESCRIPTION: The dropout rate. TYPE: int default: 0
|
||||
DROPOUT: 0
|
||||
# CHANNEL_MULT DESCRIPTION: channel multiplier for each level of the UNet. TYPE: list default: [1, 2, 4]
|
||||
CHANNEL_MULT: [1, 2, 4]
|
||||
# CONV_RESAMPLE DESCRIPTION: Use conv to resample when downsample. TYPE: bool default: True
|
||||
CONV_RESAMPLE: True
|
||||
# DIMS DESCRIPTION: The Conv dims which 2 represent Conv2D. TYPE: int default: 2
|
||||
DIMS: 2
|
||||
# NUM_CLASSES DESCRIPTION: The class num for class guided setting, also can be set as continuous. TYPE: str default: 'sequential'
|
||||
NUM_CLASSES: sequential
|
||||
# USE_CHECKPOINT DESCRIPTION: Use gradient checkpointing to reduce memory usage. TYPE: bool default: False
|
||||
USE_CHECKPOINT: False
|
||||
# NUM_HEADS DESCRIPTION: The number of attention heads in each attention layer. TYPE: int default: -1
|
||||
NUM_HEADS: -1
|
||||
# NUM_HEADS_CHANNELS DESCRIPTION: If specified, ignore num_heads and instead use a fixed channel width per attention head. TYPE: int default: 64
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
# USE_SCALE_SHIFT_NORM DESCRIPTION: The scale and shift for the outnorm of RESBLOCK, use a FiLM-like conditioning mechanism. TYPE: bool default: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
# RESBLOCK_UPDOWN DESCRIPTION: Use residual blocks for up/downsampling, if False use Conv. TYPE: bool default: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
# USE_NEW_ATTENTION_ORDER DESCRIPTION: Whether use new attention(qkv before split heads or not) or not. TYPE: bool default: True
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
# USE_SPATIAL_TRANSFORMER DESCRIPTION: Custom transformer which support the context, if context_dim is not None, the parameter must set True TYPE: bool default: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
# TRANSFORMER_DEPTH DESCRIPTION: Custom transformer's depth, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: list default: [1, 2, 10]
|
||||
TRANSFORMER_DEPTH: [1, 2, 10]
|
||||
# TRANSFORMER_DEPTH_MIDDLE DESCRIPTION: Custom transformer's depth of middle block, If set None, use TRANSFORMER_DEPTH last value. TYPE: NoneType default: None
|
||||
# TRANSFORMER_DEPTH_MIDDLE: None
|
||||
# CONTEXT_DIM DESCRIPTION: Custom context info, if set, USE_SPATIAL_TRANSFORMER also set True. TYPE: int default: 2048
|
||||
CONTEXT_DIM: 2048
|
||||
# DISABLE_SELF_ATTENTIONS DESCRIPTION: Whether disable the self-attentions on some level, should be a list, [False, True, ...] TYPE: NoneType default: None
|
||||
# DISABLE_SELF_ATTENTIONS: None
|
||||
# NUM_ATTENTION_BLOCKS DESCRIPTION: The number of attention blocks for attention layer. TYPE: NoneType default: None
|
||||
# NUM_ATTENTION_BLOCKS: None
|
||||
# DISABLE_MIDDLE_SELF_ATTN DESCRIPTION: Whether disable the self-attentions in middle blocks. TYPE: bool default: False
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
# USE_LINEAR_IN_TRANSFORMER DESCRIPTION: Custom transformer's parameter, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: bool default: True
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
# ADM_IN_CHANNELS DESCRIPTION: Used when num_classes == 'sequential' or 'timestep'. TYPE: int default: 2816
|
||||
ADM_IN_CHANNELS: 2816
|
||||
# USE_SENTENCE_EMB DESCRIPTION: Used sentence emb or not, default False. TYPE: bool default: False
|
||||
USE_SENTENCE_EMB: False
|
||||
# USE_WORD_MAPPING DESCRIPTION: Used word mapping or not, default False. TYPE: bool default: False
|
||||
USE_WORD_MAPPING: False
|
||||
# FIRST_STAGE_MODEL DESCRIPTION: TYPE: default: ''
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
# COND_STAGE_MODEL DESCRIPTION: TYPE: default: ''
|
||||
COND_STAGE_MODEL:
|
||||
# NAME DESCRIPTION: TYPE: default: 'GeneralConditioner'
|
||||
NAME: GeneralConditioner
|
||||
USE_GRAD: False
|
||||
# EMBEDDERS DESCRIPTION: TYPE: default: ''
|
||||
EMBEDDERS:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'FrozenCLIPEmbedder'
|
||||
NAME: FrozenCLIPEmbedder
|
||||
# PRETRAINED_MODEL DESCRIPTION: TYPE: str default: ''
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
# MAX_LENGTH DESCRIPTION: TYPE: int default: 77
|
||||
MAX_LENGTH: 77
|
||||
# FREEZE DESCRIPTION: TYPE: bool default: True
|
||||
FREEZE: True
|
||||
# LAYER DESCRIPTION: TYPE: str default: 'last'
|
||||
LAYER: hidden
|
||||
# LAYER_IDX DESCRIPTION: TYPE: NoneType default: None
|
||||
LAYER_IDX: 11
|
||||
# USE_FINAL_LAYER_NORM DESCRIPTION: TYPE: bool default: False
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["prompt"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'FrozenOpenCLIPEmbedder2'
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
# ARCH DESCRIPTION: TYPE: str default: 'ViT-H-14'
|
||||
ARCH: ViT-bigG-14
|
||||
# MAX_LENGTH DESCRIPTION: TYPE: int default: 77
|
||||
MAX_LENGTH: 77
|
||||
# FREEZE DESCRIPTION: TYPE: bool default: True
|
||||
FREEZE: True
|
||||
# ALWAYS_RETURN_POOLED DESCRIPTION: Whether always return pooled results or not ,default False. TYPE: bool default: False
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
# LEGACY DESCRIPTION: Whether use legacy returnd feature or not ,default True. TYPE: bool default: True
|
||||
LEGACY: False
|
||||
# LAYER DESCRIPTION: TYPE: str default: 'last'
|
||||
LAYER: penultimate
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["prompt"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["original_size_as_tuple"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["crop_coords_top_left"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["target_size_as_tuple"]
|
||||
LEGACY_UCG_VALUE:
|
||||
@@ -0,0 +1,126 @@
|
||||
NAME: SD1.5
|
||||
IS_DEFAULT: False
|
||||
DEFAULT_PARAS:
|
||||
PARAS:
|
||||
RESOLUTIONS: [[512, 512]]
|
||||
INPUT:
|
||||
IMAGE:
|
||||
PROMPT: ""
|
||||
NEGATIVE_PROMPT: ""
|
||||
TARGET_SIZE_AS_TUPLE: [512, 512]
|
||||
PROMPT_PREFIX: ""
|
||||
SAMPLE: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
OUTPUT:
|
||||
LATENT:
|
||||
IMAGES:
|
||||
SEED:
|
||||
MODULES_PARAS:
|
||||
FIRST_STAGE_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode
|
||||
DTYPE: float16
|
||||
INPUT: ["IMAGE"]
|
||||
-
|
||||
NAME: decode
|
||||
DTYPE: float16
|
||||
INPUT: ["LATENT"]
|
||||
PARAS:
|
||||
# SCALE_FACTOR DESCRIPTION: The vae embeding scale. TYPE: float default: 0.18215
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
DIFFUSION_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: forward
|
||||
DTYPE: float16
|
||||
INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE", "GUIDE_RESCALE", "DISCRETIZATION"]
|
||||
COND_STAGE_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode_text
|
||||
DTYPE: float16
|
||||
INPUT: ["PROMPT", "NEGATIVE_PROMPT"]
|
||||
|
||||
MODEL:
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-v1-5@v1-5-pruned-emaonly.safetensors
|
||||
SCHEDULE:
|
||||
PARAMETERIZATION: "eps"
|
||||
TIMESTEPS: 1000
|
||||
ZERO_TERMINAL_SNR: False
|
||||
SCHEDULE_ARGS:
|
||||
# NAME DESCRIPTION: TYPE: default: ''
|
||||
NAME: "scaled_linear"
|
||||
BETA_MIN: 0.00085
|
||||
BETA_MAX: 0.0120
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS: 8
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 768
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: False
|
||||
IGNORE_KEYS: []
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: ClipTokenizer
|
||||
PRETRAINED_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
LENGTH: 77
|
||||
CLEAN: True
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenCLIPEmbedder
|
||||
FREEZE: True
|
||||
USE_GRAD: False
|
||||
LAYER: last
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
@@ -0,0 +1,124 @@
|
||||
NAME: SD2.1
|
||||
IS_DEFAULT: False
|
||||
DEFAULT_PARAS:
|
||||
PARAS:
|
||||
RESOLUTIONS: [[768, 768]]
|
||||
INPUT:
|
||||
IMAGE:
|
||||
PROMPT: ""
|
||||
NEGATIVE_PROMPT: ""
|
||||
PROMPT_PREFIX: ""
|
||||
TARGET_SIZE_AS_TUPLE: [768, 768]
|
||||
SAMPLE: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE: 0.5
|
||||
DISCRETIZATION: trailing
|
||||
OUTPUT:
|
||||
LATENT:
|
||||
IMAGES:
|
||||
SEED:
|
||||
MODULES_PARAS:
|
||||
FIRST_STAGE_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode
|
||||
DTYPE: float16
|
||||
INPUT: ["IMAGE"]
|
||||
-
|
||||
NAME: decode
|
||||
DTYPE: float16
|
||||
INPUT: ["LATENT"]
|
||||
PARAS:
|
||||
# SCALE_FACTOR DESCRIPTION: The vae embeding scale. TYPE: float default: 0.18215
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
DIFFUSION_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: forward
|
||||
DTYPE: float16
|
||||
INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE", "GUIDE_RESCALE", "DISCRETIZATION"]
|
||||
COND_STAGE_MODEL:
|
||||
FUNCTION:
|
||||
-
|
||||
NAME: encode_text
|
||||
DTYPE: float16
|
||||
INPUT: ["PROMPT", "NEGATIVE_PROMPT"]
|
||||
|
||||
MODEL:
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-1@v2-1_768-ema-pruned.safetensors
|
||||
SCHEDULE:
|
||||
PARAMETERIZATION: "v"
|
||||
TIMESTEPS: 1000
|
||||
ZERO_TERMINAL_SNR: False
|
||||
SCHEDULE_ARGS:
|
||||
# NAME DESCRIPTION: TYPE: default: ''
|
||||
NAME: "scaled_linear"
|
||||
BETA_MIN: 0.00085
|
||||
BETA_MAX: 0.0120
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 1024
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
PRETRAINED_MODEL:
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: OpenClipTokenizer
|
||||
LENGTH: 77
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenOpenCLIPEmbedder
|
||||
ARCH: ViT-H-14
|
||||
PRETRAINED_MODEL:
|
||||
LAYER: penultimate
|
||||
@@ -0,0 +1,7 @@
|
||||
WORK_DIR: datasets
|
||||
EXPORT_DIR: export_datasets
|
||||
FILE_SYSTEM:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: ''
|
||||
NAME: LocalFs
|
||||
AUTO_CLEAN: False
|
||||
@@ -0,0 +1,85 @@
|
||||
HOST: "localhost"
|
||||
PORT: 2024
|
||||
ROOT: ""
|
||||
TITLE: SCEPTER Studio
|
||||
BANNER: |
|
||||
<style>
|
||||
.banner {
|
||||
position: relative;
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
align-items: center;
|
||||
background-color: #f0f0f0;
|
||||
color: #2a2a2a;
|
||||
padding: 35px;
|
||||
font-family: Arial, sans-serif;
|
||||
box-shadow: 0px 0px 10px rgba(0, 0, 0, 0.1);
|
||||
width: 100%;
|
||||
}
|
||||
.title {
|
||||
text-align: center;
|
||||
z-index: 1;
|
||||
}
|
||||
.qr-codes {
|
||||
position: absolute;
|
||||
right: 20px;
|
||||
display: flex;
|
||||
gap: 15px;
|
||||
}
|
||||
.qr-code-container {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
}
|
||||
.qr-codes img {
|
||||
height: 80px;
|
||||
width: 80px;
|
||||
}
|
||||
.caption {
|
||||
color: #2a2a2a;
|
||||
font-size: 0.8em;
|
||||
margin-top: 5px;
|
||||
}
|
||||
</style>
|
||||
<body>
|
||||
<div class="banner">
|
||||
<div class="title">
|
||||
<h1>🪄SCEPTER Studio</h1>
|
||||
</div>
|
||||
<div class="qr-codes">
|
||||
<div class="qr-code-container">
|
||||
<img src="https://modelscope.cn/api/v1/models/damo/scepter/repo?Revision=master&FilePath=assets/scepter_studio/ms_scepter_studio_qr.png" alt="ms_scepter_studio_qr">
|
||||
<div class="caption">Modelscope Studio</div>
|
||||
</div>
|
||||
<div class="qr-code-container">
|
||||
<img src="https://modelscope.cn/api/v1/models/damo/scepter/repo?Revision=master&FilePath=assets/scepter_studio/scepter_github_qr.png" alt="scepter_github_qr">
|
||||
<div class="caption">Github</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
WORK_DIR: "cache/scepter_ui"
|
||||
FILE_SYSTEM:
|
||||
-
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "cache/cache_data"
|
||||
-
|
||||
NAME: "HttpFs"
|
||||
TEMP_DIR: "cache/cache_data"
|
||||
INTERFACE:
|
||||
- NAME: 首页
|
||||
NAME_EN: Home
|
||||
IFID: home
|
||||
CONFIG: scepter/methods/studio/home/home.yaml
|
||||
- NAME: 数据管理
|
||||
NAME_EN: Dataset Management
|
||||
IFID: preprocess
|
||||
CONFIG: scepter/methods/studio/preprocess/preprocess.yaml
|
||||
- NAME: 训练
|
||||
NAME_EN: Train
|
||||
IFID: self_train
|
||||
CONFIG: scepter/methods/studio/self_train/self_train.yaml
|
||||
- NAME: 推理
|
||||
NAME_EN: Inference
|
||||
IFID: inference
|
||||
CONFIG: scepter/methods/studio/inference/inference.yaml
|
||||
@@ -0,0 +1,600 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
|
||||
META:
|
||||
VERSION: 'SD_XL1.0'
|
||||
DESCRIPTION: "Stable Diffusion XL1.0"
|
||||
IS_DEFAULT: True
|
||||
INFERENCE_PARAS:
|
||||
INFERENCE_BATCH_SIZE: 1
|
||||
INFERENCE_PREFIX: ""
|
||||
DEFAULT_SAMPLER: "dpmpp_2s_ancestral"
|
||||
DEFAULT_SAMPLE_STEPS: 40
|
||||
INFERENCE_N_PROMPT: ""
|
||||
RESOLUTION: 1024
|
||||
PARAS:
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 1024
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: FULL
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 1024
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: LORA
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 1024
|
||||
MEMORY: 29000
|
||||
EPOCHS: 200
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: SCE
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 1024
|
||||
MEMORY: 29000
|
||||
EPOCHS: 200
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: True
|
||||
TUNER: TEXT_SCE
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 1024
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: TEXT_LORA
|
||||
|
||||
TUNERS:
|
||||
LORA:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2)$"
|
||||
TEXT_LORA:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "(cond_stage_model.embedders.0.*(q_proj|k_proj|v_proj|out_proj|mlp.fc1|mlp.fc2))|(model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2))$"
|
||||
SCE:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [1280, 1280, 640, 640, 640, 320, 320, 320, 320]
|
||||
TARGET_MODULES: model.lsc_identity\.\d+$
|
||||
DOWN_RATIO: 1.0
|
||||
TUNER_MODE: identity
|
||||
TEXT_SCE:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [ 1280, 1280, 640, 640, 640, 320, 320, 320, 320 ]
|
||||
TARGET_MODULES: model.lsc_identity\.\d+$
|
||||
DOWN_RATIO: 1.0
|
||||
TUNER_MODE: identity
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "cond_stage_model.embedders.0.*(q_proj|k_proj|v_proj|out_proj|mlp.fc1|mlp.fc2)$"
|
||||
|
||||
MODIFY_PARAS:
|
||||
TEXT_LORA:
|
||||
TRAIN:
|
||||
SOLVER.MODEL.COND_STAGE_MODEL.USE_GRAD: True
|
||||
TEXT_SCE:
|
||||
TRAIN:
|
||||
SOLVER.MODEL.COND_STAGE_MODEL.USE_GRAD: True
|
||||
|
||||
SOLVER:
|
||||
# NAME DESCRIPTION: TYPE: default: 'LatentDiffusionSolver'
|
||||
NAME: LatentDiffusionSolver
|
||||
# MAX_STEPS DESCRIPTION: The total steps for training. TYPE: int default: 100000
|
||||
MAX_STEPS: 2000
|
||||
# USE_AMP DESCRIPTION: Use amp to surpport mix precision or not, default is False. TYPE: bool default: False
|
||||
USE_AMP: False
|
||||
# DTYPE DESCRIPTION: The precision for training. TYPE: str default: 'float32'
|
||||
DTYPE: float16
|
||||
# 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: False
|
||||
# SHARDING_STRATEGY DESCRIPTION: The shard strategy for fsdp, select from ['full_shard', 'shard_grad_op'] TYPE: str default: 'shard_grad_op'
|
||||
SHARDING_STRATEGY:
|
||||
# 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
|
||||
# CHANNELS_LAST DESCRIPTION: The channels last, default is False. TYPE: bool default: False
|
||||
CHANNELS_LAST: False
|
||||
# RESUME_FROM DESCRIPTION: Resume from some state of training! TYPE: str default: ''
|
||||
RESUME_FROM:
|
||||
# MAX_EPOCHS DESCRIPTION: Max epochs for training. TYPE: int default: 10
|
||||
MAX_EPOCHS: -1
|
||||
# NUM_FOLDS DESCRIPTION: Num folds for training. TYPE: int default: 1
|
||||
NUM_FOLDS: 1
|
||||
# WORK_DIR DESCRIPTION: Save dir of the training log or model. TYPE: str default: ''
|
||||
WORK_DIR:
|
||||
# LOG_FILE DESCRIPTION: Save log path. TYPE: str default: ''
|
||||
LOG_FILE: stg_log.txt
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
TUNER:
|
||||
# MODEL DESCRIPTION: TYPE: default: ''
|
||||
MODEL:
|
||||
# NAME DESCRIPTION:
|
||||
NAME: LatentDiffusionXL
|
||||
# PARAMETERIZATION DESCRIPTION: The prediction type, you can choose from 'eps' and 'x0' and 'v' TYPE: str default: 'v'
|
||||
PARAMETERIZATION: eps
|
||||
# TIMESTEPS DESCRIPTION: The schedule steps for diffusion. TYPE: int default: 1000
|
||||
TIMESTEPS: 1000
|
||||
# MIN_SNR_GAMMA DESCRIPTION: The minimum snr gamma, default is None. TYPE: NoneType default: None
|
||||
# MIN_SNR_GAMMA: None
|
||||
# ZERO_TERMINAL_SNR DESCRIPTION: Whether zero terminal snr, default is False. TYPE: bool default: False
|
||||
ZERO_TERMINAL_SNR: False
|
||||
# PRETRAINED_MODEL DESCRIPTION: Whole model's pretrained model path. TYPE: NoneType default: None
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-xl-base-1.0@sd_xl_base_1.0.safetensors
|
||||
# IGNORE_KEYS DESCRIPTION: The ignore keys for pretrain model loaded. TYPE: list default: []
|
||||
IGNORE_KEYS: [ ]
|
||||
# SCALE_FACTOR DESCRIPTION: The vae embeding scale. TYPE: float default: 0.18215
|
||||
SCALE_FACTOR: 0.13025
|
||||
# SIZE_FACTOR DESCRIPTION: The vae size factor. TYPE: int default: 8
|
||||
SIZE_FACTOR: 8
|
||||
# DEFAULT_N_PROMPT DESCRIPTION: The default negtive prompt. TYPE: str default: ''
|
||||
DEFAULT_N_PROMPT: ""
|
||||
# TRAIN_N_PROMPT DESCRIPTION: The negtive prompt used in train phase. TYPE: str default: ''
|
||||
TRAIN_N_PROMPT: ""
|
||||
# P_ZERO DESCRIPTION: The prob for zero or negtive prompt. TYPE: float default: 0.0
|
||||
P_ZERO: 0.1
|
||||
# USE_EMA DESCRIPTION: Use Ema or not. Default True TYPE: bool default: True
|
||||
USE_EMA: False
|
||||
LOAD_REFINER: False
|
||||
# SCHEDULE_ARGS DESCRIPTION: TYPE: default: ''
|
||||
SCHEDULE_ARGS:
|
||||
# NAME DESCRIPTION: TYPE: default: ''
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.0120
|
||||
# DIFFUSION_MODEL DESCRIPTION: TYPE: default: ''
|
||||
DIFFUSION_MODEL:
|
||||
# NAME DESCRIPTION: TYPE: default: 'DiffusionUNetXL'
|
||||
NAME: DiffusionUNetXL
|
||||
# PRETRAINED_MODEL DESCRIPTION: Whole model's pretrained model path. TYPE: NoneType default: None
|
||||
PRETRAINED_MODEL:
|
||||
# IN_CHANNELS DESCRIPTION: Unet channels for input, considering the input image's channels. TYPE: int default: 4
|
||||
IN_CHANNELS: 4
|
||||
# OUT_CHANNELS DESCRIPTION: Unet channels for output, considering the input image's channels. TYPE: int default: 4
|
||||
OUT_CHANNELS: 4
|
||||
# NUM_RES_BLOCKS DESCRIPTION: The blocks's number of res. TYPE: int default: 2
|
||||
NUM_RES_BLOCKS: 2
|
||||
# MODEL_CHANNELS DESCRIPTION: base channel count for the model. TYPE: int default: 320
|
||||
MODEL_CHANNELS: 320
|
||||
# ATTENTION_RESOLUTIONS DESCRIPTION: A collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attentio will be used. TYPE: list default: [4, 2]
|
||||
ATTENTION_RESOLUTIONS: [4, 2]
|
||||
# DROPOUT DESCRIPTION: The dropout rate. TYPE: int default: 0
|
||||
DROPOUT: 0
|
||||
# CHANNEL_MULT DESCRIPTION: channel multiplier for each level of the UNet. TYPE: list default: [1, 2, 4]
|
||||
CHANNEL_MULT: [1, 2, 4]
|
||||
# CONV_RESAMPLE DESCRIPTION: Use conv to resample when downsample. TYPE: bool default: True
|
||||
CONV_RESAMPLE: True
|
||||
# DIMS DESCRIPTION: The Conv dims which 2 represent Conv2D. TYPE: int default: 2
|
||||
DIMS: 2
|
||||
# NUM_CLASSES DESCRIPTION: The class num for class guided setting, also can be set as continuous. TYPE: str default: 'sequential'
|
||||
NUM_CLASSES: sequential
|
||||
# USE_CHECKPOINT DESCRIPTION: Use gradient checkpointing to reduce memory usage. TYPE: bool default: False
|
||||
USE_CHECKPOINT: False
|
||||
# NUM_HEADS DESCRIPTION: The number of attention heads in each attention layer. TYPE: int default: -1
|
||||
NUM_HEADS: -1
|
||||
# NUM_HEADS_CHANNELS DESCRIPTION: If specified, ignore num_heads and instead use a fixed channel width per attention head. TYPE: int default: 64
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
# USE_SCALE_SHIFT_NORM DESCRIPTION: The scale and shift for the outnorm of RESBLOCK, use a FiLM-like conditioning mechanism. TYPE: bool default: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
# RESBLOCK_UPDOWN DESCRIPTION: Use residual blocks for up/downsampling, if False use Conv. TYPE: bool default: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
# USE_NEW_ATTENTION_ORDER DESCRIPTION: Whether use new attention(qkv before split heads or not) or not. TYPE: bool default: True
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
# USE_SPATIAL_TRANSFORMER DESCRIPTION: Custom transformer which support the context, if context_dim is not None, the parameter must set True TYPE: bool default: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
# TRANSFORMER_DEPTH DESCRIPTION: Custom transformer's depth, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: list default: [1, 2, 10]
|
||||
TRANSFORMER_DEPTH: [1, 2, 10]
|
||||
# TRANSFORMER_DEPTH_MIDDLE DESCRIPTION: Custom transformer's depth of middle block, If set None, use TRANSFORMER_DEPTH last value. TYPE: NoneType default: None
|
||||
# TRANSFORMER_DEPTH_MIDDLE: None
|
||||
# CONTEXT_DIM DESCRIPTION: Custom context info, if set, USE_SPATIAL_TRANSFORMER also set True. TYPE: int default: 2048
|
||||
CONTEXT_DIM: 2048
|
||||
# DISABLE_SELF_ATTENTIONS DESCRIPTION: Whether disable the self-attentions on some level, should be a list, [False, True, ...] TYPE: NoneType default: None
|
||||
# DISABLE_SELF_ATTENTIONS: None
|
||||
# NUM_ATTENTION_BLOCKS DESCRIPTION: The number of attention blocks for attention layer. TYPE: NoneType default: None
|
||||
# NUM_ATTENTION_BLOCKS: None
|
||||
# DISABLE_MIDDLE_SELF_ATTN DESCRIPTION: Whether disable the self-attentions in middle blocks. TYPE: bool default: False
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
# USE_LINEAR_IN_TRANSFORMER DESCRIPTION: Custom transformer's parameter, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: bool default: True
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
# ADM_IN_CHANNELS DESCRIPTION: Used when num_classes == 'sequential' or 'timestep'. TYPE: int default: 2816
|
||||
ADM_IN_CHANNELS: 2816
|
||||
# USE_SENTENCE_EMB DESCRIPTION: Used sentence emb or not, default False. TYPE: bool default: False
|
||||
USE_SENTENCE_EMB: False
|
||||
# USE_WORD_MAPPING DESCRIPTION: Used word mapping or not, default False. TYPE: bool default: False
|
||||
USE_WORD_MAPPING: False
|
||||
# DIFFUSION_MODEL_EMA DESCRIPTION: TYPE: default: ''
|
||||
DIFFUSION_MODEL_EMA:
|
||||
# NAME DESCRIPTION: TYPE: default: 'DiffusionUNetXL'
|
||||
NAME: DiffusionUNetXL
|
||||
# IN_CHANNELS DESCRIPTION: Unet channels for input, considering the input image's channels. TYPE: int default: 4
|
||||
IN_CHANNELS: 4
|
||||
# OUT_CHANNELS DESCRIPTION: Unet channels for output, considering the input image's channels. TYPE: int default: 4
|
||||
OUT_CHANNELS: 4
|
||||
# NUM_RES_BLOCKS DESCRIPTION: The blocks's number of res. TYPE: int default: 2
|
||||
NUM_RES_BLOCKS: 2
|
||||
# MODEL_CHANNELS DESCRIPTION: base channel count for the model. TYPE: int default: 320
|
||||
MODEL_CHANNELS: 320
|
||||
# ATTENTION_RESOLUTIONS DESCRIPTION: A collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attentio will be used. TYPE: list default: [4, 2]
|
||||
ATTENTION_RESOLUTIONS: [4, 2]
|
||||
# DROPOUT DESCRIPTION: The dropout rate. TYPE: int default: 0
|
||||
DROPOUT: 0
|
||||
# CHANNEL_MULT DESCRIPTION: channel multiplier for each level of the UNet. TYPE: list default: [1, 2, 4]
|
||||
CHANNEL_MULT: [1, 2, 4]
|
||||
# CONV_RESAMPLE DESCRIPTION: Use conv to resample when downsample. TYPE: bool default: True
|
||||
CONV_RESAMPLE: True
|
||||
# DIMS DESCRIPTION: The Conv dims which 2 represent Conv2D. TYPE: int default: 2
|
||||
DIMS: 2
|
||||
# NUM_CLASSES DESCRIPTION: The class num for class guided setting, also can be set as continuous. TYPE: str default: 'sequential'
|
||||
NUM_CLASSES: sequential
|
||||
# USE_CHECKPOINT DESCRIPTION: Use gradient checkpointing to reduce memory usage. TYPE: bool default: False
|
||||
USE_CHECKPOINT: False
|
||||
# NUM_HEADS DESCRIPTION: The number of attention heads in each attention layer. TYPE: int default: -1
|
||||
NUM_HEADS: -1
|
||||
# NUM_HEADS_CHANNELS DESCRIPTION: If specified, ignore num_heads and instead use a fixed channel width per attention head. TYPE: int default: 64
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
# USE_SCALE_SHIFT_NORM DESCRIPTION: The scale and shift for the outnorm of RESBLOCK, use a FiLM-like conditioning mechanism. TYPE: bool default: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
# RESBLOCK_UPDOWN DESCRIPTION: Use residual blocks for up/downsampling, if False use Conv. TYPE: bool default: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
# USE_NEW_ATTENTION_ORDER DESCRIPTION: Whether use new attention(qkv before split heads or not) or not. TYPE: bool default: True
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
# USE_SPATIAL_TRANSFORMER DESCRIPTION: Custom transformer which support the context, if context_dim is not None, the parameter must set True TYPE: bool default: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
# TRANSFORMER_DEPTH DESCRIPTION: Custom transformer's depth, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: list default: [1, 2, 10]
|
||||
TRANSFORMER_DEPTH: [1, 2, 10]
|
||||
# TRANSFORMER_DEPTH_MIDDLE DESCRIPTION: Custom transformer's depth of middle block, If set None, use TRANSFORMER_DEPTH last value. TYPE: NoneType default: None
|
||||
# TRANSFORMER_DEPTH_MIDDLE: None
|
||||
# CONTEXT_DIM DESCRIPTION: Custom context info, if set, USE_SPATIAL_TRANSFORMER also set True. TYPE: int default: 2048
|
||||
CONTEXT_DIM: 2048
|
||||
# DISABLE_SELF_ATTENTIONS DESCRIPTION: Whether disable the self-attentions on some level, should be a list, [False, True, ...] TYPE: NoneType default: None
|
||||
# DISABLE_SELF_ATTENTIONS: None
|
||||
# NUM_ATTENTION_BLOCKS DESCRIPTION: The number of attention blocks for attention layer. TYPE: NoneType default: None
|
||||
# NUM_ATTENTION_BLOCKS: None
|
||||
# DISABLE_MIDDLE_SELF_ATTN DESCRIPTION: Whether disable the self-attentions in middle blocks. TYPE: bool default: False
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
# USE_LINEAR_IN_TRANSFORMER DESCRIPTION: Custom transformer's parameter, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: bool default: True
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
# ADM_IN_CHANNELS DESCRIPTION: Used when num_classes == 'sequential' or 'timestep'. TYPE: int default: 2816
|
||||
ADM_IN_CHANNELS: 2816
|
||||
# USE_SENTENCE_EMB DESCRIPTION: Used sentence emb or not, default False. TYPE: bool default: False
|
||||
USE_SENTENCE_EMB: False
|
||||
# USE_WORD_MAPPING DESCRIPTION: Used word mapping or not, default False. TYPE: bool default: False
|
||||
USE_WORD_MAPPING: False
|
||||
# FIRST_STAGE_MODEL DESCRIPTION: TYPE: default: ''
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 1
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
# COND_STAGE_MODEL DESCRIPTION: TYPE: default: ''
|
||||
COND_STAGE_MODEL:
|
||||
# NAME DESCRIPTION: TYPE: default: 'GeneralConditioner'
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
USE_GRAD: False
|
||||
# EMBEDDERS DESCRIPTION: TYPE: default: ''
|
||||
EMBEDDERS:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'FrozenCLIPEmbedder'
|
||||
NAME: FrozenCLIPEmbedder
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
TOKENIZER_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
# MAX_LENGTH DESCRIPTION: TYPE: int default: 77
|
||||
MAX_LENGTH: 77
|
||||
# FREEZE DESCRIPTION: TYPE: bool default: True
|
||||
FREEZE: True
|
||||
# LAYER DESCRIPTION: TYPE: str default: 'last'
|
||||
LAYER: hidden
|
||||
# LAYER_IDX DESCRIPTION: TYPE: NoneType default: None
|
||||
LAYER_IDX: 11
|
||||
# USE_FINAL_LAYER_NORM DESCRIPTION: TYPE: bool default: False
|
||||
USE_FINAL_LAYER_NORM: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["prompt"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'FrozenOpenCLIPEmbedder2'
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
# ARCH DESCRIPTION: TYPE: str default: 'ViT-H-14'
|
||||
ARCH: ViT-bigG-14
|
||||
# MAX_LENGTH DESCRIPTION: TYPE: int default: 77
|
||||
MAX_LENGTH: 77
|
||||
# FREEZE DESCRIPTION: TYPE: bool default: True
|
||||
FREEZE: True
|
||||
# ALWAYS_RETURN_POOLED DESCRIPTION: Whether always return pooled results or not ,default False. TYPE: bool default: False
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
# LEGACY DESCRIPTION: Whether use legacy returnd feature or not ,default True. TYPE: bool default: True
|
||||
LEGACY: False
|
||||
# LAYER DESCRIPTION: TYPE: str default: 'last'
|
||||
LAYER: penultimate
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["prompt"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["original_size_as_tuple"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["crop_coords_top_left"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["target_size_as_tuple"]
|
||||
LEGACY_UCG_VALUE:
|
||||
# APPLY REFINER
|
||||
REFINER_MODEL:
|
||||
# NAME DESCRIPTION: TYPE: default: 'DiffusionUNetXL'
|
||||
NAME: DiffusionUNetXL
|
||||
# PRETRAINED_MODEL DESCRIPTION: Whole model's pretrained model path. TYPE: NoneType default: None
|
||||
PRETRAINED_MODEL:
|
||||
# IN_CHANNELS DESCRIPTION: Unet channels for input, considering the input image's channels. TYPE: int default: 4
|
||||
IN_CHANNELS: 4
|
||||
# OUT_CHANNELS DESCRIPTION: Unet channels for output, considering the input image's channels. TYPE: int default: 4
|
||||
OUT_CHANNELS: 4
|
||||
# NUM_RES_BLOCKS DESCRIPTION: The blocks's number of res. TYPE: int default: 2
|
||||
NUM_RES_BLOCKS: 2
|
||||
# MODEL_CHANNELS DESCRIPTION: base channel count for the model. TYPE: int default: 320
|
||||
MODEL_CHANNELS: 384
|
||||
# ATTENTION_RESOLUTIONS DESCRIPTION: A collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attentio will be used. TYPE: list default: [4, 2]
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2 ]
|
||||
# DROPOUT DESCRIPTION: The dropout rate. TYPE: int default: 0
|
||||
DROPOUT: 0
|
||||
# CHANNEL_MULT DESCRIPTION: channel multiplier for each level of the UNet. TYPE: list default: [1, 2, 4]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
# CONV_RESAMPLE DESCRIPTION: Use conv to resample when downsample. TYPE: bool default: True
|
||||
CONV_RESAMPLE: True
|
||||
# DIMS DESCRIPTION: The Conv dims which 2 represent Conv2D. TYPE: int default: 2
|
||||
DIMS: 2
|
||||
# NUM_CLASSES DESCRIPTION: The class num for class guided setting, also can be set as continuous. TYPE: str default: 'sequential'
|
||||
NUM_CLASSES: sequential
|
||||
# USE_CHECKPOINT DESCRIPTION: Use gradient checkpointing to reduce memory usage. TYPE: bool default: False
|
||||
USE_CHECKPOINT: False
|
||||
# NUM_HEADS DESCRIPTION: The number of attention heads in each attention layer. TYPE: int default: -1
|
||||
NUM_HEADS: -1
|
||||
# NUM_HEADS_CHANNELS DESCRIPTION: If specified, ignore num_heads and instead use a fixed channel width per attention head. TYPE: int default: 64
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
# USE_SCALE_SHIFT_NORM DESCRIPTION: The scale and shift for the outnorm of RESBLOCK, use a FiLM-like conditioning mechanism. TYPE: bool default: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
# RESBLOCK_UPDOWN DESCRIPTION: Use residual blocks for up/downsampling, if False use Conv. TYPE: bool default: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
# USE_NEW_ATTENTION_ORDER DESCRIPTION: Whether use new attention(qkv before split heads or not) or not. TYPE: bool default: True
|
||||
USE_NEW_ATTENTION_ORDER: True
|
||||
# USE_SPATIAL_TRANSFORMER DESCRIPTION: Custom transformer which support the context, if context_dim is not None, the parameter must set True TYPE: bool default: True
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
# TRANSFORMER_DEPTH DESCRIPTION: Custom transformer's depth, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: list default: [1, 2, 10]
|
||||
TRANSFORMER_DEPTH: 4
|
||||
# TRANSFORMER_DEPTH_MIDDLE DESCRIPTION: Custom transformer's depth of middle block, If set None, use TRANSFORMER_DEPTH last value. TYPE: NoneType default: None
|
||||
# TRANSFORMER_DEPTH_MIDDLE: None
|
||||
# CONTEXT_DIM DESCRIPTION: Custom context info, if set, USE_SPATIAL_TRANSFORMER also set True. TYPE: int default: 2048
|
||||
CONTEXT_DIM: [ 1280, 1280, 1280, 1280 ]
|
||||
# DISABLE_SELF_ATTENTIONS DESCRIPTION: Whether disable the self-attentions on some level, should be a list, [False, True, ...] TYPE: NoneType default: None
|
||||
# DISABLE_SELF_ATTENTIONS: None
|
||||
# NUM_ATTENTION_BLOCKS DESCRIPTION: The number of attention blocks for attention layer. TYPE: NoneType default: None
|
||||
# NUM_ATTENTION_BLOCKS: None
|
||||
# DISABLE_MIDDLE_SELF_ATTN DESCRIPTION: Whether disable the self-attentions in middle blocks. TYPE: bool default: False
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
# USE_LINEAR_IN_TRANSFORMER DESCRIPTION: Custom transformer's parameter, valid when USE_SPATIAL_TRANSFORMER is True. TYPE: bool default: True
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
# ADM_IN_CHANNELS DESCRIPTION: Used when num_classes == 'sequential' or 'timestep'. TYPE: int default: 2816
|
||||
ADM_IN_CHANNELS: 2560
|
||||
# USE_SENTENCE_EMB DESCRIPTION: Used sentence emb or not, default False. TYPE: bool default: False
|
||||
USE_SENTENCE_EMB: False
|
||||
# USE_WORD_MAPPING DESCRIPTION: Used word mapping or not, default False. TYPE: bool default: False
|
||||
USE_WORD_MAPPING: False
|
||||
# COND_STAGE_MODEL DESCRIPTION: TYPE: default: ''
|
||||
REFINER_COND_MODEL:
|
||||
# NAME DESCRIPTION: TYPE: default: 'GeneralConditioner'
|
||||
NAME: GeneralConditioner
|
||||
PRETRAINED_MODEL:
|
||||
# EMBEDDERS DESCRIPTION: TYPE: default: ''
|
||||
EMBEDDERS:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'FrozenOpenCLIPEmbedder2'
|
||||
NAME: FrozenOpenCLIPEmbedder2
|
||||
# ARCH DESCRIPTION: TYPE: str default: 'ViT-H-14'
|
||||
ARCH: ViT-bigG-14
|
||||
# MAX_LENGTH DESCRIPTION: TYPE: int default: 77
|
||||
MAX_LENGTH: 77
|
||||
# FREEZE DESCRIPTION: TYPE: bool default: True
|
||||
FREEZE: True
|
||||
# ALWAYS_RETURN_POOLED DESCRIPTION: Whether always return pooled results or not ,default False. TYPE: bool default: False
|
||||
ALWAYS_RETURN_POOLED: True
|
||||
# LEGACY DESCRIPTION: Whether use legacy returnd feature or not ,default True. TYPE: bool default: True
|
||||
LEGACY: False
|
||||
# LAYER DESCRIPTION: TYPE: str default: 'last'
|
||||
LAYER: penultimate
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["prompt"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["original_size_as_tuple"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["crop_coords_top_left"]
|
||||
LEGACY_UCG_VALUE:
|
||||
-
|
||||
# NAME DESCRIPTION: TYPE: default: 'ConcatTimestepEmbedderND'
|
||||
NAME: ConcatTimestepEmbedderND
|
||||
# OUT_DIM DESCRIPTION: Output dim TYPE: int default: 256
|
||||
OUT_DIM: 256
|
||||
IS_TRAINABLE: False
|
||||
UCG_RATE: 0.0
|
||||
INPUT_KEYS: ["aesthetic_score"]
|
||||
LEGACY_UCG_VALUE:
|
||||
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 5.0
|
||||
GUIDE_RESCALE:
|
||||
DISCRETIZATION: linspace
|
||||
IMAGE_SIZE: [ 1024, 1024]
|
||||
RUN_TRAIN_N: False
|
||||
# OPTIMIZER DESCRIPTION: TYPE: default: ''
|
||||
OPTIMIZER:
|
||||
# NAME DESCRIPTION: TYPE: default: ''
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.0064
|
||||
EPS: 1e-8
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
#
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleResize
|
||||
INTERPOLATION: bicubic
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: FlexibleCropXL
|
||||
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: [ 'img' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Select
|
||||
KEYS: [ 'img', 'prompt', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left' ]
|
||||
META_KEYS: ['data_key', 'img_path']
|
||||
- NAME: Rename
|
||||
INPUT_KEY: [ 'img', 'img_original_size_as_tuple', 'img_target_size_as_tuple', 'img_crop_coords_top_left']
|
||||
OUTPUT_KEY: [ 'image', 'original_size_as_tuple', 'target_size_as_tuple', 'crop_coords_top_left' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
- NAME: BackwardHook
|
||||
# GRADIENT_CLIP: 1.0
|
||||
PRIORITY: 0
|
||||
- NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
SHOW_GPU_MEM: True
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
SAVE_LAST: True
|
||||
INTERVAL: 10000
|
||||
PRIORITY: 200
|
||||
@@ -0,0 +1,10 @@
|
||||
WORK_DIR: "self_train"
|
||||
SCRIPT_DIR: "scepter/studio/self_train/scripts"
|
||||
DEFAULT_FOLDER: sd_xl
|
||||
SAMPLERS:
|
||||
-
|
||||
NAME: 'ddim'
|
||||
-
|
||||
NAME: 'dpmpp_2m_sde'
|
||||
-
|
||||
NAME: 'dpmpp_2s_ancestral'
|
||||
@@ -0,0 +1,312 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
META:
|
||||
VERSION: 'SD1.5'
|
||||
DESCRIPTION: "Stable Diffusion v1.5"
|
||||
IS_DEFAULT: False
|
||||
INFERENCE_PARAS:
|
||||
INFERENCE_BATCH_SIZE: 1
|
||||
INFERENCE_PREFIX: ""
|
||||
DEFAULT_SAMPLER: "ddim"
|
||||
DEFAULT_SAMPLE_STEPS: 40
|
||||
INFERENCE_N_PROMPT: ""
|
||||
RESOLUTION: 512
|
||||
PARAS:
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 512
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: FULL
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 512
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: LORA
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 512
|
||||
MEMORY: 29000
|
||||
EPOCHS: 200
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: SCE
|
||||
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 512
|
||||
MEMORY: 29000
|
||||
EPOCHS: 200
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: True
|
||||
TUNER: TEXT_SCE
|
||||
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 4
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 512
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: TEXT_LORA
|
||||
|
||||
TUNERS:
|
||||
LORA:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2)$"
|
||||
TEXT_LORA:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "(cond_stage_model.*(q_proj|k_proj|v_proj|out_proj|mlp.fc1|mlp.fc2))|(model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2))$"
|
||||
SCE:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320]
|
||||
DOWN_RATIO: 1.0
|
||||
TARGET_MODULES: model.lsc_identity\.\d+$
|
||||
TUNER_MODE: identity
|
||||
TEXT_SCE:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [ 1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320 ]
|
||||
DOWN_RATIO: 1.0
|
||||
TARGET_MODULES: model.lsc_identity\.\d+$
|
||||
TUNER_MODE: identity
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "cond_stage_model.*(q_proj|k_proj|v_proj|out_proj|mlp.fc1|mlp.fc2)$"
|
||||
|
||||
MODIFY_PARAS:
|
||||
TEXT_LORA:
|
||||
TRAIN:
|
||||
SOLVER.MODEL.COND_STAGE_MODEL.USE_GRAD: True
|
||||
TEXT_SCE:
|
||||
TRAIN:
|
||||
SOLVER.MODEL.COND_STAGE_MODEL.USE_GRAD: True
|
||||
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 1000
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
#
|
||||
WORK_DIR:
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
#
|
||||
TUNER:
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusion
|
||||
PARAMETERIZATION: eps
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-v1-5@v1-5-pruned-emaonly.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
# DEFAULT_N_PROMPT: 'lowres, error, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS: 8
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 768
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: False
|
||||
IGNORE_KEYS: []
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
IGNORE_KEYS: []
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: ClipTokenizer
|
||||
PRETRAINED_PATH: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
LENGTH: 77
|
||||
CLEAN: True
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenCLIPEmbedder
|
||||
FREEZE: True
|
||||
USE_GRAD: False
|
||||
LAYER: last
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/clip-vit-large-patch14
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE:
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [512, 512]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 512
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 512
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ImageToTensor
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: Normalize
|
||||
MEAN: [ 0.5, 0.5, 0.5 ]
|
||||
STD: [ 0.5, 0.5, 0.5 ]
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
- NAME: BackwardHook
|
||||
# GRADIENT_CLIP: 1.0
|
||||
PRIORITY: 0
|
||||
- NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
SHOW_GPU_MEM: True
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
SAVE_LAST: True
|
||||
INTERVAL: 10000
|
||||
PRIORITY: 200
|
||||
@@ -0,0 +1,254 @@
|
||||
ENV:
|
||||
BACKEND: nccl
|
||||
META:
|
||||
VERSION: 'SD2.1'
|
||||
DESCRIPTION: "Stable Diffusion v2.1"
|
||||
IS_DEFAULT: False
|
||||
INFERENCE_PARAS:
|
||||
INFERENCE_BATCH_SIZE: 1
|
||||
INFERENCE_PREFIX: ""
|
||||
DEFAULT_SAMPLER: "ddim"
|
||||
DEFAULT_SAMPLE_STEPS: 40
|
||||
INFERENCE_N_PROMPT: ""
|
||||
RESOLUTION: 768
|
||||
PARAS:
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 768
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: FULL
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 768
|
||||
MEMORY: 29000
|
||||
EPOCHS: 50
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: False
|
||||
TUNER: LORA
|
||||
-
|
||||
TRAIN_BATCH_SIZE: 2
|
||||
TRAIN_PREFIX: ""
|
||||
TRAIN_N_PROMPT: ""
|
||||
RESOLUTION: 768
|
||||
MEMORY: 29000
|
||||
EPOCHS: 200
|
||||
SAVE_INTERVAL: 25
|
||||
EPSEC: 0.818
|
||||
LEARNING_RATE: 0.0001
|
||||
IS_DEFAULT: True
|
||||
TUNER: SCE
|
||||
|
||||
TUNERS:
|
||||
LORA:
|
||||
-
|
||||
NAME: SwiftLoRA
|
||||
R: 256
|
||||
LORA_ALPHA: 256
|
||||
LORA_DROPOUT: 0.0
|
||||
BIAS: "none"
|
||||
TARGET_MODULES: "model.*(to_q|to_k|to_v|to_out.0|net.0.proj|net.2)$"
|
||||
SCE:
|
||||
-
|
||||
NAME: SwiftSCETuning
|
||||
DIMS: [1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320]
|
||||
DOWN_RATIO: 1.0
|
||||
TARGET_MODULES: model.lsc_identity\.\d+$
|
||||
TUNER_MODE: identity
|
||||
SOLVER:
|
||||
NAME: LatentDiffusionSolver
|
||||
RESUME_FROM:
|
||||
LOAD_MODEL_ONLY: True
|
||||
USE_FSDP: False
|
||||
SHARDING_STRATEGY:
|
||||
USE_AMP: True
|
||||
DTYPE: float16
|
||||
CHANNELS_LAST: True
|
||||
MAX_STEPS: 1000
|
||||
MAX_EPOCHS: -1
|
||||
NUM_FOLDS: 1
|
||||
ACCU_STEP: 1
|
||||
#
|
||||
WORK_DIR:
|
||||
LOG_FILE: std_log.txt
|
||||
#
|
||||
FILE_SYSTEM:
|
||||
NAME: "ModelscopeFs"
|
||||
TEMP_DIR: "./cache/data"
|
||||
#
|
||||
FREEZE:
|
||||
#
|
||||
TUNER:
|
||||
#
|
||||
MODEL:
|
||||
NAME: LatentDiffusion
|
||||
PARAMETERIZATION: v
|
||||
TIMESTEPS: 1000
|
||||
MIN_SNR_GAMMA:
|
||||
ZERO_TERMINAL_SNR: False
|
||||
PRETRAINED_MODEL: ms://AI-ModelScope/stable-diffusion-2-1@v2-1_768-ema-pruned.safetensors
|
||||
IGNORE_KEYS: [ ]
|
||||
SCALE_FACTOR: 0.18215
|
||||
SIZE_FACTOR: 8
|
||||
# DEFAULT_N_PROMPT: 'lowres, error, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'
|
||||
DEFAULT_N_PROMPT:
|
||||
SCHEDULE_ARGS:
|
||||
"NAME": "scaled_linear"
|
||||
"BETA_MIN": 0.00085
|
||||
"BETA_MAX": 0.012
|
||||
USE_EMA: False
|
||||
#
|
||||
DIFFUSION_MODEL:
|
||||
NAME: DiffusionUNet
|
||||
IN_CHANNELS: 4
|
||||
OUT_CHANNELS: 4
|
||||
MODEL_CHANNELS: 320
|
||||
NUM_HEADS_CHANNELS: 64
|
||||
NUM_RES_BLOCKS: 2
|
||||
ATTENTION_RESOLUTIONS: [ 4, 2, 1 ]
|
||||
CHANNEL_MULT: [ 1, 2, 4, 4 ]
|
||||
CONV_RESAMPLE: True
|
||||
DIMS: 2
|
||||
USE_CHECKPOINT: False
|
||||
USE_SCALE_SHIFT_NORM: False
|
||||
RESBLOCK_UPDOWN: False
|
||||
USE_SPATIAL_TRANSFORMER: True
|
||||
TRANSFORMER_DEPTH: 1
|
||||
CONTEXT_DIM: 1024
|
||||
DISABLE_MIDDLE_SELF_ATTN: False
|
||||
USE_LINEAR_IN_TRANSFORMER: True
|
||||
PRETRAINED_MODEL:
|
||||
#
|
||||
FIRST_STAGE_MODEL:
|
||||
NAME: AutoencoderKL
|
||||
EMBED_DIM: 4
|
||||
PRETRAINED_MODEL:
|
||||
IGNORE_KEYS: [ ]
|
||||
BATCH_SIZE: 4
|
||||
#
|
||||
ENCODER:
|
||||
NAME: Encoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DOUBLE_Z: True
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
#
|
||||
DECODER:
|
||||
NAME: Decoder
|
||||
CH: 128
|
||||
OUT_CH: 3
|
||||
NUM_RES_BLOCKS: 2
|
||||
IN_CHANNELS: 3
|
||||
ATTN_RESOLUTIONS: [ ]
|
||||
CH_MULT: [ 1, 2, 4, 4 ]
|
||||
Z_CHANNELS: 4
|
||||
DROPOUT: 0.0
|
||||
RESAMP_WITH_CONV: True
|
||||
GIVE_PRE_END: False
|
||||
TANH_OUT: False
|
||||
#
|
||||
TOKENIZER:
|
||||
NAME: OpenClipTokenizer
|
||||
LENGTH: 77
|
||||
#
|
||||
COND_STAGE_MODEL:
|
||||
NAME: FrozenOpenCLIPEmbedder
|
||||
ARCH: ViT-H-14
|
||||
PRETRAINED_MODEL:
|
||||
USE_GRAD: False
|
||||
LAYER: penultimate
|
||||
#
|
||||
LOSS:
|
||||
NAME: ReconstructLoss
|
||||
LOSS_TYPE: l2
|
||||
#
|
||||
SAMPLE_ARGS:
|
||||
SAMPLER: ddim
|
||||
SAMPLE_STEPS: 50
|
||||
SEED: 2023
|
||||
GUIDE_SCALE: 7.5
|
||||
GUIDE_RESCALE:
|
||||
DISCRETIZATION: trailing
|
||||
IMAGE_SIZE: [768, 768]
|
||||
RUN_TRAIN_N: False
|
||||
#
|
||||
OPTIMIZER:
|
||||
NAME: AdamW
|
||||
LEARNING_RATE: 0.064
|
||||
BETAS: [ 0.9, 0.999 ]
|
||||
EPS: 1e-8
|
||||
WEIGHT_DECAY: 1e-2
|
||||
AMSGRAD: False
|
||||
#
|
||||
TRAIN_DATA:
|
||||
NAME: ImageTextPairMSDataset
|
||||
MODE: train
|
||||
MS_DATASET_NAME: style_custom_dataset
|
||||
MS_DATASET_NAMESPACE: damo
|
||||
MS_DATASET_SUBNAME: 3D
|
||||
PROMPT_PREFIX: ""
|
||||
MS_DATASET_SPLIT: train
|
||||
MS_REMAP_KEYS: { 'Image:FILE': 'Target:FILE' }
|
||||
REPLACE_STYLE: False
|
||||
PIN_MEMORY: True
|
||||
BATCH_SIZE: 1
|
||||
NUM_WORKERS: 4
|
||||
SAMPLER:
|
||||
NAME: LoopSampler
|
||||
TRANSFORMS:
|
||||
- NAME: LoadImageFromFile
|
||||
RGB_ORDER: RGB
|
||||
BACKEND: pillow
|
||||
- NAME: Resize
|
||||
SIZE: 768
|
||||
INTERPOLATION: bilinear
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: CenterCrop
|
||||
SIZE: 768
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: ImageToTensor
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'img' ]
|
||||
BACKEND: pillow
|
||||
- NAME: Normalize
|
||||
MEAN: [ 0.5, 0.5, 0.5 ]
|
||||
STD: [ 0.5, 0.5, 0.5 ]
|
||||
INPUT_KEY: [ 'img' ]
|
||||
OUTPUT_KEY: [ 'image' ]
|
||||
BACKEND: torchvision
|
||||
- NAME: Select
|
||||
KEYS: [ 'image', 'prompt' ]
|
||||
META_KEYS: [ 'data_key' ]
|
||||
#
|
||||
TRAIN_HOOKS:
|
||||
- NAME: BackwardHook
|
||||
# GRADIENT_CLIP: 1.0
|
||||
PRIORITY: 0
|
||||
- NAME: LogHook
|
||||
LOG_INTERVAL: 50
|
||||
SHOW_GPU_MEM: True
|
||||
-
|
||||
NAME: CheckpointHook
|
||||
SAVE_LAST: True
|
||||
INTERVAL: 10000
|
||||
PRIORITY: 200
|
||||
@@ -1,3 +1,4 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules import data, model, opt, solver, transform, utils
|
||||
from scepter.modules import (data, inference, model, opt, solver, transform,
|
||||
utils)
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
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.hed import HedAnnotator
|
||||
from scepter.modules.annotator.midas_op import MidasDetector
|
||||
from scepter.modules.annotator.mlsd_op import MLSDdetector
|
||||
from scepter.modules.annotator.openpose import OpenposeAnnotator
|
||||
@@ -0,0 +1,57 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from abc import ABCMeta
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.model.base_model import BaseModel
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class BaseAnnotator(BaseModel, metaclass=ABCMeta):
|
||||
para_dict = {}
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode
|
||||
def forward(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
BaseAnnotator.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class GeneralAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
anno_models = cfg.get('ANNOTATORS', [])
|
||||
self.annotators = nn.ModuleList()
|
||||
for n, anno_config in enumerate(anno_models):
|
||||
annotator = ANNOTATORS.build(anno_config, logger=logger)
|
||||
annotator.input_keys = anno_config.get('INPUT_KEYS', [])
|
||||
if isinstance(annotator.input_keys, str):
|
||||
annotator.input_keys = [annotator.input_keys]
|
||||
annotator.output_keys = anno_config.get('OUTPUT_KEYS', [])
|
||||
if isinstance(annotator.output_keys, str):
|
||||
annotator.output_keys = [annotator.output_keys]
|
||||
assert len(annotator.input_keys) == len(annotator.output_keys)
|
||||
self.annotators.append(annotator)
|
||||
|
||||
def forward(self, input_dict):
|
||||
output_dict = {}
|
||||
for annotator in self.annotators:
|
||||
for idx, in_key in enumerate(annotator.input_keys):
|
||||
if in_key in input_dict:
|
||||
image = annotator(input_dict[in_key])
|
||||
output_dict[annotator.output_keys[idx]] = image
|
||||
return output_dict
|
||||
@@ -0,0 +1,43 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from abc import ABCMeta
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class CannyAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
||||
para_dict = {}
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.low_threshold = cfg.get('LOW_THRESHOLD', 100)
|
||||
self.high_threshold = cfg.get('HIGH_THRESHOLD', 200)
|
||||
|
||||
def forward(self, image):
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
image = cv2.Canny(image, self.low_threshold, self.high_threshold)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image = image.detach().cpu().numpy()
|
||||
image = cv2.Canny(image, self.low_threshold, self.high_threshold)
|
||||
elif isinstance(image, np.ndarray):
|
||||
image = cv2.Canny(image.copy(), self.low_threshold,
|
||||
self.high_threshold)
|
||||
else:
|
||||
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
||||
assert len(image.shape) < 4
|
||||
return image[..., None].repeat(3, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
CannyAnnotator.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,44 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from abc import ABCMeta
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class ColorAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
||||
para_dict = {}
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.ratio = cfg.get('RATIO', 64)
|
||||
|
||||
def forward(self, image):
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image = image.detach().cpu().numpy()
|
||||
elif isinstance(image, np.ndarray):
|
||||
image = image.copy()
|
||||
else:
|
||||
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
||||
h, w = image.shape[:2]
|
||||
ratio = self.ratio
|
||||
image = cv2.resize(image, (w // ratio, h // ratio),
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
image = cv2.resize(image, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
assert len(image.shape) < 4
|
||||
return image
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
ColorAnnotator.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,154 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Please use this implementation in your products
|
||||
# This implementation may produce slightly different results from Saining Xie's official implementations,
|
||||
# but it generates smoother edges and is more suitable for ControlNet as well as other image-to-image translations.
|
||||
# Different from official models and other implementations, this is an RGB-input model (rather than BGR)
|
||||
# and in this way it works better for gradio's RGB protocol
|
||||
|
||||
from abc import ABCMeta
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
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
|
||||
|
||||
|
||||
def nms(x, t, s):
|
||||
x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
|
||||
|
||||
f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
|
||||
f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
|
||||
f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
|
||||
f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
|
||||
|
||||
y = np.zeros_like(x)
|
||||
|
||||
for f in [f1, f2, f3, f4]:
|
||||
np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
|
||||
|
||||
z = np.zeros_like(y, dtype=np.uint8)
|
||||
z[y > t] = 255
|
||||
return z
|
||||
|
||||
|
||||
class DoubleConvBlock(torch.nn.Module):
|
||||
def __init__(self, input_channel, output_channel, layer_number):
|
||||
super().__init__()
|
||||
self.convs = torch.nn.Sequential()
|
||||
self.convs.append(
|
||||
torch.nn.Conv2d(in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
kernel_size=(3, 3),
|
||||
stride=(1, 1),
|
||||
padding=1))
|
||||
for i in range(1, layer_number):
|
||||
self.convs.append(
|
||||
torch.nn.Conv2d(in_channels=output_channel,
|
||||
out_channels=output_channel,
|
||||
kernel_size=(3, 3),
|
||||
stride=(1, 1),
|
||||
padding=1))
|
||||
self.projection = torch.nn.Conv2d(in_channels=output_channel,
|
||||
out_channels=1,
|
||||
kernel_size=(1, 1),
|
||||
stride=(1, 1),
|
||||
padding=0)
|
||||
|
||||
def __call__(self, x, down_sampling=False):
|
||||
h = x
|
||||
if down_sampling:
|
||||
h = torch.nn.functional.max_pool2d(h,
|
||||
kernel_size=(2, 2),
|
||||
stride=(2, 2))
|
||||
for conv in self.convs:
|
||||
h = conv(h)
|
||||
h = torch.nn.functional.relu(h)
|
||||
return h, self.projection(h)
|
||||
|
||||
|
||||
class ControlNetHED_Apache2(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
|
||||
self.block1 = DoubleConvBlock(input_channel=3,
|
||||
output_channel=64,
|
||||
layer_number=2)
|
||||
self.block2 = DoubleConvBlock(input_channel=64,
|
||||
output_channel=128,
|
||||
layer_number=2)
|
||||
self.block3 = DoubleConvBlock(input_channel=128,
|
||||
output_channel=256,
|
||||
layer_number=3)
|
||||
self.block4 = DoubleConvBlock(input_channel=256,
|
||||
output_channel=512,
|
||||
layer_number=3)
|
||||
self.block5 = DoubleConvBlock(input_channel=512,
|
||||
output_channel=512,
|
||||
layer_number=3)
|
||||
|
||||
def __call__(self, x):
|
||||
h = x - self.norm
|
||||
h, projection1 = self.block1(h)
|
||||
h, projection2 = self.block2(h, down_sampling=True)
|
||||
h, projection3 = self.block3(h, down_sampling=True)
|
||||
h, projection4 = self.block4(h, down_sampling=True)
|
||||
h, projection5 = self.block5(h, down_sampling=True)
|
||||
return projection1, projection2, projection3, projection4, projection5
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class HedAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
||||
para_dict = {}
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.netNetwork = ControlNetHED_Apache2().float().eval()
|
||||
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))
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
@torch.autocast('cuda', enabled=False)
|
||||
def forward(self, image):
|
||||
if isinstance(image, torch.Tensor):
|
||||
if len(image.shape) == 3:
|
||||
image = rearrange(image, 'h w c -> 1 c h w')
|
||||
B, C, H, W = image.shape
|
||||
else:
|
||||
raise "Unsurpport input image's shape"
|
||||
elif isinstance(image, np.ndarray):
|
||||
image = torch.from_numpy(image.copy()).float()
|
||||
if len(image.shape) == 3:
|
||||
image = rearrange(image, 'h w c -> 1 c h w')
|
||||
B, C, H, W = image.shape
|
||||
else:
|
||||
raise "Unsurpport input image's shape"
|
||||
else:
|
||||
raise "Unsurpport input image's type"
|
||||
edges = self.netNetwork(image.to(we.device_id))
|
||||
edges = [
|
||||
e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges
|
||||
]
|
||||
edges = [
|
||||
cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR)
|
||||
for e in edges
|
||||
]
|
||||
edges = np.stack(edges, axis=2)
|
||||
edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
|
||||
edge = 255 - (edge * 255.0).clip(0, 255).astype(np.uint8)
|
||||
return edge[..., None].repeat(3, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
HedAnnotator.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,165 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# based on https://github.com/isl-org/MiDaS
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchvision.transforms import Compose
|
||||
|
||||
from .dpt_depth import DPTDepthModel
|
||||
from .midas_net import MidasNet
|
||||
from .midas_net_custom import MidasNet_small
|
||||
from .transforms import NormalizeImage, PrepareForNet, Resize
|
||||
|
||||
# ISL_PATHS = {
|
||||
# "dpt_large": "dpt_large-midas-2f21e586.pt",
|
||||
# "dpt_hybrid": "dpt_hybrid-midas-501f0c75.pt",
|
||||
# "midas_v21": "",
|
||||
# "midas_v21_small": "",
|
||||
# }
|
||||
|
||||
# remote_model_path =
|
||||
# "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
|
||||
|
||||
|
||||
def disabled_train(self, mode=True):
|
||||
"""Overwrite model.train with this function to make sure train/eval mode
|
||||
does not change anymore."""
|
||||
return self
|
||||
|
||||
|
||||
def load_midas_transform(model_type):
|
||||
# https://github.com/isl-org/MiDaS/blob/master/run.py
|
||||
# load transform only
|
||||
if model_type == 'dpt_large': # DPT-Large
|
||||
net_w, net_h = 384, 384
|
||||
resize_mode = 'minimal'
|
||||
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
|
||||
std=[0.5, 0.5, 0.5])
|
||||
|
||||
elif model_type == 'dpt_hybrid': # DPT-Hybrid
|
||||
net_w, net_h = 384, 384
|
||||
resize_mode = 'minimal'
|
||||
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
|
||||
std=[0.5, 0.5, 0.5])
|
||||
|
||||
elif model_type == 'midas_v21':
|
||||
net_w, net_h = 384, 384
|
||||
resize_mode = 'upper_bound'
|
||||
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225])
|
||||
|
||||
elif model_type == 'midas_v21_small':
|
||||
net_w, net_h = 256, 256
|
||||
resize_mode = 'upper_bound'
|
||||
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225])
|
||||
|
||||
else:
|
||||
assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
|
||||
|
||||
transform = Compose([
|
||||
Resize(
|
||||
net_w,
|
||||
net_h,
|
||||
resize_target=None,
|
||||
keep_aspect_ratio=True,
|
||||
ensure_multiple_of=32,
|
||||
resize_method=resize_mode,
|
||||
image_interpolation_method=cv2.INTER_CUBIC,
|
||||
),
|
||||
normalization,
|
||||
PrepareForNet(),
|
||||
])
|
||||
|
||||
return transform
|
||||
|
||||
|
||||
def load_model(model_type, model_path):
|
||||
# https://github.com/isl-org/MiDaS/blob/master/run.py
|
||||
# load network
|
||||
# model_path = ISL_PATHS[model_type]
|
||||
if model_type == 'dpt_large': # DPT-Large
|
||||
model = DPTDepthModel(
|
||||
path=model_path,
|
||||
backbone='vitl16_384',
|
||||
non_negative=True,
|
||||
)
|
||||
net_w, net_h = 384, 384
|
||||
resize_mode = 'minimal'
|
||||
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
|
||||
std=[0.5, 0.5, 0.5])
|
||||
|
||||
elif model_type == 'dpt_hybrid': # DPT-Hybrid
|
||||
model = DPTDepthModel(
|
||||
path=model_path,
|
||||
backbone='vitb_rn50_384',
|
||||
non_negative=True,
|
||||
)
|
||||
net_w, net_h = 384, 384
|
||||
resize_mode = 'minimal'
|
||||
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5],
|
||||
std=[0.5, 0.5, 0.5])
|
||||
|
||||
elif model_type == 'midas_v21':
|
||||
model = MidasNet(model_path, non_negative=True)
|
||||
net_w, net_h = 384, 384
|
||||
resize_mode = 'upper_bound'
|
||||
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225])
|
||||
|
||||
elif model_type == 'midas_v21_small':
|
||||
model = MidasNet_small(model_path,
|
||||
features=64,
|
||||
backbone='efficientnet_lite3',
|
||||
exportable=True,
|
||||
non_negative=True,
|
||||
blocks={'expand': True})
|
||||
net_w, net_h = 256, 256
|
||||
resize_mode = 'upper_bound'
|
||||
normalization = NormalizeImage(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225])
|
||||
|
||||
else:
|
||||
print(
|
||||
f"model_type '{model_type}' not implemented, use: --model_type large"
|
||||
)
|
||||
assert False
|
||||
|
||||
transform = Compose([
|
||||
Resize(
|
||||
net_w,
|
||||
net_h,
|
||||
resize_target=None,
|
||||
keep_aspect_ratio=True,
|
||||
ensure_multiple_of=32,
|
||||
resize_method=resize_mode,
|
||||
image_interpolation_method=cv2.INTER_CUBIC,
|
||||
),
|
||||
normalization,
|
||||
PrepareForNet(),
|
||||
])
|
||||
|
||||
return model.eval(), transform
|
||||
|
||||
|
||||
class MiDaSInference(nn.Module):
|
||||
MODEL_TYPES_TORCH_HUB = ['DPT_Large', 'DPT_Hybrid', 'MiDaS_small']
|
||||
MODEL_TYPES_ISL = [
|
||||
'dpt_large',
|
||||
'dpt_hybrid',
|
||||
'midas_v21',
|
||||
'midas_v21_small',
|
||||
]
|
||||
|
||||
def __init__(self, model_type, model_path):
|
||||
super().__init__()
|
||||
assert (model_type in self.MODEL_TYPES_ISL)
|
||||
model, _ = load_model(model_type, model_path)
|
||||
self.model = model
|
||||
self.model.train = disabled_train
|
||||
|
||||
def forward(self, x):
|
||||
with torch.no_grad():
|
||||
prediction = self.model(x)
|
||||
return prediction
|
||||
@@ -0,0 +1,17 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import torch
|
||||
|
||||
|
||||
class BaseModel(torch.nn.Module):
|
||||
def load(self, path):
|
||||
"""Load model from file.
|
||||
|
||||
Args:
|
||||
path (str): file path
|
||||
"""
|
||||
parameters = torch.load(path, map_location=torch.device('cpu'))
|
||||
|
||||
if 'optimizer' in parameters:
|
||||
parameters = parameters['model']
|
||||
|
||||
self.load_state_dict(parameters)
|
||||
@@ -0,0 +1,390 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .vit import (_make_pretrained_vitb16_384, _make_pretrained_vitb_rn50_384,
|
||||
_make_pretrained_vitl16_384)
|
||||
|
||||
|
||||
def _make_encoder(
|
||||
backbone,
|
||||
features,
|
||||
use_pretrained,
|
||||
groups=1,
|
||||
expand=False,
|
||||
exportable=True,
|
||||
hooks=None,
|
||||
use_vit_only=False,
|
||||
use_readout='ignore',
|
||||
):
|
||||
if backbone == 'vitl16_384':
|
||||
pretrained = _make_pretrained_vitl16_384(use_pretrained,
|
||||
hooks=hooks,
|
||||
use_readout=use_readout)
|
||||
scratch = _make_scratch(
|
||||
[256, 512, 1024, 1024], features, groups=groups,
|
||||
expand=expand) # ViT-L/16 - 85.0% Top1 (backbone)
|
||||
elif backbone == 'vitb_rn50_384':
|
||||
pretrained = _make_pretrained_vitb_rn50_384(
|
||||
use_pretrained,
|
||||
hooks=hooks,
|
||||
use_vit_only=use_vit_only,
|
||||
use_readout=use_readout,
|
||||
)
|
||||
scratch = _make_scratch(
|
||||
[256, 512, 768, 768], features, groups=groups,
|
||||
expand=expand) # ViT-H/16 - 85.0% Top1 (backbone)
|
||||
elif backbone == 'vitb16_384':
|
||||
pretrained = _make_pretrained_vitb16_384(use_pretrained,
|
||||
hooks=hooks,
|
||||
use_readout=use_readout)
|
||||
scratch = _make_scratch(
|
||||
[96, 192, 384, 768], features, groups=groups,
|
||||
expand=expand) # ViT-B/16 - 84.6% Top1 (backbone)
|
||||
elif backbone == 'resnext101_wsl':
|
||||
pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
|
||||
scratch = _make_scratch([256, 512, 1024, 2048],
|
||||
features,
|
||||
groups=groups,
|
||||
expand=expand) # efficientnet_lite3
|
||||
elif backbone == 'efficientnet_lite3':
|
||||
pretrained = _make_pretrained_efficientnet_lite3(use_pretrained,
|
||||
exportable=exportable)
|
||||
scratch = _make_scratch([32, 48, 136, 384],
|
||||
features,
|
||||
groups=groups,
|
||||
expand=expand) # efficientnet_lite3
|
||||
else:
|
||||
print(f"Backbone '{backbone}' not implemented")
|
||||
assert False
|
||||
|
||||
return pretrained, scratch
|
||||
|
||||
|
||||
def _make_scratch(in_shape, out_shape, groups=1, expand=False):
|
||||
scratch = nn.Module()
|
||||
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape
|
||||
out_shape3 = out_shape
|
||||
out_shape4 = out_shape
|
||||
if expand is True:
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape * 2
|
||||
out_shape3 = out_shape * 4
|
||||
out_shape4 = out_shape * 8
|
||||
|
||||
scratch.layer1_rn = nn.Conv2d(in_shape[0],
|
||||
out_shape1,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups)
|
||||
scratch.layer2_rn = nn.Conv2d(in_shape[1],
|
||||
out_shape2,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups)
|
||||
scratch.layer3_rn = nn.Conv2d(in_shape[2],
|
||||
out_shape3,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups)
|
||||
scratch.layer4_rn = nn.Conv2d(in_shape[3],
|
||||
out_shape4,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups)
|
||||
|
||||
return scratch
|
||||
|
||||
|
||||
def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
|
||||
efficientnet = torch.hub.load('rwightman/gen-efficientnet-pytorch',
|
||||
'tf_efficientnet_lite3',
|
||||
pretrained=use_pretrained,
|
||||
exportable=exportable)
|
||||
return _make_efficientnet_backbone(efficientnet)
|
||||
|
||||
|
||||
def _make_efficientnet_backbone(effnet):
|
||||
pretrained = nn.Module()
|
||||
|
||||
pretrained.layer1 = nn.Sequential(effnet.conv_stem, effnet.bn1,
|
||||
effnet.act1, *effnet.blocks[0:2])
|
||||
pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
|
||||
pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
|
||||
pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
|
||||
|
||||
return pretrained
|
||||
|
||||
|
||||
def _make_resnet_backbone(resnet):
|
||||
pretrained = nn.Module()
|
||||
pretrained.layer1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
|
||||
resnet.maxpool, resnet.layer1)
|
||||
|
||||
pretrained.layer2 = resnet.layer2
|
||||
pretrained.layer3 = resnet.layer3
|
||||
pretrained.layer4 = resnet.layer4
|
||||
|
||||
return pretrained
|
||||
|
||||
|
||||
def _make_pretrained_resnext101_wsl(use_pretrained):
|
||||
resnet = torch.hub.load('facebookresearch/WSL-Images',
|
||||
'resnext101_32x8d_wsl')
|
||||
return _make_resnet_backbone(resnet)
|
||||
|
||||
|
||||
class Interpolate(nn.Module):
|
||||
"""Interpolation module.
|
||||
"""
|
||||
def __init__(self, scale_factor, mode, align_corners=False):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
scale_factor (float): scaling
|
||||
mode (str): interpolation mode
|
||||
"""
|
||||
super(Interpolate, self).__init__()
|
||||
|
||||
self.interp = nn.functional.interpolate
|
||||
self.scale_factor = scale_factor
|
||||
self.mode = mode
|
||||
self.align_corners = align_corners
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input
|
||||
|
||||
Returns:
|
||||
tensor: interpolated data
|
||||
"""
|
||||
|
||||
x = self.interp(x,
|
||||
scale_factor=self.scale_factor,
|
||||
mode=self.mode,
|
||||
align_corners=self.align_corners)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class ResidualConvUnit(nn.Module):
|
||||
"""Residual convolution module.
|
||||
"""
|
||||
def __init__(self, features):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(features,
|
||||
features,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=True)
|
||||
|
||||
self.conv2 = nn.Conv2d(features,
|
||||
features,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=True)
|
||||
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
out = self.relu(x)
|
||||
out = self.conv1(out)
|
||||
out = self.relu(out)
|
||||
out = self.conv2(out)
|
||||
|
||||
return out + x
|
||||
|
||||
|
||||
class FeatureFusionBlock(nn.Module):
|
||||
"""Feature fusion block.
|
||||
"""
|
||||
def __init__(self, features):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super(FeatureFusionBlock, self).__init__()
|
||||
|
||||
self.resConfUnit1 = ResidualConvUnit(features)
|
||||
self.resConfUnit2 = ResidualConvUnit(features)
|
||||
|
||||
def forward(self, *xs):
|
||||
"""Forward pass.
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
output = xs[0]
|
||||
|
||||
if len(xs) == 2:
|
||||
output += self.resConfUnit1(xs[1])
|
||||
|
||||
output = self.resConfUnit2(output)
|
||||
|
||||
output = nn.functional.interpolate(output,
|
||||
scale_factor=2,
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class ResidualConvUnit_custom(nn.Module):
|
||||
"""Residual convolution module.
|
||||
"""
|
||||
def __init__(self, features, activation, bn):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.bn = bn
|
||||
|
||||
self.groups = 1
|
||||
|
||||
self.conv1 = nn.Conv2d(features,
|
||||
features,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=True,
|
||||
groups=self.groups)
|
||||
|
||||
self.conv2 = nn.Conv2d(features,
|
||||
features,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=True,
|
||||
groups=self.groups)
|
||||
|
||||
if self.bn is True:
|
||||
self.bn1 = nn.BatchNorm2d(features)
|
||||
self.bn2 = nn.BatchNorm2d(features)
|
||||
|
||||
self.activation = activation
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
|
||||
out = self.activation(x)
|
||||
out = self.conv1(out)
|
||||
if self.bn is True:
|
||||
out = self.bn1(out)
|
||||
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
if self.bn is True:
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.groups > 1:
|
||||
out = self.conv_merge(out)
|
||||
|
||||
return self.skip_add.add(out, x)
|
||||
|
||||
# return out + x
|
||||
|
||||
|
||||
class FeatureFusionBlock_custom(nn.Module):
|
||||
"""Feature fusion block.
|
||||
"""
|
||||
def __init__(self,
|
||||
features,
|
||||
activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=False,
|
||||
align_corners=True):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super(FeatureFusionBlock_custom, self).__init__()
|
||||
|
||||
self.deconv = deconv
|
||||
self.align_corners = align_corners
|
||||
|
||||
self.groups = 1
|
||||
|
||||
self.expand = expand
|
||||
out_features = features
|
||||
if self.expand is True:
|
||||
out_features = features // 2
|
||||
|
||||
self.out_conv = nn.Conv2d(features,
|
||||
out_features,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
bias=True,
|
||||
groups=1)
|
||||
|
||||
self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
|
||||
self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
def forward(self, *xs):
|
||||
"""Forward pass.
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
output = xs[0]
|
||||
|
||||
if len(xs) == 2:
|
||||
res = self.resConfUnit1(xs[1])
|
||||
output = self.skip_add.add(output, res)
|
||||
# output += res
|
||||
|
||||
output = self.resConfUnit2(output)
|
||||
|
||||
output = nn.functional.interpolate(output,
|
||||
scale_factor=2,
|
||||
mode='bilinear',
|
||||
align_corners=self.align_corners)
|
||||
|
||||
output = self.out_conv(output)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,106 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder
|
||||
from .vit import forward_vit
|
||||
|
||||
|
||||
def _make_fusion_block(features, use_bn):
|
||||
return FeatureFusionBlock_custom(
|
||||
features,
|
||||
nn.ReLU(False),
|
||||
deconv=False,
|
||||
bn=use_bn,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
)
|
||||
|
||||
|
||||
class DPT(BaseModel):
|
||||
def __init__(
|
||||
self,
|
||||
head,
|
||||
features=256,
|
||||
backbone='vitb_rn50_384',
|
||||
readout='project',
|
||||
channels_last=False,
|
||||
use_bn=False,
|
||||
):
|
||||
|
||||
super(DPT, self).__init__()
|
||||
|
||||
self.channels_last = channels_last
|
||||
|
||||
hooks = {
|
||||
'vitb_rn50_384': [0, 1, 8, 11],
|
||||
'vitb16_384': [2, 5, 8, 11],
|
||||
'vitl16_384': [5, 11, 17, 23],
|
||||
}
|
||||
|
||||
# Instantiate backbone and reassemble blocks
|
||||
self.pretrained, self.scratch = _make_encoder(
|
||||
backbone,
|
||||
features,
|
||||
False, # Set to true of you want to train from scratch, uses ImageNet weights
|
||||
groups=1,
|
||||
expand=False,
|
||||
exportable=False,
|
||||
hooks=hooks[backbone],
|
||||
use_readout=readout,
|
||||
)
|
||||
|
||||
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
||||
|
||||
self.scratch.output_conv = head
|
||||
|
||||
def forward(self, x):
|
||||
if self.channels_last is True:
|
||||
x.contiguous(memory_format=torch.channels_last)
|
||||
|
||||
layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn)
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
out = self.scratch.output_conv(path_1)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class DPTDepthModel(DPT):
|
||||
def __init__(self, path=None, non_negative=True, **kwargs):
|
||||
features = kwargs['features'] if 'features' in kwargs else 256
|
||||
|
||||
head = nn.Sequential(
|
||||
nn.Conv2d(features,
|
||||
features // 2,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1),
|
||||
Interpolate(scale_factor=2, mode='bilinear', align_corners=True),
|
||||
nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True) if non_negative else nn.Identity(),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
super().__init__(head, **kwargs)
|
||||
|
||||
if path is not None:
|
||||
self.load(path)
|
||||
|
||||
def forward(self, x):
|
||||
return super().forward(x).squeeze(dim=1)
|
||||
@@ -0,0 +1,79 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
|
||||
This file contains code that is adapted from
|
||||
https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
|
||||
"""
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .blocks import FeatureFusionBlock, Interpolate, _make_encoder
|
||||
|
||||
|
||||
class MidasNet(BaseModel):
|
||||
"""Network for monocular depth estimation.
|
||||
"""
|
||||
def __init__(self, path=None, features=256, non_negative=True):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
path (str, optional): Path to saved model. Defaults to None.
|
||||
features (int, optional): Number of features. Defaults to 256.
|
||||
backbone (str, optional): Backbone network for encoder. Defaults to resnet50
|
||||
"""
|
||||
print('Loading weights: ', path)
|
||||
|
||||
super(MidasNet, self).__init__()
|
||||
|
||||
use_pretrained = False if path is None else True
|
||||
|
||||
self.pretrained, self.scratch = _make_encoder(
|
||||
backbone='resnext101_wsl',
|
||||
features=features,
|
||||
use_pretrained=use_pretrained)
|
||||
|
||||
self.scratch.refinenet4 = FeatureFusionBlock(features)
|
||||
self.scratch.refinenet3 = FeatureFusionBlock(features)
|
||||
self.scratch.refinenet2 = FeatureFusionBlock(features)
|
||||
self.scratch.refinenet1 = FeatureFusionBlock(features)
|
||||
|
||||
self.scratch.output_conv = nn.Sequential(
|
||||
nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),
|
||||
Interpolate(scale_factor=2, mode='bilinear'),
|
||||
nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True) if non_negative else nn.Identity(),
|
||||
)
|
||||
|
||||
if path:
|
||||
self.load(path)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input data (image)
|
||||
|
||||
Returns:
|
||||
tensor: depth
|
||||
"""
|
||||
|
||||
layer_1 = self.pretrained.layer1(x)
|
||||
layer_2 = self.pretrained.layer2(layer_1)
|
||||
layer_3 = self.pretrained.layer3(layer_2)
|
||||
layer_4 = self.pretrained.layer4(layer_3)
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn)
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
out = self.scratch.output_conv(path_1)
|
||||
|
||||
return torch.squeeze(out, dim=1)
|
||||
@@ -0,0 +1,166 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
|
||||
This file contains code that is adapted from
|
||||
https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
|
||||
"""
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder
|
||||
|
||||
|
||||
class MidasNet_small(BaseModel):
|
||||
"""Network for monocular depth estimation.
|
||||
"""
|
||||
def __init__(self,
|
||||
path=None,
|
||||
features=64,
|
||||
backbone='efficientnet_lite3',
|
||||
non_negative=True,
|
||||
exportable=True,
|
||||
channels_last=False,
|
||||
align_corners=True,
|
||||
blocks={'expand': True}):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
path (str, optional): Path to saved model. Defaults to None.
|
||||
features (int, optional): Number of features. Defaults to 256.
|
||||
backbone (str, optional): Backbone network for encoder. Defaults to resnet50
|
||||
"""
|
||||
print('Loading weights: ', path)
|
||||
|
||||
super(MidasNet_small, self).__init__()
|
||||
|
||||
use_pretrained = False if path else True
|
||||
|
||||
self.channels_last = channels_last
|
||||
self.blocks = blocks
|
||||
self.backbone = backbone
|
||||
|
||||
self.groups = 1
|
||||
|
||||
features1 = features
|
||||
features2 = features
|
||||
features3 = features
|
||||
features4 = features
|
||||
self.expand = False
|
||||
if 'expand' in self.blocks and self.blocks['expand'] is True:
|
||||
self.expand = True
|
||||
features1 = features
|
||||
features2 = features * 2
|
||||
features3 = features * 4
|
||||
features4 = features * 8
|
||||
|
||||
self.pretrained, self.scratch = _make_encoder(self.backbone,
|
||||
features,
|
||||
use_pretrained,
|
||||
groups=self.groups,
|
||||
expand=self.expand,
|
||||
exportable=exportable)
|
||||
|
||||
self.scratch.activation = nn.ReLU(False)
|
||||
|
||||
self.scratch.refinenet4 = FeatureFusionBlock_custom(
|
||||
features4,
|
||||
self.scratch.activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=self.expand,
|
||||
align_corners=align_corners)
|
||||
self.scratch.refinenet3 = FeatureFusionBlock_custom(
|
||||
features3,
|
||||
self.scratch.activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=self.expand,
|
||||
align_corners=align_corners)
|
||||
self.scratch.refinenet2 = FeatureFusionBlock_custom(
|
||||
features2,
|
||||
self.scratch.activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=self.expand,
|
||||
align_corners=align_corners)
|
||||
self.scratch.refinenet1 = FeatureFusionBlock_custom(
|
||||
features1,
|
||||
self.scratch.activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
align_corners=align_corners)
|
||||
|
||||
self.scratch.output_conv = nn.Sequential(
|
||||
nn.Conv2d(features,
|
||||
features // 2,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
groups=self.groups),
|
||||
Interpolate(scale_factor=2, mode='bilinear'),
|
||||
nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
|
||||
self.scratch.activation,
|
||||
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True) if non_negative else nn.Identity(),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
if path:
|
||||
self.load(path)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input data (image)
|
||||
|
||||
Returns:
|
||||
tensor: depth
|
||||
"""
|
||||
if self.channels_last is True:
|
||||
print('self.channels_last = ', self.channels_last)
|
||||
x.contiguous(memory_format=torch.channels_last)
|
||||
|
||||
layer_1 = self.pretrained.layer1(x)
|
||||
layer_2 = self.pretrained.layer2(layer_1)
|
||||
layer_3 = self.pretrained.layer3(layer_2)
|
||||
layer_4 = self.pretrained.layer4(layer_3)
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn)
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
out = self.scratch.output_conv(path_1)
|
||||
|
||||
return torch.squeeze(out, dim=1)
|
||||
|
||||
|
||||
def fuse_model(m):
|
||||
prev_previous_type = nn.Identity()
|
||||
prev_previous_name = ''
|
||||
previous_type = nn.Identity()
|
||||
previous_name = ''
|
||||
for name, module in m.named_modules():
|
||||
if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(
|
||||
module) == nn.ReLU:
|
||||
# print("FUSED ", prev_previous_name, previous_name, name)
|
||||
torch.quantization.fuse_modules(
|
||||
m, [prev_previous_name, previous_name, name], inplace=True)
|
||||
elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d:
|
||||
# print("FUSED ", prev_previous_name, previous_name)
|
||||
torch.quantization.fuse_modules(
|
||||
m, [prev_previous_name, previous_name], inplace=True)
|
||||
# elif previous_type == nn.Conv2d and type(module) == nn.ReLU:
|
||||
# print("FUSED ", previous_name, name)
|
||||
# torch.quantization.fuse_modules(m, [previous_name, name], inplace=True)
|
||||
|
||||
prev_previous_type = previous_type
|
||||
prev_previous_name = previous_name
|
||||
previous_type = type(module)
|
||||
previous_name = name
|
||||
@@ -0,0 +1,230 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import math
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
|
||||
"""Rezise the sample to ensure the given size. Keeps aspect ratio.
|
||||
|
||||
Args:
|
||||
sample (dict): sample
|
||||
size (tuple): image size
|
||||
|
||||
Returns:
|
||||
tuple: new size
|
||||
"""
|
||||
shape = list(sample['disparity'].shape)
|
||||
|
||||
if shape[0] >= size[0] and shape[1] >= size[1]:
|
||||
return sample
|
||||
|
||||
scale = [0, 0]
|
||||
scale[0] = size[0] / shape[0]
|
||||
scale[1] = size[1] / shape[1]
|
||||
|
||||
scale = max(scale)
|
||||
|
||||
shape[0] = math.ceil(scale * shape[0])
|
||||
shape[1] = math.ceil(scale * shape[1])
|
||||
|
||||
# resize
|
||||
sample['image'] = cv2.resize(sample['image'],
|
||||
tuple(shape[::-1]),
|
||||
interpolation=image_interpolation_method)
|
||||
|
||||
sample['disparity'] = cv2.resize(sample['disparity'],
|
||||
tuple(shape[::-1]),
|
||||
interpolation=cv2.INTER_NEAREST)
|
||||
sample['mask'] = cv2.resize(
|
||||
sample['mask'].astype(np.float32),
|
||||
tuple(shape[::-1]),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
sample['mask'] = sample['mask'].astype(bool)
|
||||
|
||||
return tuple(shape)
|
||||
|
||||
|
||||
class Resize(object):
|
||||
"""Resize sample to given size (width, height).
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
width,
|
||||
height,
|
||||
resize_target=True,
|
||||
keep_aspect_ratio=False,
|
||||
ensure_multiple_of=1,
|
||||
resize_method='lower_bound',
|
||||
image_interpolation_method=cv2.INTER_AREA,
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
width (int): desired output width
|
||||
height (int): desired output height
|
||||
resize_target (bool, optional):
|
||||
True: Resize the full sample (image, mask, target).
|
||||
False: Resize image only.
|
||||
Defaults to True.
|
||||
keep_aspect_ratio (bool, optional):
|
||||
True: Keep the aspect ratio of the input sample.
|
||||
Output sample might not have the given width and height, and
|
||||
resize behaviour depends on the parameter 'resize_method'.
|
||||
Defaults to False.
|
||||
ensure_multiple_of (int, optional):
|
||||
Output width and height is constrained to be multiple of this parameter.
|
||||
Defaults to 1.
|
||||
resize_method (str, optional):
|
||||
"lower_bound": Output will be at least as large as the given size.
|
||||
"upper_bound": Output will be at max as large as the given size. "
|
||||
"(Output size might be smaller than given size.)"
|
||||
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
|
||||
Defaults to "lower_bound".
|
||||
"""
|
||||
self.__width = width
|
||||
self.__height = height
|
||||
|
||||
self.__resize_target = resize_target
|
||||
self.__keep_aspect_ratio = keep_aspect_ratio
|
||||
self.__multiple_of = ensure_multiple_of
|
||||
self.__resize_method = resize_method
|
||||
self.__image_interpolation_method = image_interpolation_method
|
||||
|
||||
def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
|
||||
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if max_val is not None and y > max_val:
|
||||
y = (np.floor(x / self.__multiple_of) *
|
||||
self.__multiple_of).astype(int)
|
||||
|
||||
if y < min_val:
|
||||
y = (np.ceil(x / self.__multiple_of) *
|
||||
self.__multiple_of).astype(int)
|
||||
|
||||
return y
|
||||
|
||||
def get_size(self, width, height):
|
||||
# determine new height and width
|
||||
scale_height = self.__height / height
|
||||
scale_width = self.__width / width
|
||||
|
||||
if self.__keep_aspect_ratio:
|
||||
if self.__resize_method == 'lower_bound':
|
||||
# scale such that output size is lower bound
|
||||
if scale_width > scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == 'upper_bound':
|
||||
# scale such that output size is upper bound
|
||||
if scale_width < scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == 'minimal':
|
||||
# scale as least as possbile
|
||||
if abs(1 - scale_width) < abs(1 - scale_height):
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
else:
|
||||
raise ValueError(
|
||||
f'resize_method {self.__resize_method} not implemented')
|
||||
|
||||
if self.__resize_method == 'lower_bound':
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height,
|
||||
min_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width,
|
||||
min_val=self.__width)
|
||||
elif self.__resize_method == 'upper_bound':
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height,
|
||||
max_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width,
|
||||
max_val=self.__width)
|
||||
elif self.__resize_method == 'minimal':
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width)
|
||||
else:
|
||||
raise ValueError(
|
||||
f'resize_method {self.__resize_method} not implemented')
|
||||
|
||||
return (new_width, new_height)
|
||||
|
||||
def __call__(self, sample):
|
||||
width, height = self.get_size(sample['image'].shape[1],
|
||||
sample['image'].shape[0])
|
||||
|
||||
# resize sample
|
||||
sample['image'] = cv2.resize(
|
||||
sample['image'],
|
||||
(width, height),
|
||||
interpolation=self.__image_interpolation_method,
|
||||
)
|
||||
|
||||
if self.__resize_target:
|
||||
if 'disparity' in sample:
|
||||
sample['disparity'] = cv2.resize(
|
||||
sample['disparity'],
|
||||
(width, height),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
|
||||
if 'depth' in sample:
|
||||
sample['depth'] = cv2.resize(sample['depth'], (width, height),
|
||||
interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
sample['mask'] = cv2.resize(
|
||||
sample['mask'].astype(np.float32),
|
||||
(width, height),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
sample['mask'] = sample['mask'].astype(bool)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class NormalizeImage(object):
|
||||
"""Normlize image by given mean and std.
|
||||
"""
|
||||
def __init__(self, mean, std):
|
||||
self.__mean = mean
|
||||
self.__std = std
|
||||
|
||||
def __call__(self, sample):
|
||||
sample['image'] = (sample['image'] - self.__mean) / self.__std
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class PrepareForNet(object):
|
||||
"""Prepare sample for usage as network input.
|
||||
"""
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, sample):
|
||||
image = np.transpose(sample['image'], (2, 0, 1))
|
||||
sample['image'] = np.ascontiguousarray(image).astype(np.float32)
|
||||
|
||||
if 'mask' in sample:
|
||||
sample['mask'] = sample['mask'].astype(np.float32)
|
||||
sample['mask'] = np.ascontiguousarray(sample['mask'])
|
||||
|
||||
if 'disparity' in sample:
|
||||
disparity = sample['disparity'].astype(np.float32)
|
||||
sample['disparity'] = np.ascontiguousarray(disparity)
|
||||
|
||||
if 'depth' in sample:
|
||||
depth = sample['depth'].astype(np.float32)
|
||||
sample['depth'] = np.ascontiguousarray(depth)
|
||||
|
||||
return sample
|
||||
@@ -0,0 +1,192 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Utils for monoDepth."""
|
||||
import re
|
||||
import sys
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def read_pfm(path):
|
||||
"""Read pfm file.
|
||||
|
||||
Args:
|
||||
path (str): path to file
|
||||
|
||||
Returns:
|
||||
tuple: (data, scale)
|
||||
"""
|
||||
with open(path, 'rb') as file:
|
||||
|
||||
color = None
|
||||
width = None
|
||||
height = None
|
||||
scale = None
|
||||
endian = None
|
||||
|
||||
header = file.readline().rstrip()
|
||||
if header.decode('ascii') == 'PF':
|
||||
color = True
|
||||
elif header.decode('ascii') == 'Pf':
|
||||
color = False
|
||||
else:
|
||||
raise Exception('Not a PFM file: ' + path)
|
||||
|
||||
dim_match = re.match(r'^(\d+)\s(\d+)\s$',
|
||||
file.readline().decode('ascii'))
|
||||
if dim_match:
|
||||
width, height = list(map(int, dim_match.groups()))
|
||||
else:
|
||||
raise Exception('Malformed PFM header.')
|
||||
|
||||
scale = float(file.readline().decode('ascii').rstrip())
|
||||
if scale < 0:
|
||||
# little-endian
|
||||
endian = '<'
|
||||
scale = -scale
|
||||
else:
|
||||
# big-endian
|
||||
endian = '>'
|
||||
|
||||
data = np.fromfile(file, endian + 'f')
|
||||
shape = (height, width, 3) if color else (height, width)
|
||||
|
||||
data = np.reshape(data, shape)
|
||||
data = np.flipud(data)
|
||||
|
||||
return data, scale
|
||||
|
||||
|
||||
def write_pfm(path, image, scale=1):
|
||||
"""Write pfm file.
|
||||
|
||||
Args:
|
||||
path (str): pathto file
|
||||
image (array): data
|
||||
scale (int, optional): Scale. Defaults to 1.
|
||||
"""
|
||||
|
||||
with open(path, 'wb') as file:
|
||||
color = None
|
||||
|
||||
if image.dtype.name != 'float32':
|
||||
raise Exception('Image dtype must be float32.')
|
||||
|
||||
image = np.flipud(image)
|
||||
|
||||
if len(image.shape) == 3 and image.shape[2] == 3: # color image
|
||||
color = True
|
||||
elif (len(image.shape) == 2
|
||||
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
|
||||
color = False
|
||||
else:
|
||||
raise Exception(
|
||||
'Image must have H x W x 3, H x W x 1 or H x W dimensions.')
|
||||
|
||||
file.write('PF\n' if color else 'Pf\n'.encode())
|
||||
file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0]))
|
||||
|
||||
endian = image.dtype.byteorder
|
||||
|
||||
if endian == '<' or endian == '=' and sys.byteorder == 'little':
|
||||
scale = -scale
|
||||
|
||||
file.write('%f\n'.encode() % scale)
|
||||
|
||||
image.tofile(file)
|
||||
|
||||
|
||||
def read_image(path):
|
||||
"""Read image and output RGB image (0-1).
|
||||
|
||||
Args:
|
||||
path (str): path to file
|
||||
|
||||
Returns:
|
||||
array: RGB image (0-1)
|
||||
"""
|
||||
img = cv2.imread(path)
|
||||
|
||||
if img.ndim == 2:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def resize_image(img):
|
||||
"""Resize image and make it fit for network.
|
||||
|
||||
Args:
|
||||
img (array): image
|
||||
|
||||
Returns:
|
||||
tensor: data ready for network
|
||||
"""
|
||||
height_orig = img.shape[0]
|
||||
width_orig = img.shape[1]
|
||||
|
||||
if width_orig > height_orig:
|
||||
scale = width_orig / 384
|
||||
else:
|
||||
scale = height_orig / 384
|
||||
|
||||
height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
|
||||
width = (np.ceil(width_orig / scale / 32) * 32).astype(int)
|
||||
|
||||
img_resized = cv2.resize(img, (width, height),
|
||||
interpolation=cv2.INTER_AREA)
|
||||
|
||||
img_resized = (torch.from_numpy(np.transpose(
|
||||
img_resized, (2, 0, 1))).contiguous().float())
|
||||
img_resized = img_resized.unsqueeze(0)
|
||||
|
||||
return img_resized
|
||||
|
||||
|
||||
def resize_depth(depth, width, height):
|
||||
"""Resize depth map and bring to CPU (numpy).
|
||||
|
||||
Args:
|
||||
depth (tensor): depth
|
||||
width (int): image width
|
||||
height (int): image height
|
||||
|
||||
Returns:
|
||||
array: processed depth
|
||||
"""
|
||||
depth = torch.squeeze(depth[0, :, :, :]).to('cpu')
|
||||
|
||||
depth_resized = cv2.resize(depth.numpy(), (width, height),
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
return depth_resized
|
||||
|
||||
|
||||
def write_depth(path, depth, bits=1):
|
||||
"""Write depth map to pfm and png file.
|
||||
|
||||
Args:
|
||||
path (str): filepath without extension
|
||||
depth (array): depth
|
||||
"""
|
||||
write_pfm(path + '.pfm', depth.astype(np.float32))
|
||||
|
||||
depth_min = depth.min()
|
||||
depth_max = depth.max()
|
||||
|
||||
max_val = (2**(8 * bits)) - 1
|
||||
|
||||
if depth_max - depth_min > np.finfo('float').eps:
|
||||
out = max_val * (depth - depth_min) / (depth_max - depth_min)
|
||||
else:
|
||||
out = np.zeros(depth.shape, dtype=depth.type)
|
||||
|
||||
if bits == 1:
|
||||
cv2.imwrite(path + '.png', out.astype('uint8'))
|
||||
elif bits == 2:
|
||||
cv2.imwrite(path + '.png', out.astype('uint16'))
|
||||
|
||||
return
|
||||
@@ -0,0 +1,509 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import math
|
||||
import types
|
||||
|
||||
import timm
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Slice(nn.Module):
|
||||
def __init__(self, start_index=1):
|
||||
super(Slice, self).__init__()
|
||||
self.start_index = start_index
|
||||
|
||||
def forward(self, x):
|
||||
return x[:, self.start_index:]
|
||||
|
||||
|
||||
class AddReadout(nn.Module):
|
||||
def __init__(self, start_index=1):
|
||||
super(AddReadout, self).__init__()
|
||||
self.start_index = start_index
|
||||
|
||||
def forward(self, x):
|
||||
if self.start_index == 2:
|
||||
readout = (x[:, 0] + x[:, 1]) / 2
|
||||
else:
|
||||
readout = x[:, 0]
|
||||
return x[:, self.start_index:] + readout.unsqueeze(1)
|
||||
|
||||
|
||||
class ProjectReadout(nn.Module):
|
||||
def __init__(self, in_features, start_index=1):
|
||||
super(ProjectReadout, self).__init__()
|
||||
self.start_index = start_index
|
||||
|
||||
self.project = nn.Sequential(nn.Linear(2 * in_features, in_features),
|
||||
nn.GELU())
|
||||
|
||||
def forward(self, x):
|
||||
readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index:])
|
||||
features = torch.cat((x[:, self.start_index:], readout), -1)
|
||||
|
||||
return self.project(features)
|
||||
|
||||
|
||||
class Transpose(nn.Module):
|
||||
def __init__(self, dim0, dim1):
|
||||
super(Transpose, self).__init__()
|
||||
self.dim0 = dim0
|
||||
self.dim1 = dim1
|
||||
|
||||
def forward(self, x):
|
||||
x = x.transpose(self.dim0, self.dim1)
|
||||
return x
|
||||
|
||||
|
||||
def forward_vit(pretrained, x):
|
||||
b, c, h, w = x.shape
|
||||
|
||||
_ = pretrained.model.forward_flex(x)
|
||||
|
||||
layer_1 = pretrained.activations['1']
|
||||
layer_2 = pretrained.activations['2']
|
||||
layer_3 = pretrained.activations['3']
|
||||
layer_4 = pretrained.activations['4']
|
||||
|
||||
layer_1 = pretrained.act_postprocess1[0:2](layer_1)
|
||||
layer_2 = pretrained.act_postprocess2[0:2](layer_2)
|
||||
layer_3 = pretrained.act_postprocess3[0:2](layer_3)
|
||||
layer_4 = pretrained.act_postprocess4[0:2](layer_4)
|
||||
|
||||
unflatten = nn.Sequential(
|
||||
nn.Unflatten(
|
||||
2,
|
||||
torch.Size([
|
||||
h // pretrained.model.patch_size[1],
|
||||
w // pretrained.model.patch_size[0],
|
||||
]),
|
||||
))
|
||||
|
||||
if layer_1.ndim == 3:
|
||||
layer_1 = unflatten(layer_1)
|
||||
if layer_2.ndim == 3:
|
||||
layer_2 = unflatten(layer_2)
|
||||
if layer_3.ndim == 3:
|
||||
layer_3 = unflatten(layer_3)
|
||||
if layer_4.ndim == 3:
|
||||
layer_4 = unflatten(layer_4)
|
||||
|
||||
layer_1 = pretrained.act_postprocess1[3:len(pretrained.act_postprocess1)](
|
||||
layer_1)
|
||||
layer_2 = pretrained.act_postprocess2[3:len(pretrained.act_postprocess2)](
|
||||
layer_2)
|
||||
layer_3 = pretrained.act_postprocess3[3:len(pretrained.act_postprocess3)](
|
||||
layer_3)
|
||||
layer_4 = pretrained.act_postprocess4[3:len(pretrained.act_postprocess4)](
|
||||
layer_4)
|
||||
|
||||
return layer_1, layer_2, layer_3, layer_4
|
||||
|
||||
|
||||
def _resize_pos_embed(self, posemb, gs_h, gs_w):
|
||||
posemb_tok, posemb_grid = (
|
||||
posemb[:, :self.start_index],
|
||||
posemb[0, self.start_index:],
|
||||
)
|
||||
|
||||
gs_old = int(math.sqrt(len(posemb_grid)))
|
||||
|
||||
posemb_grid = posemb_grid.reshape(1, gs_old, gs_old,
|
||||
-1).permute(0, 3, 1, 2)
|
||||
posemb_grid = F.interpolate(posemb_grid,
|
||||
size=(gs_h, gs_w),
|
||||
mode='bilinear')
|
||||
posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)
|
||||
|
||||
posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
|
||||
|
||||
return posemb
|
||||
|
||||
|
||||
def forward_flex(self, x):
|
||||
b, c, h, w = x.shape
|
||||
|
||||
pos_embed = self._resize_pos_embed(self.pos_embed, h // self.patch_size[1],
|
||||
w // self.patch_size[0])
|
||||
|
||||
B = x.shape[0]
|
||||
|
||||
if hasattr(self.patch_embed, 'backbone'):
|
||||
x = self.patch_embed.backbone(x)
|
||||
if isinstance(x, (list, tuple)):
|
||||
x = x[
|
||||
-1] # last feature if backbone outputs list/tuple of features
|
||||
|
||||
x = self.patch_embed.proj(x).flatten(2).transpose(1, 2)
|
||||
|
||||
if getattr(self, 'dist_token', None) is not None:
|
||||
cls_tokens = self.cls_token.expand(
|
||||
B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
||||
dist_token = self.dist_token.expand(B, -1, -1)
|
||||
x = torch.cat((cls_tokens, dist_token, x), dim=1)
|
||||
else:
|
||||
cls_tokens = self.cls_token.expand(
|
||||
B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
||||
x = torch.cat((cls_tokens, x), dim=1)
|
||||
|
||||
x = x + pos_embed
|
||||
x = self.pos_drop(x)
|
||||
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
x = self.norm(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
activations = {}
|
||||
|
||||
|
||||
def get_activation(name):
|
||||
def hook(model, input, output):
|
||||
activations[name] = output
|
||||
|
||||
return hook
|
||||
|
||||
|
||||
def get_readout_oper(vit_features, features, use_readout, start_index=1):
|
||||
if use_readout == 'ignore':
|
||||
readout_oper = [Slice(start_index)] * len(features)
|
||||
elif use_readout == 'add':
|
||||
readout_oper = [AddReadout(start_index)] * len(features)
|
||||
elif use_readout == 'project':
|
||||
readout_oper = [
|
||||
ProjectReadout(vit_features, start_index) for out_feat in features
|
||||
]
|
||||
else:
|
||||
assert (
|
||||
False
|
||||
), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'"
|
||||
|
||||
return readout_oper
|
||||
|
||||
|
||||
def _make_vit_b16_backbone(
|
||||
model,
|
||||
features=[96, 192, 384, 768],
|
||||
size=[384, 384],
|
||||
hooks=[2, 5, 8, 11],
|
||||
vit_features=768,
|
||||
use_readout='ignore',
|
||||
start_index=1,
|
||||
):
|
||||
pretrained = nn.Module()
|
||||
|
||||
pretrained.model = model
|
||||
pretrained.model.blocks[hooks[0]].register_forward_hook(
|
||||
get_activation('1'))
|
||||
pretrained.model.blocks[hooks[1]].register_forward_hook(
|
||||
get_activation('2'))
|
||||
pretrained.model.blocks[hooks[2]].register_forward_hook(
|
||||
get_activation('3'))
|
||||
pretrained.model.blocks[hooks[3]].register_forward_hook(
|
||||
get_activation('4'))
|
||||
|
||||
pretrained.activations = activations
|
||||
|
||||
readout_oper = get_readout_oper(vit_features, features, use_readout,
|
||||
start_index)
|
||||
|
||||
# 32, 48, 136, 384
|
||||
pretrained.act_postprocess1 = nn.Sequential(
|
||||
readout_oper[0],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[0],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=features[0],
|
||||
out_channels=features[0],
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.act_postprocess2 = nn.Sequential(
|
||||
readout_oper[1],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[1],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=features[1],
|
||||
out_channels=features[1],
|
||||
kernel_size=2,
|
||||
stride=2,
|
||||
padding=0,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.act_postprocess3 = nn.Sequential(
|
||||
readout_oper[2],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[2],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.act_postprocess4 = nn.Sequential(
|
||||
readout_oper[3],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[3],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
nn.Conv2d(
|
||||
in_channels=features[3],
|
||||
out_channels=features[3],
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.model.start_index = start_index
|
||||
pretrained.model.patch_size = [16, 16]
|
||||
|
||||
# We inject this function into the VisionTransformer instances so that
|
||||
# we can use it with interpolated position embeddings without modifying the library source.
|
||||
pretrained.model.forward_flex = types.MethodType(forward_flex,
|
||||
pretrained.model)
|
||||
pretrained.model._resize_pos_embed = types.MethodType(
|
||||
_resize_pos_embed, pretrained.model)
|
||||
|
||||
return pretrained
|
||||
|
||||
|
||||
def _make_pretrained_vitl16_384(pretrained, use_readout='ignore', hooks=None):
|
||||
model = timm.create_model('vit_large_patch16_384', pretrained=pretrained)
|
||||
|
||||
hooks = [5, 11, 17, 23] if hooks is None else hooks
|
||||
return _make_vit_b16_backbone(
|
||||
model,
|
||||
features=[256, 512, 1024, 1024],
|
||||
hooks=hooks,
|
||||
vit_features=1024,
|
||||
use_readout=use_readout,
|
||||
)
|
||||
|
||||
|
||||
def _make_pretrained_vitb16_384(pretrained, use_readout='ignore', hooks=None):
|
||||
model = timm.create_model('vit_base_patch16_384', pretrained=pretrained)
|
||||
|
||||
hooks = [2, 5, 8, 11] if hooks is None else hooks
|
||||
return _make_vit_b16_backbone(model,
|
||||
features=[96, 192, 384, 768],
|
||||
hooks=hooks,
|
||||
use_readout=use_readout)
|
||||
|
||||
|
||||
def _make_pretrained_deitb16_384(pretrained, use_readout='ignore', hooks=None):
|
||||
model = timm.create_model('vit_deit_base_patch16_384',
|
||||
pretrained=pretrained)
|
||||
|
||||
hooks = [2, 5, 8, 11] if hooks is None else hooks
|
||||
return _make_vit_b16_backbone(model,
|
||||
features=[96, 192, 384, 768],
|
||||
hooks=hooks,
|
||||
use_readout=use_readout)
|
||||
|
||||
|
||||
def _make_pretrained_deitb16_distil_384(pretrained,
|
||||
use_readout='ignore',
|
||||
hooks=None):
|
||||
model = timm.create_model('vit_deit_base_distilled_patch16_384',
|
||||
pretrained=pretrained)
|
||||
|
||||
hooks = [2, 5, 8, 11] if hooks is None else hooks
|
||||
return _make_vit_b16_backbone(
|
||||
model,
|
||||
features=[96, 192, 384, 768],
|
||||
hooks=hooks,
|
||||
use_readout=use_readout,
|
||||
start_index=2,
|
||||
)
|
||||
|
||||
|
||||
def _make_vit_b_rn50_backbone(
|
||||
model,
|
||||
features=[256, 512, 768, 768],
|
||||
size=[384, 384],
|
||||
hooks=[0, 1, 8, 11],
|
||||
vit_features=768,
|
||||
use_vit_only=False,
|
||||
use_readout='ignore',
|
||||
start_index=1,
|
||||
):
|
||||
pretrained = nn.Module()
|
||||
|
||||
pretrained.model = model
|
||||
|
||||
if use_vit_only is True:
|
||||
pretrained.model.blocks[hooks[0]].register_forward_hook(
|
||||
get_activation('1'))
|
||||
pretrained.model.blocks[hooks[1]].register_forward_hook(
|
||||
get_activation('2'))
|
||||
else:
|
||||
pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
|
||||
get_activation('1'))
|
||||
pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
|
||||
get_activation('2'))
|
||||
|
||||
pretrained.model.blocks[hooks[2]].register_forward_hook(
|
||||
get_activation('3'))
|
||||
pretrained.model.blocks[hooks[3]].register_forward_hook(
|
||||
get_activation('4'))
|
||||
|
||||
pretrained.activations = activations
|
||||
|
||||
readout_oper = get_readout_oper(vit_features, features, use_readout,
|
||||
start_index)
|
||||
|
||||
if use_vit_only is True:
|
||||
pretrained.act_postprocess1 = nn.Sequential(
|
||||
readout_oper[0],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[0],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=features[0],
|
||||
out_channels=features[0],
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.act_postprocess2 = nn.Sequential(
|
||||
readout_oper[1],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[1],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=features[1],
|
||||
out_channels=features[1],
|
||||
kernel_size=2,
|
||||
stride=2,
|
||||
padding=0,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
),
|
||||
)
|
||||
else:
|
||||
pretrained.act_postprocess1 = nn.Sequential(nn.Identity(),
|
||||
nn.Identity(),
|
||||
nn.Identity())
|
||||
pretrained.act_postprocess2 = nn.Sequential(nn.Identity(),
|
||||
nn.Identity(),
|
||||
nn.Identity())
|
||||
|
||||
pretrained.act_postprocess3 = nn.Sequential(
|
||||
readout_oper[2],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[2],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.act_postprocess4 = nn.Sequential(
|
||||
readout_oper[3],
|
||||
Transpose(1, 2),
|
||||
nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
|
||||
nn.Conv2d(
|
||||
in_channels=vit_features,
|
||||
out_channels=features[3],
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
),
|
||||
nn.Conv2d(
|
||||
in_channels=features[3],
|
||||
out_channels=features[3],
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
),
|
||||
)
|
||||
|
||||
pretrained.model.start_index = start_index
|
||||
pretrained.model.patch_size = [16, 16]
|
||||
|
||||
# We inject this function into the VisionTransformer instances so that
|
||||
# we can use it with interpolated position embeddings without modifying the library source.
|
||||
pretrained.model.forward_flex = types.MethodType(forward_flex,
|
||||
pretrained.model)
|
||||
|
||||
# We inject this function into the VisionTransformer instances so that
|
||||
# we can use it with interpolated position embeddings without modifying the library source.
|
||||
pretrained.model._resize_pos_embed = types.MethodType(
|
||||
_resize_pos_embed, pretrained.model)
|
||||
|
||||
return pretrained
|
||||
|
||||
|
||||
def _make_pretrained_vitb_rn50_384(pretrained,
|
||||
use_readout='ignore',
|
||||
hooks=None,
|
||||
use_vit_only=False):
|
||||
model = timm.create_model('vit_base_resnet50_384', pretrained=pretrained)
|
||||
|
||||
hooks = [0, 1, 8, 11] if hooks is None else hooks
|
||||
return _make_vit_b_rn50_backbone(
|
||||
model,
|
||||
features=[256, 512, 768, 768],
|
||||
size=[384, 384],
|
||||
hooks=hooks,
|
||||
use_vit_only=use_vit_only,
|
||||
use_readout=use_readout,
|
||||
)
|
||||
@@ -0,0 +1,78 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Midas Depth Estimation
|
||||
# From https://github.com/isl-org/MiDaS
|
||||
# MIT LICENSE
|
||||
from abc import ABCMeta
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
|
||||
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
||||
from scepter.modules.annotator.midas.api import MiDaSInference
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.annotator.utils import resize_image, resize_image_ori
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class MidasDetector(BaseAnnotator, metaclass=ABCMeta):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
pretrained_model = cfg.get('PRETRAINED_MODEL', None)
|
||||
if pretrained_model:
|
||||
with FS.get_from(pretrained_model, wait_finish=True) as local_path:
|
||||
self.model = MiDaSInference(model_type='dpt_hybrid',
|
||||
model_path=local_path)
|
||||
self.a = cfg.get('A', np.pi * 2.0)
|
||||
self.bg_th = cfg.get('BG_TH', 0.1)
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
@torch.autocast('cuda', enabled=False)
|
||||
def forward(self, image):
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image = image.detach().cpu().numpy()
|
||||
elif isinstance(image, np.ndarray):
|
||||
image = image.copy()
|
||||
else:
|
||||
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
||||
image_depth = image
|
||||
h, w, c = image.shape
|
||||
image_depth, k = resize_image(image_depth,
|
||||
1024 if min(h, w) > 1024 else min(h, w))
|
||||
image_depth = torch.from_numpy(image_depth).float().to(we.device_id)
|
||||
image_depth = image_depth / 127.5 - 1.0
|
||||
image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
|
||||
depth = self.model(image_depth)[0]
|
||||
|
||||
depth_pt = depth.clone()
|
||||
depth_pt -= torch.min(depth_pt)
|
||||
depth_pt /= torch.max(depth_pt)
|
||||
depth_pt = depth_pt.cpu().numpy()
|
||||
depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8)
|
||||
depth_image = depth_image[..., None].repeat(3, 2)
|
||||
|
||||
# depth_np = depth.cpu().numpy() # float16 error
|
||||
# x = cv2.Sobel(depth_np, cv2.CV_32F, 1, 0, ksize=3)
|
||||
# y = cv2.Sobel(depth_np, cv2.CV_32F, 0, 1, ksize=3)
|
||||
# z = np.ones_like(x) * self.a
|
||||
# x[depth_pt < self.bg_th] = 0
|
||||
# y[depth_pt < self.bg_th] = 0
|
||||
# normal = np.stack([x, y, z], axis=2)
|
||||
# normal /= np.sum(normal**2.0, axis=2, keepdims=True)**0.5
|
||||
# normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
|
||||
depth_image = resize_image_ori(h, w, depth_image, k)
|
||||
return depth_image
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
MidasDetector.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,303 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
class BlockTypeA(nn.Module):
|
||||
def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale=True):
|
||||
super(BlockTypeA, self).__init__()
|
||||
self.conv1 = nn.Sequential(nn.Conv2d(in_c2, out_c2, kernel_size=1),
|
||||
nn.BatchNorm2d(out_c2),
|
||||
nn.ReLU(inplace=True))
|
||||
self.conv2 = nn.Sequential(nn.Conv2d(in_c1, out_c1, kernel_size=1),
|
||||
nn.BatchNorm2d(out_c1),
|
||||
nn.ReLU(inplace=True))
|
||||
self.upscale = upscale
|
||||
|
||||
def forward(self, a, b):
|
||||
b = self.conv1(b)
|
||||
a = self.conv2(a)
|
||||
if self.upscale:
|
||||
b = F.interpolate(b,
|
||||
scale_factor=2.0,
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
return torch.cat((a, b), dim=1)
|
||||
|
||||
|
||||
class BlockTypeB(nn.Module):
|
||||
def __init__(self, in_c, out_c):
|
||||
super(BlockTypeB, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(in_c, in_c, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(in_c), nn.ReLU())
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(out_c), nn.ReLU())
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x) + x
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
||||
|
||||
class BlockTypeC(nn.Module):
|
||||
def __init__(self, in_c, out_c):
|
||||
super(BlockTypeC, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(in_c, in_c, kernel_size=3, padding=5, dilation=5),
|
||||
nn.BatchNorm2d(in_c), nn.ReLU())
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv2d(in_c, in_c, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(in_c), nn.ReLU())
|
||||
self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
x = self.conv3(x)
|
||||
return x
|
||||
|
||||
|
||||
def _make_divisible(v, divisor, min_value=None):
|
||||
"""
|
||||
This function is taken from the original tf repo.
|
||||
It ensures that all layers have a channel number that is divisible by 8
|
||||
It can be seen here:
|
||||
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
|
||||
:param v:
|
||||
:param divisor:
|
||||
:param min_value:
|
||||
:return:
|
||||
"""
|
||||
if min_value is None:
|
||||
min_value = divisor
|
||||
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
|
||||
# Make sure that round down does not go down by more than 10%.
|
||||
if new_v < 0.9 * v:
|
||||
new_v += divisor
|
||||
return new_v
|
||||
|
||||
|
||||
class ConvBNReLU(nn.Sequential):
|
||||
def __init__(self,
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
groups=1):
|
||||
self.channel_pad = out_planes - in_planes
|
||||
self.stride = stride
|
||||
# padding = (kernel_size - 1) // 2
|
||||
|
||||
# TFLite uses slightly different padding than PyTorch
|
||||
if stride == 2:
|
||||
padding = 0
|
||||
else:
|
||||
padding = (kernel_size - 1) // 2
|
||||
|
||||
super(ConvBNReLU, self).__init__(
|
||||
nn.Conv2d(in_planes,
|
||||
out_planes,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
groups=groups,
|
||||
bias=False), nn.BatchNorm2d(out_planes),
|
||||
nn.ReLU6(inplace=True))
|
||||
self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
|
||||
|
||||
def forward(self, x):
|
||||
# TFLite uses different padding
|
||||
if self.stride == 2:
|
||||
x = F.pad(x, (0, 1, 0, 1), 'constant', 0)
|
||||
# print(x.shape)
|
||||
|
||||
for module in self:
|
||||
if not isinstance(module, nn.MaxPool2d):
|
||||
x = module(x)
|
||||
return x
|
||||
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, inp, oup, stride, expand_ratio):
|
||||
super(InvertedResidual, self).__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
hidden_dim = int(round(inp * expand_ratio))
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
layers = []
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
|
||||
layers.extend([
|
||||
# dw
|
||||
ConvBNReLU(hidden_dim,
|
||||
hidden_dim,
|
||||
stride=stride,
|
||||
groups=hidden_dim),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class MobileNetV2(nn.Module):
|
||||
def __init__(self, pretrained=True):
|
||||
"""
|
||||
MobileNet V2 main class
|
||||
Args:
|
||||
num_classes (int): Number of classes
|
||||
width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
|
||||
inverted_residual_setting: Network structure
|
||||
round_nearest (int): Round the number of channels in each layer to be a multiple of this number
|
||||
Set to 1 to turn off rounding
|
||||
block: Module specifying inverted residual building block for mobilenet
|
||||
"""
|
||||
super(MobileNetV2, self).__init__()
|
||||
|
||||
block = InvertedResidual
|
||||
input_channel = 32
|
||||
last_channel = 1280
|
||||
width_mult = 1.0
|
||||
round_nearest = 8
|
||||
|
||||
inverted_residual_setting = [
|
||||
# t, c, n, s
|
||||
[1, 16, 1, 1],
|
||||
[6, 24, 2, 2],
|
||||
[6, 32, 3, 2],
|
||||
[6, 64, 4, 2],
|
||||
[6, 96, 3, 1],
|
||||
# [6, 160, 3, 2],
|
||||
# [6, 320, 1, 1],
|
||||
]
|
||||
|
||||
# only check the first element, assuming user knows t,c,n,s are required
|
||||
if len(inverted_residual_setting) == 0 or len(
|
||||
inverted_residual_setting[0]) != 4:
|
||||
raise ValueError('inverted_residual_setting should be non-empty '
|
||||
'or a 4-element list, got {}'.format(
|
||||
inverted_residual_setting))
|
||||
|
||||
# building first layer
|
||||
input_channel = _make_divisible(input_channel * width_mult,
|
||||
round_nearest)
|
||||
self.last_channel = _make_divisible(
|
||||
last_channel * max(1.0, width_mult), round_nearest)
|
||||
features = [ConvBNReLU(4, input_channel, stride=2)]
|
||||
# building inverted residual blocks
|
||||
for t, c, n, s in inverted_residual_setting:
|
||||
output_channel = _make_divisible(c * width_mult, round_nearest)
|
||||
for i in range(n):
|
||||
stride = s if i == 0 else 1
|
||||
features.append(
|
||||
block(input_channel,
|
||||
output_channel,
|
||||
stride,
|
||||
expand_ratio=t))
|
||||
input_channel = output_channel
|
||||
|
||||
self.features = nn.Sequential(*features)
|
||||
self.fpn_selected = [1, 3, 6, 10, 13]
|
||||
# weight initialization
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out')
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.ones_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
nn.init.zeros_(m.bias)
|
||||
if pretrained:
|
||||
self._load_pretrained_model()
|
||||
|
||||
def _forward_impl(self, x):
|
||||
# This exists since TorchScript doesn't support inheritance, so the superclass method
|
||||
# (this one) needs to have a name other than `forward` that can be accessed in a subclass
|
||||
fpn_features = []
|
||||
for i, f in enumerate(self.features):
|
||||
if i > self.fpn_selected[-1]:
|
||||
break
|
||||
x = f(x)
|
||||
if i in self.fpn_selected:
|
||||
fpn_features.append(x)
|
||||
|
||||
c1, c2, c3, c4, c5 = fpn_features
|
||||
return c1, c2, c3, c4, c5
|
||||
|
||||
def forward(self, x):
|
||||
return self._forward_impl(x)
|
||||
|
||||
def _load_pretrained_model(self):
|
||||
pretrain_dict = model_zoo.load_url(
|
||||
'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
|
||||
model_dict = {}
|
||||
state_dict = self.state_dict()
|
||||
for k, v in pretrain_dict.items():
|
||||
if k in state_dict:
|
||||
model_dict[k] = v
|
||||
state_dict.update(model_dict)
|
||||
self.load_state_dict(state_dict)
|
||||
|
||||
|
||||
class MobileV2_MLSD_Large(nn.Module):
|
||||
def __init__(self):
|
||||
super(MobileV2_MLSD_Large, self).__init__()
|
||||
|
||||
self.backbone = MobileNetV2(pretrained=False)
|
||||
# A, B
|
||||
self.block15 = BlockTypeA(in_c1=64,
|
||||
in_c2=96,
|
||||
out_c1=64,
|
||||
out_c2=64,
|
||||
upscale=False)
|
||||
self.block16 = BlockTypeB(128, 64)
|
||||
|
||||
# A, B
|
||||
self.block17 = BlockTypeA(in_c1=32, in_c2=64, out_c1=64, out_c2=64)
|
||||
self.block18 = BlockTypeB(128, 64)
|
||||
|
||||
# A, B
|
||||
self.block19 = BlockTypeA(in_c1=24, in_c2=64, out_c1=64, out_c2=64)
|
||||
self.block20 = BlockTypeB(128, 64)
|
||||
|
||||
# A, B, C
|
||||
self.block21 = BlockTypeA(in_c1=16, in_c2=64, out_c1=64, out_c2=64)
|
||||
self.block22 = BlockTypeB(128, 64)
|
||||
|
||||
self.block23 = BlockTypeC(64, 16)
|
||||
|
||||
def forward(self, x):
|
||||
c1, c2, c3, c4, c5 = self.backbone(x)
|
||||
|
||||
x = self.block15(c4, c5)
|
||||
x = self.block16(x)
|
||||
|
||||
x = self.block17(c3, x)
|
||||
x = self.block18(x)
|
||||
|
||||
x = self.block19(c2, x)
|
||||
x = self.block20(x)
|
||||
|
||||
x = self.block21(c1, x)
|
||||
x = self.block22(x)
|
||||
x = self.block23(x)
|
||||
x = x[:, 7:, :, :]
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,287 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
class BlockTypeA(nn.Module):
|
||||
def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale=True):
|
||||
super(BlockTypeA, self).__init__()
|
||||
self.conv1 = nn.Sequential(nn.Conv2d(in_c2, out_c2, kernel_size=1),
|
||||
nn.BatchNorm2d(out_c2),
|
||||
nn.ReLU(inplace=True))
|
||||
self.conv2 = nn.Sequential(nn.Conv2d(in_c1, out_c1, kernel_size=1),
|
||||
nn.BatchNorm2d(out_c1),
|
||||
nn.ReLU(inplace=True))
|
||||
self.upscale = upscale
|
||||
|
||||
def forward(self, a, b):
|
||||
b = self.conv1(b)
|
||||
a = self.conv2(a)
|
||||
b = F.interpolate(b,
|
||||
scale_factor=2.0,
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
return torch.cat((a, b), dim=1)
|
||||
|
||||
|
||||
class BlockTypeB(nn.Module):
|
||||
def __init__(self, in_c, out_c):
|
||||
super(BlockTypeB, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(in_c, in_c, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(in_c), nn.ReLU())
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(out_c), nn.ReLU())
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x) + x
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
||||
|
||||
class BlockTypeC(nn.Module):
|
||||
def __init__(self, in_c, out_c):
|
||||
super(BlockTypeC, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(in_c, in_c, kernel_size=3, padding=5, dilation=5),
|
||||
nn.BatchNorm2d(in_c), nn.ReLU())
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv2d(in_c, in_c, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(in_c), nn.ReLU())
|
||||
self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
x = self.conv3(x)
|
||||
return x
|
||||
|
||||
|
||||
def _make_divisible(v, divisor, min_value=None):
|
||||
"""
|
||||
This function is taken from the original tf repo.
|
||||
It ensures that all layers have a channel number that is divisible by 8
|
||||
It can be seen here:
|
||||
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
|
||||
:param v:
|
||||
:param divisor:
|
||||
:param min_value:
|
||||
:return:
|
||||
"""
|
||||
if min_value is None:
|
||||
min_value = divisor
|
||||
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
|
||||
# Make sure that round down does not go down by more than 10%.
|
||||
if new_v < 0.9 * v:
|
||||
new_v += divisor
|
||||
return new_v
|
||||
|
||||
|
||||
class ConvBNReLU(nn.Sequential):
|
||||
def __init__(self,
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
groups=1):
|
||||
self.channel_pad = out_planes - in_planes
|
||||
self.stride = stride
|
||||
# padding = (kernel_size - 1) // 2
|
||||
|
||||
# TFLite uses slightly different padding than PyTorch
|
||||
if stride == 2:
|
||||
padding = 0
|
||||
else:
|
||||
padding = (kernel_size - 1) // 2
|
||||
|
||||
super(ConvBNReLU, self).__init__(
|
||||
nn.Conv2d(in_planes,
|
||||
out_planes,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
groups=groups,
|
||||
bias=False), nn.BatchNorm2d(out_planes),
|
||||
nn.ReLU6(inplace=True))
|
||||
self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
|
||||
|
||||
def forward(self, x):
|
||||
# TFLite uses different padding
|
||||
if self.stride == 2:
|
||||
x = F.pad(x, (0, 1, 0, 1), 'constant', 0)
|
||||
# print(x.shape)
|
||||
|
||||
for module in self:
|
||||
if not isinstance(module, nn.MaxPool2d):
|
||||
x = module(x)
|
||||
return x
|
||||
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, inp, oup, stride, expand_ratio):
|
||||
super(InvertedResidual, self).__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
hidden_dim = int(round(inp * expand_ratio))
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
layers = []
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
|
||||
layers.extend([
|
||||
# dw
|
||||
ConvBNReLU(hidden_dim,
|
||||
hidden_dim,
|
||||
stride=stride,
|
||||
groups=hidden_dim),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class MobileNetV2(nn.Module):
|
||||
def __init__(self, pretrained=True):
|
||||
"""
|
||||
MobileNet V2 main class
|
||||
Args:
|
||||
num_classes (int): Number of classes
|
||||
width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
|
||||
inverted_residual_setting: Network structure
|
||||
round_nearest (int): Round the number of channels in each layer to be a multiple of this number
|
||||
Set to 1 to turn off rounding
|
||||
block: Module specifying inverted residual building block for mobilenet
|
||||
"""
|
||||
super(MobileNetV2, self).__init__()
|
||||
|
||||
block = InvertedResidual
|
||||
input_channel = 32
|
||||
last_channel = 1280
|
||||
width_mult = 1.0
|
||||
round_nearest = 8
|
||||
|
||||
inverted_residual_setting = [
|
||||
# t, c, n, s
|
||||
[1, 16, 1, 1],
|
||||
[6, 24, 2, 2],
|
||||
[6, 32, 3, 2],
|
||||
[6, 64, 4, 2],
|
||||
# [6, 96, 3, 1],
|
||||
# [6, 160, 3, 2],
|
||||
# [6, 320, 1, 1],
|
||||
]
|
||||
|
||||
# only check the first element, assuming user knows t,c,n,s are required
|
||||
if len(inverted_residual_setting) == 0 or len(
|
||||
inverted_residual_setting[0]) != 4:
|
||||
raise ValueError('inverted_residual_setting should be non-empty '
|
||||
'or a 4-element list, got {}'.format(
|
||||
inverted_residual_setting))
|
||||
|
||||
# building first layer
|
||||
input_channel = _make_divisible(input_channel * width_mult,
|
||||
round_nearest)
|
||||
self.last_channel = _make_divisible(
|
||||
last_channel * max(1.0, width_mult), round_nearest)
|
||||
features = [ConvBNReLU(4, input_channel, stride=2)]
|
||||
# building inverted residual blocks
|
||||
for t, c, n, s in inverted_residual_setting:
|
||||
output_channel = _make_divisible(c * width_mult, round_nearest)
|
||||
for i in range(n):
|
||||
stride = s if i == 0 else 1
|
||||
features.append(
|
||||
block(input_channel,
|
||||
output_channel,
|
||||
stride,
|
||||
expand_ratio=t))
|
||||
input_channel = output_channel
|
||||
self.features = nn.Sequential(*features)
|
||||
|
||||
self.fpn_selected = [3, 6, 10]
|
||||
# weight initialization
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out')
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.ones_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
# if pretrained:
|
||||
# self._load_pretrained_model()
|
||||
|
||||
def _forward_impl(self, x):
|
||||
# This exists since TorchScript doesn't support inheritance, so the superclass method
|
||||
# (this one) needs to have a name other than `forward` that can be accessed in a subclass
|
||||
fpn_features = []
|
||||
for i, f in enumerate(self.features):
|
||||
if i > self.fpn_selected[-1]:
|
||||
break
|
||||
x = f(x)
|
||||
if i in self.fpn_selected:
|
||||
fpn_features.append(x)
|
||||
|
||||
c2, c3, c4 = fpn_features
|
||||
return c2, c3, c4
|
||||
|
||||
def forward(self, x):
|
||||
return self._forward_impl(x)
|
||||
|
||||
def _load_pretrained_model(self):
|
||||
pretrain_dict = model_zoo.load_url(
|
||||
'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
|
||||
model_dict = {}
|
||||
state_dict = self.state_dict()
|
||||
for k, v in pretrain_dict.items():
|
||||
if k in state_dict:
|
||||
model_dict[k] = v
|
||||
state_dict.update(model_dict)
|
||||
self.load_state_dict(state_dict)
|
||||
|
||||
|
||||
class MobileV2_MLSD_Tiny(nn.Module):
|
||||
def __init__(self):
|
||||
super(MobileV2_MLSD_Tiny, self).__init__()
|
||||
|
||||
self.backbone = MobileNetV2(pretrained=True)
|
||||
|
||||
self.block12 = BlockTypeA(in_c1=32, in_c2=64, out_c1=64, out_c2=64)
|
||||
self.block13 = BlockTypeB(128, 64)
|
||||
|
||||
self.block14 = BlockTypeA(in_c1=24, in_c2=64, out_c1=32, out_c2=32)
|
||||
self.block15 = BlockTypeB(64, 64)
|
||||
|
||||
self.block16 = BlockTypeC(64, 16)
|
||||
|
||||
def forward(self, x):
|
||||
c2, c3, c4 = self.backbone(x)
|
||||
|
||||
x = self.block12(c3, c4)
|
||||
x = self.block13(x)
|
||||
x = self.block14(c2, x)
|
||||
x = self.block15(x)
|
||||
x = self.block16(x)
|
||||
x = x[:, 7:, :, :]
|
||||
# print(x.shape)
|
||||
x = F.interpolate(x,
|
||||
scale_factor=2.0,
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,638 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# modified by lihaoweicv
|
||||
# pytorch version
|
||||
#
|
||||
# M-LSD
|
||||
# Copyright 2021-present NAVER Corp.
|
||||
# Apache License v2.0
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def deccode_output_score_and_ptss(tpMap, topk_n=200, ksize=5):
|
||||
'''
|
||||
tpMap:
|
||||
center: tpMap[1, 0, :, :]
|
||||
displacement: tpMap[1, 1:5, :, :]
|
||||
'''
|
||||
b, c, h, w = tpMap.shape
|
||||
assert b == 1, 'only support bsize==1'
|
||||
displacement = tpMap[:, 1:5, :, :][0]
|
||||
center = tpMap[:, 0, :, :]
|
||||
heat = torch.sigmoid(center)
|
||||
hmax = F.max_pool2d(heat, (ksize, ksize),
|
||||
stride=1,
|
||||
padding=(ksize - 1) // 2)
|
||||
keep = (hmax == heat).float()
|
||||
heat = heat * keep
|
||||
heat = heat.reshape(-1, )
|
||||
|
||||
scores, indices = torch.topk(heat, topk_n, dim=-1, largest=True)
|
||||
yy = torch.floor_divide(indices, w).unsqueeze(-1)
|
||||
xx = torch.fmod(indices, w).unsqueeze(-1)
|
||||
ptss = torch.cat((yy, xx), dim=-1)
|
||||
|
||||
ptss = ptss.detach().cpu().numpy()
|
||||
scores = scores.detach().cpu().numpy()
|
||||
displacement = displacement.detach().cpu().numpy()
|
||||
displacement = displacement.transpose((1, 2, 0))
|
||||
return ptss, scores, displacement
|
||||
|
||||
|
||||
def pred_lines(image,
|
||||
model,
|
||||
input_shape=[512, 512],
|
||||
score_thr=0.10,
|
||||
dist_thr=20.0,
|
||||
device='cuda'):
|
||||
h, w, _ = image.shape
|
||||
h_ratio, w_ratio = [h / input_shape[0], w / input_shape[1]]
|
||||
|
||||
resized_image = np.concatenate([
|
||||
cv2.resize(image, (input_shape[1], input_shape[0]),
|
||||
interpolation=cv2.INTER_AREA),
|
||||
np.ones([input_shape[0], input_shape[1], 1])
|
||||
],
|
||||
axis=-1)
|
||||
|
||||
resized_image = resized_image.transpose((2, 0, 1))
|
||||
batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
|
||||
batch_image = (batch_image / 127.5) - 1.0
|
||||
|
||||
batch_image = torch.from_numpy(batch_image).float().to(device)
|
||||
outputs = model(batch_image)
|
||||
pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
|
||||
start = vmap[:, :, :2]
|
||||
end = vmap[:, :, 2:]
|
||||
dist_map = np.sqrt(np.sum((start - end)**2, axis=-1))
|
||||
|
||||
segments_list = []
|
||||
for center, score in zip(pts, pts_score):
|
||||
y, x = center
|
||||
distance = dist_map[y, x]
|
||||
if score > score_thr and distance > dist_thr:
|
||||
disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
|
||||
x_start = x + disp_x_start
|
||||
y_start = y + disp_y_start
|
||||
x_end = x + disp_x_end
|
||||
y_end = y + disp_y_end
|
||||
segments_list.append([x_start, y_start, x_end, y_end])
|
||||
|
||||
lines = 2 * np.array(segments_list) # 256 > 512
|
||||
lines[:, 0] = lines[:, 0] * w_ratio
|
||||
lines[:, 1] = lines[:, 1] * h_ratio
|
||||
lines[:, 2] = lines[:, 2] * w_ratio
|
||||
lines[:, 3] = lines[:, 3] * h_ratio
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def pred_squares(
|
||||
image,
|
||||
model,
|
||||
input_shape=[512, 512],
|
||||
device='cuda',
|
||||
params={
|
||||
'score': 0.06,
|
||||
'outside_ratio': 0.28,
|
||||
'inside_ratio': 0.45,
|
||||
'w_overlap': 0.0,
|
||||
'w_degree': 1.95,
|
||||
'w_length': 0.0,
|
||||
'w_area': 1.86,
|
||||
'w_center': 0.14
|
||||
}): # noqa
|
||||
# shape = [height, width]
|
||||
h, w, _ = image.shape
|
||||
original_shape = [h, w]
|
||||
|
||||
resized_image = np.concatenate([
|
||||
cv2.resize(image, (input_shape[0], input_shape[1]),
|
||||
interpolation=cv2.INTER_AREA),
|
||||
np.ones([input_shape[0], input_shape[1], 1])
|
||||
],
|
||||
axis=-1)
|
||||
resized_image = resized_image.transpose((2, 0, 1))
|
||||
batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
|
||||
batch_image = (batch_image / 127.5) - 1.0
|
||||
|
||||
batch_image = torch.from_numpy(batch_image).float().to(device)
|
||||
outputs = model(batch_image)
|
||||
|
||||
pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
|
||||
start = vmap[:, :, :2] # (x, y)
|
||||
end = vmap[:, :, 2:] # (x, y)
|
||||
dist_map = np.sqrt(np.sum((start - end)**2, axis=-1))
|
||||
|
||||
junc_list = []
|
||||
segments_list = []
|
||||
for junc, score in zip(pts, pts_score):
|
||||
y, x = junc
|
||||
distance = dist_map[y, x]
|
||||
if score > params['score'] and distance > 20.0:
|
||||
junc_list.append([x, y])
|
||||
disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
|
||||
d_arrow = 1.0
|
||||
x_start = x + d_arrow * disp_x_start
|
||||
y_start = y + d_arrow * disp_y_start
|
||||
x_end = x + d_arrow * disp_x_end
|
||||
y_end = y + d_arrow * disp_y_end
|
||||
segments_list.append([x_start, y_start, x_end, y_end])
|
||||
|
||||
segments = np.array(segments_list)
|
||||
|
||||
# post processing for squares
|
||||
# 1. get unique lines
|
||||
point = np.array([[0, 0]])
|
||||
point = point[0]
|
||||
start = segments[:, :2]
|
||||
end = segments[:, 2:]
|
||||
diff = start - end
|
||||
a = diff[:, 1]
|
||||
b = -diff[:, 0]
|
||||
c = a * start[:, 0] + b * start[:, 1]
|
||||
|
||||
d = np.abs(a * point[0] + b * point[1] - c) / np.sqrt(a**2 + b**2 + 1e-10)
|
||||
theta = np.arctan2(diff[:, 0], diff[:, 1]) * 180 / np.pi
|
||||
theta[theta < 0.0] += 180
|
||||
hough = np.concatenate([d[:, None], theta[:, None]], axis=-1)
|
||||
|
||||
d_quant = 1
|
||||
theta_quant = 2
|
||||
hough[:, 0] //= d_quant
|
||||
hough[:, 1] //= theta_quant
|
||||
_, indices, counts = np.unique(hough,
|
||||
axis=0,
|
||||
return_index=True,
|
||||
return_counts=True)
|
||||
|
||||
acc_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1],
|
||||
dtype='float32')
|
||||
idx_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1],
|
||||
dtype='int32') - 1
|
||||
yx_indices = hough[indices, :].astype('int32')
|
||||
acc_map[yx_indices[:, 0], yx_indices[:, 1]] = counts
|
||||
idx_map[yx_indices[:, 0], yx_indices[:, 1]] = indices
|
||||
|
||||
acc_map_np = acc_map
|
||||
# acc_map = acc_map[None, :, :, None]
|
||||
#
|
||||
# ### fast suppression using tensorflow op
|
||||
# acc_map = tf.constant(acc_map, dtype=tf.float32)
|
||||
# max_acc_map = tf.keras.layers.MaxPool2D(pool_size=(5, 5), strides=1, padding='same')(acc_map)
|
||||
# acc_map = acc_map * tf.cast(tf.math.equal(acc_map, max_acc_map), tf.float32)
|
||||
# flatten_acc_map = tf.reshape(acc_map, [1, -1])
|
||||
# topk_values, topk_indices = tf.math.top_k(flatten_acc_map, k=len(pts))
|
||||
# _, h, w, _ = acc_map.shape
|
||||
# y = tf.expand_dims(topk_indices // w, axis=-1)
|
||||
# x = tf.expand_dims(topk_indices % w, axis=-1)
|
||||
# yx = tf.concat([y, x], axis=-1)
|
||||
|
||||
# fast suppression using pytorch op
|
||||
acc_map = torch.from_numpy(acc_map_np).unsqueeze(0).unsqueeze(0)
|
||||
_, _, h, w = acc_map.shape
|
||||
max_acc_map = F.max_pool2d(acc_map, kernel_size=5, stride=1, padding=2)
|
||||
acc_map = acc_map * ((acc_map == max_acc_map).float())
|
||||
flatten_acc_map = acc_map.reshape([
|
||||
-1,
|
||||
])
|
||||
|
||||
scores, indices = torch.topk(flatten_acc_map,
|
||||
len(pts),
|
||||
dim=-1,
|
||||
largest=True)
|
||||
yy = torch.div(indices, w, rounding_mode='floor').unsqueeze(-1)
|
||||
xx = torch.fmod(indices, w).unsqueeze(-1)
|
||||
yx = torch.cat((yy, xx), dim=-1)
|
||||
|
||||
yx = yx.detach().cpu().numpy()
|
||||
|
||||
topk_values = scores.detach().cpu().numpy()
|
||||
indices = idx_map[yx[:, 0], yx[:, 1]]
|
||||
basis = 5 // 2
|
||||
|
||||
merged_segments = []
|
||||
for yx_pt, max_indice, value in zip(yx, indices, topk_values):
|
||||
y, x = yx_pt
|
||||
if max_indice == -1 or value == 0:
|
||||
continue
|
||||
segment_list = []
|
||||
for y_offset in range(-basis, basis + 1):
|
||||
for x_offset in range(-basis, basis + 1):
|
||||
indice = idx_map[y + y_offset, x + x_offset]
|
||||
cnt = int(acc_map_np[y + y_offset, x + x_offset])
|
||||
if indice != -1:
|
||||
segment_list.append(segments[indice])
|
||||
if cnt > 1:
|
||||
check_cnt = 1
|
||||
current_hough = hough[indice]
|
||||
for new_indice, new_hough in enumerate(hough):
|
||||
if (current_hough
|
||||
== new_hough).all() and indice != new_indice:
|
||||
segment_list.append(segments[new_indice])
|
||||
check_cnt += 1
|
||||
if check_cnt == cnt:
|
||||
break
|
||||
group_segments = np.array(segment_list).reshape([-1, 2])
|
||||
sorted_group_segments = np.sort(group_segments, axis=0)
|
||||
x_min, y_min = sorted_group_segments[0, :]
|
||||
x_max, y_max = sorted_group_segments[-1, :]
|
||||
|
||||
deg = theta[max_indice]
|
||||
if deg >= 90:
|
||||
merged_segments.append([x_min, y_max, x_max, y_min])
|
||||
else:
|
||||
merged_segments.append([x_min, y_min, x_max, y_max])
|
||||
|
||||
# 2. get intersections
|
||||
new_segments = np.array(merged_segments) # (x1, y1, x2, y2)
|
||||
start = new_segments[:, :2] # (x1, y1)
|
||||
end = new_segments[:, 2:] # (x2, y2)
|
||||
new_centers = (start + end) / 2.0
|
||||
diff = start - end
|
||||
dist_segments = np.sqrt(np.sum(diff**2, axis=-1))
|
||||
|
||||
# ax + by = c
|
||||
a = diff[:, 1]
|
||||
b = -diff[:, 0]
|
||||
c = a * start[:, 0] + b * start[:, 1]
|
||||
pre_det = a[:, None] * b[None, :]
|
||||
det = pre_det - np.transpose(pre_det)
|
||||
|
||||
pre_inter_y = a[:, None] * c[None, :]
|
||||
inter_y = (pre_inter_y - np.transpose(pre_inter_y)) / (det + 1e-10)
|
||||
pre_inter_x = c[:, None] * b[None, :]
|
||||
inter_x = (pre_inter_x - np.transpose(pre_inter_x)) / (det + 1e-10)
|
||||
inter_pts = np.concatenate([inter_x[:, :, None], inter_y[:, :, None]],
|
||||
axis=-1).astype('int32')
|
||||
|
||||
# 3. get corner information
|
||||
# 3.1 get distance
|
||||
'''
|
||||
dist_segments:
|
||||
| dist(0), dist(1), dist(2), ...|
|
||||
dist_inter_to_segment1:
|
||||
| dist(inter,0), dist(inter,0), dist(inter,0), ... |
|
||||
| dist(inter,1), dist(inter,1), dist(inter,1), ... |
|
||||
...
|
||||
dist_inter_to_semgnet2:
|
||||
| dist(inter,0), dist(inter,1), dist(inter,2), ... |
|
||||
| dist(inter,0), dist(inter,1), dist(inter,2), ... |
|
||||
...
|
||||
'''
|
||||
|
||||
dist_inter_to_segment1_start = np.sqrt(
|
||||
np.sum(((inter_pts - start[:, None, :])**2), axis=-1,
|
||||
keepdims=True)) # [n_batch, n_batch, 1]
|
||||
dist_inter_to_segment1_end = np.sqrt(
|
||||
np.sum(((inter_pts - end[:, None, :])**2), axis=-1,
|
||||
keepdims=True)) # [n_batch, n_batch, 1]
|
||||
dist_inter_to_segment2_start = np.sqrt(
|
||||
np.sum(((inter_pts - start[None, :, :])**2), axis=-1,
|
||||
keepdims=True)) # [n_batch, n_batch, 1]
|
||||
dist_inter_to_segment2_end = np.sqrt(
|
||||
np.sum(((inter_pts - end[None, :, :])**2), axis=-1,
|
||||
keepdims=True)) # [n_batch, n_batch, 1]
|
||||
|
||||
# sort ascending
|
||||
dist_inter_to_segment1 = np.sort(np.concatenate(
|
||||
[dist_inter_to_segment1_start, dist_inter_to_segment1_end], axis=-1),
|
||||
axis=-1) # [n_batch, n_batch, 2]
|
||||
dist_inter_to_segment2 = np.sort(np.concatenate(
|
||||
[dist_inter_to_segment2_start, dist_inter_to_segment2_end], axis=-1),
|
||||
axis=-1) # [n_batch, n_batch, 2]
|
||||
|
||||
# 3.2 get degree
|
||||
inter_to_start = new_centers[:, None, :] - inter_pts
|
||||
deg_inter_to_start = np.arctan2(inter_to_start[:, :, 1],
|
||||
inter_to_start[:, :, 0]) * 180 / np.pi
|
||||
deg_inter_to_start[deg_inter_to_start < 0.0] += 360
|
||||
inter_to_end = new_centers[None, :, :] - inter_pts
|
||||
deg_inter_to_end = np.arctan2(inter_to_end[:, :, 1],
|
||||
inter_to_end[:, :, 0]) * 180 / np.pi
|
||||
deg_inter_to_end[deg_inter_to_end < 0.0] += 360
|
||||
'''
|
||||
B -- G
|
||||
| |
|
||||
C -- R
|
||||
B : blue / G: green / C: cyan / R: red
|
||||
|
||||
0 -- 1
|
||||
| |
|
||||
3 -- 2
|
||||
'''
|
||||
# rename variables
|
||||
deg1_map, deg2_map = deg_inter_to_start, deg_inter_to_end
|
||||
# sort deg ascending
|
||||
deg_sort = np.sort(np.concatenate(
|
||||
[deg1_map[:, :, None], deg2_map[:, :, None]], axis=-1),
|
||||
axis=-1)
|
||||
|
||||
deg_diff_map = np.abs(deg1_map - deg2_map)
|
||||
# we only consider the smallest degree of intersect
|
||||
deg_diff_map[deg_diff_map > 180] = 360 - deg_diff_map[deg_diff_map > 180]
|
||||
|
||||
# define available degree range
|
||||
deg_range = [60, 120]
|
||||
|
||||
corner_dict = {corner_info: [] for corner_info in range(4)}
|
||||
inter_points = []
|
||||
for i in range(inter_pts.shape[0]):
|
||||
for j in range(i + 1, inter_pts.shape[1]):
|
||||
# i, j > line index, always i < j
|
||||
x, y = inter_pts[i, j, :]
|
||||
deg1, deg2 = deg_sort[i, j, :]
|
||||
deg_diff = deg_diff_map[i, j]
|
||||
|
||||
check_degree = deg_diff > deg_range[0] and deg_diff < deg_range[1]
|
||||
|
||||
outside_ratio = params['outside_ratio'] # over ratio >>> drop it!
|
||||
inside_ratio = params['inside_ratio'] # over ratio >>> drop it!
|
||||
check_distance = ((dist_inter_to_segment1[i, j, 1] >= dist_segments[i] and
|
||||
dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * outside_ratio) or
|
||||
(dist_inter_to_segment1[i, j, 1] <= dist_segments[i] and
|
||||
dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * inside_ratio)) and \
|
||||
((dist_inter_to_segment2[i, j, 1] >= dist_segments[j] and
|
||||
dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * outside_ratio) or
|
||||
(dist_inter_to_segment2[i, j, 1] <= dist_segments[j] and
|
||||
dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * inside_ratio))
|
||||
|
||||
if check_degree and check_distance:
|
||||
corner_info = None # noqa
|
||||
|
||||
if (deg1 >= 0 and deg1 <= 45 and deg2 >= 45 and deg2 <= 120) or \
|
||||
(deg2 >= 315 and deg1 >= 45 and deg1 <= 120):
|
||||
corner_info, color_info = 0, 'blue'
|
||||
elif (deg1 >= 45 and deg1 <= 125 and deg2 >= 125
|
||||
and deg2 <= 225):
|
||||
corner_info, color_info = 1, 'green'
|
||||
elif (deg1 >= 125 and deg1 <= 225 and deg2 >= 225
|
||||
and deg2 <= 315):
|
||||
corner_info, color_info = 2, 'black'
|
||||
elif (deg1 >= 0 and deg1 <= 45 and deg2 >= 225 and deg2 <= 315) or \
|
||||
(deg2 >= 315 and deg1 >= 225 and deg1 <= 315):
|
||||
corner_info, color_info = 3, 'cyan'
|
||||
else:
|
||||
corner_info, color_info = 4, 'red' # we don't use it # noqa
|
||||
continue
|
||||
|
||||
corner_dict[corner_info].append([x, y, i, j])
|
||||
inter_points.append([x, y])
|
||||
|
||||
square_list = []
|
||||
connect_list = []
|
||||
segments_list = []
|
||||
for corner0 in corner_dict[0]:
|
||||
for corner1 in corner_dict[1]:
|
||||
connect01 = False
|
||||
for corner0_line in corner0[2:]:
|
||||
if corner0_line in corner1[2:]:
|
||||
connect01 = True
|
||||
break
|
||||
if connect01:
|
||||
for corner2 in corner_dict[2]:
|
||||
connect12 = False
|
||||
for corner1_line in corner1[2:]:
|
||||
if corner1_line in corner2[2:]:
|
||||
connect12 = True
|
||||
break
|
||||
if connect12:
|
||||
for corner3 in corner_dict[3]:
|
||||
connect23 = False
|
||||
for corner2_line in corner2[2:]:
|
||||
if corner2_line in corner3[2:]:
|
||||
connect23 = True
|
||||
break
|
||||
if connect23:
|
||||
for corner3_line in corner3[2:]:
|
||||
if corner3_line in corner0[2:]:
|
||||
# SQUARE!!!
|
||||
'''
|
||||
0 -- 1
|
||||
| |
|
||||
3 -- 2
|
||||
square_list:
|
||||
order: 0 > 1 > 2 > 3
|
||||
| x0, y0, x1, y1, x2, y2, x3, y3 |
|
||||
| x0, y0, x1, y1, x2, y2, x3, y3 |
|
||||
...
|
||||
connect_list:
|
||||
order: 01 > 12 > 23 > 30
|
||||
| line_idx01, line_idx12, line_idx23, line_idx30 |
|
||||
| line_idx01, line_idx12, line_idx23, line_idx30 |
|
||||
...
|
||||
segments_list:
|
||||
order: 0 > 1 > 2 > 3
|
||||
| line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i,
|
||||
line_idx2_j, line_idx3_i, line_idx3_j |
|
||||
| line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i,
|
||||
line_idx2_j, line_idx3_i, line_idx3_j |
|
||||
...
|
||||
'''
|
||||
square_list.append(corner0[:2] +
|
||||
corner1[:2] +
|
||||
corner2[:2] +
|
||||
corner3[:2])
|
||||
connect_list.append([
|
||||
corner0_line, corner1_line,
|
||||
corner2_line, corner3_line
|
||||
])
|
||||
segments_list.append(corner0[2:] +
|
||||
corner1[2:] +
|
||||
corner2[2:] +
|
||||
corner3[2:])
|
||||
|
||||
def check_outside_inside(segments_info, connect_idx):
|
||||
# return 'outside or inside', min distance, cover_param, peri_param
|
||||
if connect_idx == segments_info[0]:
|
||||
check_dist_mat = dist_inter_to_segment1
|
||||
else:
|
||||
check_dist_mat = dist_inter_to_segment2
|
||||
|
||||
i, j = segments_info
|
||||
min_dist, max_dist = check_dist_mat[i, j, :]
|
||||
connect_dist = dist_segments[connect_idx]
|
||||
if max_dist > connect_dist:
|
||||
return 'outside', min_dist, 0, 1
|
||||
else:
|
||||
return 'inside', min_dist, -1, -1
|
||||
|
||||
top_square = None # noqa
|
||||
|
||||
try:
|
||||
map_size = input_shape[0] / 2
|
||||
squares = np.array(square_list).reshape([-1, 4, 2])
|
||||
score_array = []
|
||||
connect_array = np.array(connect_list)
|
||||
segments_array = np.array(segments_list).reshape([-1, 4, 2])
|
||||
|
||||
# get degree of corners:
|
||||
squares_rollup = np.roll(squares, 1, axis=1)
|
||||
squares_rolldown = np.roll(squares, -1, axis=1)
|
||||
vec1 = squares_rollup - squares
|
||||
normalized_vec1 = vec1 / (
|
||||
np.linalg.norm(vec1, axis=-1, keepdims=True) + 1e-10)
|
||||
vec2 = squares_rolldown - squares
|
||||
normalized_vec2 = vec2 / (
|
||||
np.linalg.norm(vec2, axis=-1, keepdims=True) + 1e-10)
|
||||
inner_products = np.sum(normalized_vec1 * normalized_vec2,
|
||||
axis=-1) # [n_squares, 4]
|
||||
squares_degree = np.arccos(
|
||||
inner_products) * 180 / np.pi # [n_squares, 4]
|
||||
|
||||
# get square score
|
||||
overlap_scores = []
|
||||
degree_scores = []
|
||||
length_scores = []
|
||||
|
||||
for connects, segments, square, degree in zip(connect_array,
|
||||
segments_array, squares,
|
||||
squares_degree):
|
||||
'''
|
||||
0 -- 1
|
||||
| |
|
||||
3 -- 2
|
||||
|
||||
# segments: [4, 2]
|
||||
# connects: [4]
|
||||
'''
|
||||
|
||||
# OVERLAP SCORES
|
||||
cover = 0
|
||||
perimeter = 0
|
||||
# check 0 > 1 > 2 > 3
|
||||
square_length = []
|
||||
|
||||
for start_idx in range(4):
|
||||
end_idx = (start_idx + 1) % 4
|
||||
|
||||
connect_idx = connects[start_idx] # segment idx of segment01
|
||||
start_segments = segments[start_idx]
|
||||
end_segments = segments[end_idx]
|
||||
|
||||
start_point = square[start_idx] # noqa
|
||||
end_point = square[end_idx] # noqa
|
||||
|
||||
# check whether outside or inside
|
||||
start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(
|
||||
start_segments, connect_idx)
|
||||
end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(
|
||||
end_segments, connect_idx)
|
||||
|
||||
cover += dist_segments[
|
||||
connect_idx] + start_cover_param * start_min + end_cover_param * end_min
|
||||
perimeter += dist_segments[
|
||||
connect_idx] + start_peri_param * start_min + end_peri_param * end_min
|
||||
|
||||
square_length.append(dist_segments[connect_idx] +
|
||||
start_peri_param * start_min +
|
||||
end_peri_param * end_min)
|
||||
|
||||
overlap_scores.append(cover / perimeter)
|
||||
# DEGREE SCORES
|
||||
'''
|
||||
deg0 vs deg2
|
||||
deg1 vs deg3
|
||||
'''
|
||||
deg0, deg1, deg2, deg3 = degree
|
||||
deg_ratio1 = deg0 / deg2
|
||||
if deg_ratio1 > 1.0:
|
||||
deg_ratio1 = 1 / deg_ratio1
|
||||
deg_ratio2 = deg1 / deg3
|
||||
if deg_ratio2 > 1.0:
|
||||
deg_ratio2 = 1 / deg_ratio2
|
||||
degree_scores.append((deg_ratio1 + deg_ratio2) / 2)
|
||||
# LENGTH SCORES
|
||||
'''
|
||||
len0 vs len2
|
||||
len1 vs len3
|
||||
'''
|
||||
len0, len1, len2, len3 = square_length
|
||||
len_ratio1 = len0 / len2 if len2 > len0 else len2 / len0
|
||||
len_ratio2 = len1 / len3 if len3 > len1 else len3 / len1
|
||||
length_scores.append((len_ratio1 + len_ratio2) / 2)
|
||||
|
||||
######################################
|
||||
|
||||
overlap_scores = np.array(overlap_scores)
|
||||
overlap_scores /= np.max(overlap_scores)
|
||||
|
||||
degree_scores = np.array(degree_scores)
|
||||
# degree_scores /= np.max(degree_scores)
|
||||
|
||||
length_scores = np.array(length_scores)
|
||||
|
||||
# AREA SCORES
|
||||
area_scores = np.reshape(squares, [-1, 4, 2])
|
||||
area_x = area_scores[:, :, 0]
|
||||
area_y = area_scores[:, :, 1]
|
||||
correction = area_x[:, -1] * area_y[:, 0] - area_y[:, -1] * area_x[:,
|
||||
0]
|
||||
area_scores = np.sum(area_x[:, :-1] * area_y[:, 1:], axis=-1) - np.sum(
|
||||
area_y[:, :-1] * area_x[:, 1:], axis=-1)
|
||||
area_scores = 0.5 * np.abs(area_scores + correction)
|
||||
area_scores /= (map_size * map_size) # np.max(area_scores)
|
||||
|
||||
# CENTER SCORES
|
||||
centers = np.array([[256 // 2, 256 // 2]], dtype='float32') # [1, 2]
|
||||
# squares: [n, 4, 2]
|
||||
square_centers = np.mean(squares, axis=1) # [n, 2]
|
||||
center2center = np.sqrt(np.sum((centers - square_centers)**2))
|
||||
center_scores = center2center / (map_size / np.sqrt(2.0))
|
||||
'''
|
||||
score_w = [overlap, degree, area, center, length]
|
||||
'''
|
||||
score_w = [0.0, 1.0, 10.0, 0.5, 1.0] # noqa
|
||||
score_array = (params['w_overlap'] * overlap_scores +
|
||||
params['w_degree'] * degree_scores +
|
||||
params['w_area'] * area_scores -
|
||||
params['w_center'] * center_scores +
|
||||
params['w_length'] * length_scores)
|
||||
|
||||
best_square = [] # noqa
|
||||
|
||||
sorted_idx = np.argsort(score_array)[::-1]
|
||||
score_array = score_array[sorted_idx]
|
||||
squares = squares[sorted_idx]
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
'''return list
|
||||
merged_lines, squares, scores
|
||||
'''
|
||||
|
||||
try:
|
||||
new_segments[:, 0] = new_segments[:, 0] * 2 / input_shape[
|
||||
1] * original_shape[1]
|
||||
new_segments[:, 1] = new_segments[:, 1] * 2 / input_shape[
|
||||
0] * original_shape[0]
|
||||
new_segments[:, 2] = new_segments[:, 2] * 2 / input_shape[
|
||||
1] * original_shape[1]
|
||||
new_segments[:, 3] = new_segments[:, 3] * 2 / input_shape[
|
||||
0] * original_shape[0]
|
||||
except Exception:
|
||||
new_segments = []
|
||||
|
||||
try:
|
||||
squares[:, :,
|
||||
0] = squares[:, :, 0] * 2 / input_shape[1] * original_shape[1]
|
||||
squares[:, :,
|
||||
1] = squares[:, :, 1] * 2 / input_shape[0] * original_shape[0]
|
||||
except Exception:
|
||||
squares = []
|
||||
score_array = []
|
||||
|
||||
try:
|
||||
inter_points = np.array(inter_points)
|
||||
inter_points[:, 0] = inter_points[:, 0] * 2 / input_shape[
|
||||
1] * original_shape[1]
|
||||
inter_points[:, 1] = inter_points[:, 1] * 2 / input_shape[
|
||||
0] * original_shape[0]
|
||||
except Exception:
|
||||
inter_points = []
|
||||
|
||||
return new_segments, squares, score_array, inter_points
|
||||
@@ -0,0 +1,73 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# MLSD Line Detection
|
||||
# From https://github.com/navervision/mlsd
|
||||
# Apache-2.0 license
|
||||
|
||||
import warnings
|
||||
from abc import ABCMeta
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
||||
from scepter.modules.annotator.mlsd.mbv2_mlsd_large import MobileV2_MLSD_Large
|
||||
from scepter.modules.annotator.mlsd.utils import pred_lines
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.annotator.utils import resize_image, resize_image_ori
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class MLSDdetector(BaseAnnotator, metaclass=ABCMeta):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
model = MobileV2_MLSD_Large()
|
||||
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)
|
||||
self.model = model.eval()
|
||||
self.thr_v = cfg.get('THR_V', 0.1)
|
||||
self.thr_d = cfg.get('THR_D', 0.1)
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
@torch.autocast('cuda', enabled=False)
|
||||
def forward(self, image):
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image = image.detach().cpu().numpy()
|
||||
elif isinstance(image, np.ndarray):
|
||||
image = image.copy()
|
||||
else:
|
||||
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
||||
h, w, c = image.shape
|
||||
image, k = resize_image(image, 1024 if min(h, w) > 1024 else min(h, w))
|
||||
img_output = np.zeros_like(image)
|
||||
try:
|
||||
lines = pred_lines(image,
|
||||
self.model, [image.shape[0], image.shape[1]],
|
||||
self.thr_v,
|
||||
self.thr_d,
|
||||
device=we.device_id)
|
||||
for line in lines:
|
||||
x_start, y_start, x_end, y_end = [int(val) for val in line]
|
||||
cv2.line(img_output, (x_start, y_start), (x_end, y_end),
|
||||
[255, 255, 255], 1)
|
||||
except Exception as e:
|
||||
warnings.warn(f'{e}')
|
||||
return None
|
||||
img_output = resize_image_ori(h, w, img_output, k)
|
||||
return img_output[:, :, 0]
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
MLSDdetector.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,811 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Openpose
|
||||
# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
|
||||
# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
|
||||
# The implementation is modified from 3rd Edited Version by ControlNet
|
||||
import math
|
||||
import os
|
||||
from abc import ABCMeta
|
||||
from collections import OrderedDict
|
||||
|
||||
import cv2
|
||||
import matplotlib
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from PIL import Image
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from skimage.measure import label
|
||||
|
||||
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.file_system import FS
|
||||
|
||||
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
|
||||
|
||||
|
||||
def padRightDownCorner(img, stride, padValue):
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
|
||||
pad = 4 * [None]
|
||||
pad[0] = 0 # up
|
||||
pad[1] = 0 # left
|
||||
pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
|
||||
pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
|
||||
|
||||
img_padded = img
|
||||
pad_up = np.tile(img_padded[0:1, :, :] * 0 + padValue, (pad[0], 1, 1))
|
||||
img_padded = np.concatenate((pad_up, img_padded), axis=0)
|
||||
pad_left = np.tile(img_padded[:, 0:1, :] * 0 + padValue, (1, pad[1], 1))
|
||||
img_padded = np.concatenate((pad_left, img_padded), axis=1)
|
||||
pad_down = np.tile(img_padded[-2:-1, :, :] * 0 + padValue, (pad[2], 1, 1))
|
||||
img_padded = np.concatenate((img_padded, pad_down), axis=0)
|
||||
pad_right = np.tile(img_padded[:, -2:-1, :] * 0 + padValue, (1, pad[3], 1))
|
||||
img_padded = np.concatenate((img_padded, pad_right), axis=1)
|
||||
|
||||
return img_padded, pad
|
||||
|
||||
|
||||
# transfer caffe model to pytorch which will match the layer name
|
||||
def transfer(model, model_weights):
|
||||
transfered_model_weights = {}
|
||||
for weights_name in model.state_dict().keys():
|
||||
transfered_model_weights[weights_name] = model_weights['.'.join(
|
||||
weights_name.split('.')[1:])]
|
||||
return transfered_model_weights
|
||||
|
||||
|
||||
# draw the body keypoint and lims
|
||||
def draw_bodypose(canvas, candidate, subset):
|
||||
stickwidth = 4
|
||||
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10],
|
||||
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15],
|
||||
[15, 17], [1, 16], [16, 18], [3, 17], [6, 18]]
|
||||
|
||||
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0],
|
||||
[170, 255, 0], [85, 255, 0], [0, 255, 0], [0, 255, 85],
|
||||
[0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255],
|
||||
[0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255],
|
||||
[255, 0, 170], [255, 0, 85]]
|
||||
for i in range(18):
|
||||
for n in range(len(subset)):
|
||||
index = int(subset[n][i])
|
||||
if index == -1:
|
||||
continue
|
||||
x, y = candidate[index][0:2]
|
||||
cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
|
||||
for i in range(17):
|
||||
for n in range(len(subset)):
|
||||
index = subset[n][np.array(limbSeq[i]) - 1]
|
||||
if -1 in index:
|
||||
continue
|
||||
cur_canvas = canvas.copy()
|
||||
Y = candidate[index.astype(int), 0]
|
||||
X = candidate[index.astype(int), 1]
|
||||
mX = np.mean(X)
|
||||
mY = np.mean(Y)
|
||||
length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
|
||||
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
|
||||
polygon = cv2.ellipse2Poly(
|
||||
(int(mY), int(mX)), (int(length / 2), stickwidth), int(angle),
|
||||
0, 360, 1)
|
||||
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
|
||||
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
|
||||
# plt.imsave("preview.jpg", canvas[:, :, [2, 1, 0]])
|
||||
# plt.imshow(canvas[:, :, [2, 1, 0]])
|
||||
return canvas
|
||||
|
||||
|
||||
# image drawed by opencv is not good.
|
||||
def draw_handpose(canvas, all_hand_peaks, show_number=False):
|
||||
edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8],
|
||||
[0, 9], [9, 10], [10, 11], [11, 12], [0, 13], [13, 14], [14, 15],
|
||||
[15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
|
||||
|
||||
for peaks in all_hand_peaks:
|
||||
for ie, e in enumerate(edges):
|
||||
if np.sum(np.all(peaks[e], axis=1) == 0) == 0:
|
||||
x1, y1 = peaks[e[0]]
|
||||
x2, y2 = peaks[e[1]]
|
||||
cv2.line(canvas, (x1, y1), (x2, y2),
|
||||
matplotlib.colors.hsv_to_rgb(
|
||||
[ie / float(len(edges)), 1.0, 1.0]) * 255,
|
||||
thickness=2)
|
||||
|
||||
for i, keyponit in enumerate(peaks):
|
||||
x, y = keyponit
|
||||
cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
|
||||
if show_number:
|
||||
cv2.putText(canvas,
|
||||
str(i), (x, y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.3, (0, 0, 0),
|
||||
lineType=cv2.LINE_AA)
|
||||
return canvas
|
||||
|
||||
|
||||
# detect hand according to body pose keypoints
|
||||
# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/
|
||||
# master/src/openpose/hand/handDetector.cpp
|
||||
def handDetect(candidate, subset, oriImg):
|
||||
# right hand: wrist 4, elbow 3, shoulder 2
|
||||
# left hand: wrist 7, elbow 6, shoulder 5
|
||||
ratioWristElbow = 0.33
|
||||
detect_result = []
|
||||
image_height, image_width = oriImg.shape[0:2]
|
||||
for person in subset.astype(int):
|
||||
# if any of three not detected
|
||||
has_left = np.sum(person[[5, 6, 7]] == -1) == 0
|
||||
has_right = np.sum(person[[2, 3, 4]] == -1) == 0
|
||||
if not (has_left or has_right):
|
||||
continue
|
||||
hands = []
|
||||
# left hand
|
||||
if has_left:
|
||||
left_shoulder_index, left_elbow_index, left_wrist_index = person[[
|
||||
5, 6, 7
|
||||
]]
|
||||
x1, y1 = candidate[left_shoulder_index][:2]
|
||||
x2, y2 = candidate[left_elbow_index][:2]
|
||||
x3, y3 = candidate[left_wrist_index][:2]
|
||||
hands.append([x1, y1, x2, y2, x3, y3, True])
|
||||
# right hand
|
||||
if has_right:
|
||||
right_shoulder_index, right_elbow_index, right_wrist_index = person[
|
||||
[2, 3, 4]]
|
||||
x1, y1 = candidate[right_shoulder_index][:2]
|
||||
x2, y2 = candidate[right_elbow_index][:2]
|
||||
x3, y3 = candidate[right_wrist_index][:2]
|
||||
hands.append([x1, y1, x2, y2, x3, y3, False])
|
||||
|
||||
for x1, y1, x2, y2, x3, y3, is_left in hands:
|
||||
# pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
|
||||
# handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
|
||||
# handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
|
||||
# const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
|
||||
# const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
|
||||
# handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
|
||||
x = x3 + ratioWristElbow * (x3 - x2)
|
||||
y = y3 + ratioWristElbow * (y3 - y2)
|
||||
distanceWristElbow = math.sqrt((x3 - x2)**2 + (y3 - y2)**2)
|
||||
distanceElbowShoulder = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
|
||||
width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
|
||||
# x-y refers to the center --> offset to topLeft point
|
||||
# handRectangle.x -= handRectangle.width / 2.f;
|
||||
# handRectangle.y -= handRectangle.height / 2.f;
|
||||
x -= width / 2
|
||||
y -= width / 2 # width = height
|
||||
# overflow the image
|
||||
if x < 0:
|
||||
x = 0
|
||||
if y < 0:
|
||||
y = 0
|
||||
width1 = width
|
||||
width2 = width
|
||||
if x + width > image_width:
|
||||
width1 = image_width - x
|
||||
if y + width > image_height:
|
||||
width2 = image_height - y
|
||||
width = min(width1, width2)
|
||||
# the max hand box value is 20 pixels
|
||||
if width >= 20:
|
||||
detect_result.append([int(x), int(y), int(width), is_left])
|
||||
'''
|
||||
return value: [[x, y, w, True if left hand else False]].
|
||||
width=height since the network require squared input.
|
||||
x, y is the coordinate of top left
|
||||
'''
|
||||
return detect_result
|
||||
|
||||
|
||||
# get max index of 2d array
|
||||
def npmax(array):
|
||||
arrayindex = array.argmax(1)
|
||||
arrayvalue = array.max(1)
|
||||
i = arrayvalue.argmax()
|
||||
j = arrayindex[i]
|
||||
return i, j
|
||||
|
||||
|
||||
def make_layers(block, no_relu_layers):
|
||||
layers = []
|
||||
for layer_name, v in block.items():
|
||||
if 'pool' in layer_name:
|
||||
layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2])
|
||||
layers.append((layer_name, layer))
|
||||
else:
|
||||
conv2d = nn.Conv2d(in_channels=v[0],
|
||||
out_channels=v[1],
|
||||
kernel_size=v[2],
|
||||
stride=v[3],
|
||||
padding=v[4])
|
||||
layers.append((layer_name, conv2d))
|
||||
if layer_name not in no_relu_layers:
|
||||
layers.append(('relu_' + layer_name, nn.ReLU(inplace=True)))
|
||||
|
||||
return nn.Sequential(OrderedDict(layers))
|
||||
|
||||
|
||||
class bodypose_model(nn.Module):
|
||||
def __init__(self):
|
||||
super(bodypose_model, self).__init__()
|
||||
|
||||
# these layers have no relu layer
|
||||
no_relu_layers = [
|
||||
'conv5_5_CPM_L1', 'conv5_5_CPM_L2', 'Mconv7_stage2_L1',
|
||||
'Mconv7_stage2_L2', 'Mconv7_stage3_L1', 'Mconv7_stage3_L2',
|
||||
'Mconv7_stage4_L1', 'Mconv7_stage4_L2', 'Mconv7_stage5_L1',
|
||||
'Mconv7_stage5_L2', 'Mconv7_stage6_L1', 'Mconv7_stage6_L1'
|
||||
]
|
||||
blocks = {}
|
||||
block0 = OrderedDict([('conv1_1', [3, 64, 3, 1, 1]),
|
||||
('conv1_2', [64, 64, 3, 1, 1]),
|
||||
('pool1_stage1', [2, 2, 0]),
|
||||
('conv2_1', [64, 128, 3, 1, 1]),
|
||||
('conv2_2', [128, 128, 3, 1, 1]),
|
||||
('pool2_stage1', [2, 2, 0]),
|
||||
('conv3_1', [128, 256, 3, 1, 1]),
|
||||
('conv3_2', [256, 256, 3, 1, 1]),
|
||||
('conv3_3', [256, 256, 3, 1, 1]),
|
||||
('conv3_4', [256, 256, 3, 1, 1]),
|
||||
('pool3_stage1', [2, 2, 0]),
|
||||
('conv4_1', [256, 512, 3, 1, 1]),
|
||||
('conv4_2', [512, 512, 3, 1, 1]),
|
||||
('conv4_3_CPM', [512, 256, 3, 1, 1]),
|
||||
('conv4_4_CPM', [256, 128, 3, 1, 1])])
|
||||
|
||||
# Stage 1
|
||||
block1_1 = OrderedDict([('conv5_1_CPM_L1', [128, 128, 3, 1, 1]),
|
||||
('conv5_2_CPM_L1', [128, 128, 3, 1, 1]),
|
||||
('conv5_3_CPM_L1', [128, 128, 3, 1, 1]),
|
||||
('conv5_4_CPM_L1', [128, 512, 1, 1, 0]),
|
||||
('conv5_5_CPM_L1', [512, 38, 1, 1, 0])])
|
||||
|
||||
block1_2 = OrderedDict([('conv5_1_CPM_L2', [128, 128, 3, 1, 1]),
|
||||
('conv5_2_CPM_L2', [128, 128, 3, 1, 1]),
|
||||
('conv5_3_CPM_L2', [128, 128, 3, 1, 1]),
|
||||
('conv5_4_CPM_L2', [128, 512, 1, 1, 0]),
|
||||
('conv5_5_CPM_L2', [512, 19, 1, 1, 0])])
|
||||
blocks['block1_1'] = block1_1
|
||||
blocks['block1_2'] = block1_2
|
||||
|
||||
self.model0 = make_layers(block0, no_relu_layers)
|
||||
|
||||
# Stages 2 - 6
|
||||
for i in range(2, 7):
|
||||
blocks['block%d_1' % i] = OrderedDict([
|
||||
('Mconv1_stage%d_L1' % i, [185, 128, 7, 1, 3]),
|
||||
('Mconv2_stage%d_L1' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv3_stage%d_L1' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv4_stage%d_L1' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv5_stage%d_L1' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv6_stage%d_L1' % i, [128, 128, 1, 1, 0]),
|
||||
('Mconv7_stage%d_L1' % i, [128, 38, 1, 1, 0])
|
||||
])
|
||||
|
||||
blocks['block%d_2' % i] = OrderedDict([
|
||||
('Mconv1_stage%d_L2' % i, [185, 128, 7, 1, 3]),
|
||||
('Mconv2_stage%d_L2' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv3_stage%d_L2' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv4_stage%d_L2' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv5_stage%d_L2' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv6_stage%d_L2' % i, [128, 128, 1, 1, 0]),
|
||||
('Mconv7_stage%d_L2' % i, [128, 19, 1, 1, 0])
|
||||
])
|
||||
|
||||
for k in blocks.keys():
|
||||
blocks[k] = make_layers(blocks[k], no_relu_layers)
|
||||
|
||||
self.model1_1 = blocks['block1_1']
|
||||
self.model2_1 = blocks['block2_1']
|
||||
self.model3_1 = blocks['block3_1']
|
||||
self.model4_1 = blocks['block4_1']
|
||||
self.model5_1 = blocks['block5_1']
|
||||
self.model6_1 = blocks['block6_1']
|
||||
|
||||
self.model1_2 = blocks['block1_2']
|
||||
self.model2_2 = blocks['block2_2']
|
||||
self.model3_2 = blocks['block3_2']
|
||||
self.model4_2 = blocks['block4_2']
|
||||
self.model5_2 = blocks['block5_2']
|
||||
self.model6_2 = blocks['block6_2']
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out1 = self.model0(x)
|
||||
|
||||
out1_1 = self.model1_1(out1)
|
||||
out1_2 = self.model1_2(out1)
|
||||
out2 = torch.cat([out1_1, out1_2, out1], 1)
|
||||
|
||||
out2_1 = self.model2_1(out2)
|
||||
out2_2 = self.model2_2(out2)
|
||||
out3 = torch.cat([out2_1, out2_2, out1], 1)
|
||||
|
||||
out3_1 = self.model3_1(out3)
|
||||
out3_2 = self.model3_2(out3)
|
||||
out4 = torch.cat([out3_1, out3_2, out1], 1)
|
||||
|
||||
out4_1 = self.model4_1(out4)
|
||||
out4_2 = self.model4_2(out4)
|
||||
out5 = torch.cat([out4_1, out4_2, out1], 1)
|
||||
|
||||
out5_1 = self.model5_1(out5)
|
||||
out5_2 = self.model5_2(out5)
|
||||
out6 = torch.cat([out5_1, out5_2, out1], 1)
|
||||
|
||||
out6_1 = self.model6_1(out6)
|
||||
out6_2 = self.model6_2(out6)
|
||||
|
||||
return out6_1, out6_2
|
||||
|
||||
|
||||
class handpose_model(nn.Module):
|
||||
def __init__(self):
|
||||
super(handpose_model, self).__init__()
|
||||
|
||||
# these layers have no relu layer
|
||||
no_relu_layers = [
|
||||
'conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3', 'Mconv7_stage4',
|
||||
'Mconv7_stage5', 'Mconv7_stage6'
|
||||
]
|
||||
# stage 1
|
||||
block1_0 = OrderedDict([('conv1_1', [3, 64, 3, 1, 1]),
|
||||
('conv1_2', [64, 64, 3, 1, 1]),
|
||||
('pool1_stage1', [2, 2, 0]),
|
||||
('conv2_1', [64, 128, 3, 1, 1]),
|
||||
('conv2_2', [128, 128, 3, 1, 1]),
|
||||
('pool2_stage1', [2, 2, 0]),
|
||||
('conv3_1', [128, 256, 3, 1, 1]),
|
||||
('conv3_2', [256, 256, 3, 1, 1]),
|
||||
('conv3_3', [256, 256, 3, 1, 1]),
|
||||
('conv3_4', [256, 256, 3, 1, 1]),
|
||||
('pool3_stage1', [2, 2, 0]),
|
||||
('conv4_1', [256, 512, 3, 1, 1]),
|
||||
('conv4_2', [512, 512, 3, 1, 1]),
|
||||
('conv4_3', [512, 512, 3, 1, 1]),
|
||||
('conv4_4', [512, 512, 3, 1, 1]),
|
||||
('conv5_1', [512, 512, 3, 1, 1]),
|
||||
('conv5_2', [512, 512, 3, 1, 1]),
|
||||
('conv5_3_CPM', [512, 128, 3, 1, 1])])
|
||||
|
||||
block1_1 = OrderedDict([('conv6_1_CPM', [128, 512, 1, 1, 0]),
|
||||
('conv6_2_CPM', [512, 22, 1, 1, 0])])
|
||||
|
||||
blocks = {}
|
||||
blocks['block1_0'] = block1_0
|
||||
blocks['block1_1'] = block1_1
|
||||
|
||||
# stage 2-6
|
||||
for i in range(2, 7):
|
||||
blocks['block%d' % i] = OrderedDict([
|
||||
('Mconv1_stage%d' % i, [150, 128, 7, 1, 3]),
|
||||
('Mconv2_stage%d' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv3_stage%d' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv4_stage%d' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv5_stage%d' % i, [128, 128, 7, 1, 3]),
|
||||
('Mconv6_stage%d' % i, [128, 128, 1, 1, 0]),
|
||||
('Mconv7_stage%d' % i, [128, 22, 1, 1, 0])
|
||||
])
|
||||
|
||||
for k in blocks.keys():
|
||||
blocks[k] = make_layers(blocks[k], no_relu_layers)
|
||||
|
||||
self.model1_0 = blocks['block1_0']
|
||||
self.model1_1 = blocks['block1_1']
|
||||
self.model2 = blocks['block2']
|
||||
self.model3 = blocks['block3']
|
||||
self.model4 = blocks['block4']
|
||||
self.model5 = blocks['block5']
|
||||
self.model6 = blocks['block6']
|
||||
|
||||
def forward(self, x):
|
||||
out1_0 = self.model1_0(x)
|
||||
out1_1 = self.model1_1(out1_0)
|
||||
concat_stage2 = torch.cat([out1_1, out1_0], 1)
|
||||
out_stage2 = self.model2(concat_stage2)
|
||||
concat_stage3 = torch.cat([out_stage2, out1_0], 1)
|
||||
out_stage3 = self.model3(concat_stage3)
|
||||
concat_stage4 = torch.cat([out_stage3, out1_0], 1)
|
||||
out_stage4 = self.model4(concat_stage4)
|
||||
concat_stage5 = torch.cat([out_stage4, out1_0], 1)
|
||||
out_stage5 = self.model5(concat_stage5)
|
||||
concat_stage6 = torch.cat([out_stage5, out1_0], 1)
|
||||
out_stage6 = self.model6(concat_stage6)
|
||||
return out_stage6
|
||||
|
||||
|
||||
class Hand(object):
|
||||
def __init__(self, model_path, device='cuda'):
|
||||
self.model = handpose_model()
|
||||
if torch.cuda.is_available():
|
||||
self.model = self.model.to(device)
|
||||
model_dict = transfer(self.model, torch.load(model_path))
|
||||
self.model.load_state_dict(model_dict)
|
||||
self.model.eval()
|
||||
self.device = device
|
||||
|
||||
def __call__(self, oriImg):
|
||||
scale_search = [0.5, 1.0, 1.5, 2.0]
|
||||
# scale_search = [0.5]
|
||||
boxsize = 368
|
||||
stride = 8
|
||||
padValue = 128
|
||||
thre = 0.05
|
||||
multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
|
||||
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 22))
|
||||
# paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
|
||||
|
||||
for m in range(len(multiplier)):
|
||||
scale = multiplier[m]
|
||||
imageToTest = cv2.resize(oriImg, (0, 0),
|
||||
fx=scale,
|
||||
fy=scale,
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
imageToTest_padded, pad = padRightDownCorner(
|
||||
imageToTest, stride, padValue)
|
||||
im = np.transpose(
|
||||
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
|
||||
(3, 2, 0, 1)) / 256 - 0.5
|
||||
im = np.ascontiguousarray(im)
|
||||
|
||||
data = torch.from_numpy(im).float()
|
||||
if torch.cuda.is_available():
|
||||
data = data.to(self.device)
|
||||
# data = data.permute([2, 0, 1]).unsqueeze(0).float()
|
||||
with torch.no_grad():
|
||||
output = self.model(data).cpu().numpy()
|
||||
# output = self.model(data).numpy()q
|
||||
|
||||
# extract outputs, resize, and remove padding
|
||||
heatmap = np.transpose(np.squeeze(output),
|
||||
(1, 2, 0)) # output 1 is heatmaps
|
||||
heatmap = cv2.resize(heatmap, (0, 0),
|
||||
fx=stride,
|
||||
fy=stride,
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
heatmap = heatmap[:imageToTest_padded.shape[0] -
|
||||
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
|
||||
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
heatmap_avg += heatmap / len(multiplier)
|
||||
|
||||
all_peaks = []
|
||||
for part in range(21):
|
||||
map_ori = heatmap_avg[:, :, part]
|
||||
one_heatmap = gaussian_filter(map_ori, sigma=3)
|
||||
binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
|
||||
# 全部小于阈值
|
||||
if np.sum(binary) == 0:
|
||||
all_peaks.append([0, 0])
|
||||
continue
|
||||
label_img, label_numbers = label(binary,
|
||||
return_num=True,
|
||||
connectivity=binary.ndim)
|
||||
max_index = np.argmax([
|
||||
np.sum(map_ori[label_img == i])
|
||||
for i in range(1, label_numbers + 1)
|
||||
]) + 1
|
||||
label_img[label_img != max_index] = 0
|
||||
map_ori[label_img == 0] = 0
|
||||
|
||||
y, x = npmax(map_ori)
|
||||
all_peaks.append([x, y])
|
||||
return np.array(all_peaks)
|
||||
|
||||
|
||||
class Body(object):
|
||||
def __init__(self, model_path, device='cuda'):
|
||||
self.model = bodypose_model()
|
||||
if torch.cuda.is_available():
|
||||
self.model = self.model.to(device)
|
||||
model_dict = transfer(self.model, torch.load(model_path))
|
||||
self.model.load_state_dict(model_dict)
|
||||
self.model.eval()
|
||||
self.device = device
|
||||
|
||||
def __call__(self, oriImg):
|
||||
# scale_search = [0.5, 1.0, 1.5, 2.0]
|
||||
scale_search = [0.5]
|
||||
boxsize = 368
|
||||
stride = 8
|
||||
padValue = 128
|
||||
thre1 = 0.1
|
||||
thre2 = 0.05
|
||||
multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
|
||||
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
|
||||
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
|
||||
|
||||
for m in range(len(multiplier)):
|
||||
scale = multiplier[m]
|
||||
imageToTest = cv2.resize(oriImg, (0, 0),
|
||||
fx=scale,
|
||||
fy=scale,
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
imageToTest_padded, pad = padRightDownCorner(
|
||||
imageToTest, stride, padValue)
|
||||
im = np.transpose(
|
||||
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
|
||||
(3, 2, 0, 1)) / 256 - 0.5
|
||||
im = np.ascontiguousarray(im)
|
||||
|
||||
data = torch.from_numpy(im).float()
|
||||
if torch.cuda.is_available():
|
||||
data = data.to(self.device)
|
||||
# data = data.permute([2, 0, 1]).unsqueeze(0).float()
|
||||
with torch.no_grad():
|
||||
Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
|
||||
Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
|
||||
Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
|
||||
|
||||
# extract outputs, resize, and remove padding
|
||||
# heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0))
|
||||
# output 1 is heatmaps
|
||||
heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2),
|
||||
(1, 2, 0)) # output 1 is heatmaps
|
||||
heatmap = cv2.resize(heatmap, (0, 0),
|
||||
fx=stride,
|
||||
fy=stride,
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
heatmap = heatmap[:imageToTest_padded.shape[0] -
|
||||
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
|
||||
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
# paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0)) # output 0 is PAFs
|
||||
paf = np.transpose(np.squeeze(Mconv7_stage6_L1),
|
||||
(1, 2, 0)) # output 0 is PAFs
|
||||
paf = cv2.resize(paf, (0, 0),
|
||||
fx=stride,
|
||||
fy=stride,
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
paf = paf[:imageToTest_padded.shape[0] -
|
||||
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
|
||||
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
heatmap_avg += heatmap_avg + heatmap / len(multiplier)
|
||||
paf_avg += +paf / len(multiplier)
|
||||
|
||||
all_peaks = []
|
||||
peak_counter = 0
|
||||
|
||||
for part in range(18):
|
||||
map_ori = heatmap_avg[:, :, part]
|
||||
one_heatmap = gaussian_filter(map_ori, sigma=3)
|
||||
|
||||
map_left = np.zeros(one_heatmap.shape)
|
||||
map_left[1:, :] = one_heatmap[:-1, :]
|
||||
map_right = np.zeros(one_heatmap.shape)
|
||||
map_right[:-1, :] = one_heatmap[1:, :]
|
||||
map_up = np.zeros(one_heatmap.shape)
|
||||
map_up[:, 1:] = one_heatmap[:, :-1]
|
||||
map_down = np.zeros(one_heatmap.shape)
|
||||
map_down[:, :-1] = one_heatmap[:, 1:]
|
||||
|
||||
peaks_binary = np.logical_and.reduce(
|
||||
(one_heatmap >= map_left, one_heatmap >= map_right,
|
||||
one_heatmap >= map_up, one_heatmap >= map_down,
|
||||
one_heatmap > thre1))
|
||||
peaks = list(
|
||||
zip(np.nonzero(peaks_binary)[1],
|
||||
np.nonzero(peaks_binary)[0])) # note reverse
|
||||
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
|
||||
peak_id = range(peak_counter, peak_counter + len(peaks))
|
||||
peaks_with_score_and_id = [
|
||||
peaks_with_score[i] + (peak_id[i], )
|
||||
for i in range(len(peak_id))
|
||||
]
|
||||
|
||||
all_peaks.append(peaks_with_score_and_id)
|
||||
peak_counter += len(peaks)
|
||||
|
||||
# find connection in the specified sequence, center 29 is in the position 15
|
||||
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9],
|
||||
[9, 10], [10, 11], [2, 12], [12, 13], [13, 14], [2, 1],
|
||||
[1, 15], [15, 17], [1, 16], [16, 18], [3, 17], [6, 18]]
|
||||
# the middle joints heatmap correpondence
|
||||
mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44],
|
||||
[19, 20], [21, 22], [23, 24], [25, 26], [27, 28], [29, 30],
|
||||
[47, 48], [49, 50], [53, 54], [51, 52], [55, 56], [37, 38],
|
||||
[45, 46]]
|
||||
|
||||
connection_all = []
|
||||
special_k = []
|
||||
mid_num = 10
|
||||
|
||||
for k in range(len(mapIdx)):
|
||||
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
|
||||
candA = all_peaks[limbSeq[k][0] - 1]
|
||||
candB = all_peaks[limbSeq[k][1] - 1]
|
||||
nA = len(candA)
|
||||
nB = len(candB)
|
||||
indexA, indexB = limbSeq[k]
|
||||
if (nA != 0 and nB != 0):
|
||||
connection_candidate = []
|
||||
for i in range(nA):
|
||||
for j in range(nB):
|
||||
vec = np.subtract(candB[j][:2], candA[i][:2])
|
||||
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
|
||||
norm = max(0.001, norm)
|
||||
vec = np.divide(vec, norm)
|
||||
|
||||
startend = list(
|
||||
zip(
|
||||
np.linspace(candA[i][0],
|
||||
candB[j][0],
|
||||
num=mid_num),
|
||||
np.linspace(candA[i][1],
|
||||
candB[j][1],
|
||||
num=mid_num)))
|
||||
|
||||
vec_x = np.array([
|
||||
score_mid[int(round(startend[ii][1])),
|
||||
int(round(startend[ii][0])), 0]
|
||||
for ii in range(len(startend))
|
||||
])
|
||||
vec_y = np.array([
|
||||
score_mid[int(round(startend[ii][1])),
|
||||
int(round(startend[ii][0])), 1]
|
||||
for ii in range(len(startend))
|
||||
])
|
||||
|
||||
score_midpts = np.multiply(
|
||||
vec_x, vec[0]) + np.multiply(vec_y, vec[1])
|
||||
score_with_dist_prior = sum(score_midpts) / len(
|
||||
score_midpts) + min(
|
||||
0.5 * oriImg.shape[0] / norm - 1, 0)
|
||||
criterion1 = len(np.nonzero(
|
||||
score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
|
||||
criterion2 = score_with_dist_prior > 0
|
||||
if criterion1 and criterion2:
|
||||
connection_candidate.append([
|
||||
i, j, score_with_dist_prior,
|
||||
score_with_dist_prior + candA[i][2] +
|
||||
candB[j][2]
|
||||
])
|
||||
|
||||
connection_candidate = sorted(connection_candidate,
|
||||
key=lambda x: x[2],
|
||||
reverse=True)
|
||||
connection = np.zeros((0, 5))
|
||||
for c in range(len(connection_candidate)):
|
||||
i, j, s = connection_candidate[c][0:3]
|
||||
if (i not in connection[:, 3]
|
||||
and j not in connection[:, 4]):
|
||||
connection = np.vstack(
|
||||
[connection, [candA[i][3], candB[j][3], s, i, j]])
|
||||
if (len(connection) >= min(nA, nB)):
|
||||
break
|
||||
|
||||
connection_all.append(connection)
|
||||
else:
|
||||
special_k.append(k)
|
||||
connection_all.append([])
|
||||
|
||||
# last number in each row is the total parts number of that person
|
||||
# the second last number in each row is the score of the overall configuration
|
||||
subset = -1 * np.ones((0, 20))
|
||||
candidate = np.array(
|
||||
[item for sublist in all_peaks for item in sublist])
|
||||
|
||||
for k in range(len(mapIdx)):
|
||||
if k not in special_k:
|
||||
partAs = connection_all[k][:, 0]
|
||||
partBs = connection_all[k][:, 1]
|
||||
indexA, indexB = np.array(limbSeq[k]) - 1
|
||||
|
||||
for i in range(len(connection_all[k])): # = 1:size(temp,1)
|
||||
found = 0
|
||||
subset_idx = [-1, -1]
|
||||
for j in range(len(subset)): # 1:size(subset,1):
|
||||
if subset[j][indexA] == partAs[i] or subset[j][
|
||||
indexB] == partBs[i]:
|
||||
subset_idx[found] = j
|
||||
found += 1
|
||||
|
||||
if found == 1:
|
||||
j = subset_idx[0]
|
||||
if subset[j][indexB] != partBs[i]:
|
||||
subset[j][indexB] = partBs[i]
|
||||
subset[j][-1] += 1
|
||||
subset[j][-2] += candidate[
|
||||
partBs[i].astype(int),
|
||||
2] + connection_all[k][i][2]
|
||||
elif found == 2: # if found 2 and disjoint, merge them
|
||||
j1, j2 = subset_idx
|
||||
membership = ((subset[j1] >= 0).astype(int) +
|
||||
(subset[j2] >= 0).astype(int))[:-2]
|
||||
if len(np.nonzero(membership == 2)[0]) == 0: # merge
|
||||
subset[j1][:-2] += (subset[j2][:-2] + 1)
|
||||
subset[j1][-2:] += subset[j2][-2:]
|
||||
subset[j1][-2] += connection_all[k][i][2]
|
||||
subset = np.delete(subset, j2, 0)
|
||||
else: # as like found == 1
|
||||
subset[j1][indexB] = partBs[i]
|
||||
subset[j1][-1] += 1
|
||||
subset[j1][-2] += candidate[
|
||||
partBs[i].astype(int),
|
||||
2] + connection_all[k][i][2]
|
||||
|
||||
# if find no partA in the subset, create a new subset
|
||||
elif not found and k < 17:
|
||||
row = -1 * np.ones(20)
|
||||
row[indexA] = partAs[i]
|
||||
row[indexB] = partBs[i]
|
||||
row[-1] = 2
|
||||
row[-2] = sum(
|
||||
candidate[connection_all[k][i, :2].astype(int),
|
||||
2]) + connection_all[k][i][2]
|
||||
subset = np.vstack([subset, row])
|
||||
# delete some rows of subset which has few parts occur
|
||||
deleteIdx = []
|
||||
for i in range(len(subset)):
|
||||
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
|
||||
deleteIdx.append(i)
|
||||
subset = np.delete(subset, deleteIdx, axis=0)
|
||||
|
||||
# subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
|
||||
# candidate: x, y, score, id
|
||||
return candidate, subset
|
||||
|
||||
|
||||
@ANNOTATORS.register_class()
|
||||
class OpenposeAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
||||
para_dict = {}
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
with FS.get_from(cfg.BODY_MODEL_PATH,
|
||||
wait_finish=True) as body_model_path:
|
||||
self.body_estimation = Body(body_model_path, device='cpu')
|
||||
with FS.get_from(cfg.HAND_MODEL_PATH,
|
||||
wait_finish=True) as hand_model_path:
|
||||
self.hand_estimation = Hand(hand_model_path, device='cpu')
|
||||
self.use_hand = cfg.get('USE_HAND', False)
|
||||
|
||||
def to(self, device):
|
||||
self.body_estimation.model = self.body_estimation.model.to(device)
|
||||
self.body_estimation.device = device
|
||||
self.hand_estimation.model = self.hand_estimation.model.to(device)
|
||||
self.hand_estimation.device = device
|
||||
return self
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
@torch.autocast('cuda', enabled=False)
|
||||
def forward(self, image):
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image = image.detach().cpu().numpy()
|
||||
elif isinstance(image, np.ndarray):
|
||||
image = image.copy()
|
||||
else:
|
||||
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
||||
image = image[:, :, ::-1]
|
||||
candidate, subset = self.body_estimation(image)
|
||||
canvas = np.zeros_like(image)
|
||||
canvas = draw_bodypose(canvas, candidate, subset)
|
||||
if self.use_hand:
|
||||
hands_list = handDetect(candidate, subset, image)
|
||||
all_hand_peaks = []
|
||||
for x, y, w, is_left in hands_list:
|
||||
peaks = self.hand_estimation(image[y:y + w, x:x + w, :])
|
||||
peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0],
|
||||
peaks[:, 0] + x)
|
||||
peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1],
|
||||
peaks[:, 1] + y)
|
||||
all_hand_peaks.append(peaks)
|
||||
canvas = draw_handpose(canvas, all_hand_peaks)
|
||||
return canvas
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('ANNOTATORS',
|
||||
__class__.__name__,
|
||||
OpenposeAnnotator.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,30 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.utils.config import Config
|
||||
from scepter.modules.utils.registry import Registry, build_from_config
|
||||
|
||||
|
||||
def build_annotator(cfg, registry, logger=None, *args, **kwargs):
|
||||
""" After build model, load pretrained model if exists key `pretrain`.
|
||||
|
||||
pretrain (str, dict): Describes how to load pretrained model.
|
||||
str, treat pretrain as model path;
|
||||
dict: should contains key `path`, and other parameters token by function load_pretrained();
|
||||
"""
|
||||
if not isinstance(cfg, Config):
|
||||
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)):
|
||||
raise TypeError('Pretrain parameter must be a string')
|
||||
else:
|
||||
pretrain_cfg = None
|
||||
|
||||
model = build_from_config(cfg, registry, logger=logger, *args, **kwargs)
|
||||
if pretrain_cfg is not None:
|
||||
if hasattr(model, 'load_pretrained_model'):
|
||||
model.load_pretrained_model(pretrain_cfg)
|
||||
return model
|
||||
|
||||
|
||||
ANNOTATORS = Registry('ANNOTATORS', build_func=build_annotator)
|
||||
@@ -0,0 +1,113 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def resize_image(input_image, resolution):
|
||||
H, W, C = input_image.shape
|
||||
H = float(H)
|
||||
W = float(W)
|
||||
k = float(resolution) / min(H, W)
|
||||
H *= k
|
||||
W *= k
|
||||
H = int(np.round(H / 64.0)) * 64
|
||||
W = int(np.round(W / 64.0)) * 64
|
||||
img = cv2.resize(
|
||||
input_image, (W, H),
|
||||
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
|
||||
return img, k
|
||||
|
||||
|
||||
def resize_image_ori(h, w, image, k):
|
||||
img = cv2.resize(
|
||||
image, (w, h),
|
||||
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
|
||||
return img
|
||||
|
||||
|
||||
class AnnotatorProcessor():
|
||||
canny_cfg = {
|
||||
'NAME': 'CannyAnnotator',
|
||||
'LOW_THRESHOLD': 100,
|
||||
'HIGH_THRESHOLD': 200,
|
||||
'INPUT_KEYS': ['img'],
|
||||
'OUTPUT_KEYS': ['canny']
|
||||
}
|
||||
hed_cfg = {
|
||||
'NAME': 'HedAnnotator',
|
||||
'PRETRAINED_MODEL':
|
||||
'ms://damo/scepter_scedit@annotator/ckpts/ControlNetHED.pth',
|
||||
'INPUT_KEYS': ['img'],
|
||||
'OUTPUT_KEYS': ['hed']
|
||||
}
|
||||
openpose_cfg = {
|
||||
'NAME': 'OpenposeAnnotator',
|
||||
'BODY_MODEL_PATH':
|
||||
'ms://damo/scepter_scedit@annotator/ckpts/body_pose_model.pth',
|
||||
'HAND_MODEL_PATH':
|
||||
'ms://damo/scepter_scedit@annotator/ckpts/hand_pose_model.pth',
|
||||
'INPUT_KEYS': ['img'],
|
||||
'OUTPUT_KEYS': ['openpose']
|
||||
}
|
||||
midas_cfg = {
|
||||
'NAME': 'MidasDetector',
|
||||
'PRETRAINED_MODEL':
|
||||
'ms://damo/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt',
|
||||
'INPUT_KEYS': ['img'],
|
||||
'OUTPUT_KEYS': ['depth']
|
||||
}
|
||||
mlsd_cfg = {
|
||||
'NAME': 'MLSDdetector',
|
||||
'PRETRAINED_MODEL':
|
||||
'ms://damo/scepter_scedit@annotator/ckpts/mlsd_large_512_fp32.pth',
|
||||
'INPUT_KEYS': ['img'],
|
||||
'OUTPUT_KEYS': ['mlsd']
|
||||
}
|
||||
color_cfg = {
|
||||
'NAME': 'ColorAnnotator',
|
||||
'RATIO': 64,
|
||||
'INPUT_KEYS': ['img'],
|
||||
'OUTPUT_KEYS': ['color']
|
||||
}
|
||||
|
||||
anno_type_map = {
|
||||
'canny': canny_cfg,
|
||||
'hed': hed_cfg,
|
||||
'pose': openpose_cfg,
|
||||
'depth': midas_cfg,
|
||||
'mlsd': mlsd_cfg,
|
||||
'color': color_cfg
|
||||
}
|
||||
|
||||
def __init__(self, anno_type):
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.utils.config import Config
|
||||
from scepter.modules.utils.distribute import we
|
||||
|
||||
if isinstance(anno_type, str):
|
||||
assert anno_type in self.anno_type_map.keys()
|
||||
anno_type = [anno_type]
|
||||
elif isinstance(anno_type, (list, tuple)):
|
||||
assert all(tp in self.anno_type_map.keys() for tp in anno_type)
|
||||
else:
|
||||
raise Exception(f'Error anno_type: {anno_type}')
|
||||
|
||||
general_dict = {
|
||||
'NAME': 'GeneralAnnotator',
|
||||
'ANNOTATORS': [self.anno_type_map[tp] for tp in anno_type]
|
||||
}
|
||||
general_anno = Config(cfg_dict=general_dict, load=False)
|
||||
self.general_ins = ANNOTATORS.build(general_anno).to(we.device_id)
|
||||
|
||||
def run(self, image, anno_type=None):
|
||||
output_image = self.general_ins({'img': image})
|
||||
if anno_type is not None:
|
||||
if isinstance(anno_type, str) and anno_type in output_image:
|
||||
return output_image[anno_type]
|
||||
else:
|
||||
return {
|
||||
tp: output_image[tp]
|
||||
for tp in anno_type if tp in output_image
|
||||
}
|
||||
else:
|
||||
return output_image
|
||||
@@ -6,5 +6,6 @@ from scepter.modules.data.dataset.dataset import (Image2ImageDataset,
|
||||
ImageClassifyPublicDataset,
|
||||
ImageTextPairDataset,
|
||||
Text2ImageDataset)
|
||||
from scepter.modules.data.dataset.ms_dataset import ImageTextPairMSDataset
|
||||
from scepter.modules.data.dataset.ms_dataset import (
|
||||
ImageTextPairFolderDataset, ImageTextPairMSDataset)
|
||||
from scepter.modules.data.dataset.registry import DATASETS
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import numbers
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Iterable
|
||||
|
||||
@@ -235,6 +236,7 @@ class Text2ImageDataset(BaseDataset):
|
||||
delimiter = cfg.get('DELIMITER', ',')
|
||||
fields = cfg.get('FIELDS', ['row_key', 'prompt'])
|
||||
prompt_prefix = cfg.get('PROMPT_PREFIX', '')
|
||||
path_prefix = cfg.get('PATH_PREFIX', '')
|
||||
use_num = cfg.get('USE_NUM', -1)
|
||||
|
||||
image_size = cfg.get('IMAGE_SIZE', 1024)
|
||||
@@ -257,11 +259,14 @@ class Text2ImageDataset(BaseDataset):
|
||||
if key in ['prompt', 'caption', 'text']:
|
||||
item['ori_prompt'] = value
|
||||
item['prompt'] = prompt_prefix + value
|
||||
elif key in ['oss_key', 'path', 'img_path', 'target_img_path']:
|
||||
item['meta']['img_path'] = os.path.join(path_prefix, value)
|
||||
elif key in ['width', 'height']:
|
||||
item['meta'][key] = int(value)
|
||||
elif key != 'meta':
|
||||
item[key] = value
|
||||
else:
|
||||
continue
|
||||
|
||||
self.items.append(item)
|
||||
if use_num > 0:
|
||||
self.items = self.items[:use_num]
|
||||
|
||||
@@ -9,6 +9,7 @@ from scepter.modules.data.dataset.base_dataset import BaseDataset
|
||||
from scepter.modules.data.dataset.registry import DATASETS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
|
||||
@DATASETS.register_class()
|
||||
@@ -105,14 +106,17 @@ class ImageTextPairMSDataset(BaseDataset):
|
||||
raise (
|
||||
'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)
|
||||
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 as e:
|
||||
except Exception:
|
||||
self.logger.info(
|
||||
f"Load Modelscope dataset failed with {e}, retry with download_mode='force_redownload'."
|
||||
"Load Modelscope dataset failed, retry with download_mode='force_redownload'."
|
||||
)
|
||||
try:
|
||||
self.data = MsDataset.load(
|
||||
@@ -177,3 +181,119 @@ class ImageTextPairMSDataset(BaseDataset):
|
||||
__class__.__name__,
|
||||
ImageTextPairMSDataset.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@DATASETS.register_class()
|
||||
class ImageTextPairFolderDataset(BaseDataset):
|
||||
para_dict = {
|
||||
'DATA_FOLDER': {
|
||||
'value': '',
|
||||
'description': 'Dataset folder.'
|
||||
},
|
||||
'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}>'
|
||||
},
|
||||
'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)
|
||||
data_folder = cfg.get('DATA_FOLDER', None)
|
||||
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.output_size = cfg.get('OUTPUT_SIZE', None)
|
||||
if self.output_size is not None:
|
||||
if isinstance(self.output_size, numbers.Number):
|
||||
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.')
|
||||
data_folder = FS.get_dir_to_local_dir(data_folder)
|
||||
all_lines = open(os.path.join(data_folder, 'train.csv'),
|
||||
'r').read().split('\n')
|
||||
assert all_lines[0] == 'Target:FILE,Prompt'
|
||||
self.data = []
|
||||
for line in all_lines[1:]:
|
||||
line = line.strip()
|
||||
if line == '':
|
||||
continue
|
||||
self.data.append({
|
||||
'Target:FILE':
|
||||
os.path.join(data_folder,
|
||||
line.split(',', 1)[0]),
|
||||
'Prompt':
|
||||
line.split(',', 1)[1]
|
||||
})
|
||||
self.real_number = len(self.data)
|
||||
|
||||
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)]
|
||||
# print(current_data.keys())
|
||||
image_path = current_data['Target:FILE']
|
||||
prompt = current_data['Prompt']
|
||||
style = current_data['Style'] if 'Style' in current_data else ''
|
||||
# print(prompt, style)
|
||||
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
|
||||
if we.debug:
|
||||
print(prompt, self.replace_keywords.strip())
|
||||
ret_item = {
|
||||
'meta': {
|
||||
'img_path': image_path,
|
||||
'data_key': style,
|
||||
'data_num': self.real_number
|
||||
},
|
||||
'prompt': prompt
|
||||
}
|
||||
if self.output_size is not None:
|
||||
ret_item['meta']['image_size'] = self.output_size
|
||||
return ret_item
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('DATASet',
|
||||
__class__.__name__,
|
||||
ImageTextPairMSDataset.para_dict,
|
||||
set_name=True)
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from PIL import Image
|
||||
from torch.utils.data.dataloader import default_collate
|
||||
|
||||
|
||||
def pil_collate_fn(self, batch):
|
||||
batch_data = {}
|
||||
for items in batch:
|
||||
for key, item in items.items():
|
||||
if isinstance(item, Image.Image):
|
||||
if key not in batch_data:
|
||||
batch_data[key] = []
|
||||
batch_data[key].append(item)
|
||||
else:
|
||||
if key not in batch_data:
|
||||
batch_data[key] = []
|
||||
batch_data[key].append(item)
|
||||
|
||||
for key, item in batch_data.items():
|
||||
if not all(isinstance(x, Image.Image) for x in item):
|
||||
batch_data[key] = default_collate(item)
|
||||
|
||||
return batch_data
|
||||
@@ -0,0 +1,2 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from scepter.modules.inference.diffusion_inference import DiffusionInference
|
||||
@@ -0,0 +1,848 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import os.path
|
||||
import random
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms as TT
|
||||
from peft.utils import CONFIG_NAME, SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME
|
||||
from PIL.Image import Image
|
||||
from swift import Swift, SwiftModel
|
||||
|
||||
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, TUNERS)
|
||||
from scepter.modules.utils.config import Config
|
||||
from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
|
||||
def get_model(model_tuple):
|
||||
assert 'model' in model_tuple
|
||||
return model_tuple['model']
|
||||
|
||||
|
||||
class DiffusionInference():
|
||||
'''
|
||||
define vae, unet, text-encoder, tuner, refiner components
|
||||
support to load the components dynamicly.
|
||||
create and load model when run this model at the first time.
|
||||
'''
|
||||
def __init__(self, logger=None):
|
||||
self.logger = logger
|
||||
|
||||
def init_from_cfg(self, cfg):
|
||||
self.name = cfg.NAME
|
||||
self.is_default = cfg.get('IS_DEFAULT', False)
|
||||
module_paras = self.load_default(cfg.get('DEFAULT_PARAS', None))
|
||||
assert cfg.have('MODEL')
|
||||
cfg.MODEL = self.redefine_paras(cfg.MODEL)
|
||||
self.diffusion = self.load_schedule(cfg.MODEL.SCHEDULE)
|
||||
self.diffusion_model = self.infer_model(
|
||||
cfg.MODEL.DIFFUSION_MODEL, module_paras.get(
|
||||
'DIFFUSION_MODEL',
|
||||
None)) if cfg.MODEL.have('DIFFUSION_MODEL') else None
|
||||
self.first_stage_model = self.infer_model(
|
||||
cfg.MODEL.FIRST_STAGE_MODEL,
|
||||
module_paras.get(
|
||||
'FIRST_STAGE_MODEL',
|
||||
None)) if cfg.MODEL.have('FIRST_STAGE_MODEL') else None
|
||||
self.cond_stage_model = self.infer_model(
|
||||
cfg.MODEL.COND_STAGE_MODEL,
|
||||
module_paras.get(
|
||||
'COND_STAGE_MODEL',
|
||||
None)) if cfg.MODEL.have('COND_STAGE_MODEL') else None
|
||||
self.refiner_cond_model = self.infer_model(
|
||||
cfg.MODEL.REFINER_COND_MODEL,
|
||||
module_paras.get(
|
||||
'REFINER_COND_MODEL',
|
||||
None)) if cfg.MODEL.have('REFINER_COND_MODEL') else None
|
||||
self.refiner_diffusion_model = self.infer_model(
|
||||
cfg.MODEL.REFINER_MODEL, module_paras.get(
|
||||
'REFINER_MODEL',
|
||||
None)) if cfg.MODEL.have('REFINER_MODEL') else None
|
||||
self.tokenizer = TOKENIZERS.build(
|
||||
cfg.MODEL.TOKENIZER,
|
||||
logger=self.logger) if cfg.MODEL.have('TOKENIZER') else None
|
||||
|
||||
if self.tokenizer is not None:
|
||||
self.cond_stage_model['cfg'].KWARGS = {
|
||||
'vocab_size': self.tokenizer.vocab_size
|
||||
}
|
||||
|
||||
def register_tuner(self, tuner_model_list):
|
||||
if len(tuner_model_list) < 1:
|
||||
if isinstance(self.diffusion_model['model'], SwiftModel):
|
||||
for adapter_name in self.diffusion_model['model'].adapters:
|
||||
self.diffusion_model['model'].deactivate_adapter(
|
||||
adapter_name, offload='cpu')
|
||||
if isinstance(self.cond_stage_model['model'], SwiftModel):
|
||||
for adapter_name in self.cond_stage_model['model'].adapters:
|
||||
self.cond_stage_model['model'].deactivate_adapter(
|
||||
adapter_name, offload='cpu')
|
||||
return
|
||||
all_diffusion_tuner = {}
|
||||
all_cond_tuner = {}
|
||||
save_root_dir = '.cache_tuner'
|
||||
for tuner_model in tuner_model_list:
|
||||
tunner_model_folder = tuner_model.MODEL_PATH
|
||||
local_tuner_model = FS.get_dir_to_local_dir(tunner_model_folder)
|
||||
all_tuner_datas = os.listdir(local_tuner_model)
|
||||
cur_tuner_md5 = hashlib.md5(
|
||||
tunner_model_folder.encode('utf-8')).hexdigest()
|
||||
|
||||
local_diffusion_cache = os.path.join(
|
||||
save_root_dir, cur_tuner_md5 + '_' + 'diffusion')
|
||||
local_cond_cache = os.path.join(save_root_dir,
|
||||
cur_tuner_md5 + '_' + 'cond')
|
||||
|
||||
meta_file = os.path.join(save_root_dir,
|
||||
cur_tuner_md5 + '_meta.json')
|
||||
if not os.path.exists(meta_file):
|
||||
diffusion_tuner = {}
|
||||
cond_tuner = {}
|
||||
for sub in all_tuner_datas:
|
||||
sub_file = os.path.join(local_tuner_model, sub)
|
||||
config_file = os.path.join(sub_file, CONFIG_NAME)
|
||||
safe_file = os.path.join(sub_file,
|
||||
SAFETENSORS_WEIGHTS_NAME)
|
||||
bin_file = os.path.join(sub_file, WEIGHTS_NAME)
|
||||
if os.path.isdir(sub_file) and os.path.isfile(config_file):
|
||||
# diffusion or cond
|
||||
cfg = json.load(open(config_file, 'r'))
|
||||
if 'cond_stage_model.' in cfg['target_modules']:
|
||||
cond_cfg = copy.deepcopy(cfg)
|
||||
if 'cond_stage_model.*' in cond_cfg[
|
||||
'target_modules']:
|
||||
cond_cfg['target_modules'] = cond_cfg[
|
||||
'target_modules'].replace(
|
||||
'cond_stage_model.*', '.*')
|
||||
else:
|
||||
cond_cfg['target_modules'] = cond_cfg[
|
||||
'target_modules'].replace(
|
||||
'cond_stage_model.', '')
|
||||
if cond_cfg['target_modules'].startswith('*'):
|
||||
cond_cfg['target_modules'] = '.' + cond_cfg[
|
||||
'target_modules']
|
||||
os.makedirs(local_cond_cache + '_' + sub,
|
||||
exist_ok=True)
|
||||
cond_tuner[os.path.basename(local_cond_cache) +
|
||||
'_' + sub] = hashlib.md5(
|
||||
(local_cond_cache + '_' +
|
||||
sub).encode('utf-8')).hexdigest()
|
||||
os.makedirs(local_cond_cache + '_' + sub,
|
||||
exist_ok=True)
|
||||
|
||||
json.dump(
|
||||
cond_cfg,
|
||||
open(
|
||||
os.path.join(local_cond_cache + '_' + sub,
|
||||
CONFIG_NAME), 'w'))
|
||||
if 'model.' in cfg['target_modules'].replace(
|
||||
'cond_stage_model.', ''):
|
||||
diffusion_cfg = copy.deepcopy(cfg)
|
||||
if 'model.*' in diffusion_cfg['target_modules']:
|
||||
diffusion_cfg[
|
||||
'target_modules'] = diffusion_cfg[
|
||||
'target_modules'].replace(
|
||||
'model.*', '.*')
|
||||
else:
|
||||
diffusion_cfg[
|
||||
'target_modules'] = diffusion_cfg[
|
||||
'target_modules'].replace(
|
||||
'model.', '')
|
||||
if diffusion_cfg['target_modules'].startswith('*'):
|
||||
diffusion_cfg[
|
||||
'target_modules'] = '.' + diffusion_cfg[
|
||||
'target_modules']
|
||||
os.makedirs(local_diffusion_cache + '_' + sub,
|
||||
exist_ok=True)
|
||||
diffusion_tuner[
|
||||
os.path.basename(local_diffusion_cache) + '_' +
|
||||
sub] = hashlib.md5(
|
||||
(local_diffusion_cache + '_' +
|
||||
sub).encode('utf-8')).hexdigest()
|
||||
json.dump(
|
||||
diffusion_cfg,
|
||||
open(
|
||||
os.path.join(
|
||||
local_diffusion_cache + '_' + sub,
|
||||
CONFIG_NAME), 'w'))
|
||||
|
||||
state_dict = {}
|
||||
is_bin_file = True
|
||||
if os.path.isfile(bin_file):
|
||||
state_dict = torch.load(bin_file)
|
||||
elif os.path.isfile(safe_file):
|
||||
is_bin_file = False
|
||||
from safetensors.torch import \
|
||||
load_file as safe_load_file
|
||||
state_dict = safe_load_file(
|
||||
safe_file,
|
||||
device='cuda'
|
||||
if torch.cuda.is_available() else 'cpu')
|
||||
save_diffusion_state_dict = {}
|
||||
save_cond_state_dict = {}
|
||||
for key, value in state_dict.items():
|
||||
if key.startswith('model.'):
|
||||
save_diffusion_state_dict[
|
||||
key[len('model.'):].replace(
|
||||
sub,
|
||||
os.path.basename(local_diffusion_cache)
|
||||
+ '_' + sub)] = value
|
||||
elif key.startswith('cond_stage_model.'):
|
||||
save_cond_state_dict[
|
||||
key[len('cond_stage_model.'):].replace(
|
||||
sub,
|
||||
os.path.basename(local_cond_cache) +
|
||||
'_' + sub)] = value
|
||||
|
||||
if is_bin_file:
|
||||
if len(save_diffusion_state_dict) > 0:
|
||||
torch.save(
|
||||
save_diffusion_state_dict,
|
||||
os.path.join(
|
||||
local_diffusion_cache + '_' + sub,
|
||||
WEIGHTS_NAME))
|
||||
if len(save_cond_state_dict) > 0:
|
||||
torch.save(
|
||||
save_cond_state_dict,
|
||||
os.path.join(local_cond_cache + '_' + sub,
|
||||
WEIGHTS_NAME))
|
||||
else:
|
||||
from safetensors.torch import \
|
||||
save_file as safe_save_file
|
||||
if len(save_diffusion_state_dict) > 0:
|
||||
safe_save_file(
|
||||
save_diffusion_state_dict,
|
||||
os.path.join(
|
||||
local_diffusion_cache + '_' + sub,
|
||||
SAFETENSORS_WEIGHTS_NAME),
|
||||
metadata={'format': 'pt'})
|
||||
if len(save_cond_state_dict) > 0:
|
||||
safe_save_file(
|
||||
save_cond_state_dict,
|
||||
os.path.join(local_cond_cache + '_' + sub,
|
||||
SAFETENSORS_WEIGHTS_NAME),
|
||||
metadata={'format': 'pt'})
|
||||
json.dump(
|
||||
{
|
||||
'diffusion_tuner': diffusion_tuner,
|
||||
'cond_tuner': cond_tuner
|
||||
}, open(meta_file, 'w'))
|
||||
else:
|
||||
meta_conf = json.load(open(meta_file, 'r'))
|
||||
diffusion_tuner = meta_conf['diffusion_tuner']
|
||||
cond_tuner = meta_conf['cond_tuner']
|
||||
all_diffusion_tuner.update(diffusion_tuner)
|
||||
all_cond_tuner.update(cond_tuner)
|
||||
if len(all_diffusion_tuner) > 0:
|
||||
self.load(self.diffusion_model)
|
||||
self.diffusion_model['model'] = Swift.from_pretrained(
|
||||
self.diffusion_model['model'],
|
||||
save_root_dir,
|
||||
adapter_name=all_diffusion_tuner)
|
||||
self.diffusion_model['model'].set_active_adapters(
|
||||
list(all_diffusion_tuner.values()))
|
||||
self.unload(self.diffusion_model)
|
||||
if len(all_cond_tuner) > 0:
|
||||
self.load(self.cond_stage_model)
|
||||
self.cond_stage_model['model'] = Swift.from_pretrained(
|
||||
self.cond_stage_model['model'],
|
||||
save_root_dir,
|
||||
adapter_name=all_cond_tuner)
|
||||
self.cond_stage_model['model'].set_active_adapters(
|
||||
list(all_cond_tuner.values()))
|
||||
self.unload(self.cond_stage_model)
|
||||
|
||||
def register_controllers(self, control_model_ins):
|
||||
if control_model_ins is None or control_model_ins == '':
|
||||
if isinstance(self.diffusion_model['model'], SwiftModel):
|
||||
if (hasattr(self.diffusion_model['model'].base_model,
|
||||
'control_blocks') and
|
||||
self.diffusion_model['model'].base_model.control_blocks
|
||||
): # noqa
|
||||
del self.diffusion_model['model'].base_model.control_blocks
|
||||
self.diffusion_model[
|
||||
'model'].base_model.control_blocks = None
|
||||
self.diffusion_model['model'].base_model.control_name = []
|
||||
else:
|
||||
del self.diffusion_model['model'].control_blocks
|
||||
self.diffusion_model['model'].control_blocks = None
|
||||
self.diffusion_model['model'].control_name = []
|
||||
return
|
||||
if not isinstance(control_model_ins, list):
|
||||
control_model_ins = [control_model_ins]
|
||||
control_model = nn.ModuleList([])
|
||||
control_model_folder = []
|
||||
for one_control in control_model_ins:
|
||||
one_control_model_folder = one_control.MODEL_PATH
|
||||
control_model_folder.append(one_control_model_folder)
|
||||
have_list = getattr(self.diffusion_model['model'], 'control_name',
|
||||
[])
|
||||
if one_control_model_folder in have_list:
|
||||
ind = have_list.index(one_control_model_folder)
|
||||
csc_tuners = copy.deepcopy(
|
||||
self.diffusion_model['model'].control_blocks[ind])
|
||||
else:
|
||||
one_local_control_model = FS.get_dir_to_local_dir(
|
||||
one_control_model_folder)
|
||||
control_cfg = Config(cfg_file=os.path.join(
|
||||
one_local_control_model, 'configuration.json'))
|
||||
assert hasattr(control_cfg, 'CONTROL_MODEL')
|
||||
control_cfg.CONTROL_MODEL[
|
||||
'INPUT_BLOCK_CHANS'] = self.diffusion_model[
|
||||
'model']._input_block_chans
|
||||
control_cfg.CONTROL_MODEL[
|
||||
'INPUT_DOWN_FLAG'] = self.diffusion_model[
|
||||
'model']._input_down_flag
|
||||
control_cfg.CONTROL_MODEL.PRETRAINED_MODEL = os.path.join(
|
||||
one_local_control_model, 'pytorch_model.bin')
|
||||
csc_tuners = TUNERS.build(control_cfg.CONTROL_MODEL,
|
||||
logger=self.logger)
|
||||
control_model.append(csc_tuners)
|
||||
if isinstance(self.diffusion_model['model'], SwiftModel):
|
||||
del self.diffusion_model['model'].base_model.control_blocks
|
||||
self.diffusion_model[
|
||||
'model'].base_model.control_blocks = control_model
|
||||
self.diffusion_model[
|
||||
'model'].base_model.control_name = control_model_folder
|
||||
else:
|
||||
del self.diffusion_model['model'].control_blocks
|
||||
self.diffusion_model['model'].control_blocks = control_model
|
||||
self.diffusion_model['model'].control_name = control_model_folder
|
||||
|
||||
def redefine_paras(self, cfg):
|
||||
if cfg.get('PRETRAINED_MODEL', None):
|
||||
assert FS.isfile(cfg.PRETRAINED_MODEL)
|
||||
with FS.get_from(cfg.PRETRAINED_MODEL,
|
||||
wait_finish=True) as local_path:
|
||||
if local_path.endswith('safetensors'):
|
||||
from safetensors.torch import load_file as load_safetensors
|
||||
sd = load_safetensors(local_path)
|
||||
else:
|
||||
sd = torch.load(local_path, map_location='cpu')
|
||||
first_stage_model_path = os.path.join(
|
||||
os.path.dirname(local_path), 'first_stage_model.pth')
|
||||
cond_stage_model_path = os.path.join(
|
||||
os.path.dirname(local_path), 'cond_stage_model.pth')
|
||||
diffusion_model_path = os.path.join(
|
||||
os.path.dirname(local_path), 'diffusion_model.pth')
|
||||
if (not os.path.exists(first_stage_model_path)
|
||||
or not os.path.exists(cond_stage_model_path)
|
||||
or not os.path.exists(diffusion_model_path)):
|
||||
self.logger.info(
|
||||
'Now read the whole model and rearrange the modules, it may take several mins.'
|
||||
)
|
||||
first_stage_model = OrderedDict()
|
||||
cond_stage_model = OrderedDict()
|
||||
diffusion_model = OrderedDict()
|
||||
for k, v in sd.items():
|
||||
if k.startswith('first_stage_model.'):
|
||||
first_stage_model[k.replace(
|
||||
'first_stage_model.', '')] = v
|
||||
elif k.startswith('conditioner.'):
|
||||
cond_stage_model[k.replace('conditioner.', '')] = v
|
||||
elif k.startswith('cond_stage_model.'):
|
||||
if k.startswith('cond_stage_model.model.'):
|
||||
cond_stage_model[k.replace(
|
||||
'cond_stage_model.model.', '')] = v
|
||||
else:
|
||||
cond_stage_model[k.replace(
|
||||
'cond_stage_model.', '')] = v
|
||||
elif k.startswith('model.diffusion_model.'):
|
||||
diffusion_model[k.replace('model.diffusion_model.',
|
||||
'')] = v
|
||||
else:
|
||||
continue
|
||||
if cfg.have('FIRST_STAGE_MODEL'):
|
||||
with open(first_stage_model_path + 'cache', 'wb') as f:
|
||||
torch.save(first_stage_model, f)
|
||||
os.rename(first_stage_model_path + 'cache',
|
||||
first_stage_model_path)
|
||||
self.logger.info(
|
||||
'First stage model has been processed.')
|
||||
if cfg.have('COND_STAGE_MODEL'):
|
||||
with open(cond_stage_model_path + 'cache', 'wb') as f:
|
||||
torch.save(cond_stage_model, f)
|
||||
os.rename(cond_stage_model_path + 'cache',
|
||||
cond_stage_model_path)
|
||||
self.logger.info(
|
||||
'Cond stage model has been processed.')
|
||||
if cfg.have('DIFFUSION_MODEL'):
|
||||
with open(diffusion_model_path + 'cache', 'wb') as f:
|
||||
torch.save(diffusion_model, f)
|
||||
os.rename(diffusion_model_path + 'cache',
|
||||
diffusion_model_path)
|
||||
self.logger.info('Diffusion model has been processed.')
|
||||
if not cfg.FIRST_STAGE_MODEL.get('PRETRAINED_MODEL', None):
|
||||
cfg.FIRST_STAGE_MODEL.PRETRAINED_MODEL = first_stage_model_path
|
||||
else:
|
||||
cfg.FIRST_STAGE_MODEL.RELOAD_MODEL = first_stage_model_path
|
||||
if not cfg.COND_STAGE_MODEL.get('PRETRAINED_MODEL', None):
|
||||
cfg.COND_STAGE_MODEL.PRETRAINED_MODEL = cond_stage_model_path
|
||||
else:
|
||||
cfg.COND_STAGE_MODEL.RELOAD_MODEL = cond_stage_model_path
|
||||
if not cfg.DIFFUSION_MODEL.get('PRETRAINED_MODEL', None):
|
||||
cfg.DIFFUSION_MODEL.PRETRAINED_MODEL = diffusion_model_path
|
||||
else:
|
||||
cfg.DIFFUSION_MODEL.RELOAD_MODEL = diffusion_model_path
|
||||
return cfg
|
||||
|
||||
def init_from_modules(self, modules):
|
||||
for k, v in modules.items():
|
||||
self.__setattr__(k, v)
|
||||
|
||||
def infer_model(self, cfg, module_paras=None):
|
||||
module = {
|
||||
'model': None,
|
||||
'cfg': cfg,
|
||||
'device': 'offline',
|
||||
'name': cfg.NAME,
|
||||
'function_info': {},
|
||||
'paras': {}
|
||||
}
|
||||
if module_paras is None:
|
||||
return module
|
||||
function_info = {}
|
||||
paras = {
|
||||
k.lower(): v
|
||||
for k, v in module_paras.get('PARAS', {}).items()
|
||||
}
|
||||
for function in module_paras.get('FUNCTION', []):
|
||||
input_dict = {}
|
||||
for inp in function.get('INPUT', []):
|
||||
if inp.lower() in self.input:
|
||||
input_dict[inp.lower()] = self.input[inp.lower()]
|
||||
function_info[function.NAME] = {
|
||||
'dtype': function.get('DTYPE', 'float32'),
|
||||
'input': input_dict
|
||||
}
|
||||
module['paras'] = paras
|
||||
module['function_info'] = function_info
|
||||
return module
|
||||
|
||||
def init_from_ckpt(self, path, model, ignore_keys=list()):
|
||||
if path.endswith('safetensors'):
|
||||
from safetensors.torch import load_file as load_safetensors
|
||||
sd = load_safetensors(path)
|
||||
else:
|
||||
sd = torch.load(path, map_location='cpu')
|
||||
|
||||
new_sd = OrderedDict()
|
||||
for k, v in sd.items():
|
||||
ignored = False
|
||||
for ik in ignore_keys:
|
||||
if ik in k:
|
||||
if we.rank == 0:
|
||||
self.logger.info(
|
||||
'Ignore key {} from state_dict.'.format(k))
|
||||
ignored = True
|
||||
break
|
||||
if not ignored:
|
||||
new_sd[k] = v
|
||||
|
||||
missing, unexpected = model.load_state_dict(new_sd, strict=False)
|
||||
if we.rank == 0:
|
||||
self.logger.info(
|
||||
f'Restored from {path} 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}')
|
||||
|
||||
def load(self, module):
|
||||
if module['device'] == 'offline':
|
||||
if module['cfg'].NAME in MODELS.class_map:
|
||||
model = MODELS.build(module['cfg'], logger=self.logger).eval()
|
||||
elif module['cfg'].NAME in BACKBONES.class_map:
|
||||
model = BACKBONES.build(module['cfg'],
|
||||
logger=self.logger).eval()
|
||||
elif module['cfg'].NAME in EMBEDDERS.class_map:
|
||||
model = EMBEDDERS.build(module['cfg'],
|
||||
logger=self.logger).eval()
|
||||
else:
|
||||
raise NotImplementedError
|
||||
if module['cfg'].get('RELOAD_MODEL', None):
|
||||
self.init_from_ckpt(module['cfg'].RELOAD_MODEL, model)
|
||||
module['model'] = model
|
||||
module['device'] = 'cpu'
|
||||
if module['device'] == 'cpu':
|
||||
module['device'] = we.device_id
|
||||
module['model'] = module['model'].to(we.device_id)
|
||||
return module
|
||||
|
||||
def unload(self, module):
|
||||
module['model'] = module['model'].to('cpu')
|
||||
module['device'] = 'cpu'
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
return module
|
||||
|
||||
def load_default(self, cfg):
|
||||
module_paras = {}
|
||||
if cfg is not None:
|
||||
self.paras = cfg.PARAS
|
||||
self.input = {k.lower(): 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
|
||||
|
||||
def load_schedule(self, cfg):
|
||||
parameterization = cfg.get('PARAMETERIZATION', 'eps')
|
||||
assert parameterization in [
|
||||
'eps', 'x0', 'v'
|
||||
], 'currently only supporting "eps" and "x0" and "v"'
|
||||
num_timesteps = cfg.get('TIMESTEPS', 1000)
|
||||
|
||||
schedule_args = {
|
||||
k.lower(): v
|
||||
for k, v in cfg.get('SCHEDULE_ARGS', {
|
||||
'NAME': 'logsnr_cosine_interp',
|
||||
'SCALE_MIN': 2.0,
|
||||
'SCALE_MAX': 4.0
|
||||
}).items()
|
||||
}
|
||||
|
||||
zero_terminal_snr = cfg.get('ZERO_TERMINAL_SNR', False)
|
||||
if zero_terminal_snr:
|
||||
assert parameterization == 'v', 'Now zero_terminal_snr only support v-prediction mode.'
|
||||
sigmas = noise_schedule(schedule=schedule_args.pop('name'),
|
||||
n=num_timesteps,
|
||||
zero_terminal_snr=zero_terminal_snr,
|
||||
**schedule_args)
|
||||
diffusion = GaussianDiffusion(sigmas=sigmas,
|
||||
prediction_type=parameterization)
|
||||
return diffusion
|
||||
|
||||
def get_batch(self, value_dict, num_samples=1):
|
||||
batch = {}
|
||||
batch_uc = {}
|
||||
N = num_samples
|
||||
device = we.device_id
|
||||
for key in value_dict:
|
||||
if key == 'prompt':
|
||||
if not self.tokenizer:
|
||||
batch['prompt'] = value_dict['prompt']
|
||||
batch_uc['prompt'] = value_dict['negative_prompt']
|
||||
else:
|
||||
batch['tokens'] = self.tokenizer(value_dict['prompt']).to(
|
||||
we.device_id)
|
||||
batch_uc['tokens'] = self.tokenizer(
|
||||
value_dict['negative_prompt']).to(we.device_id)
|
||||
elif key == 'original_size_as_tuple':
|
||||
batch['original_size_as_tuple'] = (torch.tensor(
|
||||
value_dict['original_size_as_tuple']).to(device).repeat(
|
||||
N, 1))
|
||||
elif key == 'crop_coords_top_left':
|
||||
batch['crop_coords_top_left'] = (torch.tensor(
|
||||
value_dict['crop_coords_top_left']).to(device).repeat(
|
||||
N, 1))
|
||||
elif key == 'aesthetic_score':
|
||||
batch['aesthetic_score'] = (torch.tensor(
|
||||
[value_dict['aesthetic_score']]).to(device).repeat(N, 1))
|
||||
batch_uc['aesthetic_score'] = (torch.tensor([
|
||||
value_dict['negative_aesthetic_score']
|
||||
]).to(device).repeat(N, 1))
|
||||
|
||||
elif key == 'target_size_as_tuple':
|
||||
batch['target_size_as_tuple'] = (torch.tensor(
|
||||
value_dict['target_size_as_tuple']).to(device).repeat(
|
||||
N, 1))
|
||||
elif key == 'image':
|
||||
batch[key] = self.load_image(value_dict[key], num_samples=N)
|
||||
else:
|
||||
batch[key] = value_dict[key]
|
||||
|
||||
for key in batch.keys():
|
||||
if key not in batch_uc and isinstance(batch[key], torch.Tensor):
|
||||
batch_uc[key] = torch.clone(batch[key])
|
||||
return batch, batch_uc
|
||||
|
||||
def load_image(self, image, num_samples=1):
|
||||
if isinstance(image, torch.Tensor):
|
||||
pass
|
||||
elif isinstance(image, Image):
|
||||
pass
|
||||
elif isinstance(image, Image):
|
||||
pass
|
||||
|
||||
def get_function_info(self, module, function_name=None):
|
||||
all_function = module['function_info']
|
||||
if function_name in all_function:
|
||||
return function_name, all_function[function_name]['dtype']
|
||||
if function_name is None and len(all_function) == 1:
|
||||
for k, v in all_function.items():
|
||||
return k, v['dtype']
|
||||
|
||||
def encode_first_stage(self, x, **kwargs):
|
||||
_, dtype = self.get_function_info(self.first_stage_model, 'encode')
|
||||
with torch.autocast('cuda',
|
||||
enabled=dtype == 'float16',
|
||||
dtype=getattr(torch, dtype)):
|
||||
z = get_model(self.first_stage_model).encode(x)
|
||||
return self.first_stage_model['paras']['scale_factor'] * z
|
||||
|
||||
def decode_first_stage(self, z):
|
||||
_, dtype = self.get_function_info(self.first_stage_model, 'encode')
|
||||
with torch.autocast('cuda',
|
||||
enabled=dtype == 'float16',
|
||||
dtype=getattr(torch, dtype)):
|
||||
z = 1. / self.first_stage_model['paras']['scale_factor'] * z
|
||||
return get_model(self.first_stage_model).decode(z)
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(self,
|
||||
input,
|
||||
num_samples=1,
|
||||
intermediate_callback=None,
|
||||
refine_strength=0,
|
||||
img_to_img_strength=0,
|
||||
cat_uc=True,
|
||||
tuner_model=None,
|
||||
control_model=None,
|
||||
**kwargs):
|
||||
|
||||
value_input = copy.deepcopy(self.input)
|
||||
value_input.update(input)
|
||||
print(value_input)
|
||||
height, width = value_input['target_size_as_tuple']
|
||||
value_output = copy.deepcopy(self.output)
|
||||
batch, batch_uc = self.get_batch(value_input, num_samples=1)
|
||||
#
|
||||
if not isinstance(tuner_model, list):
|
||||
tuner_model = [tuner_model]
|
||||
for tuner in tuner_model:
|
||||
if tuner is None or tuner == '':
|
||||
tuner_model.remove(tuner)
|
||||
self.register_tuner(tuner_model)
|
||||
# control_cond_image
|
||||
control_cond_image = kwargs.pop('control_cond_image', None)
|
||||
# crop_type = kwargs.pop('crop_type', 'center_crop')
|
||||
hints = []
|
||||
if control_cond_image and control_model:
|
||||
if not isinstance(control_model, list):
|
||||
control_model = [control_model]
|
||||
if not isinstance(control_cond_image, list):
|
||||
control_cond_image = [control_cond_image]
|
||||
assert len(control_cond_image) == len(control_model)
|
||||
for img in control_cond_image:
|
||||
if isinstance(img, Image):
|
||||
w, h = img.size
|
||||
if not h == height or not w == width:
|
||||
img = TT.Resize(min(height, width))(img)
|
||||
img = TT.CenterCrop((height, width))(img)
|
||||
hint = TT.ToTensor()(img)
|
||||
hints.append(hint)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
if len(hints) > 0:
|
||||
hints = torch.stack(hints).to(we.device_id)
|
||||
else:
|
||||
hints = None
|
||||
|
||||
# first stage encode
|
||||
image = input.pop('image', None)
|
||||
if image is not None and img_to_img_strength > 0:
|
||||
# run image2image
|
||||
b, c, ori_width, ori_height = image.shape
|
||||
if not (ori_width == width and ori_height == height):
|
||||
image = F.interpolate(image, (width, height), mode='bicubic')
|
||||
self.first_stage_model = self.load(self.first_stage_model)
|
||||
input_latent = self.encode_first_stage(image)
|
||||
self.first_stage_model = self.unload(self.first_stage_model)
|
||||
else:
|
||||
input_latent = None
|
||||
if 'input_latent' in value_output and input_latent is not None:
|
||||
value_output['input_latent'] = input_latent
|
||||
# cond stage
|
||||
self.cond_stage_model = self.load(self.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)):
|
||||
if self.tokenizer:
|
||||
context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch['tokens'])
|
||||
null_context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch_uc['tokens'])
|
||||
else:
|
||||
context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch)
|
||||
null_context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch_uc)
|
||||
self.cond_stage_model = self.unload(self.cond_stage_model)
|
||||
|
||||
if refine_strength > 0 and self.refiner_diffusion_model is not None:
|
||||
assert self.refiner_cond_model is not None
|
||||
self.refiner_cond_model = self.load(self.refiner_cond_model)
|
||||
function_name, dtype = self.get_function_info(
|
||||
self.refiner_cond_model)
|
||||
with torch.autocast('cuda',
|
||||
enabled=dtype == 'float16',
|
||||
dtype=getattr(torch, dtype)):
|
||||
if self.tokenizer:
|
||||
refine_context = getattr(
|
||||
get_model(self.refiner_cond_model),
|
||||
function_name)(batch['tokens'])
|
||||
refine_null_context = getattr(
|
||||
get_model(self.refiner_cond_model),
|
||||
function_name)(batch_uc['tokens'])
|
||||
else:
|
||||
refine_context = getattr(
|
||||
get_model(self.refiner_cond_model),
|
||||
function_name)(batch)
|
||||
refine_null_context = getattr(
|
||||
get_model(self.refiner_cond_model),
|
||||
function_name)(batch_uc)
|
||||
self.refiner_cond_model = self.unload(self.refiner_cond_model)
|
||||
self.load(self.diffusion_model)
|
||||
self.register_controllers(control_model)
|
||||
self.unload(self.diffusion_model)
|
||||
# get noise
|
||||
seed = kwargs.pop('seed', -1)
|
||||
g = torch.Generator(device=we.device_id)
|
||||
seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
|
||||
g.manual_seed(seed)
|
||||
if 'seed' in value_output:
|
||||
value_output['seed'] = seed
|
||||
for sample_id in range(num_samples):
|
||||
if self.diffusion_model is not None:
|
||||
noise = torch.empty(
|
||||
1,
|
||||
4,
|
||||
height // self.first_stage_model['paras']['size_factor'],
|
||||
width // self.first_stage_model['paras']['size_factor'],
|
||||
device=we.device_id).normal_(generator=g)
|
||||
#
|
||||
self.load(self.diffusion_model)
|
||||
# UNet use input n_prompt
|
||||
function_name, dtype = self.get_function_info(
|
||||
self.diffusion_model)
|
||||
with torch.autocast('cuda',
|
||||
enabled=dtype == 'float16',
|
||||
dtype=getattr(torch, dtype)):
|
||||
latent = self.diffusion.sample(
|
||||
noise=noise,
|
||||
x=input_latent,
|
||||
denoising_strength=img_to_img_strength
|
||||
if input_latent is not None else 1.0,
|
||||
refine_strength=refine_strength,
|
||||
solver=value_input.get('sample', 'ddim'),
|
||||
model=get_model(self.diffusion_model),
|
||||
model_kwargs=[{
|
||||
'cond': context,
|
||||
'hint': hints
|
||||
}, {
|
||||
'cond': null_context,
|
||||
'hint': hints
|
||||
}],
|
||||
steps=value_input.get('sample_steps', 50),
|
||||
guide_scale=value_input.get('guide_scale', 7.5),
|
||||
guide_rescale=value_input.get('guide_rescale', 0.5),
|
||||
discretization=value_input.get('discretization',
|
||||
'trailing'),
|
||||
show_progress=True,
|
||||
seed=seed,
|
||||
condition_fn=None,
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
t_max=None,
|
||||
t_min=None,
|
||||
discard_penultimate_step=None,
|
||||
intermediate_callback=intermediate_callback,
|
||||
cat_uc=cat_uc,
|
||||
**kwargs)
|
||||
self.diffusion_model = self.unload(self.diffusion_model)
|
||||
|
||||
# apply refiner
|
||||
if refine_strength > 0 and self.refiner_diffusion_model is not None:
|
||||
assert self.refiner_diffusion_model is not None
|
||||
# decode intermidiet latent before refine
|
||||
self.first_stage_model = self.load(self.first_stage_model)
|
||||
before_refiner_samples = self.decode_first_stage(
|
||||
latent).float()
|
||||
self.first_stage_model = self.unload(self.first_stage_model)
|
||||
|
||||
before_refiner_samples = torch.clamp(
|
||||
(before_refiner_samples + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
if 'before_refine_images' in value_output:
|
||||
if value_output['before_refine_images'] is None or (
|
||||
isinstance(value_output['before_refine_images'],
|
||||
list)
|
||||
and len(value_output['before_refine_images']) < 1):
|
||||
value_output['before_refine_images'] = []
|
||||
value_output['before_refine_images'].append(
|
||||
before_refiner_samples)
|
||||
self.refiner_model = self.load(self.refiner_diffusion_model)
|
||||
function_name, dtype = self.get_function_info(
|
||||
self.refiner_model)
|
||||
with torch.autocast('cuda',
|
||||
enabled=dtype == 'float16',
|
||||
dtype=getattr(torch, dtype)):
|
||||
latent = self.diffusion.sample(
|
||||
noise=noise,
|
||||
x=latent,
|
||||
denoising_strength=img_to_img_strength
|
||||
if input_latent is not None else 1.0,
|
||||
refine_strength=refine_strength,
|
||||
refine_stage=True,
|
||||
solver=value_input.get('refine_sample', 'ddim'),
|
||||
model=get_model(self.refiner_model),
|
||||
model_kwargs=[{
|
||||
'cond': refine_context
|
||||
}, {
|
||||
'cond': refine_null_context
|
||||
}],
|
||||
steps=value_input.get('sample_steps', 50),
|
||||
guide_scale=value_input.get('refine_guide_scale', 7.5),
|
||||
guide_rescale=value_input.get('refine_guide_rescale',
|
||||
0.5),
|
||||
discretization=value_input.get('refine_discretization',
|
||||
'trailing'),
|
||||
show_progress=True,
|
||||
seed=seed,
|
||||
condition_fn=None,
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
t_max=None,
|
||||
t_min=None,
|
||||
discard_penultimate_step=None,
|
||||
return_intermediate=None,
|
||||
intermediate_callback=intermediate_callback,
|
||||
cat_uc=cat_uc,
|
||||
**kwargs)
|
||||
self.refiner_model = self.unload(self.refiner_model)
|
||||
|
||||
if 'latent' in value_output:
|
||||
if value_output['latent'] is None or (
|
||||
isinstance(value_output['latent'], list)
|
||||
and len(value_output['latent']) < 1):
|
||||
value_output['latent'] = []
|
||||
value_output['latent'].append(latent)
|
||||
|
||||
self.first_stage_model = self.load(self.first_stage_model)
|
||||
x_samples = self.decode_first_stage(latent).float()
|
||||
self.first_stage_model = self.unload(self.first_stage_model)
|
||||
|
||||
images = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
if 'images' in value_output:
|
||||
if value_output['images'] is None or (
|
||||
isinstance(value_output['images'], list)
|
||||
and len(value_output['images']) < 1):
|
||||
value_output['images'] = []
|
||||
value_output['images'].append(images)
|
||||
|
||||
for k, v in value_output.items():
|
||||
if isinstance(v, list):
|
||||
value_output[k] = torch.cat(v, dim=0)
|
||||
if isinstance(v, torch.Tensor):
|
||||
value_output[k] = v.cpu()
|
||||
return value_output
|
||||
@@ -190,6 +190,7 @@ class DiffusionUNet(BaseModel):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.init_params(cfg)
|
||||
self.construct_network()
|
||||
self.control_blocks = None
|
||||
|
||||
def init_params(self, cfg):
|
||||
self.in_channels = cfg.IN_CHANNELS
|
||||
@@ -285,6 +286,7 @@ class DiffusionUNet(BaseModel):
|
||||
])
|
||||
self._feature_size = model_channels
|
||||
input_block_chans = [model_channels]
|
||||
input_down_flag = [False]
|
||||
ch = model_channels
|
||||
ds = 1
|
||||
for level, mult in enumerate(channel_mult):
|
||||
@@ -323,6 +325,7 @@ class DiffusionUNet(BaseModel):
|
||||
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
||||
self._feature_size += ch
|
||||
input_block_chans.append(ch)
|
||||
input_down_flag.append(False)
|
||||
if level != len(channel_mult) - 1:
|
||||
out_ch = ch
|
||||
self.input_blocks.append(
|
||||
@@ -341,9 +344,11 @@ class DiffusionUNet(BaseModel):
|
||||
)
|
||||
ch = out_ch
|
||||
input_block_chans.append(ch)
|
||||
input_down_flag.append(True)
|
||||
ds *= 2
|
||||
self._feature_size += ch
|
||||
self._input_block_chans = copy.deepcopy(input_block_chans)
|
||||
self._input_down_flag = input_down_flag
|
||||
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
@@ -479,22 +484,7 @@ class DiffusionUNet(BaseModel):
|
||||
if len(unexpected) > 0:
|
||||
self.logger.info(f'\nUnexpected Keys:\n {unexpected}')
|
||||
|
||||
def forward(self, x, t=None, cond=dict()):
|
||||
t_emb = timestep_embedding(t, self.model_channels, repeat_only=False)
|
||||
emb = self.time_embed(t_emb)
|
||||
|
||||
if isinstance(cond, dict):
|
||||
if 'y' in cond and cond['y'] is not None:
|
||||
assert self.num_classes is not None
|
||||
emb = emb + self.label_emb(cond['y'])
|
||||
if 'concat' in cond:
|
||||
c = cond['concat']
|
||||
x = torch.cat([x, c], dim=1)
|
||||
|
||||
context = cond.get('crossattn', None)
|
||||
else:
|
||||
context = cond
|
||||
|
||||
def _forward_origin(self, x, emb, context, hint=None):
|
||||
hs = []
|
||||
h = x
|
||||
for module in self.input_blocks:
|
||||
@@ -505,8 +495,76 @@ class DiffusionUNet(BaseModel):
|
||||
h = torch.cat([h, self.lsc_identity[m_id](hs.pop())], dim=1)
|
||||
target_size = hs[-1].shape[-2:] if len(hs) > 0 else None
|
||||
h = module(h, emb, context, target_size)
|
||||
out = self.out(h)
|
||||
return out
|
||||
|
||||
return self.out(h)
|
||||
def _forward_control(self, x, emb, context, hint, alpha=0.5):
|
||||
multi_csc_tuners = self.control_blocks
|
||||
# hints
|
||||
multi_hint_hs = []
|
||||
for sc_id, csc_tuners in enumerate(multi_csc_tuners):
|
||||
hint_input = hint[sc_id] if isinstance(hint, list) else hint
|
||||
hint_h = csc_tuners.pre_hint_blocks(hint_input)
|
||||
hint_hs = []
|
||||
for dsh_blk in csc_tuners.dense_hint_blocks:
|
||||
hint_h = dsh_blk(hint_h)
|
||||
hint_hs.append(hint_h)
|
||||
multi_hint_hs.append(hint_hs)
|
||||
# unet
|
||||
hs = []
|
||||
h = x
|
||||
for module in self.input_blocks:
|
||||
h = module(h, emb, context)
|
||||
hs.append(h)
|
||||
h = self.middle_block(h, emb, context)
|
||||
for m_id, module in enumerate(self.output_blocks):
|
||||
skip_h = hs.pop()
|
||||
multi_control_h = 0
|
||||
for sc_id, csc_tuners in enumerate(multi_csc_tuners):
|
||||
hint_h = multi_hint_hs[sc_id][::-1][m_id]
|
||||
control_h = csc_tuners.lsc_tuner_blocks[m_id](
|
||||
skip_h + hint_h, x_shortcut=hint_h)
|
||||
multi_control_h += csc_tuners.scale * control_h
|
||||
tuner_h = self.lsc_identity[m_id](skip_h) - skip_h
|
||||
if torch.all(
|
||||
torch.isclose(tuner_h,
|
||||
torch.zeros_like(tuner_h),
|
||||
atol=1e-6)):
|
||||
# csc-tuner
|
||||
skip_h_new = skip_h + multi_control_h
|
||||
else:
|
||||
# csc-tuner + sc-tuner
|
||||
skip_h_new = skip_h + alpha * multi_control_h + (
|
||||
1 - alpha) * tuner_h
|
||||
h = torch.cat([h, skip_h_new], dim=1)
|
||||
target_size = hs[-1].shape[-2:] if len(hs) > 0 else None
|
||||
h = module(h, emb, context, target_size)
|
||||
out = self.out(h)
|
||||
return out
|
||||
|
||||
def forward(self, x, t=None, cond=dict(), **kwargs):
|
||||
t_emb = timestep_embedding(t, self.model_channels, repeat_only=False)
|
||||
emb = self.time_embed(t_emb)
|
||||
hint = None
|
||||
if isinstance(cond, dict):
|
||||
if 'y' in cond and cond['y'] is not None:
|
||||
assert self.num_classes is not None
|
||||
emb = emb + self.label_emb(cond['y'])
|
||||
if 'concat' in cond:
|
||||
c = cond['concat']
|
||||
x = torch.cat([x, c], dim=1)
|
||||
if 'hint' in cond:
|
||||
hint = cond['hint']
|
||||
context = cond.get('crossattn', None)
|
||||
else:
|
||||
context = cond
|
||||
hint = kwargs.pop('hint', None)
|
||||
|
||||
if self.control_blocks is not None:
|
||||
out = self._forward_control(x, emb, context, hint)
|
||||
else:
|
||||
out = self._forward_origin(x, emb, context)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
@@ -602,6 +660,7 @@ class DiffusionUNetXL(DiffusionUNet):
|
||||
])
|
||||
self._feature_size = model_channels
|
||||
input_block_chans = [model_channels]
|
||||
input_down_flag = [False]
|
||||
ch = model_channels
|
||||
ds = 1
|
||||
for level, mult in enumerate(channel_mult):
|
||||
@@ -640,6 +699,7 @@ class DiffusionUNetXL(DiffusionUNet):
|
||||
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
||||
self._feature_size += ch
|
||||
input_block_chans.append(ch)
|
||||
input_down_flag.append(False)
|
||||
if level != len(channel_mult) - 1:
|
||||
out_ch = ch
|
||||
self.input_blocks.append(
|
||||
@@ -658,9 +718,11 @@ class DiffusionUNetXL(DiffusionUNet):
|
||||
)
|
||||
ch = out_ch
|
||||
input_block_chans.append(ch)
|
||||
input_down_flag.append(True)
|
||||
ds *= 2
|
||||
self._feature_size += ch
|
||||
self._input_block_chans = copy.deepcopy(input_block_chans)
|
||||
self._input_down_flag = input_down_flag
|
||||
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
@@ -760,25 +822,7 @@ class DiffusionUNetXL(DiffusionUNet):
|
||||
conv_nd(dims, model_channels, out_channels, 3, padding=1)),
|
||||
)
|
||||
|
||||
def forward(self, x, t=None, cond=dict()):
|
||||
t_emb = timestep_embedding(t,
|
||||
self.model_channels,
|
||||
repeat_only=False,
|
||||
legacy=True)
|
||||
emb = self.time_embed(t_emb)
|
||||
|
||||
if isinstance(cond, dict):
|
||||
if 'y' in cond:
|
||||
assert self.num_classes is not None
|
||||
emb = emb + self.label_emb(cond['y'])
|
||||
if 'concat' in cond:
|
||||
c = cond['concat']
|
||||
x = torch.cat([x, c], dim=1)
|
||||
|
||||
context = cond.get('crossattn', None)
|
||||
else:
|
||||
context = cond
|
||||
|
||||
def _forward_origin(self, x, emb, context, hint=None):
|
||||
hs = []
|
||||
h = x
|
||||
for module in self.input_blocks:
|
||||
@@ -789,8 +833,79 @@ class DiffusionUNetXL(DiffusionUNet):
|
||||
h = torch.cat([h, self.lsc_identity[m_id](hs.pop())], dim=1)
|
||||
target_size = hs[-1].shape[-2:] if len(hs) > 0 else None
|
||||
h = module(h, emb, context, target_size)
|
||||
out = self.out(h)
|
||||
return out
|
||||
|
||||
return self.out(h)
|
||||
def _forward_control(self, x, emb, context, hint, alpha=0.5):
|
||||
multi_csc_tuners = self.control_blocks
|
||||
# hints
|
||||
multi_hint_hs = []
|
||||
for sc_id, csc_tuners in enumerate(multi_csc_tuners):
|
||||
hint_input = hint[sc_id] if isinstance(hint, list) else hint
|
||||
hint_h = csc_tuners.pre_hint_blocks(hint_input)
|
||||
hint_hs = []
|
||||
for dsh_blk in csc_tuners.dense_hint_blocks:
|
||||
hint_h = dsh_blk(hint_h)
|
||||
hint_hs.append(hint_h)
|
||||
multi_hint_hs.append(hint_hs)
|
||||
# unet
|
||||
hs = []
|
||||
h = x
|
||||
for module in self.input_blocks:
|
||||
h = module(h, emb, context)
|
||||
hs.append(h)
|
||||
h = self.middle_block(h, emb, context)
|
||||
for m_id, module in enumerate(self.output_blocks):
|
||||
skip_h = hs.pop()
|
||||
multi_control_h = 0
|
||||
for sc_id, csc_tuners in enumerate(multi_csc_tuners):
|
||||
hint_h = multi_hint_hs[sc_id][::-1][m_id]
|
||||
control_h = csc_tuners.lsc_tuner_blocks[m_id](
|
||||
skip_h + hint_h, x_shortcut=hint_h)
|
||||
multi_control_h += csc_tuners.scale * control_h
|
||||
tuner_h = self.lsc_identity[m_id](skip_h) - skip_h
|
||||
if torch.all(
|
||||
torch.isclose(tuner_h,
|
||||
torch.zeros_like(tuner_h),
|
||||
atol=1e-6)):
|
||||
# csc-tuner
|
||||
skip_h_new = skip_h + multi_control_h
|
||||
else:
|
||||
# csc-tuner + sc-tuner
|
||||
skip_h_new = skip_h + alpha * multi_control_h + (
|
||||
1 - alpha) * tuner_h
|
||||
h = torch.cat([h, skip_h_new], dim=1)
|
||||
target_size = hs[-1].shape[-2:] if len(hs) > 0 else None
|
||||
h = module(h, emb, context, target_size)
|
||||
out = self.out(h)
|
||||
return out
|
||||
|
||||
def forward(self, x, t=None, cond=dict(), **kwargs):
|
||||
t_emb = timestep_embedding(t,
|
||||
self.model_channels,
|
||||
repeat_only=False,
|
||||
legacy=True)
|
||||
emb = self.time_embed(t_emb)
|
||||
hint = None
|
||||
if isinstance(cond, dict):
|
||||
if 'y' in cond:
|
||||
assert self.num_classes is not None
|
||||
emb = emb + self.label_emb(cond['y'])
|
||||
if 'concat' in cond:
|
||||
c = cond['concat']
|
||||
x = torch.cat([x, c], dim=1)
|
||||
if 'hint' in cond:
|
||||
hint = cond['hint']
|
||||
context = cond.get('crossattn', None)
|
||||
else:
|
||||
context = cond
|
||||
hint = kwargs.pop('hint', None)
|
||||
|
||||
if self.control_blocks is not None:
|
||||
out = self._forward_control(x, emb, context, hint)
|
||||
else:
|
||||
out = self._forward_origin(x, emb, context)
|
||||
return out
|
||||
|
||||
def convert_to_fp16(self):
|
||||
"""
|
||||
|
||||
@@ -224,16 +224,13 @@ class FrozenOpenCLIPEmbedder(BaseEmbedder):
|
||||
super().__init__(cfg, logger=logger)
|
||||
|
||||
arch = cfg.get('ARCH', 'ViT-H-14')
|
||||
if cfg.PRETRAINED_MODEL is None:
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=None)
|
||||
del model.visual
|
||||
else:
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=None)
|
||||
del model.visual
|
||||
if cfg.PRETRAINED_MODEL is not None:
|
||||
with FS.get_from(cfg.PRETRAINED_MODEL,
|
||||
wait_finish=True) as local_path:
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=local_path)
|
||||
del model.visual
|
||||
model.load_state_dict(torch.load(local_path), strict=False)
|
||||
self.model = model
|
||||
|
||||
self.use_grad = cfg.get('USE_GRAD', False)
|
||||
@@ -362,7 +359,7 @@ class FrozenOpenCLIPEmbedder2(BaseEmbedder):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
arch = cfg.get('ARCH', 'ViT-H-14')
|
||||
if cfg.PRETRAINED_MODEL is None:
|
||||
if cfg.get('PRETRAINED_MODEL', None) is None:
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=None)
|
||||
del model.visual
|
||||
|
||||
@@ -4,4 +4,4 @@ from scepter.modules.model.network.autoencoder import ae_kl
|
||||
from scepter.modules.model.network.classifier import Classifier
|
||||
from scepter.modules.model.network.diffusion import (diffusion, schedules,
|
||||
solvers)
|
||||
from scepter.modules.model.network.ldm import ldm, ldm_xl
|
||||
from scepter.modules.model.network.ldm import ldm, ldm_sce, ldm_xl
|
||||
|
||||
@@ -159,7 +159,8 @@ class GaussianDiffusion(object):
|
||||
model,
|
||||
model_kwargs={},
|
||||
reduction='mean',
|
||||
noise=None):
|
||||
noise=None,
|
||||
**kwargs):
|
||||
# hyperparams
|
||||
sigmas = _i(self.sigmas, t, x0)
|
||||
alphas = _i(self.alphas, t, x0)
|
||||
@@ -168,7 +169,7 @@ class GaussianDiffusion(object):
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x0)
|
||||
xt = self.diffuse(x0, t, noise)
|
||||
out = model(xt, t=t, **model_kwargs)
|
||||
out = model(xt, t=t, **model_kwargs, **kwargs)
|
||||
|
||||
# mse loss
|
||||
target = {
|
||||
|
||||
@@ -54,7 +54,8 @@ def sample_euler(noise,
|
||||
s_tmax=float('inf'),
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Implements Algorithm 2 (Euler steps) from Karras et al. (2022).
|
||||
"""
|
||||
@@ -87,7 +88,8 @@ def sample_euler_ancestral(noise,
|
||||
eta=1.,
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Ancestral sampling with Euler method steps.
|
||||
"""
|
||||
@@ -120,7 +122,8 @@ def sample_heun(noise,
|
||||
s_tmax=float('inf'),
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Implements Algorithm 2 (Heun steps) from Karras et al. (2022).
|
||||
"""
|
||||
@@ -165,7 +168,8 @@ def sample_dpm_2(noise,
|
||||
s_tmax=float('inf'),
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).
|
||||
"""
|
||||
@@ -211,7 +215,8 @@ def sample_dpm_2_ancestral(noise,
|
||||
eta=1.,
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Ancestral sampling with DPM-Solver second-order steps.
|
||||
"""
|
||||
@@ -253,7 +258,8 @@ def sample_dpmpp_2s_ancestral(noise,
|
||||
eta=1.,
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Ancestral sampling with DPM-Solver++ (2S) second-order steps.
|
||||
"""
|
||||
@@ -372,7 +378,8 @@ def sample_dpmpp_sde(noise,
|
||||
s_noise=1.,
|
||||
r=1 / 2,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
DPM-Solver++ (stochastic).
|
||||
"""
|
||||
@@ -429,7 +436,12 @@ def sample_dpmpp_sde(noise,
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_dpmpp_2m(noise, model, sigmas, seed=None, show_progress=True):
|
||||
def sample_dpmpp_2m(noise,
|
||||
model,
|
||||
sigmas,
|
||||
seed=None,
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
DPM-Solver++ (2M).
|
||||
"""
|
||||
@@ -475,7 +487,8 @@ def sample_dpmpp_2m_sde(noise,
|
||||
s_noise=1.,
|
||||
solver_type='midpoint',
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
DPM-Solver++ (2M) SDE.
|
||||
"""
|
||||
@@ -527,7 +540,13 @@ def sample_dpmpp_2m_sde(noise,
|
||||
|
||||
# -------------------- variation preserving (VP) solver --------------------#
|
||||
@torch.no_grad()
|
||||
def sample_ddim(noise, model, sigmas, eta=0., seed=None, show_progress=True):
|
||||
def sample_ddim(noise,
|
||||
model,
|
||||
sigmas,
|
||||
eta=0.,
|
||||
seed=None,
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
DDIM solver steps.
|
||||
"""
|
||||
@@ -556,7 +575,8 @@ def sample_img2img_euler(noise,
|
||||
s_tmax=float('inf'),
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Implements Algorithm 2 (Euler steps) from Karras et al. (2022).
|
||||
"""
|
||||
@@ -588,7 +608,8 @@ def sample_img2img_euler_ancestral(noise,
|
||||
eta=1.,
|
||||
s_noise=1.,
|
||||
seed=None,
|
||||
show_progress=True):
|
||||
show_progress=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Ancestral sampling with Euler method steps.
|
||||
"""
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.network.ldm.ldm import LatentDiffusion
|
||||
from scepter.modules.model.network.ldm.ldm_sce import (
|
||||
LatentDiffusionSCEControl, LatentDiffusionSCETuning,
|
||||
LatentDiffusionXLSCEControl, LatentDiffusionXLSCETuning)
|
||||
from scepter.modules.model.network.ldm.ldm_xl import LatentDiffusionXL
|
||||
|
||||
@@ -241,6 +241,14 @@ class LatentDiffusion(TrainModule):
|
||||
with torch.autocast(device_type='cuda', enabled=False):
|
||||
context = self.encode_condition(
|
||||
self.tokenizer(prompt).to(we.device_id))
|
||||
if 'hint' in kwargs and kwargs['hint'] is not None:
|
||||
hint = kwargs.pop('hint')
|
||||
if isinstance(context, dict):
|
||||
context['hint'] = hint
|
||||
else:
|
||||
context = {'crossattn': context, 'hint': hint}
|
||||
else:
|
||||
hint = None
|
||||
if self.min_snr_gamma is not None:
|
||||
alphas = self.diffusion.alphas.to(we.device_id)[t]
|
||||
sigmas = self.diffusion.sigmas.pow(2).to(we.device_id)[t]
|
||||
@@ -250,11 +258,13 @@ class LatentDiffusion(TrainModule):
|
||||
else:
|
||||
weights = 1
|
||||
self.register_probe({'snrs_weights': weights})
|
||||
|
||||
loss = self.diffusion.loss(x0=x_start,
|
||||
t=t,
|
||||
model=self.model,
|
||||
model_kwargs={'cond': context},
|
||||
noise=noise)
|
||||
noise=noise,
|
||||
**kwargs)
|
||||
loss = loss * weights
|
||||
loss = loss.mean()
|
||||
ret = {'loss': loss, 'probe_data': {'prompt': prompt}}
|
||||
@@ -305,7 +315,18 @@ class LatentDiffusion(TrainModule):
|
||||
null_context = self.encode_condition(self.tokenizer(n_prompt).to(
|
||||
we.device_id),
|
||||
method=method)
|
||||
|
||||
if 'hint' in kwargs and kwargs['hint'] is not None:
|
||||
hint = kwargs.pop('hint')
|
||||
if isinstance(context, dict):
|
||||
context['hint'] = hint
|
||||
else:
|
||||
context = {'crossattn': context, 'hint': hint}
|
||||
if isinstance(null_context, dict):
|
||||
null_context['hint'] = hint
|
||||
else:
|
||||
null_context = {'crossattn': null_context, 'hint': hint}
|
||||
else:
|
||||
hint = None
|
||||
if 'index' in kwargs:
|
||||
kwargs.pop('index')
|
||||
image_size = None
|
||||
@@ -317,7 +338,9 @@ class LatentDiffusion(TrainModule):
|
||||
image_size = [h, w]
|
||||
if 'image_size' in kwargs:
|
||||
image_size = kwargs.pop('image_size')
|
||||
if image_size is None or isinstance(image_size, numbers.Number):
|
||||
if isinstance(image_size, numbers.Number):
|
||||
image_size = [image_size, image_size]
|
||||
if image_size is None:
|
||||
image_size = [1024, 1024]
|
||||
height, width = image_size
|
||||
noise = self.noise_sample(num_samples, height // self.size_factor,
|
||||
@@ -387,9 +410,11 @@ class LatentDiffusion(TrainModule):
|
||||
t_x_samples = [None for _ in prompt]
|
||||
|
||||
outputs = list()
|
||||
for p, np, tnp, img, t_img in zip(prompt, n_prompt, train_n_prompt,
|
||||
x_samples, t_x_samples):
|
||||
for i, (p, np, tnp, img, t_img) in enumerate(
|
||||
zip(prompt, n_prompt, train_n_prompt, x_samples, t_x_samples)):
|
||||
one_tup = {'prompt': p, 'n_prompt': np, 'image': img}
|
||||
if hint is not None:
|
||||
one_tup.update({'hint': hint[i]})
|
||||
if t_img is not None:
|
||||
one_tup['train_n_prompt'] = tnp
|
||||
one_tup['train_n_image'] = t_img
|
||||
@@ -408,6 +433,8 @@ class LatentDiffusion(TrainModule):
|
||||
'prompt': res['prompt'],
|
||||
'n_prompt': res['n_prompt']
|
||||
}
|
||||
if 'hint' in res:
|
||||
one_tup.update({'hint': res['hint']})
|
||||
if 'train_n_prompt' in res:
|
||||
one_tup['train_n_prompt'] = res['train_n_prompt']
|
||||
one_tup['train_n_image'] = res['train_n_image']
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import copy
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as TT
|
||||
|
||||
from scepter.modules.annotator.registry import ANNOTATORS
|
||||
from scepter.modules.model.registry import MODELS, TUNERS
|
||||
from scepter.modules.utils.config import Config, dict_to_yaml
|
||||
|
||||
from .ldm import LatentDiffusion
|
||||
from .ldm_xl import LatentDiffusionXL
|
||||
|
||||
|
||||
@MODELS.register_class()
|
||||
class LatentDiffusionSCETuning(LatentDiffusion):
|
||||
para_dict = {}
|
||||
para_dict.update(LatentDiffusion.para_dict)
|
||||
|
||||
def __init__(self, cfg, logger):
|
||||
super().__init__(cfg, logger=logger)
|
||||
|
||||
def init_params(self):
|
||||
super().init_params()
|
||||
self.tuner_model_config = self.cfg.TUNER_MODEL
|
||||
|
||||
def construct_network(self):
|
||||
super().construct_network()
|
||||
input_block_channels = self.model._input_block_chans
|
||||
sc_tuner_cfg = self.tuner_model_config['SC_TUNER_CFG']
|
||||
use_layers = self.tuner_model_config.get('USE_LAYERS', None)
|
||||
lsc_tuner_blocks = nn.ModuleList([])
|
||||
for i, chan in enumerate(input_block_channels[::-1]):
|
||||
if use_layers and i not in use_layers:
|
||||
lsc_tuner_blocks.append(nn.Identity())
|
||||
continue
|
||||
tuner_cfg = copy.deepcopy(sc_tuner_cfg)
|
||||
tuner_cfg['DIM'] = chan
|
||||
tuner_cfg['TUNER_LENGTH'] = int(chan *
|
||||
tuner_cfg.get('DOWN_RATIO', 1.0))
|
||||
sc_tuner = TUNERS.build(tuner_cfg, logger=self.logger)
|
||||
lsc_tuner_blocks.append(sc_tuner)
|
||||
self.model.lsc_identity = lsc_tuner_blocks
|
||||
|
||||
def save_pretrained(self,
|
||||
*args,
|
||||
destination=None,
|
||||
prefix='',
|
||||
keep_vars=False):
|
||||
save_state = {
|
||||
key: value
|
||||
for key, value in self.state_dict().items()
|
||||
if 'lsc_identity' in key
|
||||
}
|
||||
return save_state
|
||||
|
||||
def save_pretrained_config(self):
|
||||
return copy.deepcopy(self.cfg.TUNER_MODEL.cfg_dict)
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('MODELS',
|
||||
__class__.__name__,
|
||||
LatentDiffusionSCETuning.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@MODELS.register_class()
|
||||
class LatentDiffusionXLSCETuning(LatentDiffusionSCETuning, LatentDiffusionXL):
|
||||
pass
|
||||
|
||||
|
||||
@MODELS.register_class()
|
||||
class LatentDiffusionSCEControl(LatentDiffusion):
|
||||
para_dict = {
|
||||
'CONTROL_MODEL': {},
|
||||
}
|
||||
para_dict.update(LatentDiffusion.para_dict)
|
||||
|
||||
def __init__(self, cfg, logger):
|
||||
super().__init__(cfg, logger=logger)
|
||||
|
||||
def init_params(self):
|
||||
super().init_params()
|
||||
self.control_model_config = self.cfg.CONTROL_MODEL
|
||||
self.control_anno_config = self.cfg.CONTROL_ANNO
|
||||
|
||||
def construct_network(self):
|
||||
super().construct_network()
|
||||
# anno
|
||||
self.control_processor = ANNOTATORS.build(self.control_anno_config)
|
||||
if isinstance(self.control_model_config, (dict, Config)):
|
||||
self.control_model_config = [self.control_model_config]
|
||||
control_model = nn.ModuleList([])
|
||||
for k, sub_cfg in enumerate(self.control_model_config):
|
||||
sub_cfg['INPUT_BLOCK_CHANS'] = self.model._input_block_chans
|
||||
sub_cfg['INPUT_DOWN_FLAG'] = self.model._input_down_flag
|
||||
csc_tuners = TUNERS.build(sub_cfg, logger=self.logger)
|
||||
control_model.append(csc_tuners)
|
||||
self.model.control_blocks = control_model
|
||||
|
||||
@torch.no_grad()
|
||||
def get_control_input(self, control, *args, **kwargs):
|
||||
hints = []
|
||||
for ctr in control:
|
||||
hint = self.control_processor(ctr)
|
||||
hint = TT.ToTensor()(hint)
|
||||
hints.append(hint)
|
||||
hints = torch.stack(hints).to(control.device)
|
||||
return hints
|
||||
|
||||
def forward_train(self, **kwargs):
|
||||
# if ('module' not in kwargs) or ('module' in kwargs and self.control_method not in kwargs['module']):
|
||||
# kwargs['module'] = {self.control_method: self.control_model}
|
||||
if 'hint' not in kwargs and 'image_preprocess' in kwargs:
|
||||
image_preprocess = kwargs.pop('image_preprocess')
|
||||
kwargs['hint'] = self.get_control_input(image_preprocess)
|
||||
return super().forward_train(**kwargs)
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.autocast('cuda', dtype=torch.float16)
|
||||
def forward_test(self, **kwargs):
|
||||
# if ('module' not in kwargs) or ('module' in kwargs and self.control_method not in kwargs['module']):
|
||||
# kwargs['module'] = {self.control_method: self.control_model}
|
||||
if 'hint' not in kwargs and 'image_preprocess' in kwargs:
|
||||
image_preprocess = kwargs.pop('image_preprocess')
|
||||
kwargs['hint'] = self.get_control_input(image_preprocess)
|
||||
kwargs['image_size'] = kwargs['hint'].shape[-2:]
|
||||
return super().forward_test(**kwargs)
|
||||
|
||||
def save_pretrained(self,
|
||||
*args,
|
||||
destination=None,
|
||||
prefix='',
|
||||
keep_vars=False):
|
||||
return self.model.control_blocks.state_dict(*args,
|
||||
destination=destination,
|
||||
keep_vars=keep_vars)
|
||||
|
||||
def save_pretrained_config(self):
|
||||
return copy.deepcopy(self.cfg.CONTROL_MODEL.cfg_dict)
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('MODELS',
|
||||
__class__.__name__,
|
||||
LatentDiffusionSCEControl.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@MODELS.register_class()
|
||||
class LatentDiffusionXLSCEControl(LatentDiffusionSCEControl,
|
||||
LatentDiffusionXL):
|
||||
pass
|
||||
@@ -176,7 +176,14 @@ class LatentDiffusionXL(LatentDiffusion):
|
||||
continue
|
||||
batch[key] = kwargs[key].to(we.device_id)
|
||||
context = getattr(self.cond_stage_model, 'encode')(batch)
|
||||
|
||||
if 'hint' in kwargs and kwargs['hint'] is not None:
|
||||
hint = kwargs.pop('hint')
|
||||
if isinstance(context, dict):
|
||||
context['hint'] = hint
|
||||
else:
|
||||
context = {'crossattn': context, 'hint': hint}
|
||||
else:
|
||||
hint = None
|
||||
if self.min_snr_gamma is not None:
|
||||
alphas = self.diffusion.alphas.to(we.device_id)[t]
|
||||
sigmas = self.diffusion.sigmas.pow(2).to(we.device_id)[t]
|
||||
@@ -190,7 +197,8 @@ class LatentDiffusionXL(LatentDiffusion):
|
||||
t=t,
|
||||
model=self.model,
|
||||
model_kwargs={'cond': context},
|
||||
noise=noise)
|
||||
noise=noise,
|
||||
**kwargs)
|
||||
loss = loss * weights
|
||||
loss = loss.mean()
|
||||
ret = {'loss': loss, 'probe_data': {'prompt': prompt}}
|
||||
@@ -273,6 +281,19 @@ class LatentDiffusionXL(LatentDiffusion):
|
||||
context = getattr(self.cond_stage_model, 'encode')(batch)
|
||||
null_context = getattr(self.cond_stage_model, 'encode')(batch_uc)
|
||||
|
||||
if 'hint' in kwargs and kwargs['hint'] is not None:
|
||||
hint = kwargs.pop('hint')
|
||||
if isinstance(context, dict):
|
||||
context['hint'] = hint
|
||||
else:
|
||||
context = {'crossattn': context, 'hint': hint}
|
||||
if isinstance(null_context, dict):
|
||||
null_context['hint'] = hint
|
||||
else:
|
||||
null_context = {'crossattn': null_context, 'hint': hint}
|
||||
else:
|
||||
hint = None
|
||||
|
||||
if 'index' in kwargs:
|
||||
kwargs.pop('index')
|
||||
height, width = batch['target_size_as_tuple'][0].cpu().numpy().tolist()
|
||||
@@ -482,15 +503,18 @@ class LatentDiffusionXL(LatentDiffusion):
|
||||
before_refiner_t_samples = [None for _ in prompt]
|
||||
|
||||
outputs = list()
|
||||
for p, np, tnp, img, r_img, t_img, r_t_img in zip(
|
||||
prompt, n_prompt, train_n_prompt, x_samples,
|
||||
before_refiner_samples, t_x_samples, before_refiner_t_samples):
|
||||
for i, (p, np, tnp, img, r_img, t_img, r_t_img) in enumerate(
|
||||
zip(prompt, n_prompt, train_n_prompt, x_samples,
|
||||
before_refiner_samples, t_x_samples,
|
||||
before_refiner_t_samples)):
|
||||
one_tup = {
|
||||
'prompt': p,
|
||||
'n_prompt': np,
|
||||
'image': img,
|
||||
'before_refiner_image': r_img
|
||||
}
|
||||
if hint is not None:
|
||||
one_tup.update({'hint': hint[i]})
|
||||
if t_img is not None:
|
||||
one_tup['train_n_prompt'] = tnp
|
||||
one_tup['train_n_image'] = t_img
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.tuner import sce
|
||||
from scepter.modules.model.tuner.swift_tuner import (SwiftAdapter, SwiftFull,
|
||||
SwiftLoRA)
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.tuner.sce.scetuning import CSCTuners, SCTuner
|
||||
from scepter.modules.model.tuner.sce.scetuning_component import SCEAdapter
|
||||
@@ -0,0 +1,178 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from scepter.modules.model.registry import TUNERS
|
||||
from scepter.modules.model.tuner.base_tuner import BaseTuner
|
||||
from scepter.modules.model.tuner.tuner_component import conv_nd, zero_module
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
|
||||
@TUNERS.register_class()
|
||||
class SCTuner(BaseTuner):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.logger = logger
|
||||
dim = cfg['DIM']
|
||||
tuner_length = cfg['TUNER_LENGTH']
|
||||
tuner_name = cfg.get('TUNER_NAME', 'SCEAdapter')
|
||||
self.tuner_name = tuner_name
|
||||
if tuner_name == 'SCEAdapter':
|
||||
from .scetuning_component import SCEAdapter
|
||||
self.tuner_op = SCEAdapter(dim=dim, adapter_length=tuner_length)
|
||||
else:
|
||||
raise Exception(f'Error tuner op {tuner_name}')
|
||||
|
||||
def forward(self, x, x_shortcut=None, use_shortcut=True, **kwargs):
|
||||
if self.tuner_name == 'SCEAdapter':
|
||||
out = self.tuner_op(x, x_shortcut, use_shortcut)
|
||||
else:
|
||||
out = x
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('TUNERS',
|
||||
__class__.__name__,
|
||||
SCTuner.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@TUNERS.register_class()
|
||||
class CSCTuners(BaseTuner):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.logger = logger
|
||||
input_block_channels = cfg['INPUT_BLOCK_CHANS']
|
||||
input_down_flag = cfg['INPUT_DOWN_FLAG']
|
||||
assert len(input_block_channels) == len(input_down_flag)
|
||||
pre_hint_in_channels = cfg.get('PRE_HINT_IN_CHANNELS', 3)
|
||||
pre_hint_out_channels = cfg.get('PRE_HINT_OUT_CHANNELS', 256)
|
||||
pre_hint_dim_ratio = cfg.get('PRE_HINT_DIM_RATIO', 1.0)
|
||||
dense_hint_kernal = cfg.get('DENSE_HINT_KERNAL', 3)
|
||||
sc_tuner_cfg = cfg['SC_TUNER_CFG']
|
||||
use_layers = cfg.get('USE_LAYERS', None)
|
||||
self.pretrained_model = cfg.get('PRETRAINED_MODEL', None)
|
||||
self.scale = cfg.get('SCALE', 1.0)
|
||||
self.method = 'csctuning'
|
||||
|
||||
# pre_hint
|
||||
dims = 2
|
||||
ch = pre_hint_out_channels
|
||||
self.pre_hint_blocks = nn.Sequential(
|
||||
conv_nd(dims,
|
||||
pre_hint_in_channels,
|
||||
int(16 * pre_hint_dim_ratio),
|
||||
3,
|
||||
padding=1),
|
||||
nn.SiLU(),
|
||||
conv_nd(dims,
|
||||
int(16 * pre_hint_dim_ratio),
|
||||
int(16 * pre_hint_dim_ratio),
|
||||
3,
|
||||
padding=1),
|
||||
nn.SiLU(),
|
||||
conv_nd(dims,
|
||||
int(16 * pre_hint_dim_ratio),
|
||||
int(32 * pre_hint_dim_ratio),
|
||||
3,
|
||||
padding=1,
|
||||
stride=2),
|
||||
nn.SiLU(),
|
||||
conv_nd(dims,
|
||||
int(32 * pre_hint_dim_ratio),
|
||||
int(32 * pre_hint_dim_ratio),
|
||||
3,
|
||||
padding=1),
|
||||
nn.SiLU(),
|
||||
conv_nd(dims,
|
||||
int(32 * pre_hint_dim_ratio),
|
||||
int(96 * pre_hint_dim_ratio),
|
||||
3,
|
||||
padding=1,
|
||||
stride=2),
|
||||
nn.SiLU(),
|
||||
conv_nd(dims,
|
||||
int(96 * pre_hint_dim_ratio),
|
||||
int(96 * pre_hint_dim_ratio),
|
||||
3,
|
||||
padding=1),
|
||||
nn.SiLU(),
|
||||
conv_nd(dims,
|
||||
int(96 * pre_hint_dim_ratio),
|
||||
ch,
|
||||
3,
|
||||
padding=1,
|
||||
stride=2),
|
||||
)
|
||||
# dense_hint
|
||||
self.dense_hint_blocks = nn.ModuleList([])
|
||||
stride_list = [2 if flag else 1 for flag in input_down_flag]
|
||||
for i, chan in enumerate(input_block_channels):
|
||||
if use_layers and i not in use_layers:
|
||||
self.dense_hint_blocks.append(nn.Identity())
|
||||
continue
|
||||
self.dense_hint_blocks.append(
|
||||
nn.Sequential(
|
||||
nn.SiLU(),
|
||||
zero_module(
|
||||
conv_nd(dims,
|
||||
ch,
|
||||
chan,
|
||||
dense_hint_kernal,
|
||||
padding=1,
|
||||
stride=stride_list[i]))
|
||||
if dense_hint_kernal == 3 else zero_module(
|
||||
conv_nd(dims,
|
||||
ch,
|
||||
chan,
|
||||
dense_hint_kernal,
|
||||
padding=0,
|
||||
stride=stride_list[i]))))
|
||||
ch = chan
|
||||
# tuner
|
||||
self.lsc_tuner_blocks = nn.ModuleList([])
|
||||
for i, chan in enumerate(input_block_channels[::-1]):
|
||||
if use_layers and i not in use_layers:
|
||||
self.lsc_tuner_blocks.append(nn.Identity())
|
||||
continue
|
||||
sc_tuner_cfg['DIM'] = chan
|
||||
sc_tuner_cfg['TUNER_LENGTH'] = int(chan *
|
||||
cfg.get('DOWN_RATIO', 1.0))
|
||||
sc_tuner = TUNERS.build(sc_tuner_cfg, logger=self.logger)
|
||||
self.lsc_tuner_blocks.append(sc_tuner)
|
||||
|
||||
def load_pretrained_model(self, pretrained_model):
|
||||
if self.pretrained_model:
|
||||
with FS.get_from(self.pretrained_model,
|
||||
wait_finish=True) as local_path:
|
||||
self.init_from_ckpt(local_path)
|
||||
|
||||
def init_from_ckpt(self, path):
|
||||
model_new = OrderedDict()
|
||||
model = torch.load(path, map_location='cpu')
|
||||
for k, v in model.items():
|
||||
if k.startswith('model.'):
|
||||
k = k[len('model.'):]
|
||||
if k.startswith('0.'):
|
||||
k = k[len('0.'):]
|
||||
model_new[k] = v
|
||||
missing, unexpected = self.load_state_dict(model_new, strict=False)
|
||||
print(
|
||||
f'Restored from {path} 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}')
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('TUNERS',
|
||||
__class__.__name__,
|
||||
CSCTuners.para_dict,
|
||||
set_name=True)
|
||||
@@ -0,0 +1,72 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import math
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from scepter.modules.model.tuner.tuner_utils import (choose_weight_type,
|
||||
get_weight_value)
|
||||
|
||||
|
||||
class SCEAdapter(nn.Module):
|
||||
def __init__(self,
|
||||
dim,
|
||||
adapter_length,
|
||||
adapter_type=None,
|
||||
adapter_weight=None,
|
||||
act_layer=nn.GELU,
|
||||
zero_init_last=True,
|
||||
use_bias=True):
|
||||
super(SCEAdapter, self).__init__()
|
||||
self.dim = dim
|
||||
self.adapter_length = adapter_length
|
||||
self.adapter_type = adapter_type
|
||||
self.adapter_weight = adapter_weight
|
||||
self.zero_init_last = zero_init_last
|
||||
self.use_bias = use_bias
|
||||
self.ln1 = nn.Linear(dim, adapter_length, bias=use_bias)
|
||||
self.activate = act_layer()
|
||||
self.ln2 = nn.Linear(adapter_length, dim, bias=use_bias)
|
||||
self.init_weights()
|
||||
self.init_scaling()
|
||||
|
||||
def _zero_init_weights(self, m):
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.zeros_(m.weight)
|
||||
if self.use_bias:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def _kaiming_init_weights(self, m):
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5))
|
||||
|
||||
def init_weights(self):
|
||||
self._kaiming_init_weights(self.ln1)
|
||||
if self.zero_init_last:
|
||||
self._zero_init_weights(self.ln2)
|
||||
else:
|
||||
self._kaiming_init_weights(self.ln2)
|
||||
|
||||
def init_scaling(self):
|
||||
if self.adapter_weight:
|
||||
self.scaling = choose_weight_type(self.adapter_weight, self.dim)
|
||||
else:
|
||||
self.scaling = None
|
||||
|
||||
def forward(self, x, x_shortcut=None, use_shortcut=True, **kwargs):
|
||||
if x_shortcut is None:
|
||||
x_shortcut = x
|
||||
x_shape = x.shape
|
||||
if len(x_shape) == 4:
|
||||
b, d, h, w = x_shape
|
||||
x = x.permute(0, 2, 3, 1).reshape(b, h * w, d)
|
||||
out = self.ln2(self.activate(self.ln1(x)))
|
||||
if self.adapter_weight:
|
||||
scaling = get_weight_value(self.adapter_weight, self.scaling, out)
|
||||
out = out * scaling if scaling is not None else out
|
||||
if len(x_shape) == 4:
|
||||
b, d, h, w = x_shape
|
||||
out = out.reshape(b, h, w, -1).permute(0, 3, 1, 2).contiguous()
|
||||
if use_shortcut:
|
||||
out = x_shortcut + out
|
||||
return out
|
||||
@@ -104,6 +104,7 @@ class SwiftAdapter(BaseTuner):
|
||||
SwiftAdapter.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@TUNERS.register_class()
|
||||
class SwiftSCETuning(BaseTuner):
|
||||
para_dict = {
|
||||
|
||||
@@ -6,6 +6,28 @@ import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def conv_nd(dims, *args, **kwargs):
|
||||
"""
|
||||
Create a 1D, 2D, or 3D convolution module.
|
||||
"""
|
||||
if dims == 1:
|
||||
return nn.Conv1d(*args, **kwargs)
|
||||
elif dims == 2:
|
||||
return nn.Conv2d(*args, **kwargs)
|
||||
elif dims == 3:
|
||||
return nn.Conv3d(*args, **kwargs)
|
||||
raise ValueError(f'unsupported dimensions: {dims}')
|
||||
|
||||
|
||||
def zero_module(module):
|
||||
"""
|
||||
Zero out the parameters of a module and return it.
|
||||
"""
|
||||
for p in module.parameters():
|
||||
p.detach().zero_()
|
||||
return module
|
||||
|
||||
|
||||
class Prompt(nn.Module):
|
||||
"""The implementation of vision prompt tuning method.
|
||||
|
||||
|
||||
@@ -110,7 +110,7 @@ try:
|
||||
self.solver.after_epoch(self.solver.hooks_dict['test'])
|
||||
|
||||
def setup(self, stage: str) -> None:
|
||||
self.solver.logger = get_logger(name='std_torch')
|
||||
self.solver.logger = get_logger(name='scepter')
|
||||
self.solver._prefix = FS.init_fs_client(self.solver.file_system,
|
||||
logger=self.solver.logger)
|
||||
self.solver._local_rank = self.global_rank
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import copy
|
||||
import os
|
||||
from collections import OrderedDict, defaultdict
|
||||
|
||||
@@ -97,6 +98,8 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.max_steps = cfg.MAX_STEPS
|
||||
if self.max_steps > 0:
|
||||
self.max_epochs = -1
|
||||
self.use_amp = cfg.get('USE_AMP', False)
|
||||
self.dtype = getattr(torch, cfg.DTYPE)
|
||||
self.use_fairscale = cfg.get('USE_FAIRSCALE', False)
|
||||
@@ -297,6 +300,25 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
ckpt['scaler'] = self.scaler.state_dict()
|
||||
return ckpt
|
||||
|
||||
def save_pretrained(self):
|
||||
if hasattr(self.model, 'save_pretrained'):
|
||||
ckpt = self.model.save_pretrained()
|
||||
elif hasattr(self.model, 'module') and hasattr(self.model.module,
|
||||
'save_pretrained'):
|
||||
ckpt = self.model.module.save_pretrained()
|
||||
else:
|
||||
ckpt = dict()
|
||||
if hasattr(self.model, 'save_pretrained_config'):
|
||||
cfg = self.model.save_pretrained_config()
|
||||
elif hasattr(self.model, 'module') and hasattr(
|
||||
self.model.module, 'save_pretrained_config'):
|
||||
cfg = self.model.module.save_pretrained_config()
|
||||
else:
|
||||
cfg = copy.deepcopy(self.cfg.MODEL.cfg_dict)
|
||||
if 'FILE_SYSTEM' in cfg:
|
||||
cfg.pop('FILE_SYSTEM')
|
||||
return ckpt, cfg
|
||||
|
||||
def solve(self):
|
||||
self.before_solve()
|
||||
if 'train' in self._mode_set:
|
||||
@@ -366,8 +388,17 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
log_data, log_label, ori_label = [], [], []
|
||||
for result in all_results:
|
||||
# the inference image use
|
||||
log_data.append((result['image'].permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
if 'hint' in result:
|
||||
merge_image = torch.cat([
|
||||
result['hint'][:result['image'].shape[0]], result['image']
|
||||
],
|
||||
dim=2)
|
||||
log_data.append((merge_image.permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
else:
|
||||
log_data.append(
|
||||
(result['image'].permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
log_label.append(result['prompt'] + ' NegPrompt: ' +
|
||||
result['n_prompt'])
|
||||
ori_label.append(result['prompt'])
|
||||
@@ -385,9 +416,18 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
for result in all_results:
|
||||
# the inference image use
|
||||
if 'train_n_image' in result:
|
||||
log_data.append(
|
||||
(result['train_n_image'].permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
if 'hint' in result:
|
||||
merge_image = torch.cat([
|
||||
result['hint'][:result['train_n_image'].shape[0]],
|
||||
result['train_n_image']
|
||||
],
|
||||
dim=2)
|
||||
log_data.append(
|
||||
(merge_image.permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
else:
|
||||
log_data.append((result['train_n_image'].permute(
|
||||
1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
||||
log_label.append(result['prompt'] + 'NegPrompt' +
|
||||
result['train_n_prompt'])
|
||||
ori_label.append(result['prompt'])
|
||||
@@ -578,8 +618,16 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
transfer_data_to_cuda(self.current_batch_data[self.mode]))
|
||||
log_data, log_label = [], []
|
||||
for result in outputs:
|
||||
merge_image = torch.cat([result['orig'], result['recon']],
|
||||
dim=2)
|
||||
if 'hint' in result:
|
||||
merge_image = torch.cat([
|
||||
result['orig'],
|
||||
result['hint'][:result['orig'].shape[0]],
|
||||
result['recon']
|
||||
],
|
||||
dim=2)
|
||||
else:
|
||||
merge_image = torch.cat([result['orig'], result['recon']],
|
||||
dim=2)
|
||||
log_data.append((merge_image.permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
log_label.append('recon image: ' + result['prompt'] +
|
||||
@@ -599,8 +647,16 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
log_data, log_label = [], []
|
||||
for result in outputs:
|
||||
if 'train_n_image' in result:
|
||||
merge_image = torch.cat(
|
||||
[result['orig'], result['train_n_image']], dim=2)
|
||||
if 'hint' in result:
|
||||
merge_image = torch.cat([
|
||||
result['orig'],
|
||||
result['hint'][:result['orig'].shape[0]],
|
||||
result['train_n_image']
|
||||
],
|
||||
dim=2)
|
||||
else:
|
||||
merge_image = torch.cat(
|
||||
[result['orig'], result['train_n_image']], dim=2)
|
||||
log_data.append(
|
||||
(merge_image.permute(1, 2, 0).cpu().numpy() *
|
||||
255).astype(np.uint8))
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import json
|
||||
import os
|
||||
import os.path as osp
|
||||
import shutil
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.distributed as du
|
||||
from swift import push_to_hub
|
||||
|
||||
from scepter.modules.solver.hooks.hook import Hook
|
||||
from scepter.modules.solver.hooks.registry import HOOKS
|
||||
@@ -56,6 +60,9 @@ class CheckpointHook(Hook):
|
||||
self.save_last = cfg.get('SAVE_LAST', False)
|
||||
self.save_best = cfg.get('SAVE_BEST', False)
|
||||
self.save_best_by = cfg.get('SAVE_BEST_BY', '')
|
||||
self.push_to_hub = cfg.get('PUSH_TO_HUB', False)
|
||||
self.hub_model_id = cfg.get('HUB_MODEL_ID', None)
|
||||
self.last_ckpt = None
|
||||
if self.save_best and not self.save_best_by:
|
||||
warnings.warn(
|
||||
"CheckpointHook: Parameter 'save_best_by' is not set, turn off save_best function."
|
||||
@@ -96,7 +103,6 @@ class CheckpointHook(Hook):
|
||||
if solver.total_iter != 0 and (
|
||||
(solver.total_iter + 1) % self.interval == 0
|
||||
or solver.total_iter == solver.max_steps - 1):
|
||||
checkpoint = solver.save_checkpoint()
|
||||
solver.logger.info(
|
||||
f'Saving checkpoint after {solver.total_iter + 1} steps')
|
||||
if we.rank == 0:
|
||||
@@ -106,11 +112,44 @@ class CheckpointHook(Hook):
|
||||
solver.total_iter + 1))
|
||||
with FS.put_to(save_path) as local_path:
|
||||
with open(local_path, 'wb') as f:
|
||||
checkpoint = solver.save_checkpoint()
|
||||
torch.save(checkpoint, f)
|
||||
|
||||
from swift import SwiftModel
|
||||
if isinstance(solver.model, SwiftModel):
|
||||
save_path = osp.join(
|
||||
solver.work_dir,
|
||||
'checkpoints/{}-{}'.format(self.save_name_prefix,
|
||||
solver.total_iter + 1))
|
||||
local_folder, _ = FS.map_to_local(save_path)
|
||||
solver.model.save_pretrained(local_folder)
|
||||
FS.put_dir_from_local_dir(local_folder, save_path)
|
||||
else:
|
||||
if hasattr(solver, 'save_pretrained'):
|
||||
save_path = osp.join(
|
||||
solver.work_dir, 'checkpoints/{}-{}-bin'.format(
|
||||
self.save_name_prefix, solver.total_iter + 1))
|
||||
local_folder, _ = FS.map_to_local(save_path)
|
||||
FS.make_dir(local_folder)
|
||||
ckpt, cfg = solver.save_pretrained()
|
||||
with FS.put_to(
|
||||
os.path.join(
|
||||
local_folder,
|
||||
'pytorch_model.bin')) as local_path:
|
||||
with open(local_path, 'wb') as f:
|
||||
torch.save(ckpt, f)
|
||||
with FS.put_to(
|
||||
os.path.join(
|
||||
local_folder,
|
||||
'configuration.json')) as local_path:
|
||||
json.dump(cfg, open(local_path, 'w'))
|
||||
FS.put_dir_from_local_dir(local_folder, save_path)
|
||||
|
||||
if self.save_last and solver.total_iter == solver.max_steps - 1:
|
||||
with FS.get_fs_client(save_path) as client:
|
||||
last_path = osp.join(solver.work_dir, 'checkpoint.pth')
|
||||
client.make_link(last_path, save_path)
|
||||
self.last_ckpt = save_path
|
||||
|
||||
torch.cuda.synchronize()
|
||||
if we.is_distributed:
|
||||
@@ -173,6 +212,48 @@ class CheckpointHook(Hook):
|
||||
torch.save(checkpoint['pre_state_dict'], f)
|
||||
client.put_object_from_local_file(local_file, save_path)
|
||||
|
||||
def create_or_update_model_card(self, model, output_dir: str):
|
||||
"""
|
||||
Updates or create the model card.
|
||||
"""
|
||||
if not os.path.exists(os.path.join(output_dir, 'README.md')):
|
||||
lines = []
|
||||
else:
|
||||
with open(os.path.join(output_dir, 'README.md'), 'r') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
# write the lines back to README.md
|
||||
with open(os.path.join(output_dir, 'README.md'), 'w') as f:
|
||||
f.writelines(lines)
|
||||
|
||||
def after_all_iter(self, solver):
|
||||
if we.rank == 0:
|
||||
if self.push_to_hub and self.last_ckpt:
|
||||
if os.path.isfile(self.last_ckpt):
|
||||
base_dir = os.path.dirname(self.last_ckpt)
|
||||
base_file = os.path.basename(self.last_ckpt)
|
||||
save_path = os.path.join(base_dir, 'after_all_iter')
|
||||
os.makedirs(save_path)
|
||||
self.create_or_update_model_card(solver.model, save_path)
|
||||
try:
|
||||
os.link(self.last_ckpt,
|
||||
os.path.join(save_path, base_file))
|
||||
except OSError:
|
||||
shutil.copyfile(self.last_ckpt,
|
||||
os.path.join(save_path, base_file))
|
||||
|
||||
push_to_hub(repo_name=self.hub_model_id,
|
||||
output_dir=self.last_ckpt,
|
||||
private=False)
|
||||
current_dir = os.path.dirname(__file__)
|
||||
base_path = os.sep.join(current_dir.split(os.sep)[:-4])
|
||||
base_path = os.path.join(base_path, 'config')
|
||||
file_name = self.hub_model_id.replace(os.sep, '_')
|
||||
content = {'name': self.hub_model_id}
|
||||
import json
|
||||
with open(os.path.join(base_path, file_name), 'w') as f:
|
||||
json.dump(content, f)
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('hook',
|
||||
|
||||
@@ -16,7 +16,7 @@ from scepter.modules.transform.io import (LoadCvImageFromFile,
|
||||
from scepter.modules.transform.io_video import (DecodeVideoToTensor,
|
||||
LoadVideoFromFile)
|
||||
from scepter.modules.transform.registry import TRANSFORMS, build_pipeline
|
||||
from scepter.modules.transform.tensor import Rename, Select, ToTensor
|
||||
from scepter.modules.transform.tensor import Rename, Select, ToNumpy, ToTensor
|
||||
from scepter.modules.transform.transform_xl import FlexibleCropXL
|
||||
from scepter.modules.transform.video import (AutoResizedCropVideo,
|
||||
CenterCropVideo, NormalizeVideo,
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from scepter.modules.transform.registry import TRANSFORMS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
@@ -24,6 +25,17 @@ def to_tensor(data):
|
||||
raise TypeError(f'Unsupported type {type(data)}')
|
||||
|
||||
|
||||
def to_numpy(data):
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data.detach().cpu().numpy()
|
||||
elif isinstance(data, (int, float, list, tuple, dict, Image.Image)):
|
||||
return np.array(data)
|
||||
elif isinstance(data, np.ndarray):
|
||||
return data
|
||||
else:
|
||||
raise TypeError(f'Unsupported type {type(data)}')
|
||||
|
||||
|
||||
@TRANSFORMS.register_class()
|
||||
class ToTensor(object):
|
||||
def __init__(self, cfg, logger=None):
|
||||
@@ -54,6 +66,50 @@ class ToTensor(object):
|
||||
set_name=True)
|
||||
|
||||
|
||||
@TRANSFORMS.register_class()
|
||||
class ToNumpy(object):
|
||||
def __init__(self, cfg, logger=None):
|
||||
self.input_key = cfg.get('INPUT_KEY', 'img')
|
||||
self.output_key = cfg.get('OUTPUT_KEY', 'img')
|
||||
|
||||
def __call__(self, item):
|
||||
if isinstance(self.input_key, str):
|
||||
self.input_key = [self.input_key]
|
||||
if isinstance(self.output_key, str):
|
||||
self.output_key = [self.output_key]
|
||||
for idx, key in enumerate(self.input_key):
|
||||
item[self.output_key[idx]] = to_numpy(item[key])
|
||||
return item
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
'''
|
||||
{ "ENV" :
|
||||
{ "description" : "",
|
||||
"A" : {
|
||||
"value": 1.0,
|
||||
"description": ""
|
||||
}
|
||||
}
|
||||
}
|
||||
:return:
|
||||
'''
|
||||
para_dict = [{
|
||||
'INPUT_KEY': {
|
||||
'value': [],
|
||||
'description': 'input_key'
|
||||
},
|
||||
'OUTPUT_KEY': {
|
||||
'value': [],
|
||||
'description': 'output_key'
|
||||
}
|
||||
}]
|
||||
return dict_to_yaml('TRANSFORM',
|
||||
__class__.__name__,
|
||||
para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@TRANSFORMS.register_class()
|
||||
class Select(object):
|
||||
def __init__(self, cfg, logger=None):
|
||||
@@ -107,14 +163,14 @@ class Select(object):
|
||||
@TRANSFORMS.register_class()
|
||||
class Rename(object):
|
||||
def __init__(self, cfg, logger=None):
|
||||
self.in_keys = cfg.IN_KEYS
|
||||
self.out_keys = cfg.OUT_KEYS
|
||||
self.input_key = cfg.INPUT_KEY
|
||||
self.output_key = cfg.OUTPUT_KEY
|
||||
|
||||
def __call__(self, item):
|
||||
data = {}
|
||||
for idx, key in enumerate(self.in_keys):
|
||||
data[self.out_keys[idx]] = item[key]
|
||||
have_key_set = set(self.in_keys)
|
||||
for idx, key in enumerate(self.input_key):
|
||||
data[self.output_key[idx]] = item[key]
|
||||
have_key_set = set(self.input_key)
|
||||
for k, v in item.items():
|
||||
if k not in have_key_set:
|
||||
data[k] = v
|
||||
@@ -134,12 +190,12 @@ class Rename(object):
|
||||
:return:
|
||||
'''
|
||||
para_dict = [{
|
||||
'IN_KEYS': {
|
||||
'INPUT_KEY': {
|
||||
'value': [],
|
||||
'description':
|
||||
'The keys need to rename, the other keys are outputed by default.'
|
||||
},
|
||||
'OUT_KEYS': {
|
||||
'OUTPUT_KEY': {
|
||||
'value': [],
|
||||
'description':
|
||||
'The keys need to rename, the other keys are outputed by default.'
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import numbers
|
||||
import os
|
||||
import sys
|
||||
|
||||
@@ -579,7 +580,7 @@ class Config(object):
|
||||
for key, val in cfg_dict.items():
|
||||
if isinstance(val, (Config, dict, list)):
|
||||
cfg_new[key] = Config.get_plain_cfg(val)
|
||||
else:
|
||||
elif isinstance(val, (str, numbers.Number)):
|
||||
cfg_new[key] = val
|
||||
return cfg_new
|
||||
elif isinstance(cfg, dict):
|
||||
@@ -588,7 +589,7 @@ class Config(object):
|
||||
for key, val in cfg_dict.items():
|
||||
if isinstance(val, (Config, dict, list)):
|
||||
cfg_new[key] = Config.get_plain_cfg(val)
|
||||
else:
|
||||
elif isinstance(val, (str, numbers.Number)):
|
||||
cfg_new[key] = val
|
||||
return cfg_new
|
||||
elif isinstance(cfg, list):
|
||||
@@ -597,7 +598,7 @@ class Config(object):
|
||||
for val in cfg_list:
|
||||
if isinstance(val, (Config, dict, list)):
|
||||
cfg_new.append(Config.get_plain_cfg(val))
|
||||
else:
|
||||
elif isinstance(val, (str, numbers.Number)):
|
||||
cfg_new.append(val)
|
||||
return cfg_new
|
||||
else:
|
||||
|
||||
@@ -327,7 +327,7 @@ class Workenv(object):
|
||||
self.seed = 2023
|
||||
self.debug = False
|
||||
self.use_pl = False
|
||||
self.launcher = 'spawn'
|
||||
self.launcher = 'spawn' if torch.cuda.device_count() > 1 else None
|
||||
self.data_online = False
|
||||
self.share_storage = False
|
||||
|
||||
@@ -358,7 +358,7 @@ class Workenv(object):
|
||||
fn(config)
|
||||
return
|
||||
|
||||
if (os.environ.get('WORLD_SIZE') is None or os.environ.get('WORLD_SIZE') == 1) \
|
||||
if (os.environ.get('WORLD_SIZE') is None or int(os.environ.get('WORLD_SIZE')) == 1) \
|
||||
and torch.cuda.device_count() == 1 and not self.launcher == 'dist':
|
||||
self.device_id = 0
|
||||
fn(config)
|
||||
|
||||
@@ -239,7 +239,7 @@ class LocalFs(BaseFs):
|
||||
return True
|
||||
|
||||
def walk_dir(self, file_dir, recurse=True):
|
||||
for root, dirs, files in os.walk(file_dir, topdown=True):
|
||||
for root, dirs, files in os.walk(file_dir, topdown=recurse):
|
||||
sub_files = files + dirs
|
||||
for name in sub_files:
|
||||
yield os.path.join(root, name)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user