Hugging Face's logo

Spaces:
kevinwang676
/
Diffutoon
Runtime error

App Files Community
Diffutoon
File size: 7,256 Bytes 
fb4fac3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
 
# Kolors

Kolors is a Chinese diffusion model, which is based on ChatGLM and Stable Diffusion XL. We provide training scripts here.

## Download models

The following files will be used for constructing Kolors. You can download Kolors from [huggingface](https://huggingface.co/Kwai-Kolors/Kolors) or [modelscope](https://modelscope.cn/models/Kwai-Kolors/Kolors). Due to precision overflow issues, we need to download an additional VAE model (from [huggingface](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) or [modelscope](https://modelscope.cn/models/AI-ModelScope/sdxl-vae-fp16-fix)).

```
models
β”œβ”€β”€ kolors
β”‚   └── Kolors
β”‚       β”œβ”€β”€ text_encoder
β”‚       β”‚   β”œβ”€β”€ config.json
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00001-of-00007.bin
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00002-of-00007.bin
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00003-of-00007.bin
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00004-of-00007.bin
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00005-of-00007.bin
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00006-of-00007.bin
β”‚       β”‚   β”œβ”€β”€ pytorch_model-00007-of-00007.bin
β”‚       β”‚   └── pytorch_model.bin.index.json
β”‚       β”œβ”€β”€ unet
β”‚       β”‚   └── diffusion_pytorch_model.safetensors
β”‚       └── vae
β”‚           └── diffusion_pytorch_model.safetensors
└── sdxl-vae-fp16-fix
    └── diffusion_pytorch_model.safetensors
```

You can use the following code to download these files:

```python
from diffsynth import download_models

download_models(["Kolors", "SDXL-vae-fp16-fix"])
```

## Train

### Install training dependency

```
pip install peft lightning pandas torchvision
```

### Prepare your dataset

We provide an example dataset [here](https://modelscope.cn/datasets/buptwq/lora-stable-diffusion-finetune/files). You need to manage the training images as follows:

```
data/dog/
└── train
    β”œβ”€β”€ 00.jpg
    β”œβ”€β”€ 01.jpg
    β”œβ”€β”€ 02.jpg
    β”œβ”€β”€ 03.jpg
    β”œβ”€β”€ 04.jpg
    └── metadata.csv
```

`metadata.csv`:

```
file_name,text
00.jpg,δΈ€εͺ小狗
01.jpg,δΈ€εͺ小狗
02.jpg,δΈ€εͺ小狗
03.jpg,δΈ€εͺ小狗
04.jpg,δΈ€εͺ小狗
```

### Train a LoRA model

We provide a training script `train_kolors_lora.py`. Before you run this training script, please copy it to the root directory of this project.

The following settings are recommended. 22GB VRAM is required.

```
CUDA_VISIBLE_DEVICES="0" python examples/train/kolors/train_kolors_lora.py \
  --pretrained_unet_path models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors \
  --pretrained_text_encoder_path models/kolors/Kolors/text_encoder \
  --pretrained_fp16_vae_path models/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors \
  --dataset_path data/dog \
  --output_path ./models \
  --max_epochs 10 \
  --center_crop \
  --use_gradient_checkpointing \
  --precision "16-mixed"
```

Optional arguments:
```
  -h, --help            show this help message and exit
  --pretrained_unet_path PRETRAINED_UNET_PATH
                        Path to pretrained model (UNet). For example, `models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors`.
  --pretrained_text_encoder_path PRETRAINED_TEXT_ENCODER_PATH
                        Path to pretrained model (Text Encoder). For example, `models/kolors/Kolors/text_encoder`.
  --pretrained_fp16_vae_path PRETRAINED_FP16_VAE_PATH
                        Path to pretrained model (VAE). For example, `models/kolors/Kolors/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors`.
  --dataset_path DATASET_PATH
                        The path of the Dataset.
  --output_path OUTPUT_PATH
                        Path to save the model.
  --steps_per_epoch STEPS_PER_EPOCH
                        Number of steps per epoch.
  --height HEIGHT       Image height.
  --width WIDTH         Image width.
  --center_crop         Whether to center crop the input images to the resolution. If not set, the images will be randomly cropped. The images will be resized to the resolution first before cropping.
  --random_flip         Whether to randomly flip images horizontally
  --batch_size BATCH_SIZE
                        Batch size (per device) for the training dataloader.
  --dataloader_num_workers DATALOADER_NUM_WORKERS
                        Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.
  --precision {32,16,16-mixed}
                        Training precision
  --learning_rate LEARNING_RATE
                        Learning rate.
  --lora_rank LORA_RANK
                        The dimension of the LoRA update matrices.
  --lora_alpha LORA_ALPHA
                        The weight of the LoRA update matrices.
  --use_gradient_checkpointing
                        Whether to use gradient checkpointing.
  --accumulate_grad_batches ACCUMULATE_GRAD_BATCHES
                        The number of batches in gradient accumulation.
  --training_strategy {auto,deepspeed_stage_1,deepspeed_stage_2,deepspeed_stage_3}
                        Training strategy
  --max_epochs MAX_EPOCHS
                        Number of epochs.
```

### Inference with your own LoRA model

After training, you can use your own LoRA model to generate new images. Here are some examples.

```python
from diffsynth import ModelManager, KolorsImagePipeline
from peft import LoraConfig, inject_adapter_in_model
import torch


def load_lora(model, lora_rank, lora_alpha, lora_path):
    lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_alpha,
        init_lora_weights="gaussian",
        target_modules=["to_q", "to_k", "to_v", "to_out"],
    )
    model = inject_adapter_in_model(lora_config, model)
    state_dict = torch.load(lora_path, map_location="cpu")
    model.load_state_dict(state_dict, strict=False)
    return model


# Load models
model_manager = ModelManager(torch_dtype=torch.float16, device="cuda",
                             file_path_list=[
                                 "models/kolors/Kolors/text_encoder",
                                 "models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors",
                                 "models/kolors/Kolors/vae/diffusion_pytorch_model.safetensors"
                             ])
pipe = KolorsImagePipeline.from_model_manager(model_manager)

# Generate an image with lora
pipe.unet = load_lora(
    pipe.unet,
    lora_rank=4, lora_alpha=4.0, # The two parameters should be consistent with those in your training script.
    lora_path="path/to/your/lora/model/lightning_logs/version_x/checkpoints/epoch=x-step=xxx.ckpt"
)
torch.manual_seed(0)
image = pipe(
    prompt="δΈ€εͺε°η‹—θΉ¦θΉ¦θ·³θ·³οΌŒε‘¨ε›΄ζ˜―ε§Ήη΄«ε«£ηΊ’ηš„ι²œθŠ±οΌŒθΏœε€„ζ˜―ε±±θ„‰",
    negative_prompt="",
    cfg_scale=4,
    num_inference_steps=50, height=1024, width=1024,
)
image.save("image_with_lora.jpg")
```

Prompt: δΈ€εͺε°η‹—θΉ¦θΉ¦θ·³θ·³οΌŒε‘¨ε›΄ζ˜―ε§Ήη΄«ε«£ηΊ’ηš„ι²œθŠ±οΌŒθΏœε€„ζ˜―ε±±θ„‰

|Without LoRA|With LoRA|
|-|-|
|![image_without_lora](https://github.com/modelscope/DiffSynth-Studio/assets/35051019/9d79ed7a-e8cf-4d98-800a-f182809db318)|![image_with_lora](https://github.com/modelscope/DiffSynth-Studio/assets/35051019/02f62323-6ee5-4788-97a1-549732dbe4f0)|