ShaoTengLiu commited on
Commit
edc0f20
1 Parent(s): 0fd0335
.DS_Store ADDED
Binary file (6.15 kB). View file
 
README.md CHANGED
@@ -1,12 +1,17 @@
1
- ---
2
- title: Video P2P Demo
3
- emoji: 💩
4
- colorFrom: red
5
- colorTo: pink
6
- sdk: gradio
7
- sdk_version: 3.21.0
8
- app_file: app.py
9
- pinned: false
10
- ---
11
-
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
1
+ # Video-P2P
2
+
3
+ ## Setup
4
+
5
+ All required packages are listed in the requirements file.
6
+ The code was tested on a Tesla V100 32GB but should work on other cards with at least **16GB** VRAM.
7
+
8
+ ## Quickstart
9
+
10
+ ``` bash
11
+ bash script.sh
12
+ ```
13
+
14
+ ## References
15
+ * prompt-to-prompt: https://github.com/google/prompt-to-prompt
16
+ * Tune-A-Video: https://github.com/showlab/Tune-A-Video
17
+ * diffusers: https://github.com/huggingface/diffusers
configs/.DS_Store ADDED
Binary file (6.15 kB). View file
 
configs/man-motor-tune.yaml ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "/data/stable-diffusion/stable-diffusion-v1-5"
2
+ output_dir: "./outputs/man-motor"
3
+
4
+ train_data:
5
+ video_path: "./data/motorbike"
6
+ prompt: "a man is driving a motorbike in the forest"
7
+ n_sample_frames: 8
8
+ width: 512
9
+ height: 512
10
+ sample_start_idx: 0
11
+ sample_frame_rate: 1
12
+
13
+ validation_data:
14
+ prompts:
15
+ - "a man is driving a motorbike in the forest"
16
+ - "a Spider-Man is driving a motorbike in the forest"
17
+ - "a Bat-Man is driving a motorbike in the forest"
18
+ - "an Iron-Man is driving a motorbike in the forest"
19
+ video_length: 8
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
configs/rabbit-jump-p2p.yaml ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "./outputs/rabbit-jump"
2
+ image_path: "./data/rabbit"
3
+ prompt: "a rabbit is jumping on the grass"
4
+ prompts:
5
+ - "a rabbit is jumping on the grass"
6
+ - "a origami rabbit is jumping on the grass"
7
+ blend_word:
8
+ - 'rabbit'
9
+ - 'rabbit'
10
+ eq_params:
11
+ words: "origami"
12
+ values: 2
13
+ gif_folder: "./outputs/rabbit-jump/results"
14
+ gif_name_1: "./outputs/rabbit-jump/results/original_name.gif"
15
+ gif_name_2: "./outputs/rabbit-jump/results/origami_name.gif"
16
+ IRC: False
configs/rabbit-jump-tune.yaml ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ pretrained_model_path: "/data/stable-diffusion/stable-diffusion-v1-5"
2
+ output_dir: "./outputs/rabbit-jump"
3
+
4
+ train_data:
5
+ video_path: "./data/rabbit"
6
+ prompt: "a rabbit is jumping on the grass"
7
+ n_sample_frames: 8
8
+ width: 512
9
+ height: 512
10
+ sample_start_idx: 0
11
+ sample_frame_rate: 1
12
+
13
+ validation_data:
14
+ prompts:
15
+ - "a rabbit is jumping on the grass"
16
+ - "a Lego rabbit is jumping on the grass"
17
+ - "a origami rabbit is jumping on the grass"
18
+ - "a crochet rabbit is jumping on the grass"
19
+ video_length: 8
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
data/.DS_Store ADDED
Binary file (10.2 kB). View file
 
data/motorbike/.DS_Store ADDED
Binary file (6.15 kB). View file
 
data/motorbike/1.jpg ADDED
data/motorbike/2.jpg ADDED
data/motorbike/3.jpg ADDED
data/motorbike/4.jpg ADDED
data/motorbike/5.jpg ADDED
data/motorbike/6.jpg ADDED
data/motorbike/7.jpg ADDED
data/motorbike/8.jpg ADDED
data/rabbit/1.jpg ADDED
data/rabbit/2.jpg ADDED
data/rabbit/3.jpg ADDED
data/rabbit/4.jpg ADDED
data/rabbit/5.jpg ADDED
data/rabbit/6.jpg ADDED
data/rabbit/7.jpg ADDED
data/rabbit/8.jpg ADDED
ptp_utils.py ADDED
@@ -0,0 +1,309 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import numpy as np
16
+ import torch
17
+ from PIL import Image, ImageDraw, ImageFont
18
+ import cv2
19
+ from typing import Optional, Union, Tuple, List, Callable, Dict
20
+ from IPython.display import display
21
+ from tqdm.notebook import tqdm
22
+
23
+
24
+ def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
25
+ h, w, c = image.shape
26
+ offset = int(h * .2)
27
+ img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
28
+ font = cv2.FONT_HERSHEY_SIMPLEX
29
+ img[:h] = image
30
+ textsize = cv2.getTextSize(text, font, 1, 2)[0]
31
+ text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
32
+ cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
33
+ return img
34
+
35
+
36
+ def view_images(images, num_rows=1, offset_ratio=0.02):
37
+ if type(images) is list:
38
+ num_empty = len(images) % num_rows
39
+ elif images.ndim == 4:
40
+ num_empty = images.shape[0] % num_rows
41
+ else:
42
+ images = [images]
43
+ num_empty = 0
44
+
45
+ empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
46
+ images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
47
+ num_items = len(images)
48
+
49
+ h, w, c = images[0].shape
50
+ offset = int(h * offset_ratio)
51
+ num_cols = num_items // num_rows
52
+ image_ = np.ones((h * num_rows + offset * (num_rows - 1),
53
+ w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
54
+ for i in range(num_rows):
55
+ for j in range(num_cols):
56
+ image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
57
+ i * num_cols + j]
58
+
59
+ pil_img = Image.fromarray(image_)
60
+ display(pil_img)
61
+
62
+
63
+ def diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False, simple=False):
64
+ if low_resource:
65
+ noise_pred_uncond = model.unet(latents, t, encoder_hidden_states=context[0])["sample"]
66
+ noise_prediction_text = model.unet(latents, t, encoder_hidden_states=context[1])["sample"]
67
+ else:
68
+ latents_input = torch.cat([latents] * 2)
69
+ noise_pred = model.unet(latents_input, t, encoder_hidden_states=context)["sample"]
70
+ noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
71
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
72
+ if simple:
73
+ noise_pred[0] = noise_prediction_text[0]
74
+ latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"]
75
+ # first latents: torch.Size([1, 4, 4, 64, 64])
76
+ latents = controller.step_callback(latents)
77
+ return latents
78
+
79
+
80
+ def latent2image(vae, latents):
81
+ latents = 1 / 0.18215 * latents
82
+ image = vae.decode(latents)['sample']
83
+ image = (image / 2 + 0.5).clamp(0, 1)
84
+ image = image.cpu().permute(0, 2, 3, 1).numpy()
85
+ image = (image * 255).astype(np.uint8)
86
+ return image
87
+
88
+
89
+ @torch.no_grad()
90
+ def latent2image_video(vae, latents):
91
+ latents = 1 / 0.18215 * latents
92
+ latents = latents[0].permute(1, 0, 2, 3)
93
+ image = vae.decode(latents)['sample']
94
+ image = (image / 2 + 0.5).clamp(0, 1)
95
+ image = image.cpu().permute(0, 2, 3, 1).numpy()
96
+ image = (image * 255).astype(np.uint8)
97
+ return image
98
+
99
+
100
+ def init_latent(latent, model, height, width, generator, batch_size):
101
+ if latent is None:
102
+ latent = torch.randn(
103
+ (1, model.unet.in_channels, height // 8, width // 8),
104
+ generator=generator,
105
+ )
106
+ latents = latent.expand(batch_size, model.unet.in_channels, height // 8, width // 8).to(model.device)
107
+ return latent, latents
108
+
109
+
110
+ @torch.no_grad()
111
+ def text2image_ldm(
112
+ model,
113
+ prompt: List[str],
114
+ controller,
115
+ num_inference_steps: int = 50,
116
+ guidance_scale: Optional[float] = 7.,
117
+ generator: Optional[torch.Generator] = None,
118
+ latent: Optional[torch.FloatTensor] = None,
119
+ ):
120
+ register_attention_control(model, controller)
121
+ height = width = 256
122
+ batch_size = len(prompt)
123
+
124
+ uncond_input = model.tokenizer([""] * batch_size, padding="max_length", max_length=77, return_tensors="pt")
125
+ uncond_embeddings = model.bert(uncond_input.input_ids.to(model.device))[0]
126
+
127
+ text_input = model.tokenizer(prompt, padding="max_length", max_length=77, return_tensors="pt")
128
+ text_embeddings = model.bert(text_input.input_ids.to(model.device))[0]
129
+ latent, latents = init_latent(latent, model, height, width, generator, batch_size)
130
+ context = torch.cat([uncond_embeddings, text_embeddings])
131
+
132
+ model.scheduler.set_timesteps(num_inference_steps)
133
+ for t in tqdm(model.scheduler.timesteps):
134
+ latents = diffusion_step(model, controller, latents, context, t, guidance_scale)
135
+
136
+ image = latent2image(model.vqvae, latents)
137
+
138
+ return image, latent
139
+
140
+
141
+ @torch.no_grad()
142
+ def text2image_ldm_stable(
143
+ model,
144
+ prompt: List[str],
145
+ controller,
146
+ num_inference_steps: int = 50,
147
+ guidance_scale: float = 7.5,
148
+ generator: Optional[torch.Generator] = None,
149
+ latent: Optional[torch.FloatTensor] = None,
150
+ low_resource: bool = False,
151
+ ):
152
+ register_attention_control(model, controller)
153
+ height = width = 512
154
+ batch_size = len(prompt)
155
+
156
+ text_input = model.tokenizer(
157
+ prompt,
158
+ padding="max_length",
159
+ max_length=model.tokenizer.model_max_length,
160
+ truncation=True,
161
+ return_tensors="pt",
162
+ )
163
+ text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]
164
+ max_length = text_input.input_ids.shape[-1]
165
+ uncond_input = model.tokenizer(
166
+ [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
167
+ )
168
+ uncond_embeddings = model.text_encoder(uncond_input.input_ids.to(model.device))[0]
169
+
170
+ context = [uncond_embeddings, text_embeddings]
171
+ if not low_resource:
172
+ context = torch.cat(context)
173
+
174
+ latent, latents = init_latent(latent, model, height, width, generator, batch_size)
175
+
176
+ # set timesteps
177
+ extra_set_kwargs = {"offset": 1}
178
+ model.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
179
+ for t in tqdm(model.scheduler.timesteps):
180
+ latents = diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource)
181
+
182
+ image = latent2image(model.vae, latents)
183
+
184
+ return image, latent
185
+
186
+
187
+ def register_attention_control(model, controller):
188
+ def ca_forward(self, place_in_unet):
189
+ to_out = self.to_out
190
+ if type(to_out) is torch.nn.modules.container.ModuleList:
191
+ to_out = self.to_out[0]
192
+ else:
193
+ to_out = self.to_out
194
+
195
+ def forward(x, encoder_hidden_states=None, attention_mask=None):
196
+ context = encoder_hidden_states
197
+ mask = attention_mask
198
+ batch_size, sequence_length, dim = x.shape
199
+ h = self.heads
200
+ q = self.to_q(x)
201
+ is_cross = context is not None
202
+ context = context if is_cross else x
203
+ k = self.to_k(context)
204
+ v = self.to_v(context)
205
+ q = self.reshape_heads_to_batch_dim(q)
206
+ k = self.reshape_heads_to_batch_dim(k)
207
+ v = self.reshape_heads_to_batch_dim(v)
208
+ sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale # q: torch.Size([128, 4096, 40]); k: torch.Size([64, 77, 40])
209
+
210
+ if mask is not None:
211
+ mask = mask.reshape(batch_size, -1)
212
+ max_neg_value = -torch.finfo(sim.dtype).max
213
+ mask = mask[:, None, :].repeat(h, 1, 1)
214
+ sim.masked_fill_(~mask, max_neg_value)
215
+
216
+ attn = torch.exp(sim-torch.max(sim)) / torch.sum(torch.exp(sim-torch.max(sim)), axis=-1).unsqueeze(-1)
217
+ attn = controller(attn, is_cross, place_in_unet)
218
+ out = torch.einsum("b i j, b j d -> b i d", attn, v)
219
+ out = self.reshape_batch_dim_to_heads(out)
220
+ return to_out(out)
221
+
222
+ return forward
223
+
224
+ class DummyController:
225
+
226
+ def __call__(self, *args):
227
+ return args[0]
228
+
229
+ def __init__(self):
230
+ self.num_att_layers = 0
231
+
232
+ if controller is None:
233
+ controller = DummyController()
234
+
235
+ def register_recr(net_, count, place_in_unet):
236
+ if net_.__class__.__name__ == 'CrossAttention':
237
+ net_.forward = ca_forward(net_, place_in_unet)
238
+ return count + 1
239
+ elif hasattr(net_, 'children'):
240
+ for net__ in net_.children():
241
+ count = register_recr(net__, count, place_in_unet)
242
+ return count
243
+
244
+ cross_att_count = 0
245
+ sub_nets = model.unet.named_children()
246
+ for net in sub_nets:
247
+ if "down" in net[0]:
248
+ cross_att_count += register_recr(net[1], 0, "down")
249
+ elif "up" in net[0]:
250
+ cross_att_count += register_recr(net[1], 0, "up")
251
+ elif "mid" in net[0]:
252
+ cross_att_count += register_recr(net[1], 0, "mid")
253
+
254
+ controller.num_att_layers = cross_att_count
255
+
256
+
257
+ def get_word_inds(text: str, word_place: int, tokenizer):
258
+ split_text = text.split(" ")
259
+ if type(word_place) is str:
260
+ word_place = [i for i, word in enumerate(split_text) if word_place == word]
261
+ elif type(word_place) is int:
262
+ word_place = [word_place]
263
+ out = []
264
+ if len(word_place) > 0:
265
+ words_encode = [tokenizer.decode([item]).strip("#") for item in tokenizer.encode(text)][1:-1]
266
+ cur_len, ptr = 0, 0
267
+
268
+ for i in range(len(words_encode)):
269
+ cur_len += len(words_encode[i])
270
+ if ptr in word_place:
271
+ out.append(i + 1)
272
+ if cur_len >= len(split_text[ptr]):
273
+ ptr += 1
274
+ cur_len = 0
275
+ return np.array(out)
276
+
277
+
278
+ def update_alpha_time_word(alpha, bounds: Union[float, Tuple[float, float]], prompt_ind: int,
279
+ word_inds: Optional[torch.Tensor]=None):
280
+ if type(bounds) is float:
281
+ bounds = 0, bounds
282
+ start, end = int(bounds[0] * alpha.shape[0]), int(bounds[1] * alpha.shape[0])
283
+ if word_inds is None:
284
+ word_inds = torch.arange(alpha.shape[2])
285
+ alpha[: start, prompt_ind, word_inds] = 0
286
+ alpha[start: end, prompt_ind, word_inds] = 1
287
+ alpha[end:, prompt_ind, word_inds] = 0
288
+ return alpha
289
+
290
+
291
+ def get_time_words_attention_alpha(prompts, num_steps,
292
+ cross_replace_steps: Union[float, Dict[str, Tuple[float, float]]],
293
+ tokenizer, max_num_words=77):
294
+ if type(cross_replace_steps) is not dict:
295
+ cross_replace_steps = {"default_": cross_replace_steps}
296
+ if "default_" not in cross_replace_steps:
297
+ cross_replace_steps["default_"] = (0., 1.)
298
+ alpha_time_words = torch.zeros(num_steps + 1, len(prompts) - 1, max_num_words)
299
+ for i in range(len(prompts) - 1): # 2
300
+ alpha_time_words = update_alpha_time_word(alpha_time_words, cross_replace_steps["default_"], # {'default_': 0.8}
301
+ i)
302
+ for key, item in cross_replace_steps.items():
303
+ if key != "default_":
304
+ inds = [get_word_inds(prompts[i], key, tokenizer) for i in range(1, len(prompts))]
305
+ for i, ind in enumerate(inds):
306
+ if len(ind) > 0:
307
+ alpha_time_words = update_alpha_time_word(alpha_time_words, item, i, ind)
308
+ alpha_time_words = alpha_time_words.reshape(num_steps + 1, len(prompts) - 1, 1, 1, max_num_words)
309
+ return alpha_time_words
requirements.txt ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
14
+ opencv-python
15
+ ipywidgets
run_tuning.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))
run_videop2p.py ADDED
@@ -0,0 +1,707 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from typing import Optional, Union, Tuple, List, Callable, Dict
3
+ from tqdm.notebook import tqdm
4
+ import torch
5
+ from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
6
+ import torch.nn.functional as nnf
7
+ import numpy as np
8
+ import abc
9
+ import ptp_utils
10
+ import seq_aligner
11
+ import shutil
12
+ from torch.optim.adam import Adam
13
+ from PIL import Image
14
+ from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer
15
+ from einops import rearrange
16
+
17
+ from tuneavideo.models.unet import UNet3DConditionModel
18
+ from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
19
+
20
+ import cv2
21
+ import argparse
22
+ from omegaconf import OmegaConf
23
+
24
+ scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
25
+ MY_TOKEN = ''
26
+ LOW_RESOURCE = False
27
+ NUM_DDIM_STEPS = 50
28
+ GUIDANCE_SCALE = 7.5
29
+ MAX_NUM_WORDS = 77
30
+ device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
31
+ IRC = True
32
+
33
+ # need to adjust
34
+ cross_replace_steps = {'default_': .2,}
35
+ self_replace_steps = .5
36
+ mask_th = (.3, .3)
37
+ video_len = 8
38
+
39
+ def main(
40
+ pretrained_model_path: str,
41
+ image_path: str,
42
+ prompt: str,
43
+ prompts: Tuple[str],
44
+ blend_word: Tuple[str],
45
+ eq_params: Dict,
46
+ gif_folder: str,
47
+ gif_name_1: str,
48
+ gif_name_2: str,
49
+ IRC: bool,
50
+ ):
51
+ blend_word = (((blend_word[0],), (blend_word[1],)))
52
+ eq_params["words"] = (eq_params["words"],)
53
+ eq_params["values"] = (eq_params["values"],)
54
+ eq_params = dict(eq_params)
55
+ prompts = list(prompts)
56
+ if not os.path.exists(gif_folder):
57
+ os.makedirs(gif_folder)
58
+
59
+ # Load the tokenizer
60
+ tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
61
+ # Load models and create wrapper for stable diffusion
62
+ text_encoder = CLIPTextModel.from_pretrained(
63
+ pretrained_model_path,
64
+ subfolder="text_encoder",
65
+ )
66
+ vae = AutoencoderKL.from_pretrained(
67
+ pretrained_model_path,
68
+ subfolder="vae",
69
+ )
70
+ unet = UNet3DConditionModel.from_pretrained(pretrained_model_path, subfolder="unet")
71
+ ldm_stable = TuneAVideoPipeline(
72
+ vae=vae,
73
+ text_encoder=text_encoder,
74
+ tokenizer=tokenizer,
75
+ unet=unet,
76
+ scheduler=scheduler,
77
+ ).to(device)
78
+
79
+ try:
80
+ ldm_stable.disable_xformers_memory_efficient_attention()
81
+ except AttributeError:
82
+ print("Attribute disable_xformers_memory_efficient_attention() is missing")
83
+ tokenizer = ldm_stable.tokenizer # Tokenizer of class: [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer)
84
+ # A tokenizer breaks a stream of text into tokens, usually by looking for whitespace (tabs, spaces, new lines).
85
+
86
+ class LocalBlend:
87
+
88
+ def get_mask(self, maps, alpha, use_pool): # alpha is a word map
89
+ k = 1
90
+ maps = (maps * alpha).sum(-1).mean(2) # [2, 80, 1, 16, 16, 77], [2, 1, 1, 1, 1, 77]
91
+ if use_pool:
92
+ maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
93
+ mask = nnf.interpolate(maps, size=(x_t.shape[3:]))
94
+ mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
95
+ mask = mask.gt(self.th[1-int(use_pool)])
96
+ mask = mask[:1] + mask
97
+ return mask
98
+
99
+ def __call__(self, x_t, attention_store, step):
100
+ self.counter += 1
101
+ if self.counter > self.start_blend:
102
+ # attention_store["down_cross"]: 4, attention_store["up_cross"]:6, attention_store["down_cross"][0]: torch.Size([32, 1024, 77])
103
+ maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3]
104
+ # maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, MAX_NUM_WORDS) for item in maps]
105
+ maps = [item.reshape(self.alpha_layers.shape[0], -1, 8, 16, 16, MAX_NUM_WORDS) for item in maps]
106
+ maps = torch.cat(maps, dim=2)
107
+ # self.alpha_layers: torch.Size([2, 1, 1, 1, 1, 77])
108
+ mask = self.get_mask(maps, self.alpha_layers, True)
109
+ if self.substruct_layers is not None:
110
+ maps_sub = ~self.get_mask(maps, self.substruct_layers, False)
111
+ mask = mask * maps_sub
112
+ mask = mask.float()
113
+ mask = mask.reshape(-1, 1, mask.shape[-3], mask.shape[-2], mask.shape[-1])
114
+
115
+ x_t = x_t[:1] + mask * (x_t - x_t[:1]) # line13 algorithm
116
+ return x_t
117
+
118
+ def __init__(self, prompts: List[str], words: [List[List[str]]], substruct_words=None, start_blend=0.2, th=(.3, .3)):
119
+ alpha_layers = torch.zeros(len(prompts), 1, 1, 1, 1, MAX_NUM_WORDS)
120
+ for i, (prompt, words_) in enumerate(zip(prompts, words)):
121
+ if type(words_) is str:
122
+ words_ = [words_]
123
+ for word in words_:
124
+ ind = ptp_utils.get_word_inds(prompt, word, tokenizer)
125
+ alpha_layers[i, :, :, :, :, ind] = 1
126
+
127
+ if substruct_words is not None:
128
+ substruct_layers = torch.zeros(len(prompts), 1, 1, 1, 1, MAX_NUM_WORDS)
129
+ for i, (prompt, words_) in enumerate(zip(prompts, substruct_words)):
130
+ if type(words_) is str:
131
+ words_ = [words_]
132
+ for word in words_:
133
+ ind = ptp_utils.get_word_inds(prompt, word, tokenizer)
134
+ substruct_layers[i, :, :, :, :, ind] = 1
135
+ self.substruct_layers = substruct_layers.to(device)
136
+ else:
137
+ self.substruct_layers = None
138
+ self.alpha_layers = alpha_layers.to(device)
139
+ self.start_blend = int(start_blend * NUM_DDIM_STEPS)
140
+ self.counter = 0
141
+ self.th=th
142
+
143
+
144
+
145
+
146
+ class EmptyControl:
147
+
148
+
149
+ def step_callback(self, x_t):
150
+ return x_t
151
+
152
+ def between_steps(self):
153
+ return
154
+
155
+ def __call__(self, attn, is_cross: bool, place_in_unet: str):
156
+ return attn
157
+
158
+
159
+ class AttentionControl(abc.ABC):
160
+
161
+ def step_callback(self, x_t):
162
+ return x_t
163
+
164
+ def between_steps(self):
165
+ return
166
+
167
+ @property
168
+ def num_uncond_att_layers(self):
169
+ return self.num_att_layers if LOW_RESOURCE else 0
170
+
171
+ @abc.abstractmethod
172
+ def forward (self, attn, is_cross: bool, place_in_unet: str):
173
+ raise NotImplementedError
174
+
175
+ def __call__(self, attn, is_cross: bool, place_in_unet: str):
176
+ if self.cur_att_layer >= self.num_uncond_att_layers:
177
+ if LOW_RESOURCE:
178
+ attn = self.forward(attn, is_cross, place_in_unet)
179
+ else:
180
+ h = attn.shape[0]
181
+ attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
182
+ self.cur_att_layer += 1
183
+ if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
184
+ self.cur_att_layer = 0
185
+ self.cur_step += 1
186
+ self.between_steps()
187
+ return attn
188
+
189
+ def reset(self):
190
+ self.cur_step = 0
191
+ self.cur_att_layer = 0
192
+
193
+ def __init__(self):
194
+ self.cur_step = 0
195
+ self.num_att_layers = -1
196
+ self.cur_att_layer = 0
197
+
198
+ class SpatialReplace(EmptyControl):
199
+
200
+ def step_callback(self, x_t):
201
+ if self.cur_step < self.stop_inject:
202
+ b = x_t.shape[0]
203
+ x_t = x_t[:1].expand(b, *x_t.shape[1:])
204
+ return x_t
205
+
206
+ def __init__(self, stop_inject: float):
207
+ super(SpatialReplace, self).__init__()
208
+ self.stop_inject = int((1 - stop_inject) * NUM_DDIM_STEPS)
209
+
210
+
211
+ class AttentionStore(AttentionControl):
212
+
213
+ @staticmethod
214
+ def get_empty_store():
215
+ return {"down_cross": [], "mid_cross": [], "up_cross": [],
216
+ "down_self": [], "mid_self": [], "up_self": []}
217
+
218
+ def forward(self, attn, is_cross: bool, place_in_unet: str):
219
+ key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
220
+ if attn.shape[1] <= 32 ** 2: # avoid memory overhead
221
+ self.step_store[key].append(attn) # 'down_self' torch.Size([32768, 8, 8])
222
+ return attn
223
+
224
+ def between_steps(self):
225
+ if len(self.attention_store) == 0:
226
+ self.attention_store = self.step_store
227
+ else:
228
+ for key in self.attention_store:
229
+ for i in range(len(self.attention_store[key])):
230
+ self.attention_store[key][i] += self.step_store[key][i]
231
+ self.step_store = self.get_empty_store()
232
+
233
+ def get_average_attention(self):
234
+ average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store}
235
+ return average_attention
236
+
237
+
238
+ def reset(self):
239
+ super(AttentionStore, self).reset()
240
+ self.step_store = self.get_empty_store()
241
+ self.attention_store = {}
242
+
243
+ def __init__(self):
244
+ super(AttentionStore, self).__init__()
245
+ self.step_store = self.get_empty_store()
246
+ self.attention_store = {}
247
+
248
+
249
+ class AttentionControlEdit(AttentionStore, abc.ABC):
250
+
251
+ def step_callback(self, x_t):
252
+ if self.local_blend is not None:
253
+ x_t = self.local_blend(x_t, self.attention_store, self.cur_step)
254
+ return x_t
255
+
256
+ def replace_self_attention(self, attn_base, att_replace, place_in_unet):
257
+ if att_replace.shape[2] <= 32 ** 2:
258
+ attn_base = attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
259
+ return attn_base
260
+ else:
261
+ return att_replace
262
+
263
+ @abc.abstractmethod
264
+ def replace_cross_attention(self, attn_base, att_replace):
265
+ raise NotImplementedError
266
+
267
+ def forward(self, attn, is_cross: bool, place_in_unet: str):
268
+ super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
269
+ if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):
270
+ h = attn.shape[0] // (self.batch_size)
271
+ attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
272
+ attn_base, attn_repalce = attn[0], attn[1:]
273
+ if is_cross:
274
+ alpha_words = self.cross_replace_alpha[self.cur_step]
275
+ attn_repalce_new = self.replace_cross_attention(attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce
276
+ attn[1:] = attn_repalce_new
277
+ else:
278
+ attn[1:] = self.replace_self_attention(attn_base, attn_repalce, place_in_unet)
279
+ attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
280
+ return attn
281
+
282
+ def __init__(self, prompts, num_steps: int,
283
+ cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
284
+ self_replace_steps: Union[float, Tuple[float, float]],
285
+ local_blend: Optional[LocalBlend]):
286
+ super(AttentionControlEdit, self).__init__()
287
+ self.batch_size = len(prompts)
288
+ self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, tokenizer).to(device)
289
+ if type(self_replace_steps) is float:
290
+ self_replace_steps = 0, self_replace_steps
291
+ self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])
292
+ self.local_blend = local_blend
293
+
294
+ class AttentionReplace(AttentionControlEdit):
295
+
296
+ def replace_cross_attention(self, attn_base, att_replace):
297
+ return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)
298
+
299
+ def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
300
+ local_blend: Optional[LocalBlend] = None):
301
+ super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
302
+ self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device)
303
+
304
+
305
+ class AttentionRefine(AttentionControlEdit):
306
+
307
+ def replace_cross_attention(self, attn_base, att_replace):
308
+ attn_base_replace = attn_base[:, :, self.mapper].permute(2, 0, 1, 3)
309
+ attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
310
+ return attn_replace
311
+
312
+ def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
313
+ local_blend: Optional[LocalBlend] = None):
314
+ super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
315
+ self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)
316
+ self.mapper, alphas = self.mapper.to(device), alphas.to(device)
317
+ self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])
318
+
319
+
320
+ class AttentionReweight(AttentionControlEdit):
321
+
322
+ def replace_cross_attention(self, attn_base, att_replace):
323
+ if self.prev_controller is not None:
324
+ attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace)
325
+ attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]
326
+ return attn_replace
327
+
328
+ def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer,
329
+ local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None):
330
+ super(AttentionReweight, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
331
+ self.equalizer = equalizer.to(device)
332
+ self.prev_controller = controller
333
+
334
+
335
+ def get_equalizer(text: str, word_select: Union[int, Tuple[int, ...]], values: Union[List[float],
336
+ Tuple[float, ...]]):
337
+ if type(word_select) is int or type(word_select) is str:
338
+ word_select = (word_select,)
339
+ equalizer = torch.ones(1, 77)
340
+
341
+ for word, val in zip(word_select, values):
342
+ inds = ptp_utils.get_word_inds(text, word, tokenizer)
343
+ equalizer[:, inds] = val
344
+ return equalizer
345
+
346
+ def aggregate_attention(attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
347
+ out = []
348
+ attention_maps = attention_store.get_average_attention()
349
+ num_pixels = res ** 2
350
+ for location in from_where:
351
+ for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
352
+ if item.shape[1] == num_pixels: # torch.Size([64, 256, 77]) all can pass
353
+ cross_maps = item.reshape(8, 8, res, res, item.shape[-1])
354
+ out.append(cross_maps)
355
+ out = torch.cat(out, dim=1)
356
+ out = out.sum(1) / out.shape[1]
357
+ return out.cpu()
358
+
359
+
360
+ def make_controller(prompts: List[str], is_replace_controller: bool, cross_replace_steps: Dict[str, float], self_replace_steps: float, blend_words=None, equilizer_params=None, mask_th=(.3,.3)) -> AttentionControlEdit:
361
+ if blend_words is None:
362
+ lb = None
363
+ else:
364
+ lb = LocalBlend(prompts, blend_word, th=mask_th)
365
+ if is_replace_controller:
366
+ controller = AttentionReplace(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps, self_replace_steps=self_replace_steps, local_blend=lb)
367
+ else:
368
+ controller = AttentionRefine(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps, self_replace_steps=self_replace_steps, local_blend=lb)
369
+ if equilizer_params is not None:
370
+ eq = get_equalizer(prompts[1], equilizer_params["words"], equilizer_params["values"])
371
+ controller = AttentionReweight(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps,
372
+ self_replace_steps=self_replace_steps, equalizer=eq, local_blend=lb, controller=controller)
373
+ return controller
374
+
375
+
376
+ def load_512_seq(image_path, left=0, right=0, top=0, bottom=0, n_sample_frame=video_len, sampling_rate=1):
377
+ images = []
378
+ for file in sorted(os.listdir(image_path)):
379
+ images.append(file)
380
+ n_images = len(images)
381
+ sequence_length = (n_sample_frame - 1) * sampling_rate + 1
382
+ if n_images < sequence_length:
383
+ raise ValueError
384
+ frames = []
385
+ for index in range(n_sample_frame):
386
+ p = os.path.join(image_path, images[index])
387
+ image = np.array(Image.open(p).convert("RGB"))
388
+ h, w, c = image.shape
389
+ left = min(left, w-1)
390
+ right = min(right, w - left - 1)
391
+ top = min(top, h - left - 1)
392
+ bottom = min(bottom, h - top - 1)
393
+ image = image[top:h-bottom, left:w-right]
394
+ h, w, c = image.shape
395
+ if h < w:
396
+ offset = (w - h) // 2
397
+ image = image[:, offset:offset + h]
398
+ elif w < h:
399
+ offset = (h - w) // 2
400
+ image = image[offset:offset + w]
401
+ image = np.array(Image.fromarray(image).resize((512, 512)))
402
+ frames.append(image)
403
+ return np.stack(frames)
404
+
405
+
406
+ class NullInversion:
407
+
408
+ def prev_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]):
409
+ prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
410
+ alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
411
+ alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod
412
+ beta_prod_t = 1 - alpha_prod_t
413
+ pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
414
+ pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output
415
+ prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
416
+ return prev_sample
417
+
418
+ def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]): # doing inversion (math)
419
+ timestep, next_timestep = min(timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep
420
+ alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
421
+ alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep]
422
+ beta_prod_t = 1 - alpha_prod_t
423
+ next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
424
+ next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
425
+ next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
426
+ return next_sample
427
+
428
+ def get_noise_pred_single(self, latents, t, context): # latents: torch.Size([1, 4, 64, 64]); t: tensor(1); context: torch.Size([1, 77, 768])
429
+ # formats are correct for video unet input; Tune-A-Video also predicts the residual
430
+ noise_pred = self.model.unet(latents, t, encoder_hidden_states=context)["sample"] # easy to out of mem
431
+ return noise_pred
432
+
433
+ def get_noise_pred(self, latents, t, is_forward=True, context=None):
434
+ latents_input = torch.cat([latents] * 2)
435
+ if context is None:
436
+ context = self.context
437
+ guidance_scale = 1 if is_forward else GUIDANCE_SCALE
438
+ noise_pred = self.model.unet(latents_input, t, encoder_hidden_states=context)["sample"]
439
+ noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
440
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
441
+ if is_forward:
442
+ latents = self.next_step(noise_pred, t, latents)
443
+ else:
444
+ latents = self.prev_step(noise_pred, t, latents)
445
+ return latents
446
+
447
+ @torch.no_grad()
448
+ def latent2image(self, latents, return_type='np'):
449
+ latents = 1 / 0.18215 * latents.detach()
450
+ image = self.model.vae.decode(latents)['sample']
451
+ if return_type == 'np':
452
+ image = (image / 2 + 0.5).clamp(0, 1)
453
+ image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
454
+ image = (image * 255).astype(np.uint8)
455
+ return image
456
+
457
+ @torch.no_grad()
458
+ def latent2image_video(self, latents, return_type='np'):
459
+ latents = 1 / 0.18215 * latents.detach()
460
+ latents = latents[0].permute(1, 0, 2, 3)
461
+ image = self.model.vae.decode(latents)['sample']
462
+ if return_type == 'np':
463
+ image = (image / 2 + 0.5).clamp(0, 1)
464
+ image = image.cpu().permute(0, 2, 3, 1).numpy()
465
+ image = (image * 255).astype(np.uint8)
466
+ return image
467
+
468
+ @torch.no_grad()
469
+ def image2latent(self, image):
470
+ with torch.no_grad():
471
+ if type(image) is Image:
472
+ image = np.array(image)
473
+ if type(image) is torch.Tensor and image.dim() == 4:
474
+ latents = image
475
+ else:
476
+ image = torch.from_numpy(image).float() / 127.5 - 1
477
+ image = image.permute(2, 0, 1).unsqueeze(0).to(device)
478
+ latents = self.model.vae.encode(image)['latent_dist'].mean
479
+ latents = latents * 0.18215
480
+ return latents
481
+
482
+ @torch.no_grad()
483
+ def image2latent_video(self, image):
484
+ with torch.no_grad():
485
+ image = torch.from_numpy(image).float() / 127.5 - 1
486
+ image = image.permute(0, 3, 1, 2).to(device)
487
+ latents = self.model.vae.encode(image)['latent_dist'].mean
488
+ latents = rearrange(latents, "(b f) c h w -> b c f h w", b=1)
489
+ latents = latents * 0.18215
490
+ return latents
491
+
492
+ @torch.no_grad()
493
+ def init_prompt(self, prompt: str):
494
+ uncond_input = self.model.tokenizer(
495
+ [""], padding="max_length", max_length=self.model.tokenizer.model_max_length,
496
+ return_tensors="pt"
497
+ )
498
+ uncond_embeddings = self.model.text_encoder(uncond_input.input_ids.to(self.model.device))[0] # len=2, uncond_embeddings
499
+ text_input = self.model.tokenizer(
500
+ [prompt],
501
+ padding="max_length",
502
+ max_length=self.model.tokenizer.model_max_length,
503
+ truncation=True,
504
+ return_tensors="pt",
505
+ )
506
+ text_embeddings = self.model.text_encoder(text_input.input_ids.to(self.model.device))[0]
507
+ self.context = torch.cat([uncond_embeddings, text_embeddings])
508
+ self.prompt = prompt
509
+
510
+ @torch.no_grad()
511
+ def ddim_loop(self, latent):
512
+ uncond_embeddings, cond_embeddings = self.context.chunk(2)
513
+ all_latent = [latent]
514
+ latent = latent.clone().detach()
515
+ for i in range(NUM_DDIM_STEPS):
516
+ t = self.model.scheduler.timesteps[len(self.model.scheduler.timesteps) - i - 1]
517
+ # latent: torch.Size([1, 4, 8, 16, 16])
518
+ # cond_embeddings: torch.Size([1, 77, 768])
519
+ # noise_pred: torch.Size([1, 4, 8, 16, 16])
520
+ noise_pred = self.get_noise_pred_single(latent, t, cond_embeddings) # use a unet
521
+ latent = self.next_step(noise_pred, t, latent)
522
+ all_latent.append(latent)
523
+ return all_latent
524
+
525
+ @property
526
+ def scheduler(self):
527
+ return self.model.scheduler
528
+
529
+ @torch.no_grad()
530
+ def ddim_inversion(self, image):
531
+ latent = self.image2latent_video(image)
532
+ image_rec = self.latent2image_video(latent) # image: (512, 512, 3); latent: torch.Size([1, 4, 64, 64])
533
+ ddim_latents = self.ddim_loop(latent)
534
+ return image_rec, ddim_latents
535
+
536
+ def null_optimization(self, latents, num_inner_steps, epsilon): # uncond_embeddings is what we what
537
+ uncond_embeddings, cond_embeddings = self.context.chunk(2)
538
+ uncond_embeddings_list = []
539
+ latent_cur = latents[-1]
540
+ bar = tqdm(total=num_inner_steps * NUM_DDIM_STEPS)
541
+ for i in range(NUM_DDIM_STEPS):
542
+ uncond_embeddings = uncond_embeddings.clone().detach()
543
+ uncond_embeddings.requires_grad = True
544
+ optimizer = Adam([uncond_embeddings], lr=1e-2 * (1. - i / 100.))
545
+ latent_prev = latents[len(latents) - i - 2] # GT
546
+ t = self.model.scheduler.timesteps[i]
547
+ with torch.no_grad():
548
+ noise_pred_cond = self.get_noise_pred_single(latent_cur, t, cond_embeddings)
549
+ for j in range(num_inner_steps):
550
+ noise_pred_uncond = self.get_noise_pred_single(latent_cur, t, uncond_embeddings)
551
+ noise_pred = noise_pred_uncond + GUIDANCE_SCALE * (noise_pred_cond - noise_pred_uncond)
552
+ latents_prev_rec = self.prev_step(noise_pred, t, latent_cur)
553
+ loss = nnf.mse_loss(latents_prev_rec, latent_prev)
554
+ optimizer.zero_grad()
555
+ loss.backward()
556
+ optimizer.step()
557
+ loss_item = loss.item()
558
+ bar.update()
559
+ if loss_item < epsilon + i * 2e-5:
560
+ break
561
+ for j in range(j + 1, num_inner_steps):
562
+ bar.update()
563
+ uncond_embeddings_list.append(uncond_embeddings[:1].detach())
564
+ with torch.no_grad():
565
+ context = torch.cat([uncond_embeddings, cond_embeddings])
566
+ latent_cur = self.get_noise_pred(latent_cur, t, False, context)
567
+ bar.close()
568
+ return uncond_embeddings_list
569
+
570
+ def invert(self, image_path: str, prompt: str, offsets=(0,0,0,0), num_inner_steps=10, early_stop_epsilon=1e-5, verbose=False):
571
+ self.init_prompt(prompt)
572
+ ptp_utils.register_attention_control(self.model, None)
573
+ image_gt = load_512_seq(image_path, *offsets)
574
+ if verbose:
575
+ print("DDIM inversion...")
576
+ image_rec, ddim_latents = self.ddim_inversion(image_gt) # ddim_latents is a list, like the link in Figure 3
577
+ # image_rec refers to vq-autoencoder reconstruction
578
+ if verbose:
579
+ print("Null-text optimization...")
580
+ uncond_embeddings = self.null_optimization(ddim_latents, num_inner_steps, early_stop_epsilon) # ddim_latents serve as GT; easy to out of mem
581
+ return (image_gt, image_rec), ddim_latents[-1], uncond_embeddings
582
+
583
+ def invert_(self, image_path: str, prompt: str, offsets=(0,0,0,0), num_inner_steps=10, early_stop_epsilon=1e-5, verbose=False):
584
+ self.init_prompt(prompt)
585
+ ptp_utils.register_attention_control(self.model, None)
586
+ image_gt = load_512_seq(image_path, *offsets)
587
+ if verbose:
588
+ print("DDIM inversion...")
589
+ image_rec, ddim_latents = self.ddim_inversion(image_gt) # ddim_latents is a list, like the link in Figure 3
590
+ # image_rec refers to vq-autoencoder reconstruction
591
+ if verbose:
592
+ print("Null-text optimization...")
593
+ return (image_gt, image_rec), ddim_latents[-1], None
594
+
595
+ def __init__(self, model):
596
+ scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False,
597
+ set_alpha_to_one=False)
598
+ self.model = model
599
+ self.tokenizer = self.model.tokenizer
600
+ self.model.scheduler.set_timesteps(NUM_DDIM_STEPS)
601
+ self.prompt = None
602
+ self.context = None
603
+
604
+ null_inversion = NullInversion(ldm_stable)
605
+
606
+
607
+ @torch.no_grad()
608
+ def text2image_ldm_stable(
609
+ model,
610
+ prompt: List[str],
611
+ controller,
612
+ num_inference_steps: int = 50,
613
+ guidance_scale: Optional[float] = 7.5,
614
+ generator: Optional[torch.Generator] = None,
615
+ latent: Optional[torch.FloatTensor] = None,
616
+ uncond_embeddings=None,
617
+ start_time=50,
618
+ return_type='image'
619
+ ):
620
+ batch_size = len(prompt)
621
+ ptp_utils.register_attention_control(model, controller)
622
+ height = width = 512
623
+
624
+ text_input = model.tokenizer(
625
+ prompt,
626
+ padding="max_length",
627
+ max_length=model.tokenizer.model_max_length,
628
+ truncation=True,
629
+ return_tensors="pt",
630
+ )
631
+ text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]
632
+
633
+ max_length = text_input.input_ids.shape[-1]
634
+ if uncond_embeddings is None:
635
+ uncond_input = model.tokenizer(
636
+ [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
637
+ )
638
+ uncond_embeddings_ = model.text_encoder(uncond_input.input_ids.to(model.device))[0]
639
+ else:
640
+ uncond_embeddings_ = None
641
+
642
+ model.scheduler.set_timesteps(num_inference_steps)
643
+ for i, t in enumerate(tqdm(model.scheduler.timesteps[-start_time:])):
644
+ if uncond_embeddings_ is None:
645
+ context = torch.cat([uncond_embeddings[i].expand(*text_embeddings.shape), text_embeddings])
646
+ else:
647
+ context = torch.cat([uncond_embeddings_, text_embeddings])
648
+ latents = latent
649
+ latents = ptp_utils.diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False)
650
+ if return_type == 'image':
651
+ image = ptp_utils.latent2image_video(model.vae, latents)
652
+ else:
653
+ image = latents
654
+ return image, latent
655
+
656
+
657
+ ###############
658
+ # Custom APIs:
659
+
660
+ ldm_stable.enable_xformers_memory_efficient_attention()
661
+
662
+ (image_gt, image_enc), x_t, uncond_embeddings = null_inversion.invert_(image_path, prompt, offsets=(0,0,0,0), verbose=True)
663
+
664
+ ##### load uncond #####
665
+ # uncond_embeddings_load = np.load(uncond_embeddings_path)
666
+ # uncond_embeddings = []
667
+ # for i in range(uncond_embeddings_load.shape[0]):
668
+ # uncond_embeddings.append(torch.from_numpy(uncond_embeddings_load[i]).to(device))
669
+ #######################
670
+
671
+ ##### save uncond #####
672
+ # uncond_embeddings = torch.cat(uncond_embeddings)
673
+ # uncond_embeddings = uncond_embeddings.cpu().numpy()
674
+ #######################
675
+
676
+ print("Start Video-P2P!")
677
+
678
+ controller = make_controller(prompts, IRC, cross_replace_steps, self_replace_steps, blend_word, eq_params, mask_th=mask_th)
679
+ ptp_utils.register_attention_control(ldm_stable, controller)
680
+ generator = torch.Generator(device=device)
681
+ with torch.no_grad():
682
+ sequence = ldm_stable(
683
+ prompts,
684
+ generator=generator,
685
+ latents=x_t,
686
+ uncond_embeddings_pre=uncond_embeddings,
687
+ controller = controller,
688
+ video_length=video_len,
689
+ simple=True,
690
+ ).videos
691
+ sequence1 = rearrange(sequence[0], "c t h w -> t h w c")
692
+ sequence2 = rearrange(sequence[1], "c t h w -> t h w c")
693
+ inversion = []
694
+ videop2p = []
695
+ for i in range(sequence1.shape[0]):
696
+ inversion.append( Image.fromarray((sequence1[i] * 255).numpy().astype(np.uint8)) )
697
+ videop2p.append( Image.fromarray((sequence2[i] * 255).numpy().astype(np.uint8)) )
698
+
699
+ inversion[0].save(gif_name_1.replace('name', 'inversion_simple_debug'), save_all=True, append_images=inversion[1:], optimize=False, loop=0, duration=250)
700
+ videop2p[0].save(gif_name_2.replace('name', 'p2p_simple_debug'), save_all=True, append_images=videop2p[1:], optimize=False, loop=0, duration=250)
701
+
702
+ if __name__ == "__main__":
703
+ parser = argparse.ArgumentParser()
704
+ parser.add_argument("--config", type=str, default="./configs/videop2p.yaml")
705
+ args = parser.parse_args()
706
+
707
+ main(**OmegaConf.load(args.config))
script.sh ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ # python run_tuning.py --config="configs/rabbit-jump.yaml"
2
+
3
+ python run_videop2p.py --config="configs/rabbit-jump-p2p.yaml"
4
+
5
+ # python run_tuning.py --config="configs/man-motor.yaml"
seq_aligner.py ADDED
@@ -0,0 +1,196 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ import torch
15
+ import numpy as np
16
+
17
+
18
+ class ScoreParams:
19
+
20
+ def __init__(self, gap, match, mismatch):
21
+ self.gap = gap
22
+ self.match = match
23
+ self.mismatch = mismatch
24
+
25
+ def mis_match_char(self, x, y):
26
+ if x != y:
27
+ return self.mismatch
28
+ else:
29
+ return self.match
30
+
31
+
32
+ def get_matrix(size_x, size_y, gap):
33
+ matrix = []
34
+ for i in range(len(size_x) + 1):
35
+ sub_matrix = []
36
+ for j in range(len(size_y) + 1):
37
+ sub_matrix.append(0)
38
+ matrix.append(sub_matrix)
39
+ for j in range(1, len(size_y) + 1):
40
+ matrix[0][j] = j*gap
41
+ for i in range(1, len(size_x) + 1):
42
+ matrix[i][0] = i*gap
43
+ return matrix
44
+
45
+
46
+ def get_matrix(size_x, size_y, gap):
47
+ matrix = np.zeros((size_x + 1, size_y + 1), dtype=np.int32)
48
+ matrix[0, 1:] = (np.arange(size_y) + 1) * gap
49
+ matrix[1:, 0] = (np.arange(size_x) + 1) * gap
50
+ return matrix
51
+
52
+
53
+ def get_traceback_matrix(size_x, size_y):
54
+ matrix = np.zeros((size_x + 1, size_y +1), dtype=np.int32)
55
+ matrix[0, 1:] = 1
56
+ matrix[1:, 0] = 2
57
+ matrix[0, 0] = 4
58
+ return matrix
59
+
60
+
61
+ def global_align(x, y, score):
62
+ matrix = get_matrix(len(x), len(y), score.gap)
63
+ trace_back = get_traceback_matrix(len(x), len(y))
64
+ for i in range(1, len(x) + 1):
65
+ for j in range(1, len(y) + 1):
66
+ left = matrix[i, j - 1] + score.gap
67
+ up = matrix[i - 1, j] + score.gap
68
+ diag = matrix[i - 1, j - 1] + score.mis_match_char(x[i - 1], y[j - 1])
69
+ matrix[i, j] = max(left, up, diag)
70
+ if matrix[i, j] == left:
71
+ trace_back[i, j] = 1
72
+ elif matrix[i, j] == up:
73
+ trace_back[i, j] = 2
74
+ else:
75
+ trace_back[i, j] = 3
76
+ return matrix, trace_back
77
+
78
+
79
+ def get_aligned_sequences(x, y, trace_back):
80
+ x_seq = []
81
+ y_seq = []
82
+ i = len(x)
83
+ j = len(y)
84
+ mapper_y_to_x = []
85
+ while i > 0 or j > 0:
86
+ if trace_back[i, j] == 3:
87
+ x_seq.append(x[i-1])
88
+ y_seq.append(y[j-1])
89
+ i = i-1
90
+ j = j-1
91
+ mapper_y_to_x.append((j, i))
92
+ elif trace_back[i][j] == 1:
93
+ x_seq.append('-')
94
+ y_seq.append(y[j-1])
95
+ j = j-1
96
+ mapper_y_to_x.append((j, -1))
97
+ elif trace_back[i][j] == 2:
98
+ x_seq.append(x[i-1])
99
+ y_seq.append('-')
100
+ i = i-1
101
+ elif trace_back[i][j] == 4:
102
+ break
103
+ mapper_y_to_x.reverse()
104
+ return x_seq, y_seq, torch.tensor(mapper_y_to_x, dtype=torch.int64)
105
+
106
+
107
+ def get_mapper(x: str, y: str, tokenizer, max_len=77):
108
+ x_seq = tokenizer.encode(x)
109
+ y_seq = tokenizer.encode(y)
110
+ score = ScoreParams(0, 1, -1)
111
+ matrix, trace_back = global_align(x_seq, y_seq, score)
112
+ mapper_base = get_aligned_sequences(x_seq, y_seq, trace_back)[-1]
113
+ alphas = torch.ones(max_len)
114
+ alphas[: mapper_base.shape[0]] = mapper_base[:, 1].ne(-1).float()
115
+ mapper = torch.zeros(max_len, dtype=torch.int64)
116
+ mapper[:mapper_base.shape[0]] = mapper_base[:, 1]
117
+ mapper[mapper_base.shape[0]:] = len(y_seq) + torch.arange(max_len - len(y_seq))
118
+ return mapper, alphas
119
+
120
+
121
+ def get_refinement_mapper(prompts, tokenizer, max_len=77):
122
+ x_seq = prompts[0]
123
+ mappers, alphas = [], []
124
+ for i in range(1, len(prompts)):
125
+ mapper, alpha = get_mapper(x_seq, prompts[i], tokenizer, max_len)
126
+ mappers.append(mapper)
127
+ alphas.append(alpha)
128
+ return torch.stack(mappers), torch.stack(alphas)
129
+
130
+
131
+ def get_word_inds(text: str, word_place: int, tokenizer):
132
+ split_text = text.split(" ")
133
+ if type(word_place) is str:
134
+ word_place = [i for i, word in enumerate(split_text) if word_place == word]
135
+ elif type(word_place) is int:
136
+ word_place = [word_place]
137
+ out = []
138
+ if len(word_place) > 0:
139
+ words_encode = [tokenizer.decode([item]).strip("#") for item in tokenizer.encode(text)][1:-1]
140
+ cur_len, ptr = 0, 0
141
+
142
+ for i in range(len(words_encode)):
143
+ cur_len += len(words_encode[i])
144
+ if ptr in word_place:
145
+ out.append(i + 1)
146
+ if cur_len >= len(split_text[ptr]):
147
+ ptr += 1
148
+ cur_len = 0
149
+ return np.array(out)
150
+
151
+
152
+ def get_replacement_mapper_(x: str, y: str, tokenizer, max_len=77):
153
+ words_x = x.split(' ')
154
+ words_y = y.split(' ')
155
+ if len(words_x) != len(words_y):
156
+ raise ValueError(f"attention replacement edit can only be applied on prompts with the same length"
157
+ f" but prompt A has {len(words_x)} words and prompt B has {len(words_y)} words.")
158
+ inds_replace = [i for i in range(len(words_y)) if words_y[i] != words_x[i]]
159
+ inds_source = [get_word_inds(x, i, tokenizer) for i in inds_replace]
160
+ inds_target = [get_word_inds(y, i, tokenizer) for i in inds_replace]
161
+ mapper = np.zeros((max_len, max_len))
162
+ i = j = 0
163
+ cur_inds = 0
164
+ while i < max_len and j < max_len:
165
+ if cur_inds < len(inds_source) and inds_source[cur_inds][0] == i:
166
+ inds_source_, inds_target_ = inds_source[cur_inds], inds_target[cur_inds]
167
+ if len(inds_source_) == len(inds_target_):
168
+ mapper[inds_source_, inds_target_] = 1
169
+ else:
170
+ ratio = 1 / len(inds_target_)
171
+ for i_t in inds_target_:
172
+ mapper[inds_source_, i_t] = ratio
173
+ cur_inds += 1
174
+ i += len(inds_source_)
175
+ j += len(inds_target_)
176
+ elif cur_inds < len(inds_source):
177
+ mapper[i, j] = 1
178
+ i += 1
179
+ j += 1
180
+ else:
181
+ mapper[j, j] = 1
182
+ i += 1
183
+ j += 1
184
+
185
+ return torch.from_numpy(mapper).float()
186
+
187
+
188
+
189
+ def get_replacement_mapper(prompts, tokenizer, max_len=77):
190
+ x_seq = prompts[0]
191
+ mappers = []
192
+ for i in range(1, len(prompts)):
193
+ mapper = get_replacement_mapper_(x_seq, prompts[i], tokenizer, max_len)
194
+ mappers.append(mapper)
195
+ return torch.stack(mappers)
196
+
tuneavideo/data/dataset.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import decord
2
+ decord.bridge.set_bridge('torch')
3
+
4
+ from torch.utils.data import Dataset
5
+ from einops import rearrange
6
+ import os
7
+ from PIL import Image
8
+ import numpy as np
9
+
10
+ class TuneAVideoDataset(Dataset):
11
+ def __init__(
12
+ self,
13
+ video_path: str,
14
+ prompt: str,
15
+ width: int = 512,
16
+ height: int = 512,
17
+ n_sample_frames: int = 8,
18
+ sample_start_idx: int = 0,
19
+ sample_frame_rate: int = 1,
20
+ ):
21
+ self.video_path = video_path
22
+ self.prompt = prompt
23
+ self.prompt_ids = None
24
+ self.uncond_prompt_ids = None
25
+
26
+ self.width = width
27
+ self.height = height
28
+ self.n_sample_frames = n_sample_frames
29
+ self.sample_start_idx = sample_start_idx
30
+ self.sample_frame_rate = sample_frame_rate
31
+
32
+ if 'mp4' not in self.video_path:
33
+ self.images = []
34
+ for file in sorted(os.listdir(self.video_path), key=lambda x: int(x[:-4])):
35
+ if file.endswith('jpg'):
36
+ self.images.append(np.asarray(Image.open(os.path.join(self.video_path, file)).convert('RGB').resize((self.width, self.height))))
37
+ self.images = np.stack(self.images)
38
+
39
+ def __len__(self):
40
+ return 1
41
+
42
+ def __getitem__(self, index):
43
+ # load and sample video frames
44
+ if 'mp4' in self.video_path:
45
+ vr = decord.VideoReader(self.video_path, width=self.width, height=self.height)
46
+ sample_index = list(range(self.sample_start_idx, len(vr), self.sample_frame_rate))[:self.n_sample_frames]
47
+ video = vr.get_batch(sample_index)
48
+ else:
49
+ video = self.images[:self.n_sample_frames]
50
+ video = rearrange(video, "f h w c -> f c h w")
51
+
52
+ example = {
53
+ "pixel_values": (video / 127.5 - 1.0),
54
+ "prompt_ids": self.prompt_ids,
55
+ }
56
+
57
+ return example
tuneavideo/models/attention.py ADDED
@@ -0,0 +1,388 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 = FrameAttention(
159
+ # self.attn1 = SparseCausalAttention(
160
+ query_dim=dim,
161
+ heads=num_attention_heads,
162
+ dim_head=attention_head_dim,
163
+ dropout=dropout,
164
+ bias=attention_bias,
165
+ cross_attention_dim=cross_attention_dim if only_cross_attention else None,
166
+ upcast_attention=upcast_attention,
167
+ )
168
+ self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
169
+
170
+ # Cross-Attn
171
+ if cross_attention_dim is not None:
172
+ self.attn2 = CrossAttention(
173
+ query_dim=dim,
174
+ cross_attention_dim=cross_attention_dim,
175
+ heads=num_attention_heads,
176
+ dim_head=attention_head_dim,
177
+ dropout=dropout,
178
+ bias=attention_bias,
179
+ upcast_attention=upcast_attention,
180
+ )
181
+ else:
182
+ self.attn2 = None
183
+
184
+ if cross_attention_dim is not None:
185
+ self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
186
+ else:
187
+ self.norm2 = None
188
+
189
+ # Feed-forward
190
+ self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
191
+ self.norm3 = nn.LayerNorm(dim)
192
+
193
+ # Temp-Attn
194
+ self.attn_temp = CrossAttention(
195
+ query_dim=dim,
196
+ heads=num_attention_heads,
197
+ dim_head=attention_head_dim,
198
+ dropout=dropout,
199
+ bias=attention_bias,
200
+ upcast_attention=upcast_attention,
201
+ )
202
+ nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
203
+ self.norm_temp = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
204
+
205
+ def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool):
206
+ if not is_xformers_available():
207
+ print("Here is how to install it")
208
+ raise ModuleNotFoundError(
209
+ "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
210
+ " xformers",
211
+ name="xformers",
212
+ )
213
+ elif not torch.cuda.is_available():
214
+ raise ValueError(
215
+ "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only"
216
+ " available for GPU "
217
+ )
218
+ else:
219
+ try:
220
+ # Make sure we can run the memory efficient attention
221
+ _ = xformers.ops.memory_efficient_attention(
222
+ torch.randn((1, 2, 40), device="cuda"),
223
+ torch.randn((1, 2, 40), device="cuda"),
224
+ torch.randn((1, 2, 40), device="cuda"),
225
+ )
226
+ except Exception as e:
227
+ raise e
228
+ self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
229
+ if self.attn2 is not None:
230
+ self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
231
+ # self.attn_temp._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
232
+
233
+ def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None):
234
+ # SparseCausal-Attention
235
+ norm_hidden_states = (
236
+ self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)
237
+ )
238
+
239
+ if self.only_cross_attention:
240
+ hidden_states = (
241
+ self.attn1(norm_hidden_states, encoder_hidden_states, attention_mask=attention_mask) + hidden_states
242
+ )
243
+ else:
244
+ hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states
245
+
246
+ if self.attn2 is not None:
247
+ # Cross-Attention
248
+ norm_hidden_states = (
249
+ self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
250
+ )
251
+ hidden_states = (
252
+ self.attn2(
253
+ norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
254
+ )
255
+ + hidden_states
256
+ )
257
+
258
+ # Feed-forward
259
+ hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
260
+
261
+ # Temporal-Attention
262
+ d = hidden_states.shape[1]
263
+ hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length)
264
+ norm_hidden_states = (
265
+ self.norm_temp(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_temp(hidden_states)
266
+ )
267
+ hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
268
+ hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
269
+
270
+ return hidden_states
271
+
272
+
273
+ class SparseCausalAttention(CrossAttention):
274
+ def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):
275
+ batch_size, sequence_length, _ = hidden_states.shape
276
+
277
+ encoder_hidden_states = encoder_hidden_states
278
+
279
+ if self.group_norm is not None:
280
+ hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
281
+
282
+ query = self.to_q(hidden_states)
283
+ dim = query.shape[-1]
284
+ query = self.reshape_heads_to_batch_dim(query)
285
+
286
+ if self.added_kv_proj_dim is not None:
287
+ raise NotImplementedError
288
+
289
+ encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
290
+ key = self.to_k(encoder_hidden_states)
291
+ value = self.to_v(encoder_hidden_states)
292
+
293
+ former_frame_index = torch.arange(video_length) - 1
294
+ former_frame_index[0] = 0
295
+
296
+ key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
297
+ key = torch.cat([key[:, [0] * video_length], key[:, former_frame_index]], dim=2)
298
+ key = rearrange(key, "b f d c -> (b f) d c")
299
+
300
+ value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
301
+ value = torch.cat([value[:, [0] * video_length], value[:, former_frame_index]], dim=2)
302
+ value = rearrange(value, "b f d c -> (b f) d c")
303
+
304
+ key = self.reshape_heads_to_batch_dim(key)
305
+ value = self.reshape_heads_to_batch_dim(value)
306
+
307
+ if attention_mask is not None:
308
+ if attention_mask.shape[-1] != query.shape[1]:
309
+ target_length = query.shape[1]
310
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
311
+ attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
312
+
313
+ # attention, what we cannot get enough of
314
+ if self._use_memory_efficient_attention_xformers:
315
+ hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask)
316
+ # Some versions of xformers return output in fp32, cast it back to the dtype of the input
317
+ hidden_states = hidden_states.to(query.dtype)
318
+ else:
319
+ if self._slice_size is None or query.shape[0] // self._slice_size == 1:
320
+ hidden_states = self._attention(query, key, value, attention_mask)
321
+ else:
322
+ hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask)
323
+
324
+ # linear proj
325
+ hidden_states = self.to_out[0](hidden_states)
326
+
327
+ # dropout
328
+ hidden_states = self.to_out[1](hidden_states)
329
+ return hidden_states
330
+
331
+
332
+ class FrameAttention(CrossAttention):
333
+ def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):
334
+ batch_size, sequence_length, _ = hidden_states.shape
335
+
336
+ encoder_hidden_states = encoder_hidden_states
337
+
338
+ if self.group_norm is not None:
339
+ hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
340
+
341
+ query = self.to_q(hidden_states)
342
+ dim = query.shape[-1]
343
+ query = self.reshape_heads_to_batch_dim(query)
344
+
345
+ if self.added_kv_proj_dim is not None:
346
+ raise NotImplementedError
347
+
348
+ encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
349
+ key = self.to_k(encoder_hidden_states)
350
+ value = self.to_v(encoder_hidden_states)
351
+
352
+ former_frame_index = torch.arange(video_length) - 1
353
+ former_frame_index[0] = 0
354
+
355
+ key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
356
+ key = key[:, [0] * video_length]
357
+ key = rearrange(key, "b f d c -> (b f) d c")
358
+
359
+ value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
360
+ value = value[:, [0] * video_length]
361
+ value = rearrange(value, "b f d c -> (b f) d c")
362
+
363
+ key = self.reshape_heads_to_batch_dim(key)
364
+ value = self.reshape_heads_to_batch_dim(value)
365
+
366
+ if attention_mask is not None:
367
+ if attention_mask.shape[-1] != query.shape[1]:
368
+ target_length = query.shape[1]
369
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
370
+ attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
371
+
372
+ # attention, what we cannot get enough of
373
+ if self._use_memory_efficient_attention_xformers:
374
+ hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask)
375
+ # Some versions of xformers return output in fp32, cast it back to the dtype of the input
376
+ hidden_states = hidden_states.to(query.dtype)
377
+ else:
378
+ if self._slice_size is None or query.shape[0] // self._slice_size == 1:
379
+ hidden_states = self._attention(query, key, value, attention_mask)
380
+ else:
381
+ hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask)
382
+
383
+ # linear proj
384
+ hidden_states = self.to_out[0](hidden_states)
385
+
386
+ # dropout
387
+ hidden_states = self.to_out[1](hidden_states)
388
+ return hidden_states
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))
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
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
tuneavideo/pipelines/pipeline_tuneavideo.py ADDED
@@ -0,0 +1,449 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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, repeat
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
+ bs = 4
243
+ video_list = []
244
+ for i in range(max(latents.shape[0]//bs, 1)):
245
+ video = self.vae.decode(latents[i*bs:min((i+1)*bs, latents.shape[0])]).sample
246
+ video = (video / 2 + 0.5).clamp(0, 1)
247
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
248
+ video = video.cpu().float().numpy()
249
+ video_list.append(video)
250
+ if len(video_list) > 1:
251
+ video = np.concatenate(video_list, axis=0)
252
+ else:
253
+ video = video_list[0]
254
+ video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
255
+ return video
256
+
257
+ def prepare_extra_step_kwargs(self, generator, eta):
258
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
259
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
260
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
261
+ # and should be between [0, 1]
262
+
263
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
264
+ extra_step_kwargs = {}
265
+ if accepts_eta:
266
+ extra_step_kwargs["eta"] = eta
267
+
268
+ # check if the scheduler accepts generator
269
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
270
+ if accepts_generator:
271
+ extra_step_kwargs["generator"] = generator
272
+ return extra_step_kwargs
273
+
274
+ def check_inputs(self, prompt, height, width, callback_steps):
275
+ if not isinstance(prompt, str) and not isinstance(prompt, list):
276
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
277
+
278
+ if height % 8 != 0 or width % 8 != 0:
279
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
280
+
281
+ if (callback_steps is None) or (
282
+ callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
283
+ ):
284
+ raise ValueError(
285
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
286
+ f" {type(callback_steps)}."
287
+ )
288
+
289
+ def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None):
290
+ shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor)
291
+ if isinstance(generator, list) and len(generator) != batch_size:
292
+ raise ValueError(
293
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
294
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
295
+ )
296
+
297
+ if latents is None:
298
+ rand_device = "cpu" if device.type == "mps" else device
299
+
300
+ if isinstance(generator, list):
301
+ shape = (1,) + shape[1:]
302
+ latents = [
303
+ torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)
304
+ for i in range(batch_size)
305
+ ]
306
+ latents = torch.cat(latents, dim=0).to(device)
307
+ else:
308
+ latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
309
+ else:
310
+ if latents.shape != shape:
311
+ # raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
312
+ latents = latents.expand(shape)
313
+ latents = latents.to(device)
314
+
315
+ # scale the initial noise by the standard deviation required by the scheduler
316
+ latents = latents * self.scheduler.init_noise_sigma
317
+ return latents
318
+
319
+ @torch.no_grad()
320
+ def __call__(
321
+ self,
322
+ prompt: Union[str, List[str]],
323
+ video_length: Optional[int],
324
+ height: Optional[int] = 512,
325
+ width: Optional[int] = 512,
326
+ num_inference_steps: int = 50,
327
+ guidance_scale: float = 7.5,
328
+ negative_prompt: Optional[Union[str, List[str]]] = None,
329
+ num_videos_per_prompt: Optional[int] = 1,
330
+ eta: float = 0.0,
331
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
332
+ latents: Optional[torch.FloatTensor] = None,
333
+ output_type: Optional[str] = "tensor",
334
+ return_dict: bool = True,
335
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
336
+ callback_steps: Optional[int] = 1,
337
+ uncond_embeddings_pre=None,
338
+ controller=None,
339
+ uncond2=False,
340
+ multi=False,
341
+ region=False,
342
+ simple=False,
343
+ **kwargs,
344
+ ):
345
+ # Default height and width to unet
346
+ height = height or self.unet.config.sample_size * self.vae_scale_factor
347
+ width = width or self.unet.config.sample_size * self.vae_scale_factor
348
+
349
+ # Check inputs. Raise error if not correct
350
+ self.check_inputs(prompt, height, width, callback_steps)
351
+
352
+ # Define call parameters
353
+ batch_size = 1 if isinstance(prompt, str) else len(prompt)
354
+ device = self._execution_device
355
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
356
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
357
+ # corresponds to doing no classifier free guidance.
358
+ do_classifier_free_guidance = guidance_scale > 1.0
359
+
360
+ # Encode input prompt
361
+ text_embeddings = self._encode_prompt(
362
+ prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
363
+ )
364
+ if multi:
365
+ text_embeddings = repeat(text_embeddings, 'b n c -> (b f) n c', f=video_length)
366
+
367
+ # Prepare timesteps
368
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
369
+ timesteps = self.scheduler.timesteps
370
+
371
+ # Prepare latent variables
372
+ num_channels_latents = self.unet.in_channels
373
+ latents = self.prepare_latents(
374
+ batch_size * num_videos_per_prompt,
375
+ num_channels_latents,
376
+ video_length,
377
+ height,
378
+ width,
379
+ text_embeddings.dtype,
380
+ device,
381
+ generator,
382
+ latents,
383
+ )
384
+ latents_dtype = latents.dtype
385
+
386
+ # Prepare extra step kwargs.
387
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
388
+
389
+ # Denoising loop
390
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
391
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
392
+ for i, t in enumerate(timesteps):
393
+ # expand the latents if we are doing classifier free guidance
394
+ latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
395
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) # torch.Size([2, 4, 8, 64, 64])
396
+
397
+ if uncond_embeddings_pre is not None:
398
+ if multi:
399
+ text_embeddings[:video_length] = uncond_embeddings_pre[i]
400
+ else:
401
+ text_embeddings[0] = uncond_embeddings_pre[i] # text_embeddings: torch.Size([2, 77, 768])
402
+ if region:
403
+ text_embeddings[2] = uncond_embeddings_pre[i] # text_embeddings: torch.Size([2, 77, 768])
404
+ if uncond2:
405
+ if multi:
406
+ text_embeddings[video_length: video_length*2] = uncond_embeddings_pre[i]
407
+ else:
408
+ text_embeddings[1] = uncond_embeddings_pre[i] # text_embeddings: torch.Size([2, 77, 768])
409
+
410
+ # predict the noise residual
411
+ noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample.to(dtype=latents_dtype)
412
+ if region:
413
+ mask = controller.get_mask(latents[:-1])[1:2]
414
+ noise_pred[1] = noise_pred[2]*(1-mask) + noise_pred[1]*mask
415
+ # noise_pred[1] = noise_pred[2]
416
+
417
+ # perform guidance
418
+ if do_classifier_free_guidance:
419
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
420
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
421
+ if simple:
422
+ noise_pred[0] = noise_pred_text[0]
423
+
424
+ # compute the previous noisy sample x_t -> x_t-1
425
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
426
+
427
+ if controller is not None:
428
+ latents = controller.step_callback(latents)
429
+
430
+ # call the callback, if provided
431
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
432
+ progress_bar.update()
433
+ if callback is not None and i % callback_steps == 0:
434
+ callback(i, t, latents)
435
+
436
+ if region:
437
+ latents = latents[:2]
438
+
439
+ # Post-processing
440
+ video = self.decode_latents(latents)
441
+
442
+ # Convert to tensor
443
+ if output_type == "tensor":
444
+ video = torch.from_numpy(video)
445
+
446
+ if not return_dict:
447
+ return video
448
+
449
+ return TuneAVideoPipelineOutput(videos=video)
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