ShaoTengLiu commited on
Commit
58be882
1 Parent(s): 69d3d9d
Tune-A-Video-debug/README.md ADDED
@@ -0,0 +1,227 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Tune-A-Video
2
+
3
+ This repository is the official implementation of [Tune-A-Video](https://arxiv.org/abs/2212.11565).
4
+
5
+ **[Tune-A-Video: One-Shot Tuning of Image Diffusion Models for Text-to-Video Generation](https://arxiv.org/abs/2212.11565)**
6
+ <br/>
7
+ [Jay Zhangjie Wu](https://zhangjiewu.github.io/),
8
+ [Yixiao Ge](https://geyixiao.com/),
9
+ [Xintao Wang](https://xinntao.github.io/),
10
+ [Stan Weixian Lei](),
11
+ [Yuchao Gu](https://ycgu.site/),
12
+ [Yufei Shi](),
13
+ [Wynne Hsu](https://www.comp.nus.edu.sg/~whsu/),
14
+ [Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en),
15
+ [Xiaohu Qie](https://scholar.google.com/citations?user=mk-F69UAAAAJ&hl=en),
16
+ [Mike Zheng Shou](https://sites.google.com/view/showlab)
17
+ <br/>
18
+
19
+ [![Project Website](https://img.shields.io/badge/Project-Website-orange)](https://tuneavideo.github.io/)
20
+ [![arXiv](https://img.shields.io/badge/arXiv-2212.11565-b31b1b.svg)](https://arxiv.org/abs/2212.11565)
21
+ [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Tune-A-Video-library/Tune-A-Video-Training-UI)
22
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/showlab/Tune-A-Video/blob/main/notebooks/Tune-A-Video.ipynb)
23
+
24
+
25
+ <p align="center">
26
+ <img src="https://tuneavideo.github.io/assets/overview.png" width="800px"/>
27
+ <br>
28
+ <em>Given a video-text pair as input, our method, Tune-A-Video, fine-tunes a pre-trained text-to-image diffusion model for text-to-video generation.</em>
29
+ </p>
30
+
31
+ ## News
32
+ - [02/22/2023] Improved consistency using DDIM inversion.
33
+ - [02/08/2023] [Colab demo](https://colab.research.google.com/github/showlab/Tune-A-Video/blob/main/notebooks/Tune-A-Video.ipynb) released!
34
+ - [02/03/2023] Pre-trained Tune-A-Video models are available on [Hugging Face Library](https://huggingface.co/Tune-A-Video-library)!
35
+ - [01/28/2023] New Feature: tune a video on personalized [DreamBooth](https://dreambooth.github.io/) models.
36
+ - [01/28/2023] Code released!
37
+
38
+ ## Setup
39
+
40
+ ### Requirements
41
+
42
+ ```shell
43
+ pip install -r requirements.txt
44
+ ```
45
+
46
+ Installing [xformers](https://github.com/facebookresearch/xformers) is highly recommended for more efficiency and speed on GPUs.
47
+ To enable xformers, set `enable_xformers_memory_efficient_attention=True` (default).
48
+
49
+ ### Weights
50
+
51
+ **[Stable Diffusion]** [Stable Diffusion](https://arxiv.org/abs/2112.10752) is a latent text-to-image diffusion model capable of generating photo-realistic images given any text input. The pre-trained Stable Diffusion models can be downloaded from Hugging Face (e.g., [Stable Diffusion v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4), [v2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1)). You can also use fine-tuned Stable Diffusion models trained on different styles (e.g, [Modern Disney](https://huggingface.co/nitrosocke/mo-di-diffusion), [Redshift](https://huggingface.co/nitrosocke/redshift-diffusion), etc.).
52
+
53
+ **[DreamBooth]** [DreamBooth](https://dreambooth.github.io/) is a method to personalize text-to-image models like Stable Diffusion given just a few images (3~5 images) of a subject. Tuning a video on DreamBooth models allows personalized text-to-video generation of a specific subject. There are some public DreamBooth models available on [Hugging Face](https://huggingface.co/sd-dreambooth-library) (e.g., [mr-potato-head](https://huggingface.co/sd-dreambooth-library/mr-potato-head)). You can also train your own DreamBooth model following [this training example](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth).
54
+
55
+
56
+ ## Usage
57
+
58
+ ### Training
59
+
60
+ To fine-tune the text-to-image diffusion models for text-to-video generation, run this command:
61
+
62
+ ```bash
63
+ accelerate launch train_tuneavideo.py --config="configs/man-skiing.yaml"
64
+ ```
65
+
66
+ Note: Tuning a 24-frame video usually takes `300~500` steps, about `10~15` minutes using one A100 GPU.
67
+ Reduce `n_sample_frames` if your GPU memory is limited.
68
+
69
+ ### Inference
70
+
71
+ Once the training is done, run inference:
72
+
73
+ ```python
74
+ from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
75
+ from tuneavideo.models.unet import UNet3DConditionModel
76
+ from tuneavideo.util import save_videos_grid
77
+ import torch
78
+
79
+ pretrained_model_path = "./checkpoints/stable-diffusion-v1-4"
80
+ my_model_path = "./outputs/man-skiing"
81
+ unet = UNet3DConditionModel.from_pretrained(my_model_path, subfolder='unet', torch_dtype=torch.float16).to('cuda')
82
+ pipe = TuneAVideoPipeline.from_pretrained(pretrained_model_path, unet=unet, torch_dtype=torch.float16).to("cuda")
83
+ pipe.enable_xformers_memory_efficient_attention()
84
+ pipe.enable_vae_slicing()
85
+
86
+ prompt = "spider man is skiing"
87
+ ddim_inv_latent = torch.load(f"{my_model_path}/inv_latents/ddim_latent-500.pt").to(torch.float16)
88
+ video = pipe(prompt, latents=ddim_inv_latent, video_length=24, height=512, width=512, num_inference_steps=50, guidance_scale=12.5).videos
89
+
90
+ save_videos_grid(video, f"./{prompt}.gif")
91
+ ```
92
+
93
+ ## Results
94
+
95
+ ### Pretrained T2I (Stable Diffusion)
96
+ <table class="center">
97
+ <tr>
98
+ <td style="text-align:center;"><b>Input Video</b></td>
99
+ <td style="text-align:center;" colspan="3"><b>Output Video</b></td>
100
+ </tr>
101
+ <tr>
102
+ <td><img src="https://tuneavideo.github.io/assets/data/man-skiing.gif"></td>
103
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/man-skiing/spiderman-beach.gif"></td>
104
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/man-skiing/wonder-woman.gif"></td>
105
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/man-skiing/pink-sunset.gif"></td>
106
+ </tr>
107
+ <tr>
108
+ <td width=25% style="text-align:center;color:gray;">"A man is skiing"</td>
109
+ <td width=25% style="text-align:center;">"Spider Man is skiing on the beach, cartoon style”</td>
110
+ <td width=25% style="text-align:center;">"Wonder Woman, wearing a cowboy hat, is skiing"</td>
111
+ <td width=25% style="text-align:center;">"A man, wearing pink clothes, is skiing at sunset"</td>
112
+ </tr>
113
+
114
+ <tr>
115
+ <td><img src="https://tuneavideo.github.io/assets/data/rabbit-watermelon.gif"></td>
116
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/rabbit-watermelon/rabbit.gif"></td>
117
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/rabbit-watermelon/cat.gif"></td>
118
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/rabbit-watermelon/puppy.gif"></td>
119
+ </tr>
120
+ <tr>
121
+ <td width=25% style="text-align:center;color:gray;">"A rabbit is eating a watermelon"</td>
122
+ <td width=25% style="text-align:center;">"A rabbit is <del>eating a watermelon</del> on the table"</td>
123
+ <td width=25% style="text-align:center;">"A cat with sunglasses is eating a watermelon on the beach"</td>
124
+ <td width=25% style="text-align:center;">"A puppy is eating a cheeseburger on the table, comic style"</td>
125
+ </tr>
126
+
127
+ <tr>
128
+ <td><img src="https://tuneavideo.github.io/assets/data/car-turn.gif"></td>
129
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/car-turn/porsche-beach.gif"></td>
130
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/car-turn/car-cartoon.gif"></td>
131
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/car-turn/car-snow.gif"></td>
132
+ </tr>
133
+ <tr>
134
+ <td width=25% style="text-align:center;color:gray;">"A jeep car is moving on the road"</td>
135
+ <td width=25% style="text-align:center;">"A Porsche car is moving on the beach"</td>
136
+ <td width=25% style="text-align:center;">"A car is moving on the road, cartoon style"</td>
137
+ <td width=25% style="text-align:center;">"A car is moving on the snow"</td>
138
+ </tr>
139
+
140
+ <tr>
141
+ <td><img src="https://tuneavideo.github.io/assets/data/man-basketball.gif"></td>
142
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/man-basketball/trump.gif"></td>
143
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/man-basketball/astronaut.gif"></td>
144
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/man-basketball/lego.gif"></td>
145
+ </tr>
146
+ <tr>
147
+ <td width=25% style="text-align:center;color:gray;">"A man is dribbling a basketball"</td>
148
+ <td width=25% style="text-align:center;">"Trump is dribbling a basketball"</td>
149
+ <td width=25% style="text-align:center;">"An astronaut is dribbling a basketball, cartoon style"</td>
150
+ <td width=25% style="text-align:center;">"A lego man in a black suit is dribbling a basketball"</td>
151
+ </tr>
152
+
153
+ <!-- <tr>
154
+ <td><img src="https://tuneavideo.github.io/assets/data/lion-roaring.gif"></td>
155
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/lion-roaring/tiger-roar.gif"></td>
156
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/lion-roaring/lion-vangogh.gif"></td>
157
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/lion-roaring/wolf-nyc.gif"></td>
158
+ </tr>
159
+ <tr>
160
+ <td width=25% style="text-align:center;color:gray;">"A lion is roaring"</td>
161
+ <td width=25% style="text-align:center;">"A tiger is roaring"</td>
162
+ <td width=25% style="text-align:center;">"A lion is roaring, Van Gogh style"</td>
163
+ <td width=25% style="text-align:center;">"A wolf is roaring in New York City"</td>
164
+ </tr> -->
165
+
166
+ </table>
167
+
168
+ ### Pretrained T2I (personalized DreamBooth)
169
+
170
+ <img src="https://tuneavideo.github.io/assets/results/tuneavideo/modern-disney/modern-disney.png" width="240px"/>
171
+
172
+ <table class="center">
173
+ <tr>
174
+ <td style="text-align:center;"><b>Input Video</b></td>
175
+ <td style="text-align:center;" colspan="3"><b>Output Video</b></td>
176
+ </tr>
177
+ <tr>
178
+ <td><img src="https://tuneavideo.github.io/assets/data/bear-guitar.gif"></td>
179
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/modern-disney/bear-guitar/rabbit.gif"></td>
180
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/modern-disney/bear-guitar/prince.gif"></td>
181
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/modern-disney/bear-guitar/princess.gif"></td>
182
+ </tr>
183
+ <tr>
184
+ <td width=25% style="text-align:center;color:gray;">"A bear is playing guitar"</td>
185
+ <td width=25% style="text-align:center;">"A rabbit is playing guitar, modern disney style"</td>
186
+ <td width=25% style="text-align:center;">"A handsome prince is playing guitar, modern disney style"</td>
187
+ <td width=25% style="text-align:center;">"A magic princess with sunglasses is playing guitar on the stage, modern disney style"</td>
188
+ </tr>
189
+ </table>
190
+
191
+ <img src="https://tuneavideo.github.io/assets/results/tuneavideo/mr-potato-head/mr-potato-head.png" width="240px"/>
192
+
193
+ <table class="center">
194
+ <tr>
195
+ <td style="text-align:center;"><b>Input Video</b></td>
196
+ <td style="text-align:center;" colspan="3"><b>Output Video</b></td>
197
+ </tr>
198
+ <tr>
199
+ <td><img src="https://tuneavideo.github.io/assets/data/bear-guitar.gif"></td>
200
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/mr-potato-head/bear-guitar/lego-snow.gif"></td>
201
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/mr-potato-head/bear-guitar/sunglasses-beach.gif"></td>
202
+ <td><img src="https://tuneavideo.github.io/assets/results/tuneavideo/mr-potato-head/bear-guitar/van-gogh.gif"></td>
203
+ </tr>
204
+ <tr>
205
+ <td width=25% style="text-align:center;color:gray;">"A bear is playing guitar"</td>
206
+ <td width=25% style="text-align:center;">"Mr Potato Head, made of lego, is playing guitar on the snow"</td>
207
+ <td width=25% style="text-align:center;">"Mr Potato Head, wearing sunglasses, is playing guitar on the beach"</td>
208
+ <td width=25% style="text-align:center;">"Mr Potato Head is playing guitar in the starry night, Van Gogh style"</td>
209
+ </tr>
210
+ </table>
211
+
212
+
213
+ ## Citation
214
+ If you make use of our work, please cite our paper.
215
+ ```bibtex
216
+ @article{wu2022tuneavideo,
217
+ title={Tune-A-Video: One-Shot Tuning of Image Diffusion Models for Text-to-Video Generation},
218
+ author={Wu, Jay Zhangjie and Ge, Yixiao and Wang, Xintao and Lei, Stan Weixian and Gu, Yuchao and Hsu, Wynne and Shan, Ying and Qie, Xiaohu and Shou, Mike Zheng},
219
+ journal={arXiv preprint arXiv:2212.11565},
220
+ year={2022}
221
+ }
222
+ ```
223
+
224
+ ## Shoutouts
225
+
226
+ - This code builds on [diffusers](https://github.com/huggingface/diffusers). Thanks for open-sourcing!
227
+ - Thanks [hysts](https://github.com/hysts) for the awesome [gradio demo](https://huggingface.co/spaces/Tune-A-Video-library/Tune-A-Video-Training-UI).
Tune-A-Video-debug/configs/car-turn.yaml ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "./checkpoints/stable-diffusion-v1-4"
2
+ output_dir: "./outputs/car-turn"
3
+
4
+ train_data:
5
+ video_path: "data/car-turn.mp4"
6
+ prompt: "a jeep car is moving on the road"
7
+ n_sample_frames: 24
8
+ width: 512
9
+ height: 512
10
+ sample_start_idx: 0
11
+ sample_frame_rate: 2
12
+
13
+ validation_data:
14
+ prompts:
15
+ - "a jeep car is moving on the beach"
16
+ - "a jeep car is moving on the snow"
17
+ - "a jeep car is moving on the road, cartoon style"
18
+ - "a sports car is moving on the road"
19
+ video_length: 24
20
+ width: 512
21
+ height: 512
22
+ num_inference_steps: 50
23
+ guidance_scale: 12.5
24
+ use_inv_latent: True
25
+ num_inv_steps: 50
26
+
27
+ learning_rate: 3e-5
28
+ train_batch_size: 1
29
+ max_train_steps: 500
30
+ checkpointing_steps: 1000
31
+ validation_steps: 100
32
+ trainable_modules:
33
+ - "attn1.to_q"
34
+ - "attn2.to_q"
35
+ - "attn_temp"
36
+
37
+ seed: 33
38
+ mixed_precision: fp16
39
+ use_8bit_adam: False
40
+ gradient_checkpointing: True
41
+ enable_xformers_memory_efficient_attention: True
Tune-A-Video-debug/configs/man-skiing.yaml ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "./checkpoints/stable-diffusion-v1-4"
2
+ output_dir: "./outputs/man-skiing"
3
+
4
+ train_data:
5
+ video_path: "data/man-skiing.mp4"
6
+ prompt: "a man is skiing"
7
+ n_sample_frames: 24
8
+ width: 512
9
+ height: 512
10
+ sample_start_idx: 0
11
+ sample_frame_rate: 2
12
+
13
+ validation_data:
14
+ prompts:
15
+ - "mickey mouse is skiing on the snow"
16
+ - "spider man is skiing on the beach, cartoon style"
17
+ - "wonder woman, wearing a cowboy hat, is skiing"
18
+ - "a man, wearing pink clothes, is skiing at sunset"
19
+ video_length: 24
20
+ width: 512
21
+ height: 512
22
+ num_inference_steps: 50
23
+ guidance_scale: 12.5
24
+ use_inv_latent: True
25
+ num_inv_steps: 50
26
+
27
+ learning_rate: 3e-5
28
+ train_batch_size: 1
29
+ max_train_steps: 500
30
+ checkpointing_steps: 1000
31
+ validation_steps: 100
32
+ trainable_modules:
33
+ - "attn1.to_q"
34
+ - "attn2.to_q"
35
+ - "attn_temp"
36
+
37
+ seed: 33
38
+ mixed_precision: fp16
39
+ use_8bit_adam: False
40
+ gradient_checkpointing: True
41
+ enable_xformers_memory_efficient_attention: True
Tune-A-Video-debug/configs/man-surfing.yaml ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "./checkpoints/stable-diffusion-v1-4"
2
+ output_dir: "./outputs/man-surfing"
3
+
4
+ train_data:
5
+ video_path: "data/man-surfing.mp4"
6
+ prompt: "a man is surfing"
7
+ n_sample_frames: 24
8
+ width: 512
9
+ height: 512
10
+ sample_start_idx: 0
11
+ sample_frame_rate: 1
12
+
13
+ validation_data:
14
+ prompts:
15
+ - "a panda is surfing"
16
+ - "a boy, wearing a birthday hat, is surfing"
17
+ - "a raccoon is surfing, cartoon style"
18
+ - "Iron Man is surfing in the desert"
19
+ video_length: 24
20
+ width: 512
21
+ height: 512
22
+ num_inference_steps: 50
23
+ guidance_scale: 12.5
24
+ use_inv_latent: True
25
+ num_inv_steps: 50
26
+
27
+ learning_rate: 3e-5
28
+ train_batch_size: 1
29
+ max_train_steps: 500
30
+ checkpointing_steps: 1000
31
+ validation_steps: 100
32
+ trainable_modules:
33
+ - "attn1.to_q"
34
+ - "attn2.to_q"
35
+ - "attn_temp"
36
+
37
+ seed: 33
38
+ mixed_precision: fp16
39
+ use_8bit_adam: False
40
+ gradient_checkpointing: True
41
+ enable_xformers_memory_efficient_attention: True
Tune-A-Video-debug/configs/rabbit-watermelon.yaml ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "./checkpoints/stable-diffusion-v1-4"
2
+ output_dir: "./outputs/rabbit-watermelon"
3
+
4
+ train_data:
5
+ video_path: "data/rabbit-watermelon.mp4"
6
+ prompt: "a rabbit is eating a watermelon"
7
+ n_sample_frames: 24
8
+ width: 512
9
+ height: 512
10
+ sample_start_idx: 0
11
+ sample_frame_rate: 2
12
+
13
+ validation_data:
14
+ prompts:
15
+ - "a tiger is eating a watermelon"
16
+ - "a rabbit is eating an orange"
17
+ - "a rabbit is eating a pizza"
18
+ - "a puppy is eating an orange"
19
+ video_length: 24
20
+ width: 512
21
+ height: 512
22
+ num_inference_steps: 50
23
+ guidance_scale: 12.5
24
+ use_inv_latent: True
25
+ num_inv_steps: 50
26
+
27
+ learning_rate: 3e-5
28
+ train_batch_size: 1
29
+ max_train_steps: 500
30
+ checkpointing_steps: 1000
31
+ validation_steps: 100
32
+ trainable_modules:
33
+ - "attn1.to_q"
34
+ - "attn2.to_q"
35
+ - "attn_temp"
36
+
37
+ seed: 33
38
+ mixed_precision: fp16
39
+ use_8bit_adam: False
40
+ gradient_checkpointing: True
41
+ enable_xformers_memory_efficient_attention: True
Tune-A-Video-debug/data/car-turn.mp4 ADDED
Binary file (942 kB). View file
 
Tune-A-Video-debug/data/man-skiing.mp4 ADDED
Binary file (649 kB). View file
 
Tune-A-Video-debug/data/man-surfing.mp4 ADDED
Binary file (786 kB). View file
 
Tune-A-Video-debug/data/rabbit-watermelon.mp4 ADDED
Binary file (605 kB). View file
 
Tune-A-Video-debug/notebooks/Tune-A-Video.ipynb ADDED
@@ -0,0 +1,385 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "metadata": {
6
+ "id": "fZ_xQvU70UQc"
7
+ },
8
+ "source": [
9
+ "# Tune-A-Video\n",
10
+ "\n",
11
+ "**[Tune-A-Video: One-Shot Tuning of Image Diffusion Models for Text-to-Video Generation](https://arxiv.org/abs/2212.11565)** \n",
12
+ "[Jay Zhangjie Wu](https://zhangjiewu.github.io/), \n",
13
+ "[Yixiao Ge](https://geyixiao.com/), \n",
14
+ "[Xintao Wang](https://xinntao.github.io/), \n",
15
+ "[Stan Weixian Lei](), \n",
16
+ "[Yuchao Gu](https://ycgu.site/), \n",
17
+ "[Wynne Hsu](https://www.comp.nus.edu.sg/~whsu/), \n",
18
+ "[Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en), \n",
19
+ "[Xiaohu Qie](https://scholar.google.com/citations?user=mk-F69UAAAAJ&hl=en), \n",
20
+ "[Mike Zheng Shou](https://sites.google.com/view/showlab) \n",
21
+ "\n",
22
+ "[![Project Website](https://img.shields.io/badge/Project-Website-orange)](https://tuneavideo.github.io/)\n",
23
+ "[![arXiv](https://img.shields.io/badge/arXiv-2212.11565-b31b1b.svg)](https://arxiv.org/abs/2212.11565)\n",
24
+ "[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Tune-A-Video-library/Tune-A-Video-Training-UI)\n",
25
+ "[![GitHub](https://img.shields.io/github/stars/showlab/Tune-A-Video?style=social)](https://github.com/showlab/Tune-A-Video)"
26
+ ]
27
+ },
28
+ {
29
+ "cell_type": "markdown",
30
+ "metadata": {
31
+ "id": "wnTMyW41cC1E"
32
+ },
33
+ "source": [
34
+ "## Setup"
35
+ ]
36
+ },
37
+ {
38
+ "cell_type": "code",
39
+ "execution_count": null,
40
+ "metadata": {
41
+ "id": "XU7NuMAA2drw"
42
+ },
43
+ "outputs": [],
44
+ "source": [
45
+ "#@markdown Check type of GPU and VRAM available.\n",
46
+ "!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader"
47
+ ]
48
+ },
49
+ {
50
+ "cell_type": "code",
51
+ "execution_count": null,
52
+ "metadata": {
53
+ "id": "D1PRgre3Gt5U"
54
+ },
55
+ "outputs": [],
56
+ "source": [
57
+ "#@title Install requirements\n",
58
+ "\n",
59
+ "!git clone https://github.com/showlab/Tune-A-Video.git /content/Tune-A-Video\n",
60
+ "%cd /content/Tune-A-Video \n",
61
+ "# %pip install -r requirements.txt\n",
62
+ "%pip install -q -U --pre triton\n",
63
+ "%pip install -q diffusers[torch]==0.11.1 transformers==4.26.0 bitsandbytes==0.35.4 \\\n",
64
+ "decord accelerate omegaconf einops ftfy gradio imageio-ffmpeg xformers"
65
+ ]
66
+ },
67
+ {
68
+ "cell_type": "code",
69
+ "execution_count": null,
70
+ "metadata": {
71
+ "cellView": "form",
72
+ "id": "m6I6kZNG3Inb"
73
+ },
74
+ "outputs": [],
75
+ "source": [
76
+ "#@title Download pretrained model\n",
77
+ "\n",
78
+ "#@markdown Name/Path of the initial model.\n",
79
+ "MODEL_NAME = \"CompVis/stable-diffusion-v1-4\" #@param {type:\"string\"}\n",
80
+ "\n",
81
+ "#@markdown If model should be download from a remote repo. Untick it if the model is loaded from a local path.\n",
82
+ "download_pretrained_model = True #@param {type:\"boolean\"}\n",
83
+ "if download_pretrained_model:\n",
84
+ " !git lfs install\n",
85
+ " !git clone https://huggingface.co/$MODEL_NAME checkpoints/$MODEL_NAME\n",
86
+ " MODEL_NAME = f\"./checkpoints/{MODEL_NAME}\"\n",
87
+ "print(f\"[*] MODEL_NAME={MODEL_NAME}\")"
88
+ ]
89
+ },
90
+ {
91
+ "cell_type": "markdown",
92
+ "metadata": {
93
+ "id": "qn5ILIyDJIcX"
94
+ },
95
+ "source": [
96
+ "## Usage\n"
97
+ ]
98
+ },
99
+ {
100
+ "cell_type": "markdown",
101
+ "metadata": {
102
+ "id": "REmFAHfz9Y_X"
103
+ },
104
+ "source": [
105
+ "### Training\n"
106
+ ]
107
+ },
108
+ {
109
+ "cell_type": "code",
110
+ "execution_count": null,
111
+ "metadata": {
112
+ "id": "Rxg0y5MBudmd"
113
+ },
114
+ "outputs": [],
115
+ "source": [
116
+ "#@markdown If model weights should be saved directly in google drive (takes around 4-5 GB).\n",
117
+ "save_to_gdrive = False #@param {type:\"boolean\"}\n",
118
+ "if save_to_gdrive:\n",
119
+ " from google.colab import drive\n",
120
+ " drive.mount('/content/drive')\n",
121
+ "\n",
122
+ "#@markdown Enter the directory name to save model at.\n",
123
+ "\n",
124
+ "OUTPUT_DIR = \"outputs/man-skiing\" #@param {type:\"string\"}\n",
125
+ "if save_to_gdrive:\n",
126
+ " OUTPUT_DIR = \"/content/drive/MyDrive/\" + OUTPUT_DIR\n",
127
+ "\n",
128
+ "print(f\"[*] Weights will be saved at {OUTPUT_DIR}\")\n",
129
+ "\n",
130
+ "!mkdir -p $OUTPUT_DIR\n"
131
+ ]
132
+ },
133
+ {
134
+ "cell_type": "code",
135
+ "execution_count": null,
136
+ "metadata": {
137
+ "cellView": "form",
138
+ "id": "32gYIDDR1aCp"
139
+ },
140
+ "outputs": [],
141
+ "source": [
142
+ "#@markdown Upload your video by running this cell.\n",
143
+ "\n",
144
+ "#@markdown OR\n",
145
+ "\n",
146
+ "#@markdown You can use the file manager on the left panel to upload (drag and drop) to `data` folder.\n",
147
+ "\n",
148
+ "import os\n",
149
+ "from google.colab import files\n",
150
+ "import shutil\n",
151
+ "\n",
152
+ "uploaded = files.upload()\n",
153
+ "for filename in uploaded.keys():\n",
154
+ " dst_path = os.path.join(\"data\", filename)\n",
155
+ " shutil.move(filename, dst_path)"
156
+ ]
157
+ },
158
+ {
159
+ "cell_type": "code",
160
+ "execution_count": null,
161
+ "metadata": {
162
+ "id": "wGGFFpNcR2d_"
163
+ },
164
+ "outputs": [],
165
+ "source": [
166
+ "#@markdown Train config\n",
167
+ "\n",
168
+ "from omegaconf import OmegaConf\n",
169
+ "\n",
170
+ "CONFIG_NAME = \"configs/man-skiing.yaml\" #@param {type:\"string\"}\n",
171
+ "\n",
172
+ "train_video_path = \"data/man-skiing.mp4\" #@param {type:\"string\"}\n",
173
+ "train_prompt = \"a man is skiing\" #@param {type:\"string\"}\n",
174
+ "video_length = 8 #@param {type:\"number\"}\n",
175
+ "width = 512 #@param {type:\"number\"}\n",
176
+ "height = 512 #@param {type:\"number\"}\n",
177
+ "learning_rate = 3e-5 #@param {type:\"number\"}\n",
178
+ "train_steps = 300 #@param {type:\"number\"}\n",
179
+ "\n",
180
+ "config = {\n",
181
+ " \"pretrained_model_path\": MODEL_NAME,\n",
182
+ " \"output_dir\": OUTPUT_DIR,\n",
183
+ " \"train_data\": {\n",
184
+ " \"video_path\": train_video_path,\n",
185
+ " \"prompt\": train_prompt,\n",
186
+ " \"n_sample_frames\": video_length,\n",
187
+ " \"width\": width,\n",
188
+ " \"height\": height,\n",
189
+ " \"sample_start_idx\": 0,\n",
190
+ " \"sample_frame_rate\": 2,\n",
191
+ " },\n",
192
+ " \"validation_data\": {\n",
193
+ " \"prompts\": [\n",
194
+ " \"mickey mouse is skiing on the snow\",\n",
195
+ " \"spider man is skiing on the beach, cartoon style\",\n",
196
+ " \"wonder woman, wearing a cowboy hat, is skiing\",\n",
197
+ " \"a man, wearing pink clothes, is skiing at sunset\",\n",
198
+ " ],\n",
199
+ " \"video_length\": video_length,\n",
200
+ " \"width\": width,\n",
201
+ " \"height\": height,\n",
202
+ " \"num_inference_steps\": 20,\n",
203
+ " \"guidance_scale\": 12.5,\n",
204
+ " \"use_inv_latent\": True,\n",
205
+ " \"num_inv_steps\": 50,\n",
206
+ " },\n",
207
+ " \"learning_rate\": learning_rate,\n",
208
+ " \"train_batch_size\": 1,\n",
209
+ " \"max_train_steps\": train_steps,\n",
210
+ " \"checkpointing_steps\": 1000,\n",
211
+ " \"validation_steps\": 100,\n",
212
+ " \"trainable_modules\": [\n",
213
+ " \"attn1.to_q\",\n",
214
+ " \"attn2.to_q\",\n",
215
+ " \"attn_temp\",\n",
216
+ " ],\n",
217
+ " \"seed\": 33,\n",
218
+ " \"mixed_precision\": \"fp16\",\n",
219
+ " \"use_8bit_adam\": False,\n",
220
+ " \"gradient_checkpointing\": True,\n",
221
+ " \"enable_xformers_memory_efficient_attention\": True,\n",
222
+ "}\n",
223
+ "\n",
224
+ "OmegaConf.save(config, CONFIG_NAME)"
225
+ ]
226
+ },
227
+ {
228
+ "cell_type": "code",
229
+ "execution_count": null,
230
+ "metadata": {
231
+ "id": "jjcSXTp-u-Eg"
232
+ },
233
+ "outputs": [],
234
+ "source": [
235
+ "!accelerate launch train_tuneavideo.py --config=$CONFIG_NAME"
236
+ ]
237
+ },
238
+ {
239
+ "cell_type": "markdown",
240
+ "metadata": {
241
+ "id": "ToNG4fd_dTbF"
242
+ },
243
+ "source": [
244
+ "### Inference"
245
+ ]
246
+ },
247
+ {
248
+ "cell_type": "code",
249
+ "execution_count": null,
250
+ "metadata": {
251
+ "id": "91bsSFv2Punm"
252
+ },
253
+ "outputs": [],
254
+ "source": [
255
+ "import torch\n",
256
+ "from torch import autocast\n",
257
+ "from diffusers import DDIMScheduler\n",
258
+ "from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline\n",
259
+ "from tuneavideo.models.unet import UNet3DConditionModel\n",
260
+ "from tuneavideo.util import save_videos_grid\n",
261
+ "\n",
262
+ "\n",
263
+ "unet = UNet3DConditionModel.from_pretrained(OUTPUT_DIR, subfolder='unet', torch_dtype=torch.float16).to('cuda')\n",
264
+ "scheduler = DDIMScheduler.from_pretrained(MODEL_NAME, subfolder='scheduler')\n",
265
+ "pipe = TuneAVideoPipeline.from_pretrained(MODEL_NAME, unet=unet, scheduler=scheduler, torch_dtype=torch.float16).to(\"cuda\")\n",
266
+ "pipe.enable_xformers_memory_efficient_attention()\n",
267
+ "pipe.enable_vae_slicing()\n",
268
+ "\n",
269
+ "g_cuda = None"
270
+ ]
271
+ },
272
+ {
273
+ "cell_type": "code",
274
+ "execution_count": null,
275
+ "metadata": {
276
+ "cellView": "form",
277
+ "id": "oIzkltjpVO_f"
278
+ },
279
+ "outputs": [],
280
+ "source": [
281
+ "#@markdown Can set random seed here for reproducibility.\n",
282
+ "g_cuda = torch.Generator(device='cuda')\n",
283
+ "seed = 1234 #@param {type:\"number\"}\n",
284
+ "g_cuda.manual_seed(seed)"
285
+ ]
286
+ },
287
+ {
288
+ "cell_type": "code",
289
+ "execution_count": null,
290
+ "metadata": {
291
+ "id": "K6xoHWSsbcS3",
292
+ "scrolled": false
293
+ },
294
+ "outputs": [],
295
+ "source": [
296
+ "#@markdown Run for generating videos.\n",
297
+ "\n",
298
+ "prompt = \"iron man is skiing\" #@param {type:\"string\"}\n",
299
+ "negative_prompt = \"\" #@param {type:\"string\"}\n",
300
+ "use_inv_latent = True #@param {type:\"boolean\"}\n",
301
+ "inv_latent_path = \"\" #@param {type:\"string\"}\n",
302
+ "num_samples = 1 #@param {type:\"number\"}\n",
303
+ "guidance_scale = 12.5 #@param {type:\"number\"}\n",
304
+ "num_inference_steps = 50 #@param {type:\"number\"}\n",
305
+ "video_length = 8 #@param {type:\"number\"}\n",
306
+ "height = 512 #@param {type:\"number\"}\n",
307
+ "width = 512 #@param {type:\"number\"}\n",
308
+ "\n",
309
+ "ddim_inv_latent = None\n",
310
+ "if use_inv_latent and inv_latent_path == \"\":\n",
311
+ " from natsort import natsorted\n",
312
+ " from glob import glob\n",
313
+ " import os\n",
314
+ " inv_latent_path = natsorted(glob(f\"{OUTPUT_DIR}/inv_latents/*\"))[-1]\n",
315
+ " ddim_inv_latent = torch.load(inv_latent_path).to(torch.float16)\n",
316
+ " print(f\"DDIM inversion latent loaded from {inv_latent_path}\")\n",
317
+ "\n",
318
+ "with autocast(\"cuda\"), torch.inference_mode():\n",
319
+ " videos = pipe(\n",
320
+ " prompt, \n",
321
+ " latents=ddim_inv_latent,\n",
322
+ " video_length=video_length, \n",
323
+ " height=height, \n",
324
+ " width=width, \n",
325
+ " negative_prompt=negative_prompt,\n",
326
+ " num_videos_per_prompt=num_samples,\n",
327
+ " num_inference_steps=num_inference_steps, \n",
328
+ " guidance_scale=guidance_scale,\n",
329
+ " generator=g_cuda\n",
330
+ " ).videos\n",
331
+ "\n",
332
+ "save_dir = \"./results\" #@param {type:\"string\"}\n",
333
+ "save_path = f\"{save_dir}/{prompt}.gif\"\n",
334
+ "save_videos_grid(videos, save_path)\n",
335
+ "\n",
336
+ "# display\n",
337
+ "from IPython.display import Image, display\n",
338
+ "display(Image(filename=save_path))"
339
+ ]
340
+ },
341
+ {
342
+ "cell_type": "code",
343
+ "execution_count": null,
344
+ "metadata": {
345
+ "id": "jXgi8HM4c-DA"
346
+ },
347
+ "outputs": [],
348
+ "source": [
349
+ "#@markdown Free runtime memory\n",
350
+ "exit()"
351
+ ]
352
+ }
353
+ ],
354
+ "metadata": {
355
+ "accelerator": "GPU",
356
+ "colab": {
357
+ "provenance": []
358
+ },
359
+ "gpuClass": "standard",
360
+ "kernelspec": {
361
+ "display_name": "Python 3 (ipykernel)",
362
+ "language": "python",
363
+ "name": "python3"
364
+ },
365
+ "language_info": {
366
+ "codemirror_mode": {
367
+ "name": "ipython",
368
+ "version": 3
369
+ },
370
+ "file_extension": ".py",
371
+ "mimetype": "text/x-python",
372
+ "name": "python",
373
+ "nbconvert_exporter": "python",
374
+ "pygments_lexer": "ipython3",
375
+ "version": "3.8.13-final"
376
+ },
377
+ "vscode": {
378
+ "interpreter": {
379
+ "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
380
+ }
381
+ }
382
+ },
383
+ "nbformat": 4,
384
+ "nbformat_minor": 0
385
+ }
Tune-A-Video-debug/requirements.txt ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch==1.12.1
2
+ torchvision==0.13.1
3
+ diffusers[torch]==0.11.1
4
+ transformers>=4.25.1
5
+ bitsandbytes==0.35.4
6
+ decord==0.6.0
7
+ accelerate
8
+ tensorboard
9
+ modelcards
10
+ omegaconf
11
+ einops
12
+ imageio
13
+ ftfy
Tune-A-Video-debug/train_tuneavideo.py ADDED
@@ -0,0 +1,367 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import datetime
3
+ import logging
4
+ import inspect
5
+ import math
6
+ import os
7
+ from typing import Dict, Optional, Tuple
8
+ from omegaconf import OmegaConf
9
+
10
+ import torch
11
+ import torch.nn.functional as F
12
+ import torch.utils.checkpoint
13
+
14
+ import diffusers
15
+ import transformers
16
+ from accelerate import Accelerator
17
+ from accelerate.logging import get_logger
18
+ from accelerate.utils import set_seed
19
+ from diffusers import AutoencoderKL, DDPMScheduler, DDIMScheduler
20
+ from diffusers.optimization import get_scheduler
21
+ from diffusers.utils import check_min_version
22
+ from diffusers.utils.import_utils import is_xformers_available
23
+ from tqdm.auto import tqdm
24
+ from transformers import CLIPTextModel, CLIPTokenizer
25
+
26
+ from tuneavideo.models.unet import UNet3DConditionModel
27
+ from tuneavideo.data.dataset import TuneAVideoDataset
28
+ from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
29
+ from tuneavideo.util import save_videos_grid, ddim_inversion
30
+ from einops import rearrange
31
+
32
+
33
+ # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
34
+ check_min_version("0.10.0.dev0")
35
+
36
+ logger = get_logger(__name__, log_level="INFO")
37
+
38
+
39
+ def main(
40
+ pretrained_model_path: str,
41
+ output_dir: str,
42
+ train_data: Dict,
43
+ validation_data: Dict,
44
+ validation_steps: int = 100,
45
+ trainable_modules: Tuple[str] = (
46
+ "attn1.to_q",
47
+ "attn2.to_q",
48
+ "attn_temp",
49
+ ),
50
+ train_batch_size: int = 1,
51
+ max_train_steps: int = 500,
52
+ learning_rate: float = 3e-5,
53
+ scale_lr: bool = False,
54
+ lr_scheduler: str = "constant",
55
+ lr_warmup_steps: int = 0,
56
+ adam_beta1: float = 0.9,
57
+ adam_beta2: float = 0.999,
58
+ adam_weight_decay: float = 1e-2,
59
+ adam_epsilon: float = 1e-08,
60
+ max_grad_norm: float = 1.0,
61
+ gradient_accumulation_steps: int = 1,
62
+ gradient_checkpointing: bool = True,
63
+ checkpointing_steps: int = 500,
64
+ resume_from_checkpoint: Optional[str] = None,
65
+ mixed_precision: Optional[str] = "fp16",
66
+ use_8bit_adam: bool = False,
67
+ enable_xformers_memory_efficient_attention: bool = True,
68
+ seed: Optional[int] = None,
69
+ ):
70
+ *_, config = inspect.getargvalues(inspect.currentframe())
71
+
72
+ accelerator = Accelerator(
73
+ gradient_accumulation_steps=gradient_accumulation_steps,
74
+ mixed_precision=mixed_precision,
75
+ )
76
+
77
+ # Make one log on every process with the configuration for debugging.
78
+ logging.basicConfig(
79
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
80
+ datefmt="%m/%d/%Y %H:%M:%S",
81
+ level=logging.INFO,
82
+ )
83
+ logger.info(accelerator.state, main_process_only=False)
84
+ if accelerator.is_local_main_process:
85
+ transformers.utils.logging.set_verbosity_warning()
86
+ diffusers.utils.logging.set_verbosity_info()
87
+ else:
88
+ transformers.utils.logging.set_verbosity_error()
89
+ diffusers.utils.logging.set_verbosity_error()
90
+
91
+ # If passed along, set the training seed now.
92
+ if seed is not None:
93
+ set_seed(seed)
94
+
95
+ # Handle the output folder creation
96
+ if accelerator.is_main_process:
97
+ # now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
98
+ # output_dir = os.path.join(output_dir, now)
99
+ os.makedirs(output_dir, exist_ok=True)
100
+ os.makedirs(f"{output_dir}/samples", exist_ok=True)
101
+ os.makedirs(f"{output_dir}/inv_latents", exist_ok=True)
102
+ OmegaConf.save(config, os.path.join(output_dir, 'config.yaml'))
103
+
104
+ # Load scheduler, tokenizer and models.
105
+ noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
106
+ tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
107
+ text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
108
+ vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae")
109
+ unet = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet")
110
+
111
+ # Freeze vae and text_encoder
112
+ vae.requires_grad_(False)
113
+ text_encoder.requires_grad_(False)
114
+
115
+ unet.requires_grad_(False)
116
+ for name, module in unet.named_modules():
117
+ if name.endswith(tuple(trainable_modules)):
118
+ for params in module.parameters():
119
+ params.requires_grad = True
120
+
121
+ if enable_xformers_memory_efficient_attention:
122
+ if is_xformers_available():
123
+ unet.enable_xformers_memory_efficient_attention()
124
+ else:
125
+ raise ValueError("xformers is not available. Make sure it is installed correctly")
126
+
127
+ if gradient_checkpointing:
128
+ unet.enable_gradient_checkpointing()
129
+
130
+ if scale_lr:
131
+ learning_rate = (
132
+ learning_rate * gradient_accumulation_steps * train_batch_size * accelerator.num_processes
133
+ )
134
+
135
+ # Initialize the optimizer
136
+ if use_8bit_adam:
137
+ try:
138
+ import bitsandbytes as bnb
139
+ except ImportError:
140
+ raise ImportError(
141
+ "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
142
+ )
143
+
144
+ optimizer_cls = bnb.optim.AdamW8bit
145
+ else:
146
+ optimizer_cls = torch.optim.AdamW
147
+
148
+ optimizer = optimizer_cls(
149
+ unet.parameters(),
150
+ lr=learning_rate,
151
+ betas=(adam_beta1, adam_beta2),
152
+ weight_decay=adam_weight_decay,
153
+ eps=adam_epsilon,
154
+ )
155
+
156
+ # Get the training dataset
157
+ train_dataset = TuneAVideoDataset(**train_data)
158
+
159
+ # Preprocessing the dataset
160
+ train_dataset.prompt_ids = tokenizer(
161
+ train_dataset.prompt, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
162
+ ).input_ids[0]
163
+
164
+ # DataLoaders creation:
165
+ train_dataloader = torch.utils.data.DataLoader(
166
+ train_dataset, batch_size=train_batch_size
167
+ )
168
+
169
+ # Get the validation pipeline
170
+ validation_pipeline = TuneAVideoPipeline(
171
+ vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,
172
+ scheduler=DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
173
+ )
174
+ validation_pipeline.enable_vae_slicing()
175
+ ddim_inv_scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder='scheduler')
176
+ ddim_inv_scheduler.set_timesteps(validation_data.num_inv_steps)
177
+
178
+ # Scheduler
179
+ lr_scheduler = get_scheduler(
180
+ lr_scheduler,
181
+ optimizer=optimizer,
182
+ num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
183
+ num_training_steps=max_train_steps * gradient_accumulation_steps,
184
+ )
185
+
186
+ # Prepare everything with our `accelerator`.
187
+ unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
188
+ unet, optimizer, train_dataloader, lr_scheduler
189
+ )
190
+
191
+ # For mixed precision training we cast the text_encoder and vae weights to half-precision
192
+ # as these models are only used for inference, keeping weights in full precision is not required.
193
+ weight_dtype = torch.float32
194
+ if accelerator.mixed_precision == "fp16":
195
+ weight_dtype = torch.float16
196
+ elif accelerator.mixed_precision == "bf16":
197
+ weight_dtype = torch.bfloat16
198
+
199
+ # Move text_encode and vae to gpu and cast to weight_dtype
200
+ text_encoder.to(accelerator.device, dtype=weight_dtype)
201
+ vae.to(accelerator.device, dtype=weight_dtype)
202
+
203
+ # We need to recalculate our total training steps as the size of the training dataloader may have changed.
204
+ num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
205
+ # Afterwards we recalculate our number of training epochs
206
+ num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch)
207
+
208
+ # We need to initialize the trackers we use, and also store our configuration.
209
+ # The trackers initializes automatically on the main process.
210
+ if accelerator.is_main_process:
211
+ accelerator.init_trackers("text2video-fine-tune")
212
+
213
+ # Train!
214
+ total_batch_size = train_batch_size * accelerator.num_processes * gradient_accumulation_steps
215
+
216
+ logger.info("***** Running training *****")
217
+ logger.info(f" Num examples = {len(train_dataset)}")
218
+ logger.info(f" Num Epochs = {num_train_epochs}")
219
+ logger.info(f" Instantaneous batch size per device = {train_batch_size}")
220
+ logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
221
+ logger.info(f" Gradient Accumulation steps = {gradient_accumulation_steps}")
222
+ logger.info(f" Total optimization steps = {max_train_steps}")
223
+ global_step = 0
224
+ first_epoch = 0
225
+
226
+ # Potentially load in the weights and states from a previous save
227
+ if resume_from_checkpoint:
228
+ if resume_from_checkpoint != "latest":
229
+ path = os.path.basename(resume_from_checkpoint)
230
+ else:
231
+ # Get the most recent checkpoint
232
+ dirs = os.listdir(output_dir)
233
+ dirs = [d for d in dirs if d.startswith("checkpoint")]
234
+ dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
235
+ path = dirs[-1]
236
+ accelerator.print(f"Resuming from checkpoint {path}")
237
+ accelerator.load_state(os.path.join(output_dir, path))
238
+ global_step = int(path.split("-")[1])
239
+
240
+ first_epoch = global_step // num_update_steps_per_epoch
241
+ resume_step = global_step % num_update_steps_per_epoch
242
+
243
+ # Only show the progress bar once on each machine.
244
+ progress_bar = tqdm(range(global_step, max_train_steps), disable=not accelerator.is_local_main_process)
245
+ progress_bar.set_description("Steps")
246
+
247
+ for epoch in range(first_epoch, num_train_epochs):
248
+ unet.train()
249
+ train_loss = 0.0
250
+ for step, batch in enumerate(train_dataloader):
251
+ # Skip steps until we reach the resumed step
252
+ if resume_from_checkpoint and epoch == first_epoch and step < resume_step:
253
+ if step % gradient_accumulation_steps == 0:
254
+ progress_bar.update(1)
255
+ continue
256
+
257
+ with accelerator.accumulate(unet):
258
+ # Convert videos to latent space
259
+ pixel_values = batch["pixel_values"].to(weight_dtype)
260
+ video_length = pixel_values.shape[1]
261
+ pixel_values = rearrange(pixel_values, "b f c h w -> (b f) c h w")
262
+ latents = vae.encode(pixel_values).latent_dist.sample()
263
+ latents = rearrange(latents, "(b f) c h w -> b c f h w", f=video_length)
264
+ latents = latents * 0.18215
265
+
266
+ # Sample noise that we'll add to the latents
267
+ noise = torch.randn_like(latents)
268
+ bsz = latents.shape[0]
269
+ # Sample a random timestep for each video
270
+ timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=latents.device)
271
+ timesteps = timesteps.long()
272
+
273
+ # Add noise to the latents according to the noise magnitude at each timestep
274
+ # (this is the forward diffusion process)
275
+ noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
276
+
277
+ # Get the text embedding for conditioning
278
+ encoder_hidden_states = text_encoder(batch["prompt_ids"])[0]
279
+
280
+ # Get the target for loss depending on the prediction type
281
+ if noise_scheduler.prediction_type == "epsilon":
282
+ target = noise
283
+ elif noise_scheduler.prediction_type == "v_prediction":
284
+ target = noise_scheduler.get_velocity(latents, noise, timesteps)
285
+ else:
286
+ raise ValueError(f"Unknown prediction type {noise_scheduler.prediction_type}")
287
+
288
+ # Predict the noise residual and compute loss
289
+ model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
290
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
291
+
292
+ # Gather the losses across all processes for logging (if we use distributed training).
293
+ avg_loss = accelerator.gather(loss.repeat(train_batch_size)).mean()
294
+ train_loss += avg_loss.item() / gradient_accumulation_steps
295
+
296
+ # Backpropagate
297
+ accelerator.backward(loss)
298
+ if accelerator.sync_gradients:
299
+ accelerator.clip_grad_norm_(unet.parameters(), max_grad_norm)
300
+ optimizer.step()
301
+ lr_scheduler.step()
302
+ optimizer.zero_grad()
303
+
304
+ # Checks if the accelerator has performed an optimization step behind the scenes
305
+ if accelerator.sync_gradients:
306
+ progress_bar.update(1)
307
+ global_step += 1
308
+ accelerator.log({"train_loss": train_loss}, step=global_step)
309
+ train_loss = 0.0
310
+
311
+ if global_step % checkpointing_steps == 0:
312
+ if accelerator.is_main_process:
313
+ save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
314
+ accelerator.save_state(save_path)
315
+ logger.info(f"Saved state to {save_path}")
316
+
317
+ if global_step % validation_steps == 0:
318
+ if accelerator.is_main_process:
319
+ samples = []
320
+ generator = torch.Generator(device=latents.device)
321
+ generator.manual_seed(seed)
322
+
323
+ ddim_inv_latent = None
324
+ if validation_data.use_inv_latent:
325
+ inv_latents_path = os.path.join(output_dir, f"inv_latents/ddim_latent-{global_step}.pt")
326
+ ddim_inv_latent = ddim_inversion(
327
+ validation_pipeline, ddim_inv_scheduler, video_latent=latents,
328
+ num_inv_steps=validation_data.num_inv_steps, prompt="")[-1].to(weight_dtype)
329
+ torch.save(ddim_inv_latent, inv_latents_path)
330
+
331
+ for idx, prompt in enumerate(validation_data.prompts):
332
+ sample = validation_pipeline(prompt, generator=generator, latents=ddim_inv_latent,
333
+ **validation_data).videos
334
+ save_videos_grid(sample, f"{output_dir}/samples/sample-{global_step}/{prompt}.gif")
335
+ samples.append(sample)
336
+ samples = torch.concat(samples)
337
+ save_path = f"{output_dir}/samples/sample-{global_step}.gif"
338
+ save_videos_grid(samples, save_path)
339
+ logger.info(f"Saved samples to {save_path}")
340
+
341
+ logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
342
+ progress_bar.set_postfix(**logs)
343
+
344
+ if global_step >= max_train_steps:
345
+ break
346
+
347
+ # Create the pipeline using the trained modules and save it.
348
+ accelerator.wait_for_everyone()
349
+ if accelerator.is_main_process:
350
+ unet = accelerator.unwrap_model(unet)
351
+ pipeline = TuneAVideoPipeline.from_pretrained(
352
+ pretrained_model_path,
353
+ text_encoder=text_encoder,
354
+ vae=vae,
355
+ unet=unet,
356
+ )
357
+ pipeline.save_pretrained(output_dir)
358
+
359
+ accelerator.end_training()
360
+
361
+
362
+ if __name__ == "__main__":
363
+ parser = argparse.ArgumentParser()
364
+ parser.add_argument("--config", type=str, default="./configs/tuneavideo.yaml")
365
+ args = parser.parse_args()
366
+
367
+ main(**OmegaConf.load(args.config))
Tune-A-Video-debug/tuneavideo/data/dataset.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import decord
2
+ decord.bridge.set_bridge('torch')
3
+
4
+ from torch.utils.data import Dataset
5
+ from einops import rearrange
6
+
7
+
8
+ class TuneAVideoDataset(Dataset):
9
+ def __init__(
10
+ self,
11
+ video_path: str,
12
+ prompt: str,
13
+ width: int = 512,
14
+ height: int = 512,
15
+ n_sample_frames: int = 8,
16
+ sample_start_idx: int = 0,
17
+ sample_frame_rate: int = 1,
18
+ ):
19
+ self.video_path = video_path
20
+ self.prompt = prompt
21
+ self.prompt_ids = None
22
+
23
+ self.width = width
24
+ self.height = height
25
+ self.n_sample_frames = n_sample_frames
26
+ self.sample_start_idx = sample_start_idx
27
+ self.sample_frame_rate = sample_frame_rate
28
+
29
+ def __len__(self):
30
+ return 1
31
+
32
+ def __getitem__(self, index):
33
+ # load and sample video frames
34
+ vr = decord.VideoReader(self.video_path, width=self.width, height=self.height)
35
+ sample_index = list(range(self.sample_start_idx, len(vr), self.sample_frame_rate))[:self.n_sample_frames]
36
+ video = vr.get_batch(sample_index)
37
+ video = rearrange(video, "f h w c -> f c h w")
38
+
39
+ example = {
40
+ "pixel_values": (video / 127.5 - 1.0),
41
+ "prompt_ids": self.prompt_ids
42
+ }
43
+
44
+ return example
Tune-A-Video-debug/tuneavideo/models/attention.py ADDED
@@ -0,0 +1,328 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
2
+
3
+ from dataclasses import dataclass
4
+ from typing import Optional
5
+
6
+ import torch
7
+ import torch.nn.functional as F
8
+ from torch import nn
9
+
10
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
11
+ from diffusers.modeling_utils import ModelMixin
12
+ from diffusers.utils import BaseOutput
13
+ from diffusers.utils.import_utils import is_xformers_available
14
+ from diffusers.models.attention import CrossAttention, FeedForward, AdaLayerNorm
15
+
16
+ from einops import rearrange, repeat
17
+
18
+
19
+ @dataclass
20
+ class Transformer3DModelOutput(BaseOutput):
21
+ sample: torch.FloatTensor
22
+
23
+
24
+ if is_xformers_available():
25
+ import xformers
26
+ import xformers.ops
27
+ else:
28
+ xformers = None
29
+
30
+
31
+ class Transformer3DModel(ModelMixin, ConfigMixin):
32
+ @register_to_config
33
+ def __init__(
34
+ self,
35
+ num_attention_heads: int = 16,
36
+ attention_head_dim: int = 88,
37
+ in_channels: Optional[int] = None,
38
+ num_layers: int = 1,
39
+ dropout: float = 0.0,
40
+ norm_num_groups: int = 32,
41
+ cross_attention_dim: Optional[int] = None,
42
+ attention_bias: bool = False,
43
+ activation_fn: str = "geglu",
44
+ num_embeds_ada_norm: Optional[int] = None,
45
+ use_linear_projection: bool = False,
46
+ only_cross_attention: bool = False,
47
+ upcast_attention: bool = False,
48
+ ):
49
+ super().__init__()
50
+ self.use_linear_projection = use_linear_projection
51
+ self.num_attention_heads = num_attention_heads
52
+ self.attention_head_dim = attention_head_dim
53
+ inner_dim = num_attention_heads * attention_head_dim
54
+
55
+ # Define input layers
56
+ self.in_channels = in_channels
57
+
58
+ self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
59
+ if use_linear_projection:
60
+ self.proj_in = nn.Linear(in_channels, inner_dim)
61
+ else:
62
+ self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
63
+
64
+ # Define transformers blocks
65
+ self.transformer_blocks = nn.ModuleList(
66
+ [
67
+ BasicTransformerBlock(
68
+ inner_dim,
69
+ num_attention_heads,
70
+ attention_head_dim,
71
+ dropout=dropout,
72
+ cross_attention_dim=cross_attention_dim,
73
+ activation_fn=activation_fn,
74
+ num_embeds_ada_norm=num_embeds_ada_norm,
75
+ attention_bias=attention_bias,
76
+ only_cross_attention=only_cross_attention,
77
+ upcast_attention=upcast_attention,
78
+ )
79
+ for d in range(num_layers)
80
+ ]
81
+ )
82
+
83
+ # 4. Define output layers
84
+ if use_linear_projection:
85
+ self.proj_out = nn.Linear(in_channels, inner_dim)
86
+ else:
87
+ self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
88
+
89
+ def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True):
90
+ # Input
91
+ assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
92
+ video_length = hidden_states.shape[2]
93
+ hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
94
+ encoder_hidden_states = repeat(encoder_hidden_states, 'b n c -> (b f) n c', f=video_length)
95
+
96
+ batch, channel, height, weight = hidden_states.shape
97
+ residual = hidden_states
98
+
99
+ hidden_states = self.norm(hidden_states)
100
+ if not self.use_linear_projection:
101
+ hidden_states = self.proj_in(hidden_states)
102
+ inner_dim = hidden_states.shape[1]
103
+ hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
104
+ else:
105
+ inner_dim = hidden_states.shape[1]
106
+ hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
107
+ hidden_states = self.proj_in(hidden_states)
108
+
109
+ # Blocks
110
+ for block in self.transformer_blocks:
111
+ hidden_states = block(
112
+ hidden_states,
113
+ encoder_hidden_states=encoder_hidden_states,
114
+ timestep=timestep,
115
+ video_length=video_length
116
+ )
117
+
118
+ # Output
119
+ if not self.use_linear_projection:
120
+ hidden_states = (
121
+ hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
122
+ )
123
+ hidden_states = self.proj_out(hidden_states)
124
+ else:
125
+ hidden_states = self.proj_out(hidden_states)
126
+ hidden_states = (
127
+ hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
128
+ )
129
+
130
+ output = hidden_states + residual
131
+
132
+ output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
133
+ if not return_dict:
134
+ return (output,)
135
+
136
+ return Transformer3DModelOutput(sample=output)
137
+
138
+
139
+ class BasicTransformerBlock(nn.Module):
140
+ def __init__(
141
+ self,
142
+ dim: int,
143
+ num_attention_heads: int,
144
+ attention_head_dim: int,
145
+ dropout=0.0,
146
+ cross_attention_dim: Optional[int] = None,
147
+ activation_fn: str = "geglu",
148
+ num_embeds_ada_norm: Optional[int] = None,
149
+ attention_bias: bool = False,
150
+ only_cross_attention: bool = False,
151
+ upcast_attention: bool = False,
152
+ ):
153
+ super().__init__()
154
+ self.only_cross_attention = only_cross_attention
155
+ self.use_ada_layer_norm = num_embeds_ada_norm is not None
156
+
157
+ # SC-Attn
158
+ self.attn1 = SparseCausalAttention(
159
+ query_dim=dim,
160
+ heads=num_attention_heads,
161
+ dim_head=attention_head_dim,
162
+ dropout=dropout,
163
+ bias=attention_bias,
164
+ cross_attention_dim=cross_attention_dim if only_cross_attention else None,
165
+ upcast_attention=upcast_attention,
166
+ )
167
+ self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
168
+
169
+ # Cross-Attn
170
+ if cross_attention_dim is not None:
171
+ self.attn2 = CrossAttention(
172
+ query_dim=dim,
173
+ cross_attention_dim=cross_attention_dim,
174
+ heads=num_attention_heads,
175
+ dim_head=attention_head_dim,
176
+ dropout=dropout,
177
+ bias=attention_bias,
178
+ upcast_attention=upcast_attention,
179
+ )
180
+ else:
181
+ self.attn2 = None
182
+
183
+ if cross_attention_dim is not None:
184
+ self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
185
+ else:
186
+ self.norm2 = None
187
+
188
+ # Feed-forward
189
+ self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
190
+ self.norm3 = nn.LayerNorm(dim)
191
+
192
+ # Temp-Attn
193
+ self.attn_temp = CrossAttention(
194
+ query_dim=dim,
195
+ heads=num_attention_heads,
196
+ dim_head=attention_head_dim,
197
+ dropout=dropout,
198
+ bias=attention_bias,
199
+ upcast_attention=upcast_attention,
200
+ )
201
+ nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
202
+ self.norm_temp = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
203
+
204
+ def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool):
205
+ if not is_xformers_available():
206
+ print("Here is how to install it")
207
+ raise ModuleNotFoundError(
208
+ "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
209
+ " xformers",
210
+ name="xformers",
211
+ )
212
+ elif not torch.cuda.is_available():
213
+ raise ValueError(
214
+ "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only"
215
+ " available for GPU "
216
+ )
217
+ else:
218
+ try:
219
+ # Make sure we can run the memory efficient attention
220
+ _ = xformers.ops.memory_efficient_attention(
221
+ torch.randn((1, 2, 40), device="cuda"),
222
+ torch.randn((1, 2, 40), device="cuda"),
223
+ torch.randn((1, 2, 40), device="cuda"),
224
+ )
225
+ except Exception as e:
226
+ raise e
227
+ self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
228
+ if self.attn2 is not None:
229
+ self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
230
+ # self.attn_temp._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
231
+
232
+ def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None):
233
+ # SparseCausal-Attention
234
+ norm_hidden_states = (
235
+ self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)
236
+ )
237
+
238
+ if self.only_cross_attention:
239
+ hidden_states = (
240
+ self.attn1(norm_hidden_states, encoder_hidden_states, attention_mask=attention_mask) + hidden_states
241
+ )
242
+ else:
243
+ hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states
244
+
245
+ if self.attn2 is not None:
246
+ # Cross-Attention
247
+ norm_hidden_states = (
248
+ self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
249
+ )
250
+ hidden_states = (
251
+ self.attn2(
252
+ norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
253
+ )
254
+ + hidden_states
255
+ )
256
+
257
+ # Feed-forward
258
+ hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
259
+
260
+ # Temporal-Attention
261
+ d = hidden_states.shape[1]
262
+ hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length)
263
+ norm_hidden_states = (
264
+ self.norm_temp(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_temp(hidden_states)
265
+ )
266
+ hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
267
+ hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
268
+
269
+ return hidden_states
270
+
271
+
272
+ class SparseCausalAttention(CrossAttention):
273
+ def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):
274
+ batch_size, sequence_length, _ = hidden_states.shape
275
+
276
+ encoder_hidden_states = encoder_hidden_states
277
+
278
+ if self.group_norm is not None:
279
+ hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
280
+
281
+ query = self.to_q(hidden_states)
282
+ dim = query.shape[-1]
283
+ query = self.reshape_heads_to_batch_dim(query)
284
+
285
+ if self.added_kv_proj_dim is not None:
286
+ raise NotImplementedError
287
+
288
+ encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
289
+ key = self.to_k(encoder_hidden_states)
290
+ value = self.to_v(encoder_hidden_states)
291
+
292
+ former_frame_index = torch.arange(video_length) - 1
293
+ former_frame_index[0] = 0
294
+
295
+ key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
296
+ key = torch.cat([key[:, [0] * video_length], key[:, former_frame_index]], dim=2)
297
+ key = rearrange(key, "b f d c -> (b f) d c")
298
+
299
+ value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
300
+ value = torch.cat([value[:, [0] * video_length], value[:, former_frame_index]], dim=2)
301
+ value = rearrange(value, "b f d c -> (b f) d c")
302
+
303
+ key = self.reshape_heads_to_batch_dim(key)
304
+ value = self.reshape_heads_to_batch_dim(value)
305
+
306
+ if attention_mask is not None:
307
+ if attention_mask.shape[-1] != query.shape[1]:
308
+ target_length = query.shape[1]
309
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
310
+ attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
311
+
312
+ # attention, what we cannot get enough of
313
+ if self._use_memory_efficient_attention_xformers:
314
+ hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask)
315
+ # Some versions of xformers return output in fp32, cast it back to the dtype of the input
316
+ hidden_states = hidden_states.to(query.dtype)
317
+ else:
318
+ if self._slice_size is None or query.shape[0] // self._slice_size == 1:
319
+ hidden_states = self._attention(query, key, value, attention_mask)
320
+ else:
321
+ hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask)
322
+
323
+ # linear proj
324
+ hidden_states = self.to_out[0](hidden_states)
325
+
326
+ # dropout
327
+ hidden_states = self.to_out[1](hidden_states)
328
+ return hidden_states
Tune-A-Video-debug/tuneavideo/models/resnet.py ADDED
@@ -0,0 +1,209 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+
7
+ from einops import rearrange
8
+
9
+
10
+ class InflatedConv3d(nn.Conv2d):
11
+ def forward(self, x):
12
+ video_length = x.shape[2]
13
+
14
+ x = rearrange(x, "b c f h w -> (b f) c h w")
15
+ x = super().forward(x)
16
+ x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
17
+
18
+ return x
19
+
20
+
21
+ class Upsample3D(nn.Module):
22
+ def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"):
23
+ super().__init__()
24
+ self.channels = channels
25
+ self.out_channels = out_channels or channels
26
+ self.use_conv = use_conv
27
+ self.use_conv_transpose = use_conv_transpose
28
+ self.name = name
29
+
30
+ conv = None
31
+ if use_conv_transpose:
32
+ raise NotImplementedError
33
+ elif use_conv:
34
+ conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1)
35
+
36
+ if name == "conv":
37
+ self.conv = conv
38
+ else:
39
+ self.Conv2d_0 = conv
40
+
41
+ def forward(self, hidden_states, output_size=None):
42
+ assert hidden_states.shape[1] == self.channels
43
+
44
+ if self.use_conv_transpose:
45
+ raise NotImplementedError
46
+
47
+ # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
48
+ dtype = hidden_states.dtype
49
+ if dtype == torch.bfloat16:
50
+ hidden_states = hidden_states.to(torch.float32)
51
+
52
+ # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
53
+ if hidden_states.shape[0] >= 64:
54
+ hidden_states = hidden_states.contiguous()
55
+
56
+ # if `output_size` is passed we force the interpolation output
57
+ # size and do not make use of `scale_factor=2`
58
+ if output_size is None:
59
+ hidden_states = F.interpolate(hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest")
60
+ else:
61
+ hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
62
+
63
+ # If the input is bfloat16, we cast back to bfloat16
64
+ if dtype == torch.bfloat16:
65
+ hidden_states = hidden_states.to(dtype)
66
+
67
+ if self.use_conv:
68
+ if self.name == "conv":
69
+ hidden_states = self.conv(hidden_states)
70
+ else:
71
+ hidden_states = self.Conv2d_0(hidden_states)
72
+
73
+ return hidden_states
74
+
75
+
76
+ class Downsample3D(nn.Module):
77
+ def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"):
78
+ super().__init__()
79
+ self.channels = channels
80
+ self.out_channels = out_channels or channels
81
+ self.use_conv = use_conv
82
+ self.padding = padding
83
+ stride = 2
84
+ self.name = name
85
+
86
+ if use_conv:
87
+ conv = InflatedConv3d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
88
+ else:
89
+ raise NotImplementedError
90
+
91
+ if name == "conv":
92
+ self.Conv2d_0 = conv
93
+ self.conv = conv
94
+ elif name == "Conv2d_0":
95
+ self.conv = conv
96
+ else:
97
+ self.conv = conv
98
+
99
+ def forward(self, hidden_states):
100
+ assert hidden_states.shape[1] == self.channels
101
+ if self.use_conv and self.padding == 0:
102
+ raise NotImplementedError
103
+
104
+ assert hidden_states.shape[1] == self.channels
105
+ hidden_states = self.conv(hidden_states)
106
+
107
+ return hidden_states
108
+
109
+
110
+ class ResnetBlock3D(nn.Module):
111
+ def __init__(
112
+ self,
113
+ *,
114
+ in_channels,
115
+ out_channels=None,
116
+ conv_shortcut=False,
117
+ dropout=0.0,
118
+ temb_channels=512,
119
+ groups=32,
120
+ groups_out=None,
121
+ pre_norm=True,
122
+ eps=1e-6,
123
+ non_linearity="swish",
124
+ time_embedding_norm="default",
125
+ output_scale_factor=1.0,
126
+ use_in_shortcut=None,
127
+ ):
128
+ super().__init__()
129
+ self.pre_norm = pre_norm
130
+ self.pre_norm = True
131
+ self.in_channels = in_channels
132
+ out_channels = in_channels if out_channels is None else out_channels
133
+ self.out_channels = out_channels
134
+ self.use_conv_shortcut = conv_shortcut
135
+ self.time_embedding_norm = time_embedding_norm
136
+ self.output_scale_factor = output_scale_factor
137
+
138
+ if groups_out is None:
139
+ groups_out = groups
140
+
141
+ self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
142
+
143
+ self.conv1 = InflatedConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
144
+
145
+ if temb_channels is not None:
146
+ if self.time_embedding_norm == "default":
147
+ time_emb_proj_out_channels = out_channels
148
+ elif self.time_embedding_norm == "scale_shift":
149
+ time_emb_proj_out_channels = out_channels * 2
150
+ else:
151
+ raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")
152
+
153
+ self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)
154
+ else:
155
+ self.time_emb_proj = None
156
+
157
+ self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
158
+ self.dropout = torch.nn.Dropout(dropout)
159
+ self.conv2 = InflatedConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
160
+
161
+ if non_linearity == "swish":
162
+ self.nonlinearity = lambda x: F.silu(x)
163
+ elif non_linearity == "mish":
164
+ self.nonlinearity = Mish()
165
+ elif non_linearity == "silu":
166
+ self.nonlinearity = nn.SiLU()
167
+
168
+ self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut
169
+
170
+ self.conv_shortcut = None
171
+ if self.use_in_shortcut:
172
+ self.conv_shortcut = InflatedConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
173
+
174
+ def forward(self, input_tensor, temb):
175
+ hidden_states = input_tensor
176
+
177
+ hidden_states = self.norm1(hidden_states)
178
+ hidden_states = self.nonlinearity(hidden_states)
179
+
180
+ hidden_states = self.conv1(hidden_states)
181
+
182
+ if temb is not None:
183
+ temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None]
184
+
185
+ if temb is not None and self.time_embedding_norm == "default":
186
+ hidden_states = hidden_states + temb
187
+
188
+ hidden_states = self.norm2(hidden_states)
189
+
190
+ if temb is not None and self.time_embedding_norm == "scale_shift":
191
+ scale, shift = torch.chunk(temb, 2, dim=1)
192
+ hidden_states = hidden_states * (1 + scale) + shift
193
+
194
+ hidden_states = self.nonlinearity(hidden_states)
195
+
196
+ hidden_states = self.dropout(hidden_states)
197
+ hidden_states = self.conv2(hidden_states)
198
+
199
+ if self.conv_shortcut is not None:
200
+ input_tensor = self.conv_shortcut(input_tensor)
201
+
202
+ output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
203
+
204
+ return output_tensor
205
+
206
+
207
+ class Mish(torch.nn.Module):
208
+ def forward(self, hidden_states):
209
+ return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))
Tune-A-Video-debug/tuneavideo/models/unet.py ADDED
@@ -0,0 +1,450 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_condition.py
2
+
3
+ from dataclasses import dataclass
4
+ from typing import List, Optional, Tuple, Union
5
+
6
+ import os
7
+ import json
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.utils.checkpoint
12
+
13
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
14
+ from diffusers.modeling_utils import ModelMixin
15
+ from diffusers.utils import BaseOutput, logging
16
+ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
17
+ from .unet_blocks import (
18
+ CrossAttnDownBlock3D,
19
+ CrossAttnUpBlock3D,
20
+ DownBlock3D,
21
+ UNetMidBlock3DCrossAttn,
22
+ UpBlock3D,
23
+ get_down_block,
24
+ get_up_block,
25
+ )
26
+ from .resnet import InflatedConv3d
27
+
28
+
29
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
30
+
31
+
32
+ @dataclass
33
+ class UNet3DConditionOutput(BaseOutput):
34
+ sample: torch.FloatTensor
35
+
36
+
37
+ class UNet3DConditionModel(ModelMixin, ConfigMixin):
38
+ _supports_gradient_checkpointing = True
39
+
40
+ @register_to_config
41
+ def __init__(
42
+ self,
43
+ sample_size: Optional[int] = None,
44
+ in_channels: int = 4,
45
+ out_channels: int = 4,
46
+ center_input_sample: bool = False,
47
+ flip_sin_to_cos: bool = True,
48
+ freq_shift: int = 0,
49
+ down_block_types: Tuple[str] = (
50
+ "CrossAttnDownBlock3D",
51
+ "CrossAttnDownBlock3D",
52
+ "CrossAttnDownBlock3D",
53
+ "DownBlock3D",
54
+ ),
55
+ mid_block_type: str = "UNetMidBlock3DCrossAttn",
56
+ up_block_types: Tuple[str] = (
57
+ "UpBlock3D",
58
+ "CrossAttnUpBlock3D",
59
+ "CrossAttnUpBlock3D",
60
+ "CrossAttnUpBlock3D"
61
+ ),
62
+ only_cross_attention: Union[bool, Tuple[bool]] = False,
63
+ block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
64
+ layers_per_block: int = 2,
65
+ downsample_padding: int = 1,
66
+ mid_block_scale_factor: float = 1,
67
+ act_fn: str = "silu",
68
+ norm_num_groups: int = 32,
69
+ norm_eps: float = 1e-5,
70
+ cross_attention_dim: int = 1280,
71
+ attention_head_dim: Union[int, Tuple[int]] = 8,
72
+ dual_cross_attention: bool = False,
73
+ use_linear_projection: bool = False,
74
+ class_embed_type: Optional[str] = None,
75
+ num_class_embeds: Optional[int] = None,
76
+ upcast_attention: bool = False,
77
+ resnet_time_scale_shift: str = "default",
78
+ ):
79
+ super().__init__()
80
+
81
+ self.sample_size = sample_size
82
+ time_embed_dim = block_out_channels[0] * 4
83
+
84
+ # input
85
+ self.conv_in = InflatedConv3d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))
86
+
87
+ # time
88
+ self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
89
+ timestep_input_dim = block_out_channels[0]
90
+
91
+ self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
92
+
93
+ # class embedding
94
+ if class_embed_type is None and num_class_embeds is not None:
95
+ self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
96
+ elif class_embed_type == "timestep":
97
+ self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
98
+ elif class_embed_type == "identity":
99
+ self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
100
+ else:
101
+ self.class_embedding = None
102
+
103
+ self.down_blocks = nn.ModuleList([])
104
+ self.mid_block = None
105
+ self.up_blocks = nn.ModuleList([])
106
+
107
+ if isinstance(only_cross_attention, bool):
108
+ only_cross_attention = [only_cross_attention] * len(down_block_types)
109
+
110
+ if isinstance(attention_head_dim, int):
111
+ attention_head_dim = (attention_head_dim,) * len(down_block_types)
112
+
113
+ # down
114
+ output_channel = block_out_channels[0]
115
+ for i, down_block_type in enumerate(down_block_types):
116
+ input_channel = output_channel
117
+ output_channel = block_out_channels[i]
118
+ is_final_block = i == len(block_out_channels) - 1
119
+
120
+ down_block = get_down_block(
121
+ down_block_type,
122
+ num_layers=layers_per_block,
123
+ in_channels=input_channel,
124
+ out_channels=output_channel,
125
+ temb_channels=time_embed_dim,
126
+ add_downsample=not is_final_block,
127
+ resnet_eps=norm_eps,
128
+ resnet_act_fn=act_fn,
129
+ resnet_groups=norm_num_groups,
130
+ cross_attention_dim=cross_attention_dim,
131
+ attn_num_head_channels=attention_head_dim[i],
132
+ downsample_padding=downsample_padding,
133
+ dual_cross_attention=dual_cross_attention,
134
+ use_linear_projection=use_linear_projection,
135
+ only_cross_attention=only_cross_attention[i],
136
+ upcast_attention=upcast_attention,
137
+ resnet_time_scale_shift=resnet_time_scale_shift,
138
+ )
139
+ self.down_blocks.append(down_block)
140
+
141
+ # mid
142
+ if mid_block_type == "UNetMidBlock3DCrossAttn":
143
+ self.mid_block = UNetMidBlock3DCrossAttn(
144
+ in_channels=block_out_channels[-1],
145
+ temb_channels=time_embed_dim,
146
+ resnet_eps=norm_eps,
147
+ resnet_act_fn=act_fn,
148
+ output_scale_factor=mid_block_scale_factor,
149
+ resnet_time_scale_shift=resnet_time_scale_shift,
150
+ cross_attention_dim=cross_attention_dim,
151
+ attn_num_head_channels=attention_head_dim[-1],
152
+ resnet_groups=norm_num_groups,
153
+ dual_cross_attention=dual_cross_attention,
154
+ use_linear_projection=use_linear_projection,
155
+ upcast_attention=upcast_attention,
156
+ )
157
+ else:
158
+ raise ValueError(f"unknown mid_block_type : {mid_block_type}")
159
+
160
+ # count how many layers upsample the videos
161
+ self.num_upsamplers = 0
162
+
163
+ # up
164
+ reversed_block_out_channels = list(reversed(block_out_channels))
165
+ reversed_attention_head_dim = list(reversed(attention_head_dim))
166
+ only_cross_attention = list(reversed(only_cross_attention))
167
+ output_channel = reversed_block_out_channels[0]
168
+ for i, up_block_type in enumerate(up_block_types):
169
+ is_final_block = i == len(block_out_channels) - 1
170
+
171
+ prev_output_channel = output_channel
172
+ output_channel = reversed_block_out_channels[i]
173
+ input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
174
+
175
+ # add upsample block for all BUT final layer
176
+ if not is_final_block:
177
+ add_upsample = True
178
+ self.num_upsamplers += 1
179
+ else:
180
+ add_upsample = False
181
+
182
+ up_block = get_up_block(
183
+ up_block_type,
184
+ num_layers=layers_per_block + 1,
185
+ in_channels=input_channel,
186
+ out_channels=output_channel,
187
+ prev_output_channel=prev_output_channel,
188
+ temb_channels=time_embed_dim,
189
+ add_upsample=add_upsample,
190
+ resnet_eps=norm_eps,
191
+ resnet_act_fn=act_fn,
192
+ resnet_groups=norm_num_groups,
193
+ cross_attention_dim=cross_attention_dim,
194
+ attn_num_head_channels=reversed_attention_head_dim[i],
195
+ dual_cross_attention=dual_cross_attention,
196
+ use_linear_projection=use_linear_projection,
197
+ only_cross_attention=only_cross_attention[i],
198
+ upcast_attention=upcast_attention,
199
+ resnet_time_scale_shift=resnet_time_scale_shift,
200
+ )
201
+ self.up_blocks.append(up_block)
202
+ prev_output_channel = output_channel
203
+
204
+ # out
205
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps)
206
+ self.conv_act = nn.SiLU()
207
+ self.conv_out = InflatedConv3d(block_out_channels[0], out_channels, kernel_size=3, padding=1)
208
+
209
+ def set_attention_slice(self, slice_size):
210
+ r"""
211
+ Enable sliced attention computation.
212
+
213
+ When this option is enabled, the attention module will split the input tensor in slices, to compute attention
214
+ in several steps. This is useful to save some memory in exchange for a small speed decrease.
215
+
216
+ Args:
217
+ slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
218
+ When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
219
+ `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
220
+ provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
221
+ must be a multiple of `slice_size`.
222
+ """
223
+ sliceable_head_dims = []
224
+
225
+ def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
226
+ if hasattr(module, "set_attention_slice"):
227
+ sliceable_head_dims.append(module.sliceable_head_dim)
228
+
229
+ for child in module.children():
230
+ fn_recursive_retrieve_slicable_dims(child)
231
+
232
+ # retrieve number of attention layers
233
+ for module in self.children():
234
+ fn_recursive_retrieve_slicable_dims(module)
235
+
236
+ num_slicable_layers = len(sliceable_head_dims)
237
+
238
+ if slice_size == "auto":
239
+ # half the attention head size is usually a good trade-off between
240
+ # speed and memory
241
+ slice_size = [dim // 2 for dim in sliceable_head_dims]
242
+ elif slice_size == "max":
243
+ # make smallest slice possible
244
+ slice_size = num_slicable_layers * [1]
245
+
246
+ slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
247
+
248
+ if len(slice_size) != len(sliceable_head_dims):
249
+ raise ValueError(
250
+ f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
251
+ f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
252
+ )
253
+
254
+ for i in range(len(slice_size)):
255
+ size = slice_size[i]
256
+ dim = sliceable_head_dims[i]
257
+ if size is not None and size > dim:
258
+ raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
259
+
260
+ # Recursively walk through all the children.
261
+ # Any children which exposes the set_attention_slice method
262
+ # gets the message
263
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
264
+ if hasattr(module, "set_attention_slice"):
265
+ module.set_attention_slice(slice_size.pop())
266
+
267
+ for child in module.children():
268
+ fn_recursive_set_attention_slice(child, slice_size)
269
+
270
+ reversed_slice_size = list(reversed(slice_size))
271
+ for module in self.children():
272
+ fn_recursive_set_attention_slice(module, reversed_slice_size)
273
+
274
+ def _set_gradient_checkpointing(self, module, value=False):
275
+ if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):
276
+ module.gradient_checkpointing = value
277
+
278
+ def forward(
279
+ self,
280
+ sample: torch.FloatTensor,
281
+ timestep: Union[torch.Tensor, float, int],
282
+ encoder_hidden_states: torch.Tensor,
283
+ class_labels: Optional[torch.Tensor] = None,
284
+ attention_mask: Optional[torch.Tensor] = None,
285
+ return_dict: bool = True,
286
+ ) -> Union[UNet3DConditionOutput, Tuple]:
287
+ r"""
288
+ Args:
289
+ sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
290
+ timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
291
+ encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
292
+ return_dict (`bool`, *optional*, defaults to `True`):
293
+ Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
294
+
295
+ Returns:
296
+ [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
297
+ [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
298
+ returning a tuple, the first element is the sample tensor.
299
+ """
300
+ # By default samples have to be AT least a multiple of the overall upsampling factor.
301
+ # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
302
+ # However, the upsampling interpolation output size can be forced to fit any upsampling size
303
+ # on the fly if necessary.
304
+ default_overall_up_factor = 2**self.num_upsamplers
305
+
306
+ # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
307
+ forward_upsample_size = False
308
+ upsample_size = None
309
+
310
+ if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
311
+ logger.info("Forward upsample size to force interpolation output size.")
312
+ forward_upsample_size = True
313
+
314
+ # prepare attention_mask
315
+ if attention_mask is not None:
316
+ attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
317
+ attention_mask = attention_mask.unsqueeze(1)
318
+
319
+ # center input if necessary
320
+ if self.config.center_input_sample:
321
+ sample = 2 * sample - 1.0
322
+
323
+ # time
324
+ timesteps = timestep
325
+ if not torch.is_tensor(timesteps):
326
+ # This would be a good case for the `match` statement (Python 3.10+)
327
+ is_mps = sample.device.type == "mps"
328
+ if isinstance(timestep, float):
329
+ dtype = torch.float32 if is_mps else torch.float64
330
+ else:
331
+ dtype = torch.int32 if is_mps else torch.int64
332
+ timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
333
+ elif len(timesteps.shape) == 0:
334
+ timesteps = timesteps[None].to(sample.device)
335
+
336
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
337
+ timesteps = timesteps.expand(sample.shape[0])
338
+
339
+ t_emb = self.time_proj(timesteps)
340
+
341
+ # timesteps does not contain any weights and will always return f32 tensors
342
+ # but time_embedding might actually be running in fp16. so we need to cast here.
343
+ # there might be better ways to encapsulate this.
344
+ t_emb = t_emb.to(dtype=self.dtype)
345
+ emb = self.time_embedding(t_emb)
346
+
347
+ if self.class_embedding is not None:
348
+ if class_labels is None:
349
+ raise ValueError("class_labels should be provided when num_class_embeds > 0")
350
+
351
+ if self.config.class_embed_type == "timestep":
352
+ class_labels = self.time_proj(class_labels)
353
+
354
+ class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
355
+ emb = emb + class_emb
356
+
357
+ # pre-process
358
+ sample = self.conv_in(sample)
359
+
360
+ # down
361
+ down_block_res_samples = (sample,)
362
+ for downsample_block in self.down_blocks:
363
+ if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
364
+ sample, res_samples = downsample_block(
365
+ hidden_states=sample,
366
+ temb=emb,
367
+ encoder_hidden_states=encoder_hidden_states,
368
+ attention_mask=attention_mask,
369
+ )
370
+ else:
371
+ sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
372
+
373
+ down_block_res_samples += res_samples
374
+
375
+ # mid
376
+ sample = self.mid_block(
377
+ sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
378
+ )
379
+
380
+ # up
381
+ for i, upsample_block in enumerate(self.up_blocks):
382
+ is_final_block = i == len(self.up_blocks) - 1
383
+
384
+ res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
385
+ down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
386
+
387
+ # if we have not reached the final block and need to forward the
388
+ # upsample size, we do it here
389
+ if not is_final_block and forward_upsample_size:
390
+ upsample_size = down_block_res_samples[-1].shape[2:]
391
+
392
+ if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
393
+ sample = upsample_block(
394
+ hidden_states=sample,
395
+ temb=emb,
396
+ res_hidden_states_tuple=res_samples,
397
+ encoder_hidden_states=encoder_hidden_states,
398
+ upsample_size=upsample_size,
399
+ attention_mask=attention_mask,
400
+ )
401
+ else:
402
+ sample = upsample_block(
403
+ hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
404
+ )
405
+ # post-process
406
+ sample = self.conv_norm_out(sample)
407
+ sample = self.conv_act(sample)
408
+ sample = self.conv_out(sample)
409
+
410
+ if not return_dict:
411
+ return (sample,)
412
+
413
+ return UNet3DConditionOutput(sample=sample)
414
+
415
+ @classmethod
416
+ def from_pretrained_2d(cls, pretrained_model_path, subfolder=None):
417
+ if subfolder is not None:
418
+ pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
419
+
420
+ config_file = os.path.join(pretrained_model_path, 'config.json')
421
+ if not os.path.isfile(config_file):
422
+ raise RuntimeError(f"{config_file} does not exist")
423
+ with open(config_file, "r") as f:
424
+ config = json.load(f)
425
+ config["_class_name"] = cls.__name__
426
+ config["down_block_types"] = [
427
+ "CrossAttnDownBlock3D",
428
+ "CrossAttnDownBlock3D",
429
+ "CrossAttnDownBlock3D",
430
+ "DownBlock3D"
431
+ ]
432
+ config["up_block_types"] = [
433
+ "UpBlock3D",
434
+ "CrossAttnUpBlock3D",
435
+ "CrossAttnUpBlock3D",
436
+ "CrossAttnUpBlock3D"
437
+ ]
438
+
439
+ from diffusers.utils import WEIGHTS_NAME
440
+ model = cls.from_config(config)
441
+ model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
442
+ if not os.path.isfile(model_file):
443
+ raise RuntimeError(f"{model_file} does not exist")
444
+ state_dict = torch.load(model_file, map_location="cpu")
445
+ for k, v in model.state_dict().items():
446
+ if '_temp.' in k:
447
+ state_dict.update({k: v})
448
+ model.load_state_dict(state_dict)
449
+
450
+ return model
Tune-A-Video-debug/tuneavideo/models/unet_blocks.py ADDED
@@ -0,0 +1,588 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py
2
+
3
+ import torch
4
+ from torch import nn
5
+
6
+ from .attention import Transformer3DModel
7
+ from .resnet import Downsample3D, ResnetBlock3D, Upsample3D
8
+
9
+
10
+ def get_down_block(
11
+ down_block_type,
12
+ num_layers,
13
+ in_channels,
14
+ out_channels,
15
+ temb_channels,
16
+ add_downsample,
17
+ resnet_eps,
18
+ resnet_act_fn,
19
+ attn_num_head_channels,
20
+ resnet_groups=None,
21
+ cross_attention_dim=None,
22
+ downsample_padding=None,
23
+ dual_cross_attention=False,
24
+ use_linear_projection=False,
25
+ only_cross_attention=False,
26
+ upcast_attention=False,
27
+ resnet_time_scale_shift="default",
28
+ ):
29
+ down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
30
+ if down_block_type == "DownBlock3D":
31
+ return DownBlock3D(
32
+ num_layers=num_layers,
33
+ in_channels=in_channels,
34
+ out_channels=out_channels,
35
+ temb_channels=temb_channels,
36
+ add_downsample=add_downsample,
37
+ resnet_eps=resnet_eps,
38
+ resnet_act_fn=resnet_act_fn,
39
+ resnet_groups=resnet_groups,
40
+ downsample_padding=downsample_padding,
41
+ resnet_time_scale_shift=resnet_time_scale_shift,
42
+ )
43
+ elif down_block_type == "CrossAttnDownBlock3D":
44
+ if cross_attention_dim is None:
45
+ raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D")
46
+ return CrossAttnDownBlock3D(
47
+ num_layers=num_layers,
48
+ in_channels=in_channels,
49
+ out_channels=out_channels,
50
+ temb_channels=temb_channels,
51
+ add_downsample=add_downsample,
52
+ resnet_eps=resnet_eps,
53
+ resnet_act_fn=resnet_act_fn,
54
+ resnet_groups=resnet_groups,
55
+ downsample_padding=downsample_padding,
56
+ cross_attention_dim=cross_attention_dim,
57
+ attn_num_head_channels=attn_num_head_channels,
58
+ dual_cross_attention=dual_cross_attention,
59
+ use_linear_projection=use_linear_projection,
60
+ only_cross_attention=only_cross_attention,
61
+ upcast_attention=upcast_attention,
62
+ resnet_time_scale_shift=resnet_time_scale_shift,
63
+ )
64
+ raise ValueError(f"{down_block_type} does not exist.")
65
+
66
+
67
+ def get_up_block(
68
+ up_block_type,
69
+ num_layers,
70
+ in_channels,
71
+ out_channels,
72
+ prev_output_channel,
73
+ temb_channels,
74
+ add_upsample,
75
+ resnet_eps,
76
+ resnet_act_fn,
77
+ attn_num_head_channels,
78
+ resnet_groups=None,
79
+ cross_attention_dim=None,
80
+ dual_cross_attention=False,
81
+ use_linear_projection=False,
82
+ only_cross_attention=False,
83
+ upcast_attention=False,
84
+ resnet_time_scale_shift="default",
85
+ ):
86
+ up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
87
+ if up_block_type == "UpBlock3D":
88
+ return UpBlock3D(
89
+ num_layers=num_layers,
90
+ in_channels=in_channels,
91
+ out_channels=out_channels,
92
+ prev_output_channel=prev_output_channel,
93
+ temb_channels=temb_channels,
94
+ add_upsample=add_upsample,
95
+ resnet_eps=resnet_eps,
96
+ resnet_act_fn=resnet_act_fn,
97
+ resnet_groups=resnet_groups,
98
+ resnet_time_scale_shift=resnet_time_scale_shift,
99
+ )
100
+ elif up_block_type == "CrossAttnUpBlock3D":
101
+ if cross_attention_dim is None:
102
+ raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D")
103
+ return CrossAttnUpBlock3D(
104
+ num_layers=num_layers,
105
+ in_channels=in_channels,
106
+ out_channels=out_channels,
107
+ prev_output_channel=prev_output_channel,
108
+ temb_channels=temb_channels,
109
+ add_upsample=add_upsample,
110
+ resnet_eps=resnet_eps,
111
+ resnet_act_fn=resnet_act_fn,
112
+ resnet_groups=resnet_groups,
113
+ cross_attention_dim=cross_attention_dim,
114
+ attn_num_head_channels=attn_num_head_channels,
115
+ dual_cross_attention=dual_cross_attention,
116
+ use_linear_projection=use_linear_projection,
117
+ only_cross_attention=only_cross_attention,
118
+ upcast_attention=upcast_attention,
119
+ resnet_time_scale_shift=resnet_time_scale_shift,
120
+ )
121
+ raise ValueError(f"{up_block_type} does not exist.")
122
+
123
+
124
+ class UNetMidBlock3DCrossAttn(nn.Module):
125
+ def __init__(
126
+ self,
127
+ in_channels: int,
128
+ temb_channels: int,
129
+ dropout: float = 0.0,
130
+ num_layers: int = 1,
131
+ resnet_eps: float = 1e-6,
132
+ resnet_time_scale_shift: str = "default",
133
+ resnet_act_fn: str = "swish",
134
+ resnet_groups: int = 32,
135
+ resnet_pre_norm: bool = True,
136
+ attn_num_head_channels=1,
137
+ output_scale_factor=1.0,
138
+ cross_attention_dim=1280,
139
+ dual_cross_attention=False,
140
+ use_linear_projection=False,
141
+ upcast_attention=False,
142
+ ):
143
+ super().__init__()
144
+
145
+ self.has_cross_attention = True
146
+ self.attn_num_head_channels = attn_num_head_channels
147
+ resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
148
+
149
+ # there is always at least one resnet
150
+ resnets = [
151
+ ResnetBlock3D(
152
+ in_channels=in_channels,
153
+ out_channels=in_channels,
154
+ temb_channels=temb_channels,
155
+ eps=resnet_eps,
156
+ groups=resnet_groups,
157
+ dropout=dropout,
158
+ time_embedding_norm=resnet_time_scale_shift,
159
+ non_linearity=resnet_act_fn,
160
+ output_scale_factor=output_scale_factor,
161
+ pre_norm=resnet_pre_norm,
162
+ )
163
+ ]
164
+ attentions = []
165
+
166
+ for _ in range(num_layers):
167
+ if dual_cross_attention:
168
+ raise NotImplementedError
169
+ attentions.append(
170
+ Transformer3DModel(
171
+ attn_num_head_channels,
172
+ in_channels // attn_num_head_channels,
173
+ in_channels=in_channels,
174
+ num_layers=1,
175
+ cross_attention_dim=cross_attention_dim,
176
+ norm_num_groups=resnet_groups,
177
+ use_linear_projection=use_linear_projection,
178
+ upcast_attention=upcast_attention,
179
+ )
180
+ )
181
+ resnets.append(
182
+ ResnetBlock3D(
183
+ in_channels=in_channels,
184
+ out_channels=in_channels,
185
+ temb_channels=temb_channels,
186
+ eps=resnet_eps,
187
+ groups=resnet_groups,
188
+ dropout=dropout,
189
+ time_embedding_norm=resnet_time_scale_shift,
190
+ non_linearity=resnet_act_fn,
191
+ output_scale_factor=output_scale_factor,
192
+ pre_norm=resnet_pre_norm,
193
+ )
194
+ )
195
+
196
+ self.attentions = nn.ModuleList(attentions)
197
+ self.resnets = nn.ModuleList(resnets)
198
+
199
+ def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
200
+ hidden_states = self.resnets[0](hidden_states, temb)
201
+ for attn, resnet in zip(self.attentions, self.resnets[1:]):
202
+ hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
203
+ hidden_states = resnet(hidden_states, temb)
204
+
205
+ return hidden_states
206
+
207
+
208
+ class CrossAttnDownBlock3D(nn.Module):
209
+ def __init__(
210
+ self,
211
+ in_channels: int,
212
+ out_channels: int,
213
+ temb_channels: int,
214
+ dropout: float = 0.0,
215
+ num_layers: int = 1,
216
+ resnet_eps: float = 1e-6,
217
+ resnet_time_scale_shift: str = "default",
218
+ resnet_act_fn: str = "swish",
219
+ resnet_groups: int = 32,
220
+ resnet_pre_norm: bool = True,
221
+ attn_num_head_channels=1,
222
+ cross_attention_dim=1280,
223
+ output_scale_factor=1.0,
224
+ downsample_padding=1,
225
+ add_downsample=True,
226
+ dual_cross_attention=False,
227
+ use_linear_projection=False,
228
+ only_cross_attention=False,
229
+ upcast_attention=False,
230
+ ):
231
+ super().__init__()
232
+ resnets = []
233
+ attentions = []
234
+
235
+ self.has_cross_attention = True
236
+ self.attn_num_head_channels = attn_num_head_channels
237
+
238
+ for i in range(num_layers):
239
+ in_channels = in_channels if i == 0 else out_channels
240
+ resnets.append(
241
+ ResnetBlock3D(
242
+ in_channels=in_channels,
243
+ out_channels=out_channels,
244
+ temb_channels=temb_channels,
245
+ eps=resnet_eps,
246
+ groups=resnet_groups,
247
+ dropout=dropout,
248
+ time_embedding_norm=resnet_time_scale_shift,
249
+ non_linearity=resnet_act_fn,
250
+ output_scale_factor=output_scale_factor,
251
+ pre_norm=resnet_pre_norm,
252
+ )
253
+ )
254
+ if dual_cross_attention:
255
+ raise NotImplementedError
256
+ attentions.append(
257
+ Transformer3DModel(
258
+ attn_num_head_channels,
259
+ out_channels // attn_num_head_channels,
260
+ in_channels=out_channels,
261
+ num_layers=1,
262
+ cross_attention_dim=cross_attention_dim,
263
+ norm_num_groups=resnet_groups,
264
+ use_linear_projection=use_linear_projection,
265
+ only_cross_attention=only_cross_attention,
266
+ upcast_attention=upcast_attention,
267
+ )
268
+ )
269
+ self.attentions = nn.ModuleList(attentions)
270
+ self.resnets = nn.ModuleList(resnets)
271
+
272
+ if add_downsample:
273
+ self.downsamplers = nn.ModuleList(
274
+ [
275
+ Downsample3D(
276
+ out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
277
+ )
278
+ ]
279
+ )
280
+ else:
281
+ self.downsamplers = None
282
+
283
+ self.gradient_checkpointing = False
284
+
285
+ def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
286
+ output_states = ()
287
+
288
+ for resnet, attn in zip(self.resnets, self.attentions):
289
+ if self.training and self.gradient_checkpointing:
290
+
291
+ def create_custom_forward(module, return_dict=None):
292
+ def custom_forward(*inputs):
293
+ if return_dict is not None:
294
+ return module(*inputs, return_dict=return_dict)
295
+ else:
296
+ return module(*inputs)
297
+
298
+ return custom_forward
299
+
300
+ hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
301
+ hidden_states = torch.utils.checkpoint.checkpoint(
302
+ create_custom_forward(attn, return_dict=False),
303
+ hidden_states,
304
+ encoder_hidden_states,
305
+ )[0]
306
+ else:
307
+ hidden_states = resnet(hidden_states, temb)
308
+ hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
309
+
310
+ output_states += (hidden_states,)
311
+
312
+ if self.downsamplers is not None:
313
+ for downsampler in self.downsamplers:
314
+ hidden_states = downsampler(hidden_states)
315
+
316
+ output_states += (hidden_states,)
317
+
318
+ return hidden_states, output_states
319
+
320
+
321
+ class DownBlock3D(nn.Module):
322
+ def __init__(
323
+ self,
324
+ in_channels: int,
325
+ out_channels: int,
326
+ temb_channels: int,
327
+ dropout: float = 0.0,
328
+ num_layers: int = 1,
329
+ resnet_eps: float = 1e-6,
330
+ resnet_time_scale_shift: str = "default",
331
+ resnet_act_fn: str = "swish",
332
+ resnet_groups: int = 32,
333
+ resnet_pre_norm: bool = True,
334
+ output_scale_factor=1.0,
335
+ add_downsample=True,
336
+ downsample_padding=1,
337
+ ):
338
+ super().__init__()
339
+ resnets = []
340
+
341
+ for i in range(num_layers):
342
+ in_channels = in_channels if i == 0 else out_channels
343
+ resnets.append(
344
+ ResnetBlock3D(
345
+ in_channels=in_channels,
346
+ out_channels=out_channels,
347
+ temb_channels=temb_channels,
348
+ eps=resnet_eps,
349
+ groups=resnet_groups,
350
+ dropout=dropout,
351
+ time_embedding_norm=resnet_time_scale_shift,
352
+ non_linearity=resnet_act_fn,
353
+ output_scale_factor=output_scale_factor,
354
+ pre_norm=resnet_pre_norm,
355
+ )
356
+ )
357
+
358
+ self.resnets = nn.ModuleList(resnets)
359
+
360
+ if add_downsample:
361
+ self.downsamplers = nn.ModuleList(
362
+ [
363
+ Downsample3D(
364
+ out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
365
+ )
366
+ ]
367
+ )
368
+ else:
369
+ self.downsamplers = None
370
+
371
+ self.gradient_checkpointing = False
372
+
373
+ def forward(self, hidden_states, temb=None):
374
+ output_states = ()
375
+
376
+ for resnet in self.resnets:
377
+ if self.training and self.gradient_checkpointing:
378
+
379
+ def create_custom_forward(module):
380
+ def custom_forward(*inputs):
381
+ return module(*inputs)
382
+
383
+ return custom_forward
384
+
385
+ hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
386
+ else:
387
+ hidden_states = resnet(hidden_states, temb)
388
+
389
+ output_states += (hidden_states,)
390
+
391
+ if self.downsamplers is not None:
392
+ for downsampler in self.downsamplers:
393
+ hidden_states = downsampler(hidden_states)
394
+
395
+ output_states += (hidden_states,)
396
+
397
+ return hidden_states, output_states
398
+
399
+
400
+ class CrossAttnUpBlock3D(nn.Module):
401
+ def __init__(
402
+ self,
403
+ in_channels: int,
404
+ out_channels: int,
405
+ prev_output_channel: int,
406
+ temb_channels: int,
407
+ dropout: float = 0.0,
408
+ num_layers: int = 1,
409
+ resnet_eps: float = 1e-6,
410
+ resnet_time_scale_shift: str = "default",
411
+ resnet_act_fn: str = "swish",
412
+ resnet_groups: int = 32,
413
+ resnet_pre_norm: bool = True,
414
+ attn_num_head_channels=1,
415
+ cross_attention_dim=1280,
416
+ output_scale_factor=1.0,
417
+ add_upsample=True,
418
+ dual_cross_attention=False,
419
+ use_linear_projection=False,
420
+ only_cross_attention=False,
421
+ upcast_attention=False,
422
+ ):
423
+ super().__init__()
424
+ resnets = []
425
+ attentions = []
426
+
427
+ self.has_cross_attention = True
428
+ self.attn_num_head_channels = attn_num_head_channels
429
+
430
+ for i in range(num_layers):
431
+ res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
432
+ resnet_in_channels = prev_output_channel if i == 0 else out_channels
433
+
434
+ resnets.append(
435
+ ResnetBlock3D(
436
+ in_channels=resnet_in_channels + res_skip_channels,
437
+ out_channels=out_channels,
438
+ temb_channels=temb_channels,
439
+ eps=resnet_eps,
440
+ groups=resnet_groups,
441
+ dropout=dropout,
442
+ time_embedding_norm=resnet_time_scale_shift,
443
+ non_linearity=resnet_act_fn,
444
+ output_scale_factor=output_scale_factor,
445
+ pre_norm=resnet_pre_norm,
446
+ )
447
+ )
448
+ if dual_cross_attention:
449
+ raise NotImplementedError
450
+ attentions.append(
451
+ Transformer3DModel(
452
+ attn_num_head_channels,
453
+ out_channels // attn_num_head_channels,
454
+ in_channels=out_channels,
455
+ num_layers=1,
456
+ cross_attention_dim=cross_attention_dim,
457
+ norm_num_groups=resnet_groups,
458
+ use_linear_projection=use_linear_projection,
459
+ only_cross_attention=only_cross_attention,
460
+ upcast_attention=upcast_attention,
461
+ )
462
+ )
463
+
464
+ self.attentions = nn.ModuleList(attentions)
465
+ self.resnets = nn.ModuleList(resnets)
466
+
467
+ if add_upsample:
468
+ self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)])
469
+ else:
470
+ self.upsamplers = None
471
+
472
+ self.gradient_checkpointing = False
473
+
474
+ def forward(
475
+ self,
476
+ hidden_states,
477
+ res_hidden_states_tuple,
478
+ temb=None,
479
+ encoder_hidden_states=None,
480
+ upsample_size=None,
481
+ attention_mask=None,
482
+ ):
483
+ for resnet, attn in zip(self.resnets, self.attentions):
484
+ # pop res hidden states
485
+ res_hidden_states = res_hidden_states_tuple[-1]
486
+ res_hidden_states_tuple = res_hidden_states_tuple[:-1]
487
+ hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
488
+
489
+ if self.training and self.gradient_checkpointing:
490
+
491
+ def create_custom_forward(module, return_dict=None):
492
+ def custom_forward(*inputs):
493
+ if return_dict is not None:
494
+ return module(*inputs, return_dict=return_dict)
495
+ else:
496
+ return module(*inputs)
497
+
498
+ return custom_forward
499
+
500
+ hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
501
+ hidden_states = torch.utils.checkpoint.checkpoint(
502
+ create_custom_forward(attn, return_dict=False),
503
+ hidden_states,
504
+ encoder_hidden_states,
505
+ )[0]
506
+ else:
507
+ hidden_states = resnet(hidden_states, temb)
508
+ hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
509
+
510
+ if self.upsamplers is not None:
511
+ for upsampler in self.upsamplers:
512
+ hidden_states = upsampler(hidden_states, upsample_size)
513
+
514
+ return hidden_states
515
+
516
+
517
+ class UpBlock3D(nn.Module):
518
+ def __init__(
519
+ self,
520
+ in_channels: int,
521
+ prev_output_channel: int,
522
+ out_channels: int,
523
+ temb_channels: int,
524
+ dropout: float = 0.0,
525
+ num_layers: int = 1,
526
+ resnet_eps: float = 1e-6,
527
+ resnet_time_scale_shift: str = "default",
528
+ resnet_act_fn: str = "swish",
529
+ resnet_groups: int = 32,
530
+ resnet_pre_norm: bool = True,
531
+ output_scale_factor=1.0,
532
+ add_upsample=True,
533
+ ):
534
+ super().__init__()
535
+ resnets = []
536
+
537
+ for i in range(num_layers):
538
+ res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
539
+ resnet_in_channels = prev_output_channel if i == 0 else out_channels
540
+
541
+ resnets.append(
542
+ ResnetBlock3D(
543
+ in_channels=resnet_in_channels + res_skip_channels,
544
+ out_channels=out_channels,
545
+ temb_channels=temb_channels,
546
+ eps=resnet_eps,
547
+ groups=resnet_groups,
548
+ dropout=dropout,
549
+ time_embedding_norm=resnet_time_scale_shift,
550
+ non_linearity=resnet_act_fn,
551
+ output_scale_factor=output_scale_factor,
552
+ pre_norm=resnet_pre_norm,
553
+ )
554
+ )
555
+
556
+ self.resnets = nn.ModuleList(resnets)
557
+
558
+ if add_upsample:
559
+ self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)])
560
+ else:
561
+ self.upsamplers = None
562
+
563
+ self.gradient_checkpointing = False
564
+
565
+ def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
566
+ for resnet in self.resnets:
567
+ # pop res hidden states
568
+ res_hidden_states = res_hidden_states_tuple[-1]
569
+ res_hidden_states_tuple = res_hidden_states_tuple[:-1]
570
+ hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
571
+
572
+ if self.training and self.gradient_checkpointing:
573
+
574
+ def create_custom_forward(module):
575
+ def custom_forward(*inputs):
576
+ return module(*inputs)
577
+
578
+ return custom_forward
579
+
580
+ hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
581
+ else:
582
+ hidden_states = resnet(hidden_states, temb)
583
+
584
+ if self.upsamplers is not None:
585
+ for upsampler in self.upsamplers:
586
+ hidden_states = upsampler(hidden_states, upsample_size)
587
+
588
+ return hidden_states
Tune-A-Video-debug/tuneavideo/pipelines/pipeline_tuneavideo.py ADDED
@@ -0,0 +1,407 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
2
+
3
+ import inspect
4
+ from typing import Callable, List, Optional, Union
5
+ from dataclasses import dataclass
6
+
7
+ import numpy as np
8
+ import torch
9
+
10
+ from diffusers.utils import is_accelerate_available
11
+ from packaging import version
12
+ from transformers import CLIPTextModel, CLIPTokenizer
13
+
14
+ from diffusers.configuration_utils import FrozenDict
15
+ from diffusers.models import AutoencoderKL
16
+ from diffusers.pipeline_utils import DiffusionPipeline
17
+ from diffusers.schedulers import (
18
+ DDIMScheduler,
19
+ DPMSolverMultistepScheduler,
20
+ EulerAncestralDiscreteScheduler,
21
+ EulerDiscreteScheduler,
22
+ LMSDiscreteScheduler,
23
+ PNDMScheduler,
24
+ )
25
+ from diffusers.utils import deprecate, logging, BaseOutput
26
+
27
+ from einops import rearrange
28
+
29
+ from ..models.unet import UNet3DConditionModel
30
+
31
+
32
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
33
+
34
+
35
+ @dataclass
36
+ class TuneAVideoPipelineOutput(BaseOutput):
37
+ videos: Union[torch.Tensor, np.ndarray]
38
+
39
+
40
+ class TuneAVideoPipeline(DiffusionPipeline):
41
+ _optional_components = []
42
+
43
+ def __init__(
44
+ self,
45
+ vae: AutoencoderKL,
46
+ text_encoder: CLIPTextModel,
47
+ tokenizer: CLIPTokenizer,
48
+ unet: UNet3DConditionModel,
49
+ scheduler: Union[
50
+ DDIMScheduler,
51
+ PNDMScheduler,
52
+ LMSDiscreteScheduler,
53
+ EulerDiscreteScheduler,
54
+ EulerAncestralDiscreteScheduler,
55
+ DPMSolverMultistepScheduler,
56
+ ],
57
+ ):
58
+ super().__init__()
59
+
60
+ if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
61
+ deprecation_message = (
62
+ f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
63
+ f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
64
+ "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
65
+ " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
66
+ " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
67
+ " file"
68
+ )
69
+ deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
70
+ new_config = dict(scheduler.config)
71
+ new_config["steps_offset"] = 1
72
+ scheduler._internal_dict = FrozenDict(new_config)
73
+
74
+ if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
75
+ deprecation_message = (
76
+ f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
77
+ " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
78
+ " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
79
+ " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
80
+ " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
81
+ )
82
+ deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
83
+ new_config = dict(scheduler.config)
84
+ new_config["clip_sample"] = False
85
+ scheduler._internal_dict = FrozenDict(new_config)
86
+
87
+ is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
88
+ version.parse(unet.config._diffusers_version).base_version
89
+ ) < version.parse("0.9.0.dev0")
90
+ is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
91
+ if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
92
+ deprecation_message = (
93
+ "The configuration file of the unet has set the default `sample_size` to smaller than"
94
+ " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
95
+ " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
96
+ " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
97
+ " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
98
+ " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
99
+ " in the config might lead to incorrect results in future versions. If you have downloaded this"
100
+ " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
101
+ " the `unet/config.json` file"
102
+ )
103
+ deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
104
+ new_config = dict(unet.config)
105
+ new_config["sample_size"] = 64
106
+ unet._internal_dict = FrozenDict(new_config)
107
+
108
+ self.register_modules(
109
+ vae=vae,
110
+ text_encoder=text_encoder,
111
+ tokenizer=tokenizer,
112
+ unet=unet,
113
+ scheduler=scheduler,
114
+ )
115
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
116
+
117
+ def enable_vae_slicing(self):
118
+ self.vae.enable_slicing()
119
+
120
+ def disable_vae_slicing(self):
121
+ self.vae.disable_slicing()
122
+
123
+ def enable_sequential_cpu_offload(self, gpu_id=0):
124
+ if is_accelerate_available():
125
+ from accelerate import cpu_offload
126
+ else:
127
+ raise ImportError("Please install accelerate via `pip install accelerate`")
128
+
129
+ device = torch.device(f"cuda:{gpu_id}")
130
+
131
+ for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
132
+ if cpu_offloaded_model is not None:
133
+ cpu_offload(cpu_offloaded_model, device)
134
+
135
+
136
+ @property
137
+ def _execution_device(self):
138
+ if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
139
+ return self.device
140
+ for module in self.unet.modules():
141
+ if (
142
+ hasattr(module, "_hf_hook")
143
+ and hasattr(module._hf_hook, "execution_device")
144
+ and module._hf_hook.execution_device is not None
145
+ ):
146
+ return torch.device(module._hf_hook.execution_device)
147
+ return self.device
148
+
149
+ def _encode_prompt(self, prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt):
150
+ batch_size = len(prompt) if isinstance(prompt, list) else 1
151
+
152
+ text_inputs = self.tokenizer(
153
+ prompt,
154
+ padding="max_length",
155
+ max_length=self.tokenizer.model_max_length,
156
+ truncation=True,
157
+ return_tensors="pt",
158
+ )
159
+ text_input_ids = text_inputs.input_ids
160
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
161
+
162
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
163
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])
164
+ logger.warning(
165
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
166
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
167
+ )
168
+
169
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
170
+ attention_mask = text_inputs.attention_mask.to(device)
171
+ else:
172
+ attention_mask = None
173
+
174
+ text_embeddings = self.text_encoder(
175
+ text_input_ids.to(device),
176
+ attention_mask=attention_mask,
177
+ )
178
+ text_embeddings = text_embeddings[0]
179
+
180
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
181
+ bs_embed, seq_len, _ = text_embeddings.shape
182
+ text_embeddings = text_embeddings.repeat(1, num_videos_per_prompt, 1)
183
+ text_embeddings = text_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
184
+
185
+ # get unconditional embeddings for classifier free guidance
186
+ if do_classifier_free_guidance:
187
+ uncond_tokens: List[str]
188
+ if negative_prompt is None:
189
+ uncond_tokens = [""] * batch_size
190
+ elif type(prompt) is not type(negative_prompt):
191
+ raise TypeError(
192
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
193
+ f" {type(prompt)}."
194
+ )
195
+ elif isinstance(negative_prompt, str):
196
+ uncond_tokens = [negative_prompt]
197
+ elif batch_size != len(negative_prompt):
198
+ raise ValueError(
199
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
200
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
201
+ " the batch size of `prompt`."
202
+ )
203
+ else:
204
+ uncond_tokens = negative_prompt
205
+
206
+ max_length = text_input_ids.shape[-1]
207
+ uncond_input = self.tokenizer(
208
+ uncond_tokens,
209
+ padding="max_length",
210
+ max_length=max_length,
211
+ truncation=True,
212
+ return_tensors="pt",
213
+ )
214
+
215
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
216
+ attention_mask = uncond_input.attention_mask.to(device)
217
+ else:
218
+ attention_mask = None
219
+
220
+ uncond_embeddings = self.text_encoder(
221
+ uncond_input.input_ids.to(device),
222
+ attention_mask=attention_mask,
223
+ )
224
+ uncond_embeddings = uncond_embeddings[0]
225
+
226
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
227
+ seq_len = uncond_embeddings.shape[1]
228
+ uncond_embeddings = uncond_embeddings.repeat(1, num_videos_per_prompt, 1)
229
+ uncond_embeddings = uncond_embeddings.view(batch_size * num_videos_per_prompt, seq_len, -1)
230
+
231
+ # For classifier free guidance, we need to do two forward passes.
232
+ # Here we concatenate the unconditional and text embeddings into a single batch
233
+ # to avoid doing two forward passes
234
+ text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
235
+
236
+ return text_embeddings
237
+
238
+ def decode_latents(self, latents):
239
+ video_length = latents.shape[2]
240
+ latents = 1 / 0.18215 * latents
241
+ latents = rearrange(latents, "b c f h w -> (b f) c h w")
242
+ video = self.vae.decode(latents).sample
243
+ video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
244
+ video = (video / 2 + 0.5).clamp(0, 1)
245
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
246
+ video = video.cpu().float().numpy()
247
+ return video
248
+
249
+ def prepare_extra_step_kwargs(self, generator, eta):
250
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
251
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
252
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
253
+ # and should be between [0, 1]
254
+
255
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
256
+ extra_step_kwargs = {}
257
+ if accepts_eta:
258
+ extra_step_kwargs["eta"] = eta
259
+
260
+ # check if the scheduler accepts generator
261
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
262
+ if accepts_generator:
263
+ extra_step_kwargs["generator"] = generator
264
+ return extra_step_kwargs
265
+
266
+ def check_inputs(self, prompt, height, width, callback_steps):
267
+ if not isinstance(prompt, str) and not isinstance(prompt, list):
268
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
269
+
270
+ if height % 8 != 0 or width % 8 != 0:
271
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
272
+
273
+ if (callback_steps is None) or (
274
+ callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
275
+ ):
276
+ raise ValueError(
277
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
278
+ f" {type(callback_steps)}."
279
+ )
280
+
281
+ def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None):
282
+ shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor)
283
+ if isinstance(generator, list) and len(generator) != batch_size:
284
+ raise ValueError(
285
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
286
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
287
+ )
288
+
289
+ if latents is None:
290
+ rand_device = "cpu" if device.type == "mps" else device
291
+
292
+ if isinstance(generator, list):
293
+ shape = (1,) + shape[1:]
294
+ latents = [
295
+ torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)
296
+ for i in range(batch_size)
297
+ ]
298
+ latents = torch.cat(latents, dim=0).to(device)
299
+ else:
300
+ latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
301
+ else:
302
+ if latents.shape != shape:
303
+ raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
304
+ latents = latents.to(device)
305
+
306
+ # scale the initial noise by the standard deviation required by the scheduler
307
+ latents = latents * self.scheduler.init_noise_sigma
308
+ return latents
309
+
310
+ @torch.no_grad()
311
+ def __call__(
312
+ self,
313
+ prompt: Union[str, List[str]],
314
+ video_length: Optional[int],
315
+ height: Optional[int] = None,
316
+ width: Optional[int] = None,
317
+ num_inference_steps: int = 50,
318
+ guidance_scale: float = 7.5,
319
+ negative_prompt: Optional[Union[str, List[str]]] = None,
320
+ num_videos_per_prompt: Optional[int] = 1,
321
+ eta: float = 0.0,
322
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
323
+ latents: Optional[torch.FloatTensor] = None,
324
+ output_type: Optional[str] = "tensor",
325
+ return_dict: bool = True,
326
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
327
+ callback_steps: Optional[int] = 1,
328
+ **kwargs,
329
+ ):
330
+ # Default height and width to unet
331
+ height = height or self.unet.config.sample_size * self.vae_scale_factor
332
+ width = width or self.unet.config.sample_size * self.vae_scale_factor
333
+
334
+ # Check inputs. Raise error if not correct
335
+ self.check_inputs(prompt, height, width, callback_steps)
336
+
337
+ # Define call parameters
338
+ batch_size = 1 if isinstance(prompt, str) else len(prompt)
339
+ device = self._execution_device
340
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
341
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
342
+ # corresponds to doing no classifier free guidance.
343
+ do_classifier_free_guidance = guidance_scale > 1.0
344
+
345
+ # Encode input prompt
346
+ text_embeddings = self._encode_prompt(
347
+ prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
348
+ )
349
+
350
+ # Prepare timesteps
351
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
352
+ timesteps = self.scheduler.timesteps
353
+
354
+ # Prepare latent variables
355
+ num_channels_latents = self.unet.in_channels
356
+ latents = self.prepare_latents(
357
+ batch_size * num_videos_per_prompt,
358
+ num_channels_latents,
359
+ video_length,
360
+ height,
361
+ width,
362
+ text_embeddings.dtype,
363
+ device,
364
+ generator,
365
+ latents,
366
+ )
367
+ latents_dtype = latents.dtype
368
+
369
+ # Prepare extra step kwargs.
370
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
371
+
372
+ # Denoising loop
373
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
374
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
375
+ for i, t in enumerate(timesteps):
376
+ # expand the latents if we are doing classifier free guidance
377
+ latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
378
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
379
+
380
+ # predict the noise residual
381
+ noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample.to(dtype=latents_dtype)
382
+
383
+ # perform guidance
384
+ if do_classifier_free_guidance:
385
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
386
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
387
+
388
+ # compute the previous noisy sample x_t -> x_t-1
389
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
390
+
391
+ # call the callback, if provided
392
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
393
+ progress_bar.update()
394
+ if callback is not None and i % callback_steps == 0:
395
+ callback(i, t, latents)
396
+
397
+ # Post-processing
398
+ video = self.decode_latents(latents)
399
+
400
+ # Convert to tensor
401
+ if output_type == "tensor":
402
+ video = torch.from_numpy(video)
403
+
404
+ if not return_dict:
405
+ return video
406
+
407
+ return TuneAVideoPipelineOutput(videos=video)
Tune-A-Video-debug/tuneavideo/util.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import imageio
3
+ import numpy as np
4
+ from typing import Union
5
+
6
+ import torch
7
+ import torchvision
8
+
9
+ from tqdm import tqdm
10
+ from einops import rearrange
11
+
12
+
13
+ def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=4, fps=8):
14
+ videos = rearrange(videos, "b c t h w -> t b c h w")
15
+ outputs = []
16
+ for x in videos:
17
+ x = torchvision.utils.make_grid(x, nrow=n_rows)
18
+ x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
19
+ if rescale:
20
+ x = (x + 1.0) / 2.0 # -1,1 -> 0,1
21
+ x = (x * 255).numpy().astype(np.uint8)
22
+ outputs.append(x)
23
+
24
+ os.makedirs(os.path.dirname(path), exist_ok=True)
25
+ imageio.mimsave(path, outputs, fps=fps)
26
+
27
+
28
+ # DDIM Inversion
29
+ @torch.no_grad()
30
+ def init_prompt(prompt, pipeline):
31
+ uncond_input = pipeline.tokenizer(
32
+ [""], padding="max_length", max_length=pipeline.tokenizer.model_max_length,
33
+ return_tensors="pt"
34
+ )
35
+ uncond_embeddings = pipeline.text_encoder(uncond_input.input_ids.to(pipeline.device))[0]
36
+ text_input = pipeline.tokenizer(
37
+ [prompt],
38
+ padding="max_length",
39
+ max_length=pipeline.tokenizer.model_max_length,
40
+ truncation=True,
41
+ return_tensors="pt",
42
+ )
43
+ text_embeddings = pipeline.text_encoder(text_input.input_ids.to(pipeline.device))[0]
44
+ context = torch.cat([uncond_embeddings, text_embeddings])
45
+
46
+ return context
47
+
48
+
49
+ def next_step(model_output: Union[torch.FloatTensor, np.ndarray], timestep: int,
50
+ sample: Union[torch.FloatTensor, np.ndarray], ddim_scheduler):
51
+ timestep, next_timestep = min(
52
+ timestep - ddim_scheduler.config.num_train_timesteps // ddim_scheduler.num_inference_steps, 999), timestep
53
+ alpha_prod_t = ddim_scheduler.alphas_cumprod[timestep] if timestep >= 0 else ddim_scheduler.final_alpha_cumprod
54
+ alpha_prod_t_next = ddim_scheduler.alphas_cumprod[next_timestep]
55
+ beta_prod_t = 1 - alpha_prod_t
56
+ next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
57
+ next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
58
+ next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
59
+ return next_sample
60
+
61
+
62
+ def get_noise_pred_single(latents, t, context, unet):
63
+ noise_pred = unet(latents, t, encoder_hidden_states=context)["sample"]
64
+ return noise_pred
65
+
66
+
67
+ @torch.no_grad()
68
+ def ddim_loop(pipeline, ddim_scheduler, latent, num_inv_steps, prompt):
69
+ context = init_prompt(prompt, pipeline)
70
+ uncond_embeddings, cond_embeddings = context.chunk(2)
71
+ all_latent = [latent]
72
+ latent = latent.clone().detach()
73
+ for i in tqdm(range(num_inv_steps)):
74
+ t = ddim_scheduler.timesteps[len(ddim_scheduler.timesteps) - i - 1]
75
+ noise_pred = get_noise_pred_single(latent, t, cond_embeddings, pipeline.unet)
76
+ latent = next_step(noise_pred, t, latent, ddim_scheduler)
77
+ all_latent.append(latent)
78
+ return all_latent
79
+
80
+
81
+ @torch.no_grad()
82
+ def ddim_inversion(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt=""):
83
+ ddim_latents = ddim_loop(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt)
84
+ return ddim_latents
trainer.py CHANGED
@@ -19,7 +19,7 @@ from utils import save_model_card
19
 
20
  sys.path.append('Tune-A-Video')
21
 
22
- URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Video-library/share/YjTcaNJmKyeHFpMBioHhzBcTzCYddVErEk'
23
  ORIGINAL_SPACE_ID = 'video-p2p-library/Video-P2P-Demo'
24
  SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
25
 
@@ -103,7 +103,8 @@ class Trainer:
103
  self.join_model_library_org(
104
  self.hf_token if self.hf_token else input_token)
105
 
106
- config = OmegaConf.load('Video-P2P/configs/man-surfing-tune.yaml')
 
107
  config.pretrained_model_path = self.download_base_model(base_model)
108
  config.output_dir = output_dir.as_posix()
109
  config.train_data.video_path = training_video.name # type: ignore
@@ -133,7 +134,8 @@ class Trainer:
133
  with open(config_path, 'w') as f:
134
  OmegaConf.save(config, f)
135
 
136
- command = f'accelerate launch Video-P2P/run_tuning.py --config {config_path}'
 
137
  subprocess.run(shlex.split(command))
138
  save_model_card(save_dir=output_dir,
139
  base_model=base_model,
 
19
 
20
  sys.path.append('Tune-A-Video')
21
 
22
+ URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/video-p2p-library/share/pZwQaStCpdmMCGLURsMhMkEpvIlsdMdnkk'
23
  ORIGINAL_SPACE_ID = 'video-p2p-library/Video-P2P-Demo'
24
  SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
25
 
 
103
  self.join_model_library_org(
104
  self.hf_token if self.hf_token else input_token)
105
 
106
+ # config = OmegaConf.load('Video-P2P/configs/man-surfing-tune.yaml')
107
+ config = OmegaConf.load('Video-P2P/configs/man-surfing.yaml')
108
  config.pretrained_model_path = self.download_base_model(base_model)
109
  config.output_dir = output_dir.as_posix()
110
  config.train_data.video_path = training_video.name # type: ignore
 
134
  with open(config_path, 'w') as f:
135
  OmegaConf.save(config, f)
136
 
137
+ # command = f'accelerate launch Video-P2P/run_tuning.py --config {config_path}'
138
+ command = f'accelerate launch Tune-A-Video-debug/train_tuneavideo.py --config {config_path}'
139
  subprocess.run(shlex.split(command))
140
  save_model_card(save_dir=output_dir,
141
  base_model=base_model,