Skip to content

Commit 82188ce

Browse files
zRzRzRzRzRzRzROleehyOyiyixuxu
authored
CogView4 Control Block (#10809)
* cogview4 control training --------- Co-authored-by: OleehyO <leehy0357@gmail.com> Co-authored-by: yiyixuxu <yixu310@gmail.com>
1 parent cc19726 commit 82188ce

13 files changed

+2287
-37
lines changed

‎examples/cogview4-control/README.md

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
1+
# Training CogView4 Control
2+
3+
This (experimental) example shows how to train Control LoRAs with [CogView4](https://huggingface.co/THUDM/CogView4-6B) by conditioning it with additional structural controls (like depth maps, poses, etc.). We provide a script for full fine-tuning, too, refer to [this section](#full-fine-tuning). To know more about CogView4 Control family, refer to the following resources:
4+
5+
To incorporate additional condition latents, we expand the input features of CogView-4 from 64 to 128. The first 64 channels correspond to the original input latents to be denoised, while the latter 64 channels correspond to control latents. This expansion happens on the `patch_embed` layer, where the combined latents are projected to the expected feature dimension of rest of the network. Inference is performed using the `CogView4ControlPipeline`.
6+
7+
> [!NOTE]
8+
> **Gated model**
9+
>
10+
> As the model is gated, before using it with diffusers you first need to go to the [CogView4 Hugging Face page](https://huggingface.co/THUDM/CogView4-6B), fill in the form and accept the gate. Once you are in, you need to log in so that your system knows you’ve accepted the gate. Use the command below to log in:
11+
12+
```bash
13+
huggingface-cli login
14+
```
15+
16+
The example command below shows how to launch fine-tuning for pose conditions. The dataset ([`raulc0399/open_pose_controlnet`](https://huggingface.co/datasets/raulc0399/open_pose_controlnet)) being used here already has the pose conditions of the original images, so we don't have to compute them.
17+
18+
```bash
19+
accelerate launch train_control_lora_cogview4.py \
20+
--pretrained_model_name_or_path="THUDM/CogView4-6B" \
21+
--dataset_name="raulc0399/open_pose_controlnet" \
22+
--output_dir="pose-control-lora" \
23+
--mixed_precision="bf16" \
24+
--train_batch_size=1 \
25+
--rank=64 \
26+
--gradient_accumulation_steps=4 \
27+
--gradient_checkpointing \
28+
--use_8bit_adam \
29+
--learning_rate=1e-4 \
30+
--report_to="wandb" \
31+
--lr_scheduler="constant" \
32+
--lr_warmup_steps=0 \
33+
--max_train_steps=5000 \
34+
--validation_image="openpose.png" \
35+
--validation_prompt="A couple, 4k photo, highly detailed" \
36+
--offload \
37+
--seed="0" \
38+
--push_to_hub
39+
```
40+
41+
`openpose.png` comes from [here](https://huggingface.co/Adapter/t2iadapter/resolve/main/openpose.png).
42+
43+
You need to install `diffusers` from the branch of [this PR](https://github.com/huggingface/diffusers/pull/9999). When it's merged, you should install `diffusers` from the `main`.
44+
45+
The training script exposes additional CLI args that might be useful to experiment with:
46+
47+
* `use_lora_bias`: When set, additionally trains the biases of the `lora_B` layer.
48+
* `train_norm_layers`: When set, additionally trains the normalization scales. Takes care of saving and loading.
49+
* `lora_layers`: Specify the layers you want to apply LoRA to. If you specify "all-linear", all the linear layers will be LoRA-attached.
50+
51+
### Training with DeepSpeed
52+
53+
It's possible to train with [DeepSpeed](https://github.com/microsoft/DeepSpeed), specifically leveraging the Zero2 system optimization. To use it, save the following config to an YAML file (feel free to modify as needed):
54+
55+
```yaml
56+
compute_environment: LOCAL_MACHINE
57+
debug: false
58+
deepspeed_config:
59+
gradient_accumulation_steps: 1
60+
gradient_clipping: 1.0
61+
offload_optimizer_device: cpu
62+
offload_param_device: cpu
63+
zero3_init_flag: false
64+
zero_stage: 2
65+
distributed_type: DEEPSPEED
66+
downcast_bf16: 'no'
67+
enable_cpu_affinity: false
68+
machine_rank: 0
69+
main_training_function: main
70+
mixed_precision: bf16
71+
num_machines: 1
72+
num_processes: 1
73+
rdzv_backend: static
74+
same_network: true
75+
tpu_env: []
76+
tpu_use_cluster: false
77+
tpu_use_sudo: false
78+
use_cpu: false
79+
```
80+
81+
And then while launching training, pass the config file:
82+
83+
```bash
84+
accelerate launch --config_file=CONFIG_FILE.yaml ...
85+
```
86+
87+
### Inference
88+
89+
The pose images in our dataset were computed using the [`controlnet_aux`](https://github.com/huggingface/controlnet_aux) library. Let's install it first:
90+
91+
```bash
92+
pip install controlnet_aux
93+
```
94+
95+
And then we are ready:
96+
97+
```py
98+
from controlnet_aux import OpenposeDetector
99+
from diffusers import CogView4ControlPipeline
100+
from diffusers.utils import load_image
101+
from PIL import Image
102+
import numpy as np
103+
import torch
104+
105+
pipe = CogView4ControlPipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16).to("cuda")
106+
pipe.load_lora_weights("...") # change this.
107+
108+
open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
109+
110+
# prepare pose condition.
111+
url = "https://huggingface.co/Adapter/t2iadapter/resolve/main/people.jpg"
112+
image = load_image(url)
113+
image = open_pose(image, detect_resolution=512, image_resolution=1024)
114+
image = np.array(image)[:, :, ::-1]
115+
image = Image.fromarray(np.uint8(image))
116+
117+
prompt = "A couple, 4k photo, highly detailed"
118+
119+
gen_images = pipe(
120+
prompt=prompt,
121+
control_image=image,
122+
num_inference_steps=50,
123+
joint_attention_kwargs={"scale": 0.9},
124+
guidance_scale=25.,
125+
).images[0]
126+
gen_images.save("output.png")
127+
```
128+
129+
## Full fine-tuning
130+
131+
We provide a non-LoRA version of the training script `train_control_cogview4.py`. Here is an example command:
132+
133+
```bash
134+
accelerate launch --config_file=accelerate_ds2.yaml train_control_cogview4.py \
135+
--pretrained_model_name_or_path="THUDM/CogView4-6B" \
136+
--dataset_name="raulc0399/open_pose_controlnet" \
137+
--output_dir="pose-control" \
138+
--mixed_precision="bf16" \
139+
--train_batch_size=2 \
140+
--dataloader_num_workers=4 \
141+
--gradient_accumulation_steps=4 \
142+
--gradient_checkpointing \
143+
--use_8bit_adam \
144+
--proportion_empty_prompts=0.2 \
145+
--learning_rate=5e-5 \
146+
--adam_weight_decay=1e-4 \
147+
--report_to="wandb" \
148+
--lr_scheduler="cosine" \
149+
--lr_warmup_steps=1000 \
150+
--checkpointing_steps=1000 \
151+
--max_train_steps=10000 \
152+
--validation_steps=200 \
153+
--validation_image "2_pose_1024.jpg" "3_pose_1024.jpg" \
154+
--validation_prompt "two friends sitting by each other enjoying a day at the park, full hd, cinematic" "person enjoying a day at the park, full hd, cinematic" \
155+
--offload \
156+
--seed="0" \
157+
--push_to_hub
158+
```
159+
160+
Change the `validation_image` and `validation_prompt` as needed.
161+
162+
For inference, this time, we will run:
163+
164+
```py
165+
from controlnet_aux import OpenposeDetector
166+
from diffusers import CogView4ControlPipeline, CogView4Transformer2DModel
167+
from diffusers.utils import load_image
168+
from PIL import Image
169+
import numpy as np
170+
import torch
171+
172+
transformer = CogView4Transformer2DModel.from_pretrained("...") # change this.
173+
pipe = CogView4ControlPipeline.from_pretrained(
174+
"THUDM/CogView4-6B", transformer=transformer, torch_dtype=torch.bfloat16
175+
).to("cuda")
176+
177+
open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
178+
179+
# prepare pose condition.
180+
url = "https://huggingface.co/Adapter/t2iadapter/resolve/main/people.jpg"
181+
image = load_image(url)
182+
image = open_pose(image, detect_resolution=512, image_resolution=1024)
183+
image = np.array(image)[:, :, ::-1]
184+
image = Image.fromarray(np.uint8(image))
185+
186+
prompt = "A couple, 4k photo, highly detailed"
187+
188+
gen_images = pipe(
189+
prompt=prompt,
190+
control_image=image,
191+
num_inference_steps=50,
192+
guidance_scale=25.,
193+
).images[0]
194+
gen_images.save("output.png")
195+
```
196+
197+
## Things to note
198+
199+
* The scripts provided in this directory are experimental and educational. This means we may have to tweak things around to get good results on a given condition. We believe this is best done with the community 🤗
200+
* The scripts are not memory-optimized but we offload the VAE and the text encoders to CPU when they are not used if `--offload` is specified.
201+
* We can extract LoRAs from the fully fine-tuned model. While we currently don't provide any utilities for that, users are welcome to refer to [this script](https://github.com/Stability-AI/stability-ComfyUI-nodes/blob/master/control_lora_create.py) that provides a similar functionality.

‎examples/cogview4-control/requirements.txt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
transformers==4.47.0
2+
wandb
3+
torch
4+
torchvision
5+
accelerate==1.2.0
6+
peft>=0.14.0

0 commit comments

Comments
 (0)