lev1 commited on
Commit
ccfd1d5
1 Parent(s): 92ffe41

Text to video files

Browse files
models/cldm_v15.yaml ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ target: cldm.cldm.ControlLDM
3
+ params:
4
+ linear_start: 0.00085
5
+ linear_end: 0.0120
6
+ num_timesteps_cond: 1
7
+ log_every_t: 200
8
+ timesteps: 1000
9
+ first_stage_key: "jpg"
10
+ cond_stage_key: "txt"
11
+ control_key: "hint"
12
+ image_size: 64
13
+ channels: 4
14
+ cond_stage_trainable: false
15
+ conditioning_key: crossattn
16
+ monitor: val/loss_simple_ema
17
+ scale_factor: 0.18215
18
+ use_ema: False
19
+ only_mid_control: False
20
+
21
+ control_stage_config:
22
+ target: cldm.cldm.ControlNet
23
+ params:
24
+ image_size: 32 # unused
25
+ in_channels: 4
26
+ hint_channels: 3
27
+ model_channels: 320
28
+ attention_resolutions: [ 4, 2, 1 ]
29
+ num_res_blocks: 2
30
+ channel_mult: [ 1, 2, 4, 4 ]
31
+ num_heads: 8
32
+ use_spatial_transformer: True
33
+ transformer_depth: 1
34
+ context_dim: 768
35
+ use_checkpoint: True
36
+ use_cf_attn: True
37
+ legacy: False
38
+
39
+ unet_config:
40
+ target: cldm.cldm.ControlledUnetModel
41
+ params:
42
+ image_size: 32 # unused
43
+ in_channels: 4
44
+ out_channels: 4
45
+ model_channels: 320
46
+ attention_resolutions: [ 4, 2, 1 ]
47
+ num_res_blocks: 2
48
+ channel_mult: [ 1, 2, 4, 4 ]
49
+ num_heads: 8
50
+ use_spatial_transformer: True
51
+ transformer_depth: 1
52
+ context_dim: 768
53
+ use_checkpoint: True
54
+ legacy: False
55
+ use_cf_attn: True
56
+
57
+ first_stage_config:
58
+ target: ldm.models.autoencoder.AutoencoderKL
59
+ params:
60
+ embed_dim: 4
61
+ monitor: val/rec_loss
62
+ ddconfig:
63
+ double_z: true
64
+ z_channels: 4
65
+ resolution: 256
66
+ in_channels: 3
67
+ out_ch: 3
68
+ ch: 128
69
+ ch_mult:
70
+ - 1
71
+ - 2
72
+ - 4
73
+ - 4
74
+ num_res_blocks: 2
75
+ attn_resolutions: []
76
+ dropout: 0.0
77
+ lossconfig:
78
+ target: torch.nn.Identity
79
+
80
+ cond_stage_config:
81
+ target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
models/cldm_v15_no_cf_attn.yaml ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ target: cldm.cldm.ControlLDM
3
+ params:
4
+ linear_start: 0.00085
5
+ linear_end: 0.0120
6
+ num_timesteps_cond: 1
7
+ log_every_t: 200
8
+ timesteps: 1000
9
+ first_stage_key: "jpg"
10
+ cond_stage_key: "txt"
11
+ control_key: "hint"
12
+ image_size: 64
13
+ channels: 4
14
+ cond_stage_trainable: false
15
+ conditioning_key: crossattn
16
+ monitor: val/loss_simple_ema
17
+ scale_factor: 0.18215
18
+ use_ema: False
19
+ only_mid_control: False
20
+
21
+ control_stage_config:
22
+ target: cldm.cldm.ControlNet
23
+ params:
24
+ image_size: 32 # unused
25
+ in_channels: 4
26
+ hint_channels: 3
27
+ model_channels: 320
28
+ attention_resolutions: [ 4, 2, 1 ]
29
+ num_res_blocks: 2
30
+ channel_mult: [ 1, 2, 4, 4 ]
31
+ num_heads: 8
32
+ use_spatial_transformer: True
33
+ transformer_depth: 1
34
+ context_dim: 768
35
+ use_checkpoint: True
36
+ use_cf_attn: False
37
+ legacy: False
38
+
39
+ unet_config:
40
+ target: cldm.cldm.ControlledUnetModel
41
+ params:
42
+ image_size: 32 # unused
43
+ in_channels: 4
44
+ out_channels: 4
45
+ model_channels: 320
46
+ attention_resolutions: [ 4, 2, 1 ]
47
+ num_res_blocks: 2
48
+ channel_mult: [ 1, 2, 4, 4 ]
49
+ num_heads: 8
50
+ use_spatial_transformer: True
51
+ transformer_depth: 1
52
+ context_dim: 768
53
+ use_checkpoint: True
54
+ legacy: False
55
+ use_cf_attn: False
56
+
57
+ first_stage_config:
58
+ target: ldm.models.autoencoder.AutoencoderKL
59
+ params:
60
+ embed_dim: 4
61
+ monitor: val/rec_loss
62
+ ddconfig:
63
+ double_z: true
64
+ z_channels: 4
65
+ resolution: 256
66
+ in_channels: 3
67
+ out_ch: 3
68
+ ch: 128
69
+ ch_mult:
70
+ - 1
71
+ - 2
72
+ - 4
73
+ - 4
74
+ num_res_blocks: 2
75
+ attn_resolutions: []
76
+ dropout: 0.0
77
+ lossconfig:
78
+ target: torch.nn.Identity
79
+
80
+ cond_stage_config:
81
+ target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
models/cldm_v21.yaml ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ target: cldm.cldm.ControlLDM
3
+ params:
4
+ linear_start: 0.00085
5
+ linear_end: 0.0120
6
+ num_timesteps_cond: 1
7
+ log_every_t: 200
8
+ timesteps: 1000
9
+ first_stage_key: "jpg"
10
+ cond_stage_key: "txt"
11
+ control_key: "hint"
12
+ image_size: 64
13
+ channels: 4
14
+ cond_stage_trainable: false
15
+ conditioning_key: crossattn
16
+ monitor: val/loss_simple_ema
17
+ scale_factor: 0.18215
18
+ use_ema: False
19
+ only_mid_control: False
20
+
21
+ control_stage_config:
22
+ target: cldm.cldm.ControlNet
23
+ params:
24
+ use_checkpoint: True
25
+ image_size: 32 # unused
26
+ in_channels: 4
27
+ hint_channels: 3
28
+ model_channels: 320
29
+ attention_resolutions: [ 4, 2, 1 ]
30
+ num_res_blocks: 2
31
+ channel_mult: [ 1, 2, 4, 4 ]
32
+ num_head_channels: 64 # need to fix for flash-attn
33
+ use_spatial_transformer: True
34
+ use_linear_in_transformer: True
35
+ transformer_depth: 1
36
+ context_dim: 1024
37
+ legacy: False
38
+
39
+ unet_config:
40
+ target: cldm.cldm.ControlledUnetModel
41
+ params:
42
+ use_checkpoint: True
43
+ image_size: 32 # unused
44
+ in_channels: 4
45
+ out_channels: 4
46
+ model_channels: 320
47
+ attention_resolutions: [ 4, 2, 1 ]
48
+ num_res_blocks: 2
49
+ channel_mult: [ 1, 2, 4, 4 ]
50
+ num_head_channels: 64 # need to fix for flash-attn
51
+ use_spatial_transformer: True
52
+ use_linear_in_transformer: True
53
+ transformer_depth: 1
54
+ context_dim: 1024
55
+ legacy: False
56
+
57
+ first_stage_config:
58
+ target: ldm.models.autoencoder.AutoencoderKL
59
+ params:
60
+ embed_dim: 4
61
+ monitor: val/rec_loss
62
+ ddconfig:
63
+ #attn_type: "vanilla-xformers"
64
+ double_z: true
65
+ z_channels: 4
66
+ resolution: 256
67
+ in_channels: 3
68
+ out_ch: 3
69
+ ch: 128
70
+ ch_mult:
71
+ - 1
72
+ - 2
73
+ - 4
74
+ - 4
75
+ num_res_blocks: 2
76
+ attn_resolutions: []
77
+ dropout: 0.0
78
+ lossconfig:
79
+ target: torch.nn.Identity
80
+
81
+ cond_stage_config:
82
+ target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
83
+ params:
84
+ freeze: True
85
+ layer: "penultimate"
text_to_video/text_to_video_generator.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from text_to_video.tuneavideo.pipelines.pipeline_text_to_video import TuneAVideoPipeline
2
+ from text_to_video.tuneavideo.models.unet import UNet3DConditionModel
3
+ import torch
4
+ from diffusers import AutoencoderKL, DDIMScheduler
5
+ from transformers import CLIPTextModel, CLIPTokenizer
6
+
7
+
8
+ class TextToVideo():
9
+
10
+
11
+ def __init__(self,sd_path = None,motion_field_strength = 12, video_length = 8,t0 = 881, t1=941,use_cf_attn=True,use_motion_field=True) -> None:
12
+ g = torch.Generator(device='cuda')
13
+ g.manual_seed(22)
14
+ self.g = g
15
+
16
+ print(f"Loading model SD-Net model file from {sd_path}")
17
+
18
+ self.dtype = torch.float16
19
+ noise_scheduler = DDIMScheduler.from_pretrained(
20
+ sd_path, subfolder="scheduler")
21
+ tokenizer = CLIPTokenizer.from_pretrained(
22
+ sd_path, subfolder="tokenizer")
23
+ text_encoder = CLIPTextModel.from_pretrained(
24
+ sd_path, subfolder="text_encoder")
25
+ vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae")
26
+
27
+
28
+ unet = UNet3DConditionModel.from_pretrained_2d(
29
+ sd_path, subfolder="unet", use_cf_attn=use_cf_attn)
30
+ self.pipe = TuneAVideoPipeline(
31
+ vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,
32
+ scheduler=DDIMScheduler.from_pretrained(
33
+ sd_path, subfolder="scheduler")
34
+ ).to('cuda').to(self.dtype)
35
+
36
+ noise_scheduler.set_timesteps(50, device='cuda')
37
+
38
+ # t0 parameter (DDIM backward from noise until t0)
39
+ self.t0 = t0
40
+
41
+
42
+ # from t0 apply DDPM forward until t1
43
+ self.t1 = t1
44
+
45
+ self.use_foreground_motion_field = False # apply motion field on forground object (not used)
46
+
47
+ # strength of motion field (delta_x = delta_y in Sect 3.3.1)
48
+ self.motion_field_strength = motion_field_strength
49
+ self.use_motion_field = use_motion_field # apply general motion field
50
+ self.smooth_bg = False # temporally smooth background
51
+ self.smooth_bg_strength = 0.4 # alpha = (1-self.smooth_bg_strength) in Eq (9)
52
+
53
+
54
+ self.video_length = video_length
55
+
56
+ def inference(self, prompt):
57
+
58
+ prompt_compute = [prompt]
59
+ xT = torch.randn((1, 4, 1, 64, 64), dtype=self.dtype, device="cuda")
60
+ result = self.pipe(prompt_compute,
61
+ video_length=self.video_length,
62
+ height=512,
63
+ width=512,
64
+ num_inference_steps=50,
65
+ guidance_scale=7.5,
66
+ guidance_stop_step=1.0,
67
+ t0=self.t0,
68
+ t1=self.t1,
69
+ xT=xT,
70
+ use_foreground_motion_field=self.use_foreground_motion_field,
71
+ motion_field_strength=self.motion_field_strength,
72
+ use_motion_field=self.use_motion_field,
73
+ smooth_bg=self.smooth_bg,
74
+ smooth_bg_strength=self.smooth_bg_strength,
75
+ generator=self.g)
76
+
77
+ return result.videos[0]
text_to_video/text_to_video_pipeline.py ADDED
@@ -0,0 +1,550 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import StableDiffusionPipeline
2
+ import torch
3
+ from dataclasses import dataclass
4
+ from typing import Callable, List, Optional, Union
5
+ import numpy as np
6
+ from diffusers.utils import deprecate, logging, BaseOutput
7
+ from einops import rearrange, repeat
8
+ from torch.nn.functional import grid_sample
9
+ import torchvision.transforms as T
10
+ from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
11
+ from diffusers.models import AutoencoderKL, UNet2DConditionModel
12
+ from diffusers.schedulers import KarrasDiffusionSchedulers
13
+ from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
14
+
15
+ @dataclass
16
+ class TextToVideoPipelineOutput(BaseOutput):
17
+ videos: Union[torch.Tensor, np.ndarray]
18
+ code: Union[torch.Tensor, np.ndarray]
19
+
20
+
21
+
22
+ def coords_grid(batch, ht, wd, device):
23
+ # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py
24
+ coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
25
+ coords = torch.stack(coords[::-1], dim=0).float()
26
+ return coords[None].repeat(batch, 1, 1, 1)
27
+
28
+
29
+
30
+ class TextToVideoPipeline(StableDiffusionPipeline):
31
+ def __init__(
32
+ self,
33
+ vae: AutoencoderKL,
34
+ text_encoder: CLIPTextModel,
35
+ tokenizer: CLIPTokenizer,
36
+ unet: UNet2DConditionModel,
37
+ scheduler: KarrasDiffusionSchedulers,
38
+ safety_checker: StableDiffusionSafetyChecker,
39
+ feature_extractor: CLIPFeatureExtractor,
40
+ requires_safety_checker: bool = True,
41
+ ):
42
+ #super().__init__(*args,**kwargs)
43
+ super().__init__(vae,text_encoder,tokenizer,unet,scheduler,safety_checker,feature_extractor,requires_safety_checker)
44
+
45
+
46
+ def DDPM_forward(self, x0, t0, tMax, generator, device, shape, text_embeddings):
47
+ rand_device = "cpu" if device.type == "mps" else device
48
+
49
+ if x0 is None:
50
+ return torch.randn(shape, generator=generator, device=rand_device, dtype=text_embeddings.dtype).to(device)
51
+ else:
52
+ eps = torch.randn_like(x0, dtype=text_embeddings.dtype).to(device)
53
+ alpha_vec = torch.prod(self.scheduler.alphas[t0:tMax])
54
+ xt = torch.sqrt(alpha_vec) * x0 + \
55
+ torch.sqrt(1-alpha_vec) * eps
56
+ return xt
57
+
58
+
59
+ def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None):
60
+ shape = (batch_size, num_channels_latents, video_length, height //
61
+ self.vae_scale_factor, width // self.vae_scale_factor)
62
+ if isinstance(generator, list) and len(generator) != batch_size:
63
+ raise ValueError(
64
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
65
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
66
+ )
67
+
68
+ if latents is None:
69
+ rand_device = "cpu" if device.type == "mps" else device
70
+
71
+ if isinstance(generator, list):
72
+ shape = (1,) + shape[1:]
73
+ latents = [
74
+ torch.randn(
75
+ shape, generator=generator[i], device=rand_device, dtype=dtype)
76
+ for i in range(batch_size)
77
+ ]
78
+ latents = torch.cat(latents, dim=0).to(device)
79
+ else:
80
+ latents = torch.randn(
81
+ shape, generator=generator, device=rand_device, dtype=dtype).to(device)
82
+ else:
83
+ latents = latents.to(device)
84
+
85
+ # scale the initial noise by the standard deviation required by the scheduler
86
+ latents = latents * self.scheduler.init_noise_sigma
87
+ return latents
88
+
89
+
90
+
91
+ def warp_latents(self, latents, reference_flow):
92
+ _, _, H, W = reference_flow.size()
93
+ b, c, f, h, w = latents.size()
94
+ coords0 = coords_grid(f, H, W, device=latents.device).to(latents.dtype)
95
+ coords_t0 = coords0 + reference_flow
96
+ coords_t0[:, 0] /= W
97
+ coords_t0[:, 1] /= H
98
+ coords_t0 = coords_t0 * 2.0 - 1.0
99
+ coords_t0 = T.Resize((h, w))(coords_t0)
100
+ coords_t0 = rearrange(coords_t0, 'f c h w -> f h w c')
101
+ latents_0 = latents[:, :, 0]
102
+ latents_0 = latents_0.repeat(f, 1, 1, 1)
103
+ warped = grid_sample(latents_0, coords_t0,
104
+ mode='nearest', padding_mode='reflection')
105
+ warped = rearrange(warped, '(b f) c h w -> b c f h w', f=f)
106
+ return warped
107
+
108
+ def warp_latents_independently(self, latents, reference_flow):
109
+ _, _, H, W = reference_flow.size()
110
+ b, c, f, h, w = latents.size()
111
+ assert b == 1
112
+ coords0 = coords_grid(f, H, W, device=latents.device).to(latents.dtype)
113
+ coords_t0 = coords0 + reference_flow
114
+
115
+ coords_t0[:, 0] /= W
116
+ coords_t0[:, 1] /= H
117
+ coords_t0 = coords_t0 * 2.0 - 1.0
118
+
119
+ coords_t0 = T.Resize((h, w))(coords_t0)
120
+
121
+ coords_t0 = rearrange(coords_t0, 'f c h w -> f h w c')
122
+
123
+ latents_0 = rearrange(latents[0], 'c f h w -> f c h w')
124
+
125
+ warped = grid_sample(latents_0, coords_t0,
126
+ mode='nearest', padding_mode='reflection')
127
+ warped = rearrange(warped, '(b f) c h w -> b c f h w', f=f)
128
+ return warped
129
+
130
+ def DDIM_backward(self, num_inference_steps, timesteps, skip_t, t0, t1, do_classifier_free_guidance, null_embs, text_embeddings, latents_local, latents_dtype, guidance_scale, guidance_stop_step, callback, callback_steps, extra_step_kwargs, num_warmup_steps):
131
+ entered = False
132
+
133
+ f = latents_local.shape[2]
134
+ latents_local = rearrange(latents_local,"b c f w h -> (b f) c w h")
135
+
136
+ latents = latents_local.detach().clone()
137
+ x_t0_1 = None
138
+ x_t1_1 = None
139
+
140
+
141
+
142
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
143
+ for i, t in enumerate(timesteps):
144
+ if t > skip_t:
145
+ # print("Skipping frame!")
146
+ continue
147
+ else:
148
+ if not entered:
149
+ print(
150
+ f"Continue DDIM with i = {i}, t = {t}, latent = {latents.shape}, device = {latents.device}, type = {latents.dtype}")
151
+ entered = True
152
+
153
+ latents = latents.detach()
154
+ # expand the latents if we are doing classifier free guidance
155
+ latent_model_input = torch.cat(
156
+ [latents] * 2) if do_classifier_free_guidance else latents
157
+ latent_model_input = self.scheduler.scale_model_input(
158
+ latent_model_input, t)
159
+
160
+ # predict the noise residual
161
+ with torch.no_grad():
162
+ if null_embs is not None:
163
+ text_embeddings[0] = null_embs[i][0]
164
+ te = torch.cat([repeat(text_embeddings[0,:,:], "c k -> f c k",f=f),repeat(text_embeddings[1,:,:], "c k -> f c k",f=f)])
165
+ noise_pred = self.unet(
166
+ latent_model_input, t, encoder_hidden_states=te).sample.to(dtype=latents_dtype)
167
+
168
+ # perform guidance
169
+ if do_classifier_free_guidance:
170
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(
171
+ 2)
172
+ noise_pred = noise_pred_uncond + guidance_scale * \
173
+ (noise_pred_text - noise_pred_uncond)
174
+
175
+ if i >= guidance_stop_step * len(timesteps):
176
+ alpha = 0
177
+ # compute the previous noisy sample x_t -> x_t-1
178
+ latents = self.scheduler.step(
179
+ noise_pred, t, latents, **extra_step_kwargs).prev_sample
180
+ # latents = latents - alpha * grads / (torch.norm(grads) + 1e-10)
181
+ # call the callback, if provided
182
+
183
+ if i < len(timesteps)-1 and timesteps[i+1] == t0:
184
+ x_t0_1 = latents.detach().clone()
185
+ print(f"latent t0 found at i = {i}, t = {t}")
186
+ elif i < len(timesteps)-1 and timesteps[i+1] == t1:
187
+ x_t1_1 = latents.detach().clone()
188
+ print(f"latent t1 found at i={i}, t = {t}")
189
+
190
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
191
+ progress_bar.update()
192
+ if callback is not None and i % callback_steps == 0:
193
+ callback(i, t, latents)
194
+
195
+
196
+ latents = rearrange(latents,"(b f) c w h -> b c f w h",f = f)
197
+
198
+
199
+
200
+ res = {"x0": latents.detach().clone()}
201
+ if x_t0_1 is not None:
202
+ x_t0_1 = rearrange(x_t0_1,"(b f) c w h -> b c f w h",f = f)
203
+ res["x_t0_1"] = x_t0_1.detach().clone()
204
+ if x_t1_1 is not None:
205
+ x_t1_1 = rearrange(x_t1_1,"(b f) c w h -> b c f w h",f = f)
206
+ res["x_t1_1"] = x_t1_1.detach().clone()
207
+ return res
208
+
209
+ def decode_latents(self, latents):
210
+ video_length = latents.shape[2]
211
+ latents = 1 / 0.18215 * latents
212
+ latents = rearrange(latents, "b c f h w -> (b f) c h w")
213
+ video = self.vae.decode(latents).sample
214
+ video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
215
+ video = (video / 2 + 0.5).clamp(0, 1)
216
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
217
+ return video
218
+
219
+
220
+
221
+ @torch.no_grad()
222
+ def __call__(
223
+ self,
224
+ prompt: Union[str, List[str]],
225
+ video_length: Optional[int],
226
+ height: Optional[int] = None,
227
+ width: Optional[int] = None,
228
+ num_inference_steps: int = 50,
229
+ guidance_scale: float = 7.5,
230
+ guidance_stop_step: float = 0.5,
231
+ negative_prompt: Optional[Union[str, List[str]]] = None,
232
+ num_videos_per_prompt: Optional[int] = 1,
233
+ eta: float = 0.0,
234
+ generator: Optional[Union[torch.Generator,
235
+ List[torch.Generator]]] = None,
236
+ xT: Optional[torch.FloatTensor] = None,
237
+ null_embs: Optional[torch.FloatTensor] = None,
238
+ motion_field_strength: float = 12,
239
+ output_type: Optional[str] = "tensor",
240
+ return_dict: bool = True,
241
+ callback: Optional[Callable[[
242
+ int, int, torch.FloatTensor], None]] = None,
243
+ callback_steps: Optional[int] = 1,
244
+ use_foreground_motion_field: bool = True,
245
+ use_motion_field: bool = True,
246
+ smooth_bg: bool = True,
247
+ smooth_bg_strength: float = 0.4,
248
+ **kwargs,
249
+ ):
250
+
251
+ print(f" Use: Motion field = {use_motion_field}")
252
+ print(f" Use: Background smoothing = {smooth_bg}")
253
+ # Default height and width to unet
254
+ height = height or self.unet.config.sample_size * self.vae_scale_factor
255
+ width = width or self.unet.config.sample_size * self.vae_scale_factor
256
+
257
+ # Check inputs. Raise error if not correct
258
+ self.check_inputs(prompt, height, width, callback_steps)
259
+
260
+ # Define call parameters
261
+ batch_size = 1 if isinstance(prompt, str) else len(prompt)
262
+ device = self._execution_device
263
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
264
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
265
+ # corresponds to doing no classifier free guidance.
266
+ do_classifier_free_guidance = guidance_scale > 1.0
267
+
268
+ # Encode input prompt
269
+ text_embeddings = self._encode_prompt(
270
+ prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
271
+ )
272
+
273
+ # Prepare timesteps
274
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
275
+ timesteps = self.scheduler.timesteps
276
+
277
+ # print(f" Latent shape = {latents.shape}")
278
+
279
+ # Prepare latent variables
280
+ num_channels_latents = self.unet.in_channels
281
+
282
+ xT = self.prepare_latents(
283
+ batch_size * num_videos_per_prompt,
284
+ num_channels_latents,
285
+ video_length,
286
+ height,
287
+ width,
288
+ text_embeddings.dtype,
289
+ device,
290
+ generator,
291
+ xT,
292
+ )
293
+ dtype = xT.dtype
294
+
295
+ # when motion field is not used, augment with random latent codes
296
+ if use_motion_field:
297
+ xT = xT[:, :, :1]
298
+ else:
299
+ if xT.shape[2] < video_length:
300
+ xT_missing = self.prepare_latents(
301
+ batch_size * num_videos_per_prompt,
302
+ num_channels_latents,
303
+ video_length-xT.shape[2],
304
+ height,
305
+ width,
306
+ text_embeddings.dtype,
307
+ device,
308
+ generator,
309
+ None,
310
+ )
311
+ xT = torch.cat([xT, xT_missing], dim=2)
312
+
313
+
314
+ xInit = xT.clone()
315
+ t0 = kwargs["t0"]
316
+ t1 = kwargs["t1"]
317
+ x_t1_1 = None
318
+
319
+
320
+ # Prepare extra step kwargs.
321
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
322
+ # Denoising loop
323
+ num_warmup_steps = len(timesteps) - \
324
+ num_inference_steps * self.scheduler.order
325
+
326
+
327
+
328
+ ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=1000, t0=t0, t1=t1, do_classifier_free_guidance=do_classifier_free_guidance,
329
+ null_embs=null_embs, text_embeddings=text_embeddings, latents_local=xT, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
330
+
331
+ x0 = ddim_res["x0"].detach()
332
+
333
+ if "x_t0_1" in ddim_res:
334
+ x_t0_1 = ddim_res["x_t0_1"].detach()
335
+ if "x_t1_1" in ddim_res:
336
+ x_t1_1 = ddim_res["x_t1_1"].detach()
337
+ del ddim_res
338
+ del xT
339
+
340
+ if use_motion_field:
341
+ del x0
342
+ shape = (batch_size, num_channels_latents, 1, height //
343
+ self.vae_scale_factor, width // self.vae_scale_factor)
344
+
345
+
346
+ x_t0_k = x_t0_1[:, :, :1, :, :].repeat(1, 1, video_length-1, 1, 1)
347
+
348
+
349
+ reference_flow = torch.zeros(
350
+ (video_length-1, 2, 512, 512), device=x_t0_1.device, dtype=x_t0_1.dtype)
351
+ for fr_idx in range(video_length-1):
352
+ reference_flow[fr_idx, :, :, :] = motion_field_strength*(fr_idx+1)
353
+
354
+ for idx, latent in enumerate(x_t0_k):
355
+ x_t0_k[idx] = self.warp_latents_independently(
356
+ latent[None], reference_flow)
357
+
358
+ # assuming t0=t1=1000, if t0 = 1000
359
+ if t1 > t0:
360
+ x_t1_k = self.DDPM_forward(
361
+ x0=x_t0_k, t0=t0, tMax=t1, device=device, shape=shape, text_embeddings=text_embeddings, generator=generator)
362
+ else:
363
+ x_t1_k = x_t0_k
364
+
365
+ if x_t1_1 is None:
366
+ raise Exception
367
+
368
+ x_t1 = torch.cat([x_t1_1, x_t1_k], dim=2).clone().detach()
369
+
370
+ ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=t1, t0=-1, t1=-1, do_classifier_free_guidance=do_classifier_free_guidance,
371
+ null_embs=null_embs, text_embeddings=text_embeddings, latents_local=x_t1, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
372
+
373
+ x0 = ddim_res["x0"].detach()
374
+ del ddim_res
375
+ else:
376
+ x_t1 = x_t1_1.clone()
377
+ x_t1_1 = x_t1_1[:,:,:1,:,:].clone()
378
+ x_t1_k = x_t1_1[:,:,1:,:,:].clone()
379
+ x_t0_k = x_t0_1[:, :, 1:, :, :].clone()
380
+ x_t0_1 = x_t0_1[:,:,:1,:,:].clone()
381
+
382
+
383
+ move_object = use_foreground_motion_field
384
+ if move_object:
385
+ h, w = x0.shape[3], x0.shape[4]
386
+ # Move object
387
+ # reference_flow = torch.zeros(
388
+ # (video_length-1, 2, 512, 512), device=x_t0_1.device, dtype=x_t0_1.dtype)
389
+ reference_flow_obj = torch.zeros(
390
+ (batch_size, video_length, 2, 512, 512), device=x_t0_1.device, dtype=x_t0_1.dtype)
391
+
392
+ for batch_idx, x0_b in enumerate(x0):
393
+ tmp = x0_b[None]
394
+ z0_b = []
395
+ for fr_split in range(tmp.shape[2]):
396
+ z0_b.append(self.decode_latents(
397
+ tmp[:, :, fr_split, None]).detach())
398
+ z0_b = torch.cat(z0_b, dim=2)
399
+ z0_b = rearrange(z0_b[0], "c f h w -> f h w c")
400
+ shift = (-5 - 5) * torch.rand(2,
401
+ device=x0.device, dtype=x0.dtype) + 5
402
+ for frame_idx, z0_f in enumerate(z0_b):
403
+ if frame_idx > 0:
404
+
405
+ z0_f = torch.round(
406
+ z0_f * 255).cpu().numpy().astype(np.uint8)
407
+
408
+ # apply SOD detection to obtain mask of foreground object
409
+ m_f = torch.tensor(self.sod_model.process_data(
410
+ z0_f), device=x0.device).to(x0.dtype)
411
+ kernel = torch.ones(
412
+ 5, 5, device=x0.device, dtype=x0.dtype)
413
+ mask = dilation(
414
+ m_f[None, None].to(x0.device), kernel)[0]
415
+ for coord_idx in range(2):
416
+ reference_flow_obj[batch_idx, frame_idx,
417
+ coord_idx, :, :] = (1+frame_idx) * shift[coord_idx] * mask
418
+
419
+
420
+
421
+ for idx, x_t0_k_b in enumerate(x_t0_k):
422
+ x_t0_k[idx] = self.warp_latents_independently(
423
+ x_t0_k_b[None], reference_flow_obj[idx, 1:])
424
+
425
+ x_t1_k = self.DDPM_forward(
426
+ x0=x_t0_k, t0=t0, tMax=t1, device=device, shape=shape, text_embeddings=text_embeddings, generator=generator)
427
+
428
+ if x_t1_1 is None:
429
+ raise Exception
430
+ x_t1 = torch.cat([x_t1_1, x_t1_k], dim=2)
431
+
432
+ # del latent
433
+ ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=t1, t0=-1, t1=-1, do_classifier_free_guidance=do_classifier_free_guidance,
434
+ null_embs=null_embs, text_embeddings=text_embeddings, latents_local=x_t1, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
435
+ x0 = ddim_res["x0"].detach()
436
+ del ddim_res
437
+
438
+
439
+ # smooth background
440
+ if smooth_bg:
441
+ h, w = x0.shape[3], x0.shape[4]
442
+ M_FG = torch.zeros((batch_size, video_length, h, w),
443
+ device=x0.device).to(x0.dtype)
444
+ for batch_idx, x0_b in enumerate(x0):
445
+ z0_b = self.decode_latents(x0_b[None]).detach()
446
+ z0_b = rearrange(z0_b[0], "c f h w -> f h w c")
447
+ for frame_idx, z0_f in enumerate(z0_b):
448
+ z0_f = torch.round(
449
+ z0_f * 255).cpu().numpy().astype(np.uint8)
450
+ # apply SOD detection
451
+ m_f = torch.tensor(self.sod_model.process_data(
452
+ z0_f), device=x0.device).to(x0.dtype)
453
+ mask = T.Resize(
454
+ size=(h, w), interpolation=T.InterpolationMode.NEAREST)(m_f[None])
455
+ kernel = torch.ones(5, 5, device=x0.device, dtype=x0.dtype)
456
+ mask = dilation(mask[None].to(x0.device), kernel)[0]
457
+ M_FG[batch_idx, frame_idx, :, :] = mask
458
+
459
+
460
+ x_t1_1_fg_masked = x_t1_1 * \
461
+ (1 - repeat(M_FG[:, 0, :, :],
462
+ "b w h -> b c 1 w h", c=x_t1_1.shape[1]))
463
+
464
+
465
+ x_t1_1_fg_masked_moved = []
466
+ for batch_idx, x_t1_1_fg_masked_b in enumerate(x_t1_1_fg_masked):
467
+ x_t1_fg_masked_b = x_t1_1_fg_masked_b.clone()
468
+
469
+ x_t1_fg_masked_b = x_t1_fg_masked_b.repeat(
470
+ 1, video_length-1, 1, 1)
471
+ if use_motion_field:
472
+ x_t1_fg_masked_b = x_t1_fg_masked_b[None]
473
+ x_t1_fg_masked_b = self.warp_latents_independently(
474
+ x_t1_fg_masked_b, reference_flow)
475
+ else:
476
+ x_t1_fg_masked_b = x_t1_fg_masked_b[None]
477
+ if move_object:
478
+ x_t1_fg_masked_b = self.warp_latents_independently(
479
+ x_t1_fg_masked_b, reference_flow_obj[batch_idx, 1:])
480
+
481
+ x_t1_fg_masked_b = torch.cat(
482
+ [x_t1_1_fg_masked_b[None], x_t1_fg_masked_b], dim=2)
483
+ x_t1_1_fg_masked_moved.append(x_t1_fg_masked_b)
484
+
485
+ x_t1_1_fg_masked_moved = torch.cat(x_t1_1_fg_masked_moved, dim=0)
486
+
487
+ M_FG_1 = M_FG[:, :1, :, :]
488
+
489
+ M_FG_warped = []
490
+ for batch_idx, m_fg_1_b in enumerate(M_FG_1):
491
+ m_fg_1_b = m_fg_1_b[None, None]
492
+ m_fg_b = m_fg_1_b.repeat(1, 1, video_length-1, 1, 1)
493
+ if use_motion_field:
494
+ m_fg_b = self.warp_latents_independently(
495
+ m_fg_b.clone(), reference_flow)
496
+ if move_object:
497
+ m_fg_b = self.warp_latents_independently(
498
+ m_fg_b, reference_flow_obj[batch_idx, 1:])
499
+ M_FG_warped.append(
500
+ torch.cat([m_fg_1_b[:1, 0], m_fg_b[:1, 0]], dim=1))
501
+
502
+ M_FG_warped = torch.cat(M_FG_warped, dim=0)
503
+
504
+ channels = x0.shape[1]
505
+
506
+ M_BG = (1-M_FG) * (1 - M_FG_warped)
507
+ M_BG = repeat(M_BG, "b f h w -> b c f h w", c=channels)
508
+ a_convex = smooth_bg_strength
509
+
510
+ x_t1_blending = (1-M_BG) * x_t1 + M_BG * (a_convex *
511
+ x_t1 + (1-a_convex) * x_t1_1_fg_masked_moved)
512
+
513
+ '''
514
+ x_t1_blending = self.DDPM_forward(
515
+ x0=x_t1_blending, t0=t1, tMax=961, device=device, shape=shape, text_embeddings=text_embeddings, generator=generator)
516
+ t1 = 961
517
+ '''
518
+ latents = x_t1_blending
519
+
520
+ ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=t1, t0=-1, t1=-1, do_classifier_free_guidance=do_classifier_free_guidance,
521
+ null_embs=null_embs, text_embeddings=text_embeddings, latents_local=latents, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
522
+ x0 = ddim_res["x0"].detach()
523
+ del ddim_res
524
+
525
+
526
+ # Post-processing
527
+ video_list = []
528
+ for latent in x0:
529
+ tmp = latent[None]
530
+ print("Frame spit shape", tmp.shape)
531
+ frames = []
532
+ for fr_split in range(tmp.shape[2]):
533
+ print("frame decoding")
534
+ frames.append(self.decode_latents(
535
+ tmp[:, :, fr_split, None]).detach())
536
+
537
+ video_list.append(torch.cat(frames, dim=2).cpu().float().numpy())
538
+
539
+ # Convert to tensor
540
+ videos = []
541
+ if output_type == "tensor":
542
+ for video in video_list:
543
+ videos.append(torch.from_numpy(video))
544
+ if output_type == 'numpy':
545
+ for video in video_list:
546
+ videos.append(rearrange(video, 'b c f h w -> (b f) h w c'))
547
+ if not return_dict:
548
+ return video
549
+
550
+ return TextToVideoPipelineOutput(videos=videos, code=torch.split(xInit.detach().cpu(), 1, dim=0))