Issues should be raised directly in the repository. For professional support and recommendations please welcomedenk@gmail.com.
100 | -------------------------------------------------------------------------------- /cogvideo_controlnet.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional, Tuple, Union 2 | 3 | import torch 4 | from torch import nn 5 | from einops import rearrange 6 | import torch.nn.functional as F 7 | from diffusers.models.transformers.cogvideox_transformer_3d import Transformer2DModelOutput, CogVideoXBlock 8 | from diffusers.utils import is_torch_version 9 | from diffusers.loaders import PeftAdapterMixin 10 | from diffusers.utils.torch_utils import maybe_allow_in_graph 11 | from diffusers.models.embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps, get_3d_sincos_pos_embed 12 | from diffusers.models.modeling_utils import ModelMixin 13 | from diffusers.models.attention import Attention, FeedForward 14 | from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor2_0 15 | from diffusers.models.normalization import AdaLayerNorm, CogVideoXLayerNormZero, AdaLayerNormZeroSingle 16 | from diffusers.configuration_utils import ConfigMixin, register_to_config 17 | 18 | 19 | class CogVideoXControlnet(ModelMixin, ConfigMixin, PeftAdapterMixin): 20 | _supports_gradient_checkpointing = True 21 | 22 | @register_to_config 23 | def __init__( 24 | self, 25 | num_attention_heads: int = 30, 26 | attention_head_dim: int = 64, 27 | vae_channels: int = 16, 28 | in_channels: int = 3, 29 | downscale_coef: int = 8, 30 | flip_sin_to_cos: bool = True, 31 | freq_shift: int = 0, 32 | time_embed_dim: int = 512, 33 | num_layers: int = 8, 34 | dropout: float = 0.0, 35 | attention_bias: bool = True, 36 | sample_width: int = 90, 37 | sample_height: int = 60, 38 | sample_frames: int = 49, 39 | patch_size: int = 2, 40 | temporal_compression_ratio: int = 4, 41 | max_text_seq_length: int = 226, 42 | activation_fn: str = "gelu-approximate", 43 | timestep_activation_fn: str = "silu", 44 | norm_elementwise_affine: bool = True, 45 | norm_eps: float = 1e-5, 46 | spatial_interpolation_scale: float = 1.875, 47 | temporal_interpolation_scale: float = 1.0, 48 | use_rotary_positional_embeddings: bool = False, 49 | use_learned_positional_embeddings: bool = False, 50 | out_proj_dim = None, 51 | ): 52 | super().__init__() 53 | inner_dim = num_attention_heads * attention_head_dim 54 | 55 | if not use_rotary_positional_embeddings and use_learned_positional_embeddings: 56 | raise ValueError( 57 | "There are no CogVideoX checkpoints available with disable rotary embeddings and learned positional " 58 | "embeddings. If you're using a custom model and/or believe this should be supported, please open an " 59 | "issue at https://github.com/huggingface/diffusers/issues." 60 | ) 61 | 62 | start_channels = in_channels * (downscale_coef ** 2) 63 | input_channels = [start_channels, start_channels // 2, start_channels // 4] 64 | self.unshuffle = nn.PixelUnshuffle(downscale_coef) 65 | 66 | self.controlnet_encode_first = nn.Sequential( 67 | nn.Conv2d(input_channels[0], input_channels[1], kernel_size=1, stride=1, padding=0), 68 | nn.GroupNorm(2, input_channels[1]), 69 | nn.ReLU(), 70 | ) 71 | 72 | self.controlnet_encode_second = nn.Sequential( 73 | nn.Conv2d(input_channels[1], input_channels[2], kernel_size=1, stride=1, padding=0), 74 | nn.GroupNorm(2, input_channels[2]), 75 | nn.ReLU(), 76 | ) 77 | 78 | # 1. Patch embedding 79 | self.patch_embed = CogVideoXPatchEmbed( 80 | patch_size=patch_size, 81 | in_channels=vae_channels + input_channels[2], 82 | embed_dim=inner_dim, 83 | bias=True, 84 | sample_width=sample_width, 85 | sample_height=sample_height, 86 | sample_frames=sample_frames, 87 | temporal_compression_ratio=temporal_compression_ratio, 88 | spatial_interpolation_scale=spatial_interpolation_scale, 89 | temporal_interpolation_scale=temporal_interpolation_scale, 90 | use_positional_embeddings=not use_rotary_positional_embeddings, 91 | use_learned_positional_embeddings=use_learned_positional_embeddings, 92 | ) 93 | 94 | self.embedding_dropout = nn.Dropout(dropout) 95 | 96 | # 2. Time embeddings 97 | self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift) 98 | self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn) 99 | 100 | # 3. Define spatio-temporal transformers blocks 101 | self.transformer_blocks = nn.ModuleList( 102 | [ 103 | CogVideoXBlock( 104 | dim=inner_dim, 105 | num_attention_heads=num_attention_heads, 106 | attention_head_dim=attention_head_dim, 107 | time_embed_dim=time_embed_dim, 108 | dropout=dropout, 109 | activation_fn=activation_fn, 110 | attention_bias=attention_bias, 111 | norm_elementwise_affine=norm_elementwise_affine, 112 | norm_eps=norm_eps, 113 | ) 114 | for _ in range(num_layers) 115 | ] 116 | ) 117 | 118 | self.out_projectors = None 119 | if out_proj_dim is not None: 120 | self.out_projectors = nn.ModuleList( 121 | [nn.Linear(inner_dim, out_proj_dim) for _ in range(num_layers)] 122 | ) 123 | 124 | self.gradient_checkpointing = False 125 | 126 | def _set_gradient_checkpointing(self, module, value=False): 127 | self.gradient_checkpointing = value 128 | 129 | def compress_time(self, x, num_frames): 130 | x = rearrange(x, '(b f) c h w -> b f c h w', f=num_frames) 131 | batch_size, frames, channels, height, width = x.shape 132 | x = rearrange(x, 'b f c h w -> (b h w) c f') 133 | 134 | if x.shape[-1] % 2 == 1: 135 | x_first, x_rest = x[..., 0], x[..., 1:] 136 | if x_rest.shape[-1] > 0: 137 | x_rest = F.avg_pool1d(x_rest, kernel_size=2, stride=2) 138 | 139 | x = torch.cat([x_first[..., None], x_rest], dim=-1) 140 | else: 141 | x = F.avg_pool1d(x, kernel_size=2, stride=2) 142 | x = rearrange(x, '(b h w) c f -> (b f) c h w', b=batch_size, h=height, w=width) 143 | return x 144 | 145 | def forward( 146 | self, 147 | hidden_states: torch.Tensor, 148 | encoder_hidden_states: torch.Tensor, 149 | controlnet_states: torch.Tensor, 150 | timestep: Union[int, float, torch.LongTensor], 151 | image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, 152 | timestep_cond: Optional[torch.Tensor] = None, 153 | return_dict: bool = True, 154 | ): 155 | batch_size, num_frames, channels, height, width = controlnet_states.shape 156 | # 0. Controlnet encoder 157 | controlnet_states = rearrange(controlnet_states, 'b f c h w -> (b f) c h w') 158 | controlnet_states = self.unshuffle(controlnet_states) 159 | controlnet_states = self.controlnet_encode_first(controlnet_states) 160 | controlnet_states = self.compress_time(controlnet_states, num_frames=num_frames) 161 | num_frames = controlnet_states.shape[0] // batch_size 162 | 163 | controlnet_states = self.controlnet_encode_second(controlnet_states) 164 | controlnet_states = self.compress_time(controlnet_states, num_frames=num_frames) 165 | controlnet_states = rearrange(controlnet_states, '(b f) c h w -> b f c h w', b=batch_size) 166 | 167 | hidden_states = torch.cat([hidden_states, controlnet_states], dim=2) 168 | # controlnet_states = self.controlnext_encoder(controlnet_states, timestep=timestep) 169 | # 1. Time embedding 170 | timesteps = timestep 171 | t_emb = self.time_proj(timesteps) 172 | 173 | # timesteps does not contain any weights and will always return f32 tensors 174 | # but time_embedding might actually be running in fp16. so we need to cast here. 175 | # there might be better ways to encapsulate this. 176 | t_emb = t_emb.to(dtype=hidden_states.dtype) 177 | emb = self.time_embedding(t_emb, timestep_cond) 178 | 179 | hidden_states = self.patch_embed(encoder_hidden_states, hidden_states) 180 | hidden_states = self.embedding_dropout(hidden_states) 181 | 182 | 183 | text_seq_length = encoder_hidden_states.shape[1] 184 | encoder_hidden_states = hidden_states[:, :text_seq_length] 185 | hidden_states = hidden_states[:, text_seq_length:] 186 | 187 | 188 | controlnet_hidden_states = () 189 | # 3. Transformer blocks 190 | for i, block in enumerate(self.transformer_blocks): 191 | if self.training and self.gradient_checkpointing: 192 | 193 | def create_custom_forward(module): 194 | def custom_forward(*inputs): 195 | return module(*inputs) 196 | 197 | return custom_forward 198 | 199 | ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} 200 | hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( 201 | create_custom_forward(block), 202 | hidden_states, 203 | encoder_hidden_states, 204 | emb, 205 | image_rotary_emb, 206 | **ckpt_kwargs, 207 | ) 208 | else: 209 | hidden_states, encoder_hidden_states = block( 210 | hidden_states=hidden_states, 211 | encoder_hidden_states=encoder_hidden_states, 212 | temb=emb, 213 | image_rotary_emb=image_rotary_emb, 214 | ) 215 | 216 | if self.out_projectors is not None: 217 | controlnet_hidden_states += (self.out_projectors[i](hidden_states),) 218 | else: 219 | controlnet_hidden_states += (hidden_states,) 220 | 221 | if not return_dict: 222 | return (controlnet_hidden_states,) 223 | return Transformer2DModelOutput(sample=controlnet_hidden_states) -------------------------------------------------------------------------------- /cogvideo_transformer.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional, Tuple, Union 2 | 3 | import torch 4 | import numpy as np 5 | from diffusers.utils import is_torch_version 6 | from diffusers.models.transformers.cogvideox_transformer_3d import CogVideoXTransformer3DModel, Transformer2DModelOutput 7 | 8 | 9 | class CustomCogVideoXTransformer3DModel(CogVideoXTransformer3DModel): 10 | def forward( 11 | self, 12 | hidden_states: torch.Tensor, 13 | encoder_hidden_states: torch.Tensor, 14 | timestep: Union[int, float, torch.LongTensor], 15 | start_frame = None, 16 | timestep_cond: Optional[torch.Tensor] = None, 17 | image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, 18 | controlnet_states: torch.Tensor = None, 19 | controlnet_weights: Optional[Union[float, int, list, np.ndarray, torch.FloatTensor]] = 1.0, 20 | return_dict: bool = True, 21 | ): 22 | batch_size, num_frames, channels, height, width = hidden_states.shape 23 | 24 | if start_frame is not None: 25 | hidden_states = torch.cat([start_frame, hidden_states], dim=2) 26 | # 1. Time embedding 27 | timesteps = timestep 28 | t_emb = self.time_proj(timesteps) 29 | 30 | # timesteps does not contain any weights and will always return f32 tensors 31 | # but time_embedding might actually be running in fp16. so we need to cast here. 32 | # there might be better ways to encapsulate this. 33 | t_emb = t_emb.to(dtype=hidden_states.dtype) 34 | emb = self.time_embedding(t_emb, timestep_cond) 35 | 36 | # 2. Patch embedding 37 | hidden_states = self.patch_embed(encoder_hidden_states, hidden_states) 38 | hidden_states = self.embedding_dropout(hidden_states) 39 | 40 | text_seq_length = encoder_hidden_states.shape[1] 41 | encoder_hidden_states = hidden_states[:, :text_seq_length] 42 | hidden_states = hidden_states[:, text_seq_length:] 43 | 44 | # 3. Transformer blocks 45 | for i, block in enumerate(self.transformer_blocks): 46 | if self.training and self.gradient_checkpointing: 47 | 48 | def create_custom_forward(module): 49 | def custom_forward(*inputs): 50 | return module(*inputs) 51 | 52 | return custom_forward 53 | 54 | ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} 55 | hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( 56 | create_custom_forward(block), 57 | hidden_states, 58 | encoder_hidden_states, 59 | emb, 60 | image_rotary_emb, 61 | **ckpt_kwargs, 62 | ) 63 | else: 64 | hidden_states, encoder_hidden_states = block( 65 | hidden_states=hidden_states, 66 | encoder_hidden_states=encoder_hidden_states, 67 | temb=emb, 68 | image_rotary_emb=image_rotary_emb, 69 | ) 70 | 71 | if (controlnet_states is not None) and (i < len(controlnet_states)): 72 | controlnet_states_block = controlnet_states[i] 73 | controlnet_block_weight = 1.0 74 | if isinstance(controlnet_weights, (list, np.ndarray)) or torch.is_tensor(controlnet_weights): 75 | controlnet_block_weight = controlnet_weights[i] 76 | elif isinstance(controlnet_weights, (float, int)): 77 | controlnet_block_weight = controlnet_weights 78 | 79 | hidden_states = hidden_states + controlnet_states_block * controlnet_block_weight 80 | 81 | if not self.config.use_rotary_positional_embeddings: 82 | # CogVideoX-2B 83 | hidden_states = self.norm_final(hidden_states) 84 | else: 85 | # CogVideoX-5B 86 | hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) 87 | hidden_states = self.norm_final(hidden_states) 88 | hidden_states = hidden_states[:, text_seq_length:] 89 | 90 | # 4. Final block 91 | hidden_states = self.norm_out(hidden_states, temb=emb) 92 | hidden_states = self.proj_out(hidden_states) 93 | 94 | # 5. Unpatchify 95 | p = self.config.patch_size 96 | p_t = self.config.patch_size_t 97 | 98 | if p_t is None: 99 | output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p) 100 | output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) 101 | else: 102 | output = hidden_states.reshape( 103 | batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p 104 | ) 105 | output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2) 106 | 107 | if not return_dict: 108 | return (output,) 109 | return Transformer2DModelOutput(sample=output) -------------------------------------------------------------------------------- /controlnet_img2vid_pipeline.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import math 3 | from typing import Any, Callable, Dict, List, Optional, Tuple, Union 4 | 5 | import torch 6 | import numpy as np 7 | import PIL 8 | from PIL import Image 9 | from torchvision import transforms 10 | from einops import rearrange, repeat 11 | from transformers import T5EncoderModel, T5Tokenizer 12 | from diffusers.video_processor import VideoProcessor 13 | from diffusers.utils.torch_utils import randn_tensor 14 | from diffusers.models.embeddings import get_3d_rotary_pos_embed 15 | from diffusers.pipelines.pipeline_utils import DiffusionPipeline 16 | from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel 17 | from diffusers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler, CogVideoXImageToVideoPipeline 18 | from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback 19 | from diffusers.pipelines.cogvideo.pipeline_cogvideox import CogVideoXPipelineOutput, CogVideoXLoraLoaderMixin 20 | 21 | from cogvideo_controlnet import CogVideoXControlnet 22 | 23 | 24 | def resize_for_crop(image, crop_h, crop_w): 25 | img_h, img_w = image.shape[-2:] 26 | if img_h >= crop_h and img_w >= crop_w: 27 | coef = max(crop_h / img_h, crop_w / img_w) 28 | elif img_h <= crop_h and img_w <= crop_w: 29 | coef = max(crop_h / img_h, crop_w / img_w) 30 | else: 31 | coef = crop_h / img_h if crop_h > img_h else crop_w / img_w 32 | out_h, out_w = int(img_h * coef), int(img_w * coef) 33 | resized_image = transforms.functional.resize(image, (out_h, out_w), antialias=True) 34 | return resized_image 35 | 36 | 37 | def prepare_frames(input_images, video_size, do_resize=True, do_crop=True): 38 | input_images = np.stack([np.array(x) for x in input_images]) 39 | images_tensor = torch.from_numpy(input_images).permute(0, 3, 1, 2) / 127.5 - 1 40 | if do_resize: 41 | images_tensor = [resize_for_crop(x, crop_h=video_size[0], crop_w=video_size[1]) for x in images_tensor] 42 | if do_crop: 43 | images_tensor = [transforms.functional.center_crop(x, video_size) for x in images_tensor] 44 | if isinstance(images_tensor, list): 45 | images_tensor = torch.stack(images_tensor) 46 | return images_tensor.unsqueeze(0) 47 | 48 | 49 | def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): 50 | tw = tgt_width 51 | th = tgt_height 52 | h, w = src 53 | r = h / w 54 | if r > (th / tw): 55 | resize_height = th 56 | resize_width = int(round(th / h * w)) 57 | else: 58 | resize_width = tw 59 | resize_height = int(round(tw / w * h)) 60 | 61 | crop_top = int(round((th - resize_height) / 2.0)) 62 | crop_left = int(round((tw - resize_width) / 2.0)) 63 | 64 | return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) 65 | 66 | 67 | # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps 68 | def retrieve_timesteps( 69 | scheduler, 70 | num_inference_steps: Optional[int] = None, 71 | device: Optional[Union[str, torch.device]] = None, 72 | timesteps: Optional[List[int]] = None, 73 | sigmas: Optional[List[float]] = None, 74 | **kwargs, 75 | ): 76 | """ 77 | Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles 78 | custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. 79 | 80 | Args: 81 | scheduler (`SchedulerMixin`): 82 | The scheduler to get timesteps from. 83 | num_inference_steps (`int`): 84 | The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` 85 | must be `None`. 86 | device (`str` or `torch.device`, *optional*): 87 | The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. 88 | timesteps (`List[int]`, *optional*): 89 | Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, 90 | `num_inference_steps` and `sigmas` must be `None`. 91 | sigmas (`List[float]`, *optional*): 92 | Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, 93 | `num_inference_steps` and `timesteps` must be `None`. 94 | 95 | Returns: 96 | `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the 97 | second element is the number of inference steps. 98 | """ 99 | if timesteps is not None and sigmas is not None: 100 | raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") 101 | if timesteps is not None: 102 | accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) 103 | if not accepts_timesteps: 104 | raise ValueError( 105 | f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" 106 | f" timestep schedules. Please check whether you are using the correct scheduler." 107 | ) 108 | scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) 109 | timesteps = scheduler.timesteps 110 | num_inference_steps = len(timesteps) 111 | elif sigmas is not None: 112 | accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) 113 | if not accept_sigmas: 114 | raise ValueError( 115 | f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" 116 | f" sigmas schedules. Please check whether you are using the correct scheduler." 117 | ) 118 | scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) 119 | timesteps = scheduler.timesteps 120 | num_inference_steps = len(timesteps) 121 | else: 122 | scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) 123 | timesteps = scheduler.timesteps 124 | return timesteps, num_inference_steps 125 | 126 | 127 | def retrieve_latents( 128 | encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" 129 | ): 130 | if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": 131 | return encoder_output.latent_dist.sample(generator) 132 | elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": 133 | return encoder_output.latent_dist.mode() 134 | elif hasattr(encoder_output, "latents"): 135 | return encoder_output.latents 136 | else: 137 | raise AttributeError("Could not access latents of provided encoder_output") 138 | 139 | 140 | class CogVideoXImageToVideoControlnetPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): 141 | r""" 142 | Pipeline for image-to-video generation using CogVideoX. 143 | 144 | This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the 145 | library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) 146 | 147 | Args: 148 | vae ([`AutoencoderKL`]): 149 | Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. 150 | text_encoder ([`T5EncoderModel`]): 151 | Frozen text-encoder. CogVideoX uses 152 | [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the 153 | [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. 154 | tokenizer (`T5Tokenizer`): 155 | Tokenizer of class 156 | [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). 157 | transformer ([`CogVideoXTransformer3DModel`]): 158 | A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents. 159 | scheduler ([`SchedulerMixin`]): 160 | A scheduler to be used in combination with `transformer` to denoise the encoded video latents. 161 | """ 162 | 163 | _optional_components = [] 164 | model_cpu_offload_seq = "text_encoder->transformer->vae" 165 | 166 | _callback_tensor_inputs = [ 167 | "latents", 168 | "prompt_embeds", 169 | "negative_prompt_embeds", 170 | ] 171 | 172 | def __init__( 173 | self, 174 | tokenizer: T5Tokenizer, 175 | text_encoder: T5EncoderModel, 176 | vae: AutoencoderKLCogVideoX, 177 | transformer: CogVideoXTransformer3DModel, 178 | controlnet: CogVideoXControlnet, 179 | scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], 180 | ): 181 | super().__init__() 182 | 183 | self.register_modules( 184 | tokenizer=tokenizer, 185 | text_encoder=text_encoder, 186 | vae=vae, 187 | transformer=transformer, 188 | controlnet=controlnet, 189 | scheduler=scheduler, 190 | ) 191 | self.vae_scale_factor_spatial = ( 192 | 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 193 | ) 194 | self.vae_scale_factor_temporal = ( 195 | self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4 196 | ) 197 | self.vae_scaling_factor_image = ( 198 | self.vae.config.scaling_factor if hasattr(self, "vae") and self.vae is not None else 0.7 199 | ) 200 | 201 | self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) 202 | 203 | # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds 204 | def _get_t5_prompt_embeds( 205 | self, 206 | prompt: Union[str, List[str]] = None, 207 | num_videos_per_prompt: int = 1, 208 | max_sequence_length: int = 226, 209 | device: Optional[torch.device] = None, 210 | dtype: Optional[torch.dtype] = None, 211 | ): 212 | device = device or self._execution_device 213 | dtype = dtype or self.text_encoder.dtype 214 | 215 | prompt = [prompt] if isinstance(prompt, str) else prompt 216 | batch_size = len(prompt) 217 | 218 | text_inputs = self.tokenizer( 219 | prompt, 220 | padding="max_length", 221 | max_length=max_sequence_length, 222 | truncation=True, 223 | add_special_tokens=True, 224 | return_tensors="pt", 225 | ) 226 | text_input_ids = text_inputs.input_ids 227 | untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids 228 | 229 | if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): 230 | removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) 231 | logger.warning( 232 | "The following part of your input was truncated because `max_sequence_length` is set to " 233 | f" {max_sequence_length} tokens: {removed_text}" 234 | ) 235 | 236 | prompt_embeds = self.text_encoder(text_input_ids.to(device))[0] 237 | prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) 238 | 239 | # duplicate text embeddings for each generation per prompt, using mps friendly method 240 | _, seq_len, _ = prompt_embeds.shape 241 | prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) 242 | prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) 243 | 244 | return prompt_embeds 245 | 246 | # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt 247 | def encode_prompt( 248 | self, 249 | prompt: Union[str, List[str]], 250 | negative_prompt: Optional[Union[str, List[str]]] = None, 251 | do_classifier_free_guidance: bool = True, 252 | num_videos_per_prompt: int = 1, 253 | prompt_embeds: Optional[torch.Tensor] = None, 254 | negative_prompt_embeds: Optional[torch.Tensor] = None, 255 | max_sequence_length: int = 226, 256 | device: Optional[torch.device] = None, 257 | dtype: Optional[torch.dtype] = None, 258 | ): 259 | r""" 260 | Encodes the prompt into text encoder hidden states. 261 | 262 | Args: 263 | prompt (`str` or `List[str]`, *optional*): 264 | prompt to be encoded 265 | negative_prompt (`str` or `List[str]`, *optional*): 266 | The prompt or prompts not to guide the image generation. If not defined, one has to pass 267 | `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is 268 | less than `1`). 269 | do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): 270 | Whether to use classifier free guidance or not. 271 | num_videos_per_prompt (`int`, *optional*, defaults to 1): 272 | Number of videos that should be generated per prompt. torch device to place the resulting embeddings on 273 | prompt_embeds (`torch.Tensor`, *optional*): 274 | Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not 275 | provided, text embeddings will be generated from `prompt` input argument. 276 | negative_prompt_embeds (`torch.Tensor`, *optional*): 277 | Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt 278 | weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input 279 | argument. 280 | device: (`torch.device`, *optional*): 281 | torch device 282 | dtype: (`torch.dtype`, *optional*): 283 | torch dtype 284 | """ 285 | device = device or self._execution_device 286 | 287 | prompt = [prompt] if isinstance(prompt, str) else prompt 288 | if prompt is not None: 289 | batch_size = len(prompt) 290 | else: 291 | batch_size = prompt_embeds.shape[0] 292 | 293 | if prompt_embeds is None: 294 | prompt_embeds = self._get_t5_prompt_embeds( 295 | prompt=prompt, 296 | num_videos_per_prompt=num_videos_per_prompt, 297 | max_sequence_length=max_sequence_length, 298 | device=device, 299 | dtype=dtype, 300 | ) 301 | 302 | if do_classifier_free_guidance and negative_prompt_embeds is None: 303 | negative_prompt = negative_prompt or "" 304 | negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt 305 | 306 | if prompt is not None and type(prompt) is not type(negative_prompt): 307 | raise TypeError( 308 | f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" 309 | f" {type(prompt)}." 310 | ) 311 | elif batch_size != len(negative_prompt): 312 | raise ValueError( 313 | f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" 314 | f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" 315 | " the batch size of `prompt`." 316 | ) 317 | 318 | negative_prompt_embeds = self._get_t5_prompt_embeds( 319 | prompt=negative_prompt, 320 | num_videos_per_prompt=num_videos_per_prompt, 321 | max_sequence_length=max_sequence_length, 322 | device=device, 323 | dtype=dtype, 324 | ) 325 | 326 | return prompt_embeds, negative_prompt_embeds 327 | 328 | def prepare_latents( 329 | self, 330 | image: torch.Tensor, 331 | batch_size: int = 1, 332 | num_channels_latents: int = 16, 333 | num_frames: int = 13, 334 | height: int = 60, 335 | width: int = 90, 336 | dtype: Optional[torch.dtype] = None, 337 | device: Optional[torch.device] = None, 338 | generator: Optional[torch.Generator] = None, 339 | latents: Optional[torch.Tensor] = None, 340 | ): 341 | if isinstance(generator, list) and len(generator) != batch_size: 342 | raise ValueError( 343 | f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" 344 | f" size of {batch_size}. Make sure the batch size matches the length of the generators." 345 | ) 346 | 347 | num_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 348 | shape = ( 349 | batch_size, 350 | num_frames, 351 | num_channels_latents, 352 | height // self.vae_scale_factor_spatial, 353 | width // self.vae_scale_factor_spatial, 354 | ) 355 | 356 | # For CogVideoX1.5, the latent should add 1 for padding (Not use) 357 | if self.transformer.config.patch_size_t is not None: 358 | shape = shape[:1] + (shape[1] + shape[1] % self.transformer.config.patch_size_t,) + shape[2:] 359 | 360 | image = image.unsqueeze(2) # [B, C, F, H, W] 361 | 362 | if isinstance(generator, list): 363 | image_latents = [ 364 | retrieve_latents(self.vae.encode(image[i].unsqueeze(0)), generator[i]) for i in range(batch_size) 365 | ] 366 | else: 367 | image_latents = [retrieve_latents(self.vae.encode(img.unsqueeze(0)), generator) for img in image] 368 | 369 | image_latents = torch.cat(image_latents, dim=0).to(dtype).permute(0, 2, 1, 3, 4) # [B, F, C, H, W] 370 | 371 | if not self.vae.config.invert_scale_latents: 372 | image_latents = self.vae_scaling_factor_image * image_latents 373 | else: 374 | # This is awkward but required because the CogVideoX team forgot to multiply the 375 | # scaling factor during training :) 376 | image_latents = 1 / self.vae_scaling_factor_image * image_latents 377 | 378 | padding_shape = ( 379 | batch_size, 380 | num_frames - 1, 381 | num_channels_latents, 382 | height // self.vae_scale_factor_spatial, 383 | width // self.vae_scale_factor_spatial, 384 | ) 385 | 386 | latent_padding = torch.zeros(padding_shape, device=device, dtype=dtype) 387 | image_latents = torch.cat([image_latents, latent_padding], dim=1) 388 | 389 | # Select the first frame along the second dimension 390 | if self.transformer.config.patch_size_t is not None: 391 | first_frame = image_latents[:, : image_latents.size(1) % self.transformer.config.patch_size_t, ...] 392 | image_latents = torch.cat([first_frame, image_latents], dim=1) 393 | 394 | if latents is None: 395 | latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) 396 | else: 397 | latents = latents.to(device) 398 | 399 | # scale the initial noise by the standard deviation required by the scheduler 400 | latents = latents * self.scheduler.init_noise_sigma 401 | return latents, image_latents 402 | 403 | # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.decode_latents 404 | def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: 405 | latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] 406 | latents = 1 / self.vae_scaling_factor_image * latents 407 | 408 | frames = self.vae.decode(latents).sample 409 | return frames 410 | 411 | # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps 412 | def get_timesteps(self, num_inference_steps, timesteps, strength, device): 413 | # get the original timestep using init_timestep 414 | init_timestep = min(int(num_inference_steps * strength), num_inference_steps) 415 | 416 | t_start = max(num_inference_steps - init_timestep, 0) 417 | timesteps = timesteps[t_start * self.scheduler.order :] 418 | 419 | return timesteps, num_inference_steps - t_start 420 | 421 | # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs 422 | def prepare_extra_step_kwargs(self, generator, eta): 423 | # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature 424 | # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. 425 | # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 426 | # and should be between [0, 1] 427 | 428 | accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) 429 | extra_step_kwargs = {} 430 | if accepts_eta: 431 | extra_step_kwargs["eta"] = eta 432 | 433 | # check if the scheduler accepts generator 434 | accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) 435 | if accepts_generator: 436 | extra_step_kwargs["generator"] = generator 437 | return extra_step_kwargs 438 | 439 | def check_inputs( 440 | self, 441 | image, 442 | prompt, 443 | height, 444 | width, 445 | negative_prompt, 446 | callback_on_step_end_tensor_inputs, 447 | latents=None, 448 | prompt_embeds=None, 449 | negative_prompt_embeds=None, 450 | ): 451 | if ( 452 | not isinstance(image, torch.Tensor) 453 | and not isinstance(image, PIL.Image.Image) 454 | and not isinstance(image, list) 455 | ): 456 | raise ValueError( 457 | "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" 458 | f" {type(image)}" 459 | ) 460 | 461 | if height % 8 != 0 or width % 8 != 0: 462 | raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") 463 | 464 | if callback_on_step_end_tensor_inputs is not None and not all( 465 | k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs 466 | ): 467 | raise ValueError( 468 | f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" 469 | ) 470 | if prompt is not None and prompt_embeds is not None: 471 | raise ValueError( 472 | f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" 473 | " only forward one of the two." 474 | ) 475 | elif prompt is None and prompt_embeds is None: 476 | raise ValueError( 477 | "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." 478 | ) 479 | elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): 480 | raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") 481 | 482 | if prompt is not None and negative_prompt_embeds is not None: 483 | raise ValueError( 484 | f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" 485 | f" {negative_prompt_embeds}. Please make sure to only forward one of the two." 486 | ) 487 | 488 | if negative_prompt is not None and negative_prompt_embeds is not None: 489 | raise ValueError( 490 | f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" 491 | f" {negative_prompt_embeds}. Please make sure to only forward one of the two." 492 | ) 493 | 494 | if prompt_embeds is not None and negative_prompt_embeds is not None: 495 | if prompt_embeds.shape != negative_prompt_embeds.shape: 496 | raise ValueError( 497 | "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" 498 | f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" 499 | f" {negative_prompt_embeds.shape}." 500 | ) 501 | 502 | # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.fuse_qkv_projections 503 | def fuse_qkv_projections(self) -> None: 504 | r"""Enables fused QKV projections.""" 505 | self.fusing_transformer = True 506 | self.transformer.fuse_qkv_projections() 507 | 508 | # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.unfuse_qkv_projections 509 | def unfuse_qkv_projections(self) -> None: 510 | r"""Disable QKV projection fusion if enabled.""" 511 | if not self.fusing_transformer: 512 | logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.") 513 | else: 514 | self.transformer.unfuse_qkv_projections() 515 | self.fusing_transformer = False 516 | 517 | # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._prepare_rotary_positional_embeddings 518 | def _prepare_rotary_positional_embeddings( 519 | self, 520 | height: int, 521 | width: int, 522 | num_frames: int, 523 | device: torch.device, 524 | ) -> Tuple[torch.Tensor, torch.Tensor]: 525 | grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) 526 | grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) 527 | 528 | p = self.transformer.config.patch_size 529 | p_t = self.transformer.config.patch_size_t 530 | 531 | base_size_width = self.transformer.config.sample_width // p 532 | base_size_height = self.transformer.config.sample_height // p 533 | 534 | if p_t is None: 535 | # CogVideoX 1.0 536 | grid_crops_coords = get_resize_crop_region_for_grid( 537 | (grid_height, grid_width), base_size_width, base_size_height 538 | ) 539 | freqs_cos, freqs_sin = get_3d_rotary_pos_embed( 540 | embed_dim=self.transformer.config.attention_head_dim, 541 | crops_coords=grid_crops_coords, 542 | grid_size=(grid_height, grid_width), 543 | temporal_size=num_frames, 544 | ) 545 | else: 546 | # CogVideoX 1.5 547 | base_num_frames = (num_frames + p_t - 1) // p_t 548 | 549 | freqs_cos, freqs_sin = get_3d_rotary_pos_embed( 550 | embed_dim=self.transformer.config.attention_head_dim, 551 | crops_coords=None, 552 | grid_size=(grid_height, grid_width), 553 | temporal_size=base_num_frames, 554 | grid_type="slice", 555 | max_size=(base_size_height, base_size_width), 556 | ) 557 | 558 | freqs_cos = freqs_cos.to(device=device) 559 | freqs_sin = freqs_sin.to(device=device) 560 | return freqs_cos, freqs_sin 561 | 562 | def prepare_controlnet_frames(self, controlnet_frames, height, width, do_classifier_free_guidance): 563 | prepared_frames = prepare_frames(controlnet_frames, (height, width)) 564 | controlnet_encoded_frames = prepared_frames.to(dtype=self.vae.dtype, device='cuda') 565 | controlnet_encoded_frames = torch.cat([controlnet_encoded_frames] * 2) if do_classifier_free_guidance else controlnet_encoded_frames 566 | return controlnet_encoded_frames.contiguous() 567 | 568 | @property 569 | def guidance_scale(self): 570 | return self._guidance_scale 571 | 572 | @property 573 | def num_timesteps(self): 574 | return self._num_timesteps 575 | 576 | @property 577 | def attention_kwargs(self): 578 | return self._attention_kwargs 579 | 580 | @property 581 | def interrupt(self): 582 | return self._interrupt 583 | 584 | @torch.no_grad() 585 | def __call__( 586 | self, 587 | image, 588 | controlnet_frames: List[Image.Image] = None, 589 | prompt: Optional[Union[str, List[str]]] = None, 590 | negative_prompt: Optional[Union[str, List[str]]] = None, 591 | height: Optional[int] = None, 592 | width: Optional[int] = None, 593 | num_frames: int = 49, 594 | num_inference_steps: int = 50, 595 | timesteps: Optional[List[int]] = None, 596 | guidance_scale: float = 6, 597 | use_dynamic_cfg: bool = False, 598 | num_videos_per_prompt: int = 1, 599 | eta: float = 0.0, 600 | generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, 601 | latents: Optional[torch.FloatTensor] = None, 602 | prompt_embeds: Optional[torch.FloatTensor] = None, 603 | negative_prompt_embeds: Optional[torch.FloatTensor] = None, 604 | controlnet_latents: Optional[torch.FloatTensor] = None, 605 | output_type: str = "pil", 606 | return_dict: bool = True, 607 | attention_kwargs: Optional[Dict[str, Any]] = None, 608 | callback_on_step_end: Optional[ 609 | Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] 610 | ] = None, 611 | callback_on_step_end_tensor_inputs: List[str] = ["latents"], 612 | max_sequence_length: int = 226, 613 | controlnet_weights: Optional[Union[float, list, np.ndarray, torch.FloatTensor]] = 1.0, 614 | controlnet_guidance_start: float = 0.0, 615 | controlnet_guidance_end: float = 1.0, 616 | ) -> Union[CogVideoXPipelineOutput, Tuple]: 617 | """ 618 | Function invoked when calling the pipeline for generation. 619 | 620 | Args: 621 | image (`PipelineImageInput`): 622 | The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. 623 | prompt (`str` or `List[str]`, *optional*): 624 | The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. 625 | instead. 626 | negative_prompt (`str` or `List[str]`, *optional*): 627 | The prompt or prompts not to guide the image generation. If not defined, one has to pass 628 | `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is 629 | less than `1`). 630 | height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): 631 | The height in pixels of the generated image. This is set to 480 by default for the best results. 632 | width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): 633 | The width in pixels of the generated image. This is set to 720 by default for the best results. 634 | num_frames (`int`, defaults to `48`): 635 | Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will 636 | contain 1 extra frame because CogVideoX is conditioned with (num_seconds * fps + 1) frames where 637 | num_seconds is 6 and fps is 8. However, since videos can be saved at any fps, the only condition that 638 | needs to be satisfied is that of divisibility mentioned above. 639 | num_inference_steps (`int`, *optional*, defaults to 50): 640 | The number of denoising steps. More denoising steps usually lead to a higher quality image at the 641 | expense of slower inference. 642 | timesteps (`List[int]`, *optional*): 643 | Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument 644 | in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is 645 | passed will be used. Must be in descending order. 646 | guidance_scale (`float`, *optional*, defaults to 7.0): 647 | Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). 648 | `guidance_scale` is defined as `w` of equation 2. of [Imagen 649 | Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > 650 | 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, 651 | usually at the expense of lower image quality. 652 | num_videos_per_prompt (`int`, *optional*, defaults to 1): 653 | The number of videos to generate per prompt. 654 | generator (`torch.Generator` or `List[torch.Generator]`, *optional*): 655 | One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) 656 | to make generation deterministic. 657 | latents (`torch.FloatTensor`, *optional*): 658 | Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image 659 | generation. Can be used to tweak the same generation with different prompts. If not provided, a latents 660 | tensor will ge generated by sampling using the supplied random `generator`. 661 | prompt_embeds (`torch.FloatTensor`, *optional*): 662 | Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not 663 | provided, text embeddings will be generated from `prompt` input argument. 664 | negative_prompt_embeds (`torch.FloatTensor`, *optional*): 665 | Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt 666 | weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input 667 | argument. 668 | output_type (`str`, *optional*, defaults to `"pil"`): 669 | The output format of the generate image. Choose between 670 | [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. 671 | return_dict (`bool`, *optional*, defaults to `True`): 672 | Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead 673 | of a plain tuple. 674 | attention_kwargs (`dict`, *optional*): 675 | A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under 676 | `self.processor` in 677 | [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). 678 | callback_on_step_end (`Callable`, *optional*): 679 | A function that calls at the end of each denoising steps during the inference. The function is called 680 | with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, 681 | callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by 682 | `callback_on_step_end_tensor_inputs`. 683 | callback_on_step_end_tensor_inputs (`List`, *optional*): 684 | The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list 685 | will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the 686 | `._callback_tensor_inputs` attribute of your pipeline class. 687 | max_sequence_length (`int`, defaults to `226`): 688 | Maximum sequence length in encoded prompt. Must be consistent with 689 | `self.transformer.config.max_text_seq_length` otherwise may lead to poor results. 690 | 691 | Examples: 692 | 693 | Returns: 694 | [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] or `tuple`: 695 | [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a 696 | `tuple`. When returning a tuple, the first element is a list with the generated images. 697 | """ 698 | 699 | if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): 700 | callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs 701 | 702 | height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial 703 | width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial 704 | num_frames = num_frames or self.transformer.config.sample_frames 705 | 706 | num_videos_per_prompt = 1 707 | 708 | # 1. Check inputs. Raise error if not correct 709 | self.check_inputs( 710 | image=image, 711 | prompt=prompt, 712 | height=height, 713 | width=width, 714 | negative_prompt=negative_prompt, 715 | callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, 716 | latents=latents, 717 | prompt_embeds=prompt_embeds, 718 | negative_prompt_embeds=negative_prompt_embeds, 719 | ) 720 | self._guidance_scale = guidance_scale 721 | self._attention_kwargs = attention_kwargs 722 | self._interrupt = False 723 | 724 | # 2. Default call parameters 725 | if prompt is not None and isinstance(prompt, str): 726 | batch_size = 1 727 | elif prompt is not None and isinstance(prompt, list): 728 | batch_size = len(prompt) 729 | else: 730 | batch_size = prompt_embeds.shape[0] 731 | 732 | device = self._execution_device 733 | 734 | # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) 735 | # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` 736 | # corresponds to doing no classifier free guidance. 737 | do_classifier_free_guidance = guidance_scale > 1.0 738 | 739 | # 3. Encode input prompt 740 | prompt_embeds, negative_prompt_embeds = self.encode_prompt( 741 | prompt=prompt, 742 | negative_prompt=negative_prompt, 743 | do_classifier_free_guidance=do_classifier_free_guidance, 744 | num_videos_per_prompt=num_videos_per_prompt, 745 | prompt_embeds=prompt_embeds, 746 | negative_prompt_embeds=negative_prompt_embeds, 747 | max_sequence_length=max_sequence_length, 748 | device=device, 749 | ) 750 | if do_classifier_free_guidance: 751 | prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) 752 | 753 | # 4. Prepare timesteps 754 | timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) 755 | self._num_timesteps = len(timesteps) 756 | 757 | # 5. Prepare latents 758 | latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 759 | 760 | # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t 761 | patch_size_t = self.transformer.config.patch_size_t 762 | additional_frames = 0 763 | if patch_size_t is not None and latent_frames % patch_size_t != 0: 764 | additional_frames = patch_size_t - latent_frames % patch_size_t 765 | num_frames += additional_frames * self.vae_scale_factor_temporal 766 | 767 | image = self.video_processor.preprocess(image, height=height, width=width).to( 768 | device, dtype=prompt_embeds.dtype 769 | ) 770 | 771 | latent_channels = self.transformer.config.in_channels // 2 772 | latents, image_latents = self.prepare_latents( 773 | image, 774 | batch_size * num_videos_per_prompt, 775 | latent_channels, 776 | num_frames, 777 | height, 778 | width, 779 | prompt_embeds.dtype, 780 | device, 781 | generator, 782 | latents, 783 | ) 784 | 785 | # 6. Encode controlnet frames 786 | if controlnet_latents is None: 787 | controlnet_latents = self.prepare_controlnet_frames( 788 | controlnet_frames, 789 | height, 790 | width, 791 | do_classifier_free_guidance, 792 | ) 793 | 794 | # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline 795 | extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) 796 | 797 | # 8. Create rotary embeds if required 798 | image_rotary_emb = ( 799 | self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) 800 | if self.transformer.config.use_rotary_positional_embeddings 801 | else None 802 | ) 803 | 804 | # 9. Create ofs embeds if required 805 | ofs_emb = None if self.transformer.config.ofs_embed_dim is None else latents.new_full((1,), fill_value=2.0) 806 | 807 | # 10. Denoising loop 808 | num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) 809 | 810 | with self.progress_bar(total=num_inference_steps) as progress_bar: 811 | # for DPM-solver++ 812 | old_pred_original_sample = None 813 | for i, t in enumerate(timesteps): 814 | if self.interrupt: 815 | continue 816 | 817 | latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents 818 | latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) 819 | latent_image_input = torch.cat([image_latents] * 2) if do_classifier_free_guidance else image_latents 820 | latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=2) 821 | 822 | # broadcast to batch dimension in a way that's compatible with ONNX/Core ML 823 | timestep = t.expand(latent_model_input.shape[0]) 824 | 825 | current_sampling_percent = i / len(timesteps) 826 | 827 | controlnet_states = [] 828 | if (controlnet_guidance_start <= current_sampling_percent < controlnet_guidance_end): 829 | # extract controlnet hidden state 830 | controlnet_states = self.controlnet( 831 | hidden_states=latent_model_input[:, :, :16, :, :], 832 | encoder_hidden_states=prompt_embeds, 833 | image_rotary_emb=image_rotary_emb, 834 | controlnet_states=controlnet_latents, 835 | timestep=timestep, 836 | return_dict=False, 837 | )[0] 838 | if isinstance(controlnet_states, (tuple, list)): 839 | controlnet_states = [x.to(dtype=self.transformer.dtype) for x in controlnet_states] 840 | else: 841 | controlnet_states = controlnet_states.to(dtype=self.transformer.dtype) 842 | 843 | # predict noise model_output 844 | noise_pred = self.transformer( 845 | hidden_states=latent_model_input, 846 | encoder_hidden_states=prompt_embeds, 847 | timestep=timestep, 848 | # ofs=ofs_emb, 849 | image_rotary_emb=image_rotary_emb, 850 | # attention_kwargs=attention_kwargs, 851 | controlnet_states=controlnet_states, 852 | controlnet_weights=controlnet_weights, 853 | return_dict=False, 854 | )[0] 855 | noise_pred = noise_pred.float() 856 | 857 | # perform guidance 858 | if use_dynamic_cfg: 859 | self._guidance_scale = 1 + guidance_scale * ( 860 | (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2 861 | ) 862 | if do_classifier_free_guidance: 863 | noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) 864 | noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) 865 | 866 | # compute the previous noisy sample x_t -> x_t-1 867 | if not isinstance(self.scheduler, CogVideoXDPMScheduler): 868 | latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] 869 | else: 870 | latents, old_pred_original_sample = self.scheduler.step( 871 | noise_pred, 872 | old_pred_original_sample, 873 | t, 874 | timesteps[i - 1] if i > 0 else None, 875 | latents, 876 | **extra_step_kwargs, 877 | return_dict=False, 878 | ) 879 | latents = latents.to(prompt_embeds.dtype) 880 | 881 | # call the callback, if provided 882 | if callback_on_step_end is not None: 883 | callback_kwargs = {} 884 | for k in callback_on_step_end_tensor_inputs: 885 | callback_kwargs[k] = locals()[k] 886 | callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) 887 | 888 | latents = callback_outputs.pop("latents", latents) 889 | prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) 890 | negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) 891 | 892 | if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): 893 | progress_bar.update() 894 | 895 | if not output_type == "latent": 896 | # Discard any padding frames that were added for CogVideoX 1.5 897 | latents = latents[:, additional_frames:] 898 | video = self.decode_latents(latents) 899 | video = self.video_processor.postprocess_video(video=video, output_type=output_type) 900 | else: 901 | video = latents 902 | 903 | # Offload all models 904 | self.maybe_free_model_hooks() 905 | 906 | if not return_dict: 907 | return (video,) 908 | 909 | return CogVideoXPipelineOutput(frames=video) -------------------------------------------------------------------------------- /controlnet_pipeline.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import math 3 | from typing import Callable, Dict, List, Optional, Tuple, Union 4 | 5 | import torch 6 | import numpy as np 7 | from PIL import Image 8 | from torchvision import transforms 9 | from einops import rearrange, repeat 10 | from transformers import T5EncoderModel, T5Tokenizer 11 | from diffusers.video_processor import VideoProcessor 12 | from diffusers.utils.torch_utils import randn_tensor 13 | from diffusers.models.embeddings import get_3d_rotary_pos_embed 14 | from diffusers.pipelines.pipeline_utils import DiffusionPipeline 15 | from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel 16 | from diffusers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler 17 | from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback 18 | from diffusers.pipelines.cogvideo.pipeline_cogvideox import CogVideoXPipelineOutput, CogVideoXLoraLoaderMixin 19 | 20 | from cogvideo_controlnet import CogVideoXControlnet 21 | 22 | 23 | def resize_for_crop(image, crop_h, crop_w): 24 | img_h, img_w = image.shape[-2:] 25 | if img_h >= crop_h and img_w >= crop_w: 26 | coef = max(crop_h / img_h, crop_w / img_w) 27 | elif img_h <= crop_h and img_w <= crop_w: 28 | coef = max(crop_h / img_h, crop_w / img_w) 29 | else: 30 | coef = crop_h / img_h if crop_h > img_h else crop_w / img_w 31 | out_h, out_w = int(img_h * coef), int(img_w * coef) 32 | resized_image = transforms.functional.resize(image, (out_h, out_w), antialias=True) 33 | return resized_image 34 | 35 | 36 | def prepare_frames(input_images, video_size, do_resize=True, do_crop=True): 37 | input_images = np.stack([np.array(x) for x in input_images]) 38 | images_tensor = torch.from_numpy(input_images).permute(0, 3, 1, 2) / 127.5 - 1 39 | if do_resize: 40 | images_tensor = [resize_for_crop(x, crop_h=video_size[0], crop_w=video_size[1]) for x in images_tensor] 41 | if do_crop: 42 | images_tensor = [transforms.functional.center_crop(x, video_size) for x in images_tensor] 43 | if isinstance(images_tensor, list): 44 | images_tensor = torch.stack(images_tensor) 45 | return images_tensor.unsqueeze(0) 46 | 47 | 48 | def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): 49 | tw = tgt_width 50 | th = tgt_height 51 | h, w = src 52 | r = h / w 53 | if r > (th / tw): 54 | resize_height = th 55 | resize_width = int(round(th / h * w)) 56 | else: 57 | resize_width = tw 58 | resize_height = int(round(tw / w * h)) 59 | 60 | crop_top = int(round((th - resize_height) / 2.0)) 61 | crop_left = int(round((tw - resize_width) / 2.0)) 62 | 63 | return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) 64 | 65 | 66 | # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps 67 | def retrieve_timesteps( 68 | scheduler, 69 | num_inference_steps: Optional[int] = None, 70 | device: Optional[Union[str, torch.device]] = None, 71 | timesteps: Optional[List[int]] = None, 72 | sigmas: Optional[List[float]] = None, 73 | **kwargs, 74 | ): 75 | """ 76 | Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles 77 | custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. 78 | 79 | Args: 80 | scheduler (`SchedulerMixin`): 81 | The scheduler to get timesteps from. 82 | num_inference_steps (`int`): 83 | The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` 84 | must be `None`. 85 | device (`str` or `torch.device`, *optional*): 86 | The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. 87 | timesteps (`List[int]`, *optional*): 88 | Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, 89 | `num_inference_steps` and `sigmas` must be `None`. 90 | sigmas (`List[float]`, *optional*): 91 | Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, 92 | `num_inference_steps` and `timesteps` must be `None`. 93 | 94 | Returns: 95 | `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the 96 | second element is the number of inference steps. 97 | """ 98 | if timesteps is not None and sigmas is not None: 99 | raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") 100 | if timesteps is not None: 101 | accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) 102 | if not accepts_timesteps: 103 | raise ValueError( 104 | f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" 105 | f" timestep schedules. Please check whether you are using the correct scheduler." 106 | ) 107 | scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) 108 | timesteps = scheduler.timesteps 109 | num_inference_steps = len(timesteps) 110 | elif sigmas is not None: 111 | accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) 112 | if not accept_sigmas: 113 | raise ValueError( 114 | f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" 115 | f" sigmas schedules. Please check whether you are using the correct scheduler." 116 | ) 117 | scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) 118 | timesteps = scheduler.timesteps 119 | num_inference_steps = len(timesteps) 120 | else: 121 | scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) 122 | timesteps = scheduler.timesteps 123 | return timesteps, num_inference_steps 124 | 125 | 126 | class ControlnetCogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): 127 | _optional_components = [] 128 | model_cpu_offload_seq = "text_encoder->transformer->vae" 129 | 130 | _callback_tensor_inputs = [ 131 | "latents", 132 | "prompt_embeds", 133 | "negative_prompt_embeds", 134 | ] 135 | 136 | def __init__( 137 | self, 138 | tokenizer: T5Tokenizer, 139 | text_encoder: T5EncoderModel, 140 | vae: AutoencoderKLCogVideoX, 141 | transformer: CogVideoXTransformer3DModel, 142 | scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], 143 | controlnet: CogVideoXControlnet, 144 | ): 145 | super().__init__() 146 | 147 | self.register_modules( 148 | tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, controlnet=controlnet, scheduler=scheduler 149 | ) 150 | self.vae_scale_factor_spatial = ( 151 | 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 152 | ) 153 | self.vae_scale_factor_temporal = ( 154 | self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4 155 | ) 156 | 157 | self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) 158 | 159 | def prepare_controlnet_frames(self, controlnet_frames, height, width, do_classifier_free_guidance): 160 | prepared_frames = prepare_frames(controlnet_frames, (height, width)) 161 | controlnet_encoded_frames = prepared_frames.to(dtype=self.vae.dtype, device='cuda') 162 | controlnet_encoded_frames = torch.cat([controlnet_encoded_frames] * 2) if do_classifier_free_guidance else controlnet_encoded_frames 163 | return controlnet_encoded_frames.contiguous() 164 | 165 | def _get_t5_prompt_embeds( 166 | self, 167 | prompt: Union[str, List[str]] = None, 168 | num_videos_per_prompt: int = 1, 169 | max_sequence_length: int = 226, 170 | device: Optional[torch.device] = None, 171 | dtype: Optional[torch.dtype] = None, 172 | ): 173 | device = device or self._execution_device 174 | dtype = dtype or self.text_encoder.dtype 175 | 176 | prompt = [prompt] if isinstance(prompt, str) else prompt 177 | batch_size = len(prompt) 178 | 179 | text_inputs = self.tokenizer( 180 | prompt, 181 | padding="max_length", 182 | max_length=max_sequence_length, 183 | truncation=True, 184 | add_special_tokens=True, 185 | return_tensors="pt", 186 | ) 187 | text_input_ids = text_inputs.input_ids 188 | untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids 189 | 190 | if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): 191 | removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) 192 | logger.warning( 193 | "The following part of your input was truncated because `max_sequence_length` is set to " 194 | f" {max_sequence_length} tokens: {removed_text}" 195 | ) 196 | 197 | prompt_embeds = self.text_encoder(text_input_ids.to(device))[0] 198 | prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) 199 | 200 | # duplicate text embeddings for each generation per prompt, using mps friendly method 201 | _, seq_len, _ = prompt_embeds.shape 202 | prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) 203 | prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) 204 | 205 | return prompt_embeds 206 | 207 | def encode_prompt( 208 | self, 209 | prompt: Union[str, List[str]], 210 | negative_prompt: Optional[Union[str, List[str]]] = None, 211 | do_classifier_free_guidance: bool = True, 212 | num_videos_per_prompt: int = 1, 213 | prompt_embeds: Optional[torch.Tensor] = None, 214 | negative_prompt_embeds: Optional[torch.Tensor] = None, 215 | max_sequence_length: int = 226, 216 | device: Optional[torch.device] = None, 217 | dtype: Optional[torch.dtype] = None, 218 | ): 219 | r""" 220 | Encodes the prompt into text encoder hidden states. 221 | 222 | Args: 223 | prompt (`str` or `List[str]`, *optional*): 224 | prompt to be encoded 225 | negative_prompt (`str` or `List[str]`, *optional*): 226 | The prompt or prompts not to guide the image generation. If not defined, one has to pass 227 | `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is 228 | less than `1`). 229 | do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): 230 | Whether to use classifier free guidance or not. 231 | num_videos_per_prompt (`int`, *optional*, defaults to 1): 232 | Number of videos that should be generated per prompt. torch device to place the resulting embeddings on 233 | prompt_embeds (`torch.Tensor`, *optional*): 234 | Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not 235 | provided, text embeddings will be generated from `prompt` input argument. 236 | negative_prompt_embeds (`torch.Tensor`, *optional*): 237 | Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt 238 | weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input 239 | argument. 240 | device: (`torch.device`, *optional*): 241 | torch device 242 | dtype: (`torch.dtype`, *optional*): 243 | torch dtype 244 | """ 245 | device = device or self._execution_device 246 | 247 | prompt = [prompt] if isinstance(prompt, str) else prompt 248 | if prompt is not None: 249 | batch_size = len(prompt) 250 | else: 251 | batch_size = prompt_embeds.shape[0] 252 | 253 | if prompt_embeds is None: 254 | prompt_embeds = self._get_t5_prompt_embeds( 255 | prompt=prompt, 256 | num_videos_per_prompt=num_videos_per_prompt, 257 | max_sequence_length=max_sequence_length, 258 | device=device, 259 | dtype=dtype, 260 | ) 261 | 262 | if do_classifier_free_guidance and negative_prompt_embeds is None: 263 | negative_prompt = negative_prompt or "" 264 | negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt 265 | 266 | if prompt is not None and type(prompt) is not type(negative_prompt): 267 | raise TypeError( 268 | f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" 269 | f" {type(prompt)}." 270 | ) 271 | elif batch_size != len(negative_prompt): 272 | raise ValueError( 273 | f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" 274 | f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" 275 | " the batch size of `prompt`." 276 | ) 277 | 278 | negative_prompt_embeds = self._get_t5_prompt_embeds( 279 | prompt=negative_prompt, 280 | num_videos_per_prompt=num_videos_per_prompt, 281 | max_sequence_length=max_sequence_length, 282 | device=device, 283 | dtype=dtype, 284 | ) 285 | 286 | return prompt_embeds, negative_prompt_embeds 287 | 288 | def prepare_latents( 289 | self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None 290 | ): 291 | shape = ( 292 | batch_size, 293 | (num_frames - 1) // self.vae_scale_factor_temporal + 1, 294 | num_channels_latents, 295 | height // self.vae_scale_factor_spatial, 296 | width // self.vae_scale_factor_spatial, 297 | ) 298 | if isinstance(generator, list) and len(generator) != batch_size: 299 | raise ValueError( 300 | f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" 301 | f" size of {batch_size}. Make sure the batch size matches the length of the generators." 302 | ) 303 | 304 | if latents is None: 305 | latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) 306 | else: 307 | latents = latents.to(device) 308 | 309 | # scale the initial noise by the standard deviation required by the scheduler 310 | latents = latents * self.scheduler.init_noise_sigma 311 | return latents 312 | 313 | def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: 314 | latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] 315 | latents = 1 / self.vae.config.scaling_factor * latents 316 | 317 | frames = self.vae.decode(latents).sample 318 | return frames 319 | 320 | # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs 321 | def prepare_extra_step_kwargs(self, generator, eta): 322 | # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature 323 | # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. 324 | # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 325 | # and should be between [0, 1] 326 | 327 | accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) 328 | extra_step_kwargs = {} 329 | if accepts_eta: 330 | extra_step_kwargs["eta"] = eta 331 | 332 | # check if the scheduler accepts generator 333 | accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) 334 | if accepts_generator: 335 | extra_step_kwargs["generator"] = generator 336 | return extra_step_kwargs 337 | 338 | # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs 339 | def check_inputs( 340 | self, 341 | prompt, 342 | height, 343 | width, 344 | negative_prompt, 345 | callback_on_step_end_tensor_inputs, 346 | controlnet_frames=None, 347 | controlnet_latents=None, 348 | prompt_embeds=None, 349 | negative_prompt_embeds=None, 350 | ): 351 | if height % 8 != 0 or width % 8 != 0: 352 | raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") 353 | 354 | if callback_on_step_end_tensor_inputs is not None and not all( 355 | k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs 356 | ): 357 | raise ValueError( 358 | f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" 359 | ) 360 | if prompt is not None and prompt_embeds is not None: 361 | raise ValueError( 362 | f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" 363 | " only forward one of the two." 364 | ) 365 | elif prompt is None and prompt_embeds is None: 366 | raise ValueError( 367 | "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." 368 | ) 369 | elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): 370 | raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") 371 | 372 | if prompt is not None and negative_prompt_embeds is not None: 373 | raise ValueError( 374 | f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" 375 | f" {negative_prompt_embeds}. Please make sure to only forward one of the two." 376 | ) 377 | 378 | if negative_prompt is not None and negative_prompt_embeds is not None: 379 | raise ValueError( 380 | f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" 381 | f" {negative_prompt_embeds}. Please make sure to only forward one of the two." 382 | ) 383 | 384 | if prompt_embeds is not None and negative_prompt_embeds is not None: 385 | if prompt_embeds.shape != negative_prompt_embeds.shape: 386 | raise ValueError( 387 | "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" 388 | f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" 389 | f" {negative_prompt_embeds.shape}." 390 | ) 391 | if controlnet_frames is not None and controlnet_latents is not None: 392 | raise ValueError("Only one of `controlnet_frames` or `controlnet_latents` should be provided") 393 | 394 | def fuse_qkv_projections(self) -> None: 395 | r"""Enables fused QKV projections.""" 396 | self.fusing_transformer = True 397 | self.transformer.fuse_qkv_projections() 398 | 399 | def unfuse_qkv_projections(self) -> None: 400 | r"""Disable QKV projection fusion if enabled.""" 401 | if not self.fusing_transformer: 402 | logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.") 403 | else: 404 | self.transformer.unfuse_qkv_projections() 405 | self.fusing_transformer = False 406 | 407 | def _prepare_rotary_positional_embeddings( 408 | self, 409 | height: int, 410 | width: int, 411 | num_frames: int, 412 | device: torch.device, 413 | ) -> Tuple[torch.Tensor, torch.Tensor]: 414 | grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) 415 | grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) 416 | base_size_width = 720 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) 417 | base_size_height = 480 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) 418 | 419 | grid_crops_coords = get_resize_crop_region_for_grid( 420 | (grid_height, grid_width), base_size_width, base_size_height 421 | ) 422 | freqs_cos, freqs_sin = get_3d_rotary_pos_embed( 423 | embed_dim=self.transformer.config.attention_head_dim, 424 | crops_coords=grid_crops_coords, 425 | grid_size=(grid_height, grid_width), 426 | temporal_size=num_frames, 427 | ) 428 | 429 | freqs_cos = freqs_cos.to(device=device) 430 | freqs_sin = freqs_sin.to(device=device) 431 | return freqs_cos, freqs_sin 432 | 433 | @property 434 | def guidance_scale(self): 435 | return self._guidance_scale 436 | 437 | @property 438 | def num_timesteps(self): 439 | return self._num_timesteps 440 | 441 | @property 442 | def attention_kwargs(self): 443 | return self._attention_kwargs 444 | 445 | @property 446 | def interrupt(self): 447 | return self._interrupt 448 | 449 | @torch.no_grad() 450 | def __call__( 451 | self, 452 | controlnet_frames: List[Image.Image] = None, 453 | prompt: Optional[Union[str, List[str]]] = None, 454 | negative_prompt: Optional[Union[str, List[str]]] = None, 455 | height: int = 480, 456 | width: int = 720, 457 | num_frames: int = 49, 458 | num_inference_steps: int = 50, 459 | timesteps: Optional[List[int]] = None, 460 | guidance_scale: float = 6, 461 | use_dynamic_cfg: bool = False, 462 | num_videos_per_prompt: int = 1, 463 | eta: float = 0.0, 464 | generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, 465 | latents: Optional[torch.FloatTensor] = None, 466 | controlnet_latents: Optional[torch.FloatTensor] = None, 467 | prompt_embeds: Optional[torch.FloatTensor] = None, 468 | negative_prompt_embeds: Optional[torch.FloatTensor] = None, 469 | output_type: str = "pil", 470 | return_dict: bool = True, 471 | callback_on_step_end: Optional[ 472 | Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] 473 | ] = None, 474 | callback_on_step_end_tensor_inputs: List[str] = ["latents"], 475 | max_sequence_length: int = 226, 476 | controlnet_weights: Optional[Union[float, list, np.ndarray, torch.FloatTensor]] = 1.0, 477 | controlnet_guidance_start: float = 0.0, 478 | controlnet_guidance_end: float = 1.0, 479 | ) -> Union[CogVideoXPipelineOutput, Tuple]: 480 | if num_frames > 49: 481 | raise ValueError( 482 | "The number of frames must be less than 49 for now due to static positional embeddings. This will be updated in the future to remove this limitation." 483 | ) 484 | 485 | if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): 486 | callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs 487 | 488 | height = height or self.transformer.config.sample_size * self.vae_scale_factor_spatial 489 | width = width or self.transformer.config.sample_size * self.vae_scale_factor_spatial 490 | num_videos_per_prompt = 1 491 | 492 | # 1. Check inputs. Raise error if not correct 493 | self.check_inputs( 494 | prompt, 495 | height, 496 | width, 497 | negative_prompt, 498 | callback_on_step_end_tensor_inputs, 499 | controlnet_frames, 500 | controlnet_latents, 501 | prompt_embeds, 502 | negative_prompt_embeds, 503 | ) 504 | self._guidance_scale = guidance_scale 505 | self._interrupt = False 506 | 507 | # 2. Default call parameters 508 | if prompt is not None and isinstance(prompt, str): 509 | batch_size = 1 510 | elif prompt is not None and isinstance(prompt, list): 511 | batch_size = len(prompt) 512 | else: 513 | batch_size = prompt_embeds.shape[0] 514 | 515 | device = self._execution_device 516 | 517 | # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) 518 | # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` 519 | # corresponds to doing no classifier free guidance. 520 | do_classifier_free_guidance = guidance_scale > 1.0 521 | 522 | # 3. Encode input prompt 523 | prompt_embeds, negative_prompt_embeds = self.encode_prompt( 524 | prompt, 525 | negative_prompt, 526 | do_classifier_free_guidance, 527 | num_videos_per_prompt=num_videos_per_prompt, 528 | prompt_embeds=prompt_embeds, 529 | negative_prompt_embeds=negative_prompt_embeds, 530 | max_sequence_length=max_sequence_length, 531 | device=device, 532 | ) 533 | if do_classifier_free_guidance: 534 | prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) 535 | 536 | # 4. Prepare timesteps 537 | timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) 538 | self._num_timesteps = len(timesteps) 539 | 540 | # 5. Prepare latents. 541 | latent_channels = 16 #self.transformer.config.in_channels 542 | latents = self.prepare_latents( 543 | batch_size * num_videos_per_prompt, 544 | latent_channels, 545 | num_frames, 546 | height, 547 | width, 548 | prompt_embeds.dtype, 549 | device, 550 | generator, 551 | latents, 552 | ) 553 | 554 | # 6. Encode controlnet frames 555 | if controlnet_latents is None: 556 | duplicate_frames_count = num_frames - len(controlnet_frames) 557 | if duplicate_frames_count > 0: 558 | # Simple duplicate first frame 559 | # controlnet_frames = [controlnet_frames[0]] * duplicate_frames_count + controlnet_frames 560 | # Or reversed duplicate frames ? 561 | reversed_controlnet_frames = list(reversed(controlnet_frames)) 562 | controlnet_sum_frames = controlnet_frames + reversed_controlnet_frames 563 | reversed_chunks_count = num_frames // len(controlnet_sum_frames) 564 | controlnet_frames = [*controlnet_sum_frames] 565 | for _ in range(reversed_chunks_count): 566 | controlnet_frames += controlnet_sum_frames 567 | 568 | # If controlnet frames count greater than num_frames parameter 569 | controlnet_frames = controlnet_frames[:num_frames] 570 | 571 | controlnet_latents = self.prepare_controlnet_frames( 572 | controlnet_frames, 573 | height, 574 | width, 575 | do_classifier_free_guidance, 576 | ) 577 | 578 | # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline 579 | extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) 580 | 581 | # 8. Create rotary embeds if required 582 | image_rotary_emb = ( 583 | self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) 584 | if self.transformer.config.use_rotary_positional_embeddings 585 | else None 586 | ) 587 | 588 | # 9. Denoising loop 589 | num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) 590 | 591 | with self.progress_bar(total=num_inference_steps) as progress_bar: 592 | # for DPM-solver++ 593 | old_pred_original_sample = None 594 | for i, t in enumerate(timesteps): 595 | if self.interrupt: 596 | continue 597 | 598 | latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents 599 | latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) 600 | 601 | timestep = t.expand(latent_model_input.shape[0]) 602 | 603 | current_sampling_percent = i / len(timesteps) 604 | 605 | latent_model_input = latent_model_input.to(dtype=self.transformer.dtype) 606 | prompt_embeds = prompt_embeds.to(dtype=self.transformer.dtype) 607 | 608 | controlnet_states = None 609 | if (controlnet_guidance_start <= current_sampling_percent <= controlnet_guidance_end): 610 | # extract controlnet hidden state 611 | controlnet_states = self.controlnet( 612 | hidden_states=latent_model_input, 613 | encoder_hidden_states=prompt_embeds, 614 | image_rotary_emb=image_rotary_emb, 615 | controlnet_states=controlnet_latents, 616 | timestep=timestep, 617 | return_dict=False, 618 | )[0] 619 | if isinstance(controlnet_states, (tuple, list)): 620 | controlnet_states = [x.to(dtype=self.transformer.dtype) for x in controlnet_states] 621 | else: 622 | controlnet_states = controlnet_states.to(dtype=self.transformer.dtype) 623 | 624 | # predict noise model_output 625 | noise_pred = self.transformer( 626 | hidden_states=latent_model_input, 627 | encoder_hidden_states=prompt_embeds, 628 | timestep=timestep, 629 | image_rotary_emb=image_rotary_emb, 630 | controlnet_states=controlnet_states, 631 | controlnet_weights=controlnet_weights, 632 | return_dict=False, 633 | )[0] 634 | noise_pred = noise_pred.float() 635 | 636 | # perform guidance 637 | if use_dynamic_cfg: 638 | self._guidance_scale = 1 + guidance_scale * ( 639 | (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2 640 | ) 641 | if do_classifier_free_guidance: 642 | noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) 643 | noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) 644 | 645 | # compute the previous noisy sample x_t -> x_t-1 646 | if not isinstance(self.scheduler, CogVideoXDPMScheduler): 647 | latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] 648 | else: 649 | latents, old_pred_original_sample = self.scheduler.step( 650 | noise_pred, 651 | old_pred_original_sample, 652 | t, 653 | timesteps[i - 1] if i > 0 else None, 654 | latents, 655 | **extra_step_kwargs, 656 | return_dict=False, 657 | ) 658 | latents = latents.to(prompt_embeds.dtype) 659 | 660 | # call the callback, if provided 661 | if callback_on_step_end is not None: 662 | callback_kwargs = {} 663 | for k in callback_on_step_end_tensor_inputs: 664 | callback_kwargs[k] = locals()[k] 665 | callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) 666 | 667 | latents = callback_outputs.pop("latents", latents) 668 | prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) 669 | negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) 670 | 671 | if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): 672 | progress_bar.update() 673 | 674 | if not output_type == "latent": 675 | video = self.decode_latents(latents) 676 | video = self.video_processor.postprocess_video(video=video, output_type=output_type) 677 | else: 678 | video = latents 679 | 680 | # Offload all models 681 | self.maybe_free_model_hooks() 682 | 683 | if not return_dict: 684 | return (video,) 685 | 686 | return CogVideoXPipelineOutput(frames=video) 687 | -------------------------------------------------------------------------------- /inference/cli_demo.py: -------------------------------------------------------------------------------- 1 | """ 2 | Running the Script: 3 | To run the script, use the following command with appropriate arguments: 4 | 5 | ```bash 6 | $ python inference/cli_demo.py \ 7 | --video_path "test_video/car.mp4" \ 8 | --prompt "the car is driving on a mountain road" \ 9 | --controlnet_type "hed" \ 10 | --model_path THUDM/CogVideoX-5b \ 11 | --controlnet_path TheDenk/cogvideox-5b-controlnet-hed-v1 12 | ``` 13 | 14 | Additional options are available to specify the guidance scale, number of inference steps, video generation type, and output paths. 15 | """ 16 | import sys 17 | sys.path.append('..') 18 | import argparse 19 | 20 | import torch 21 | from transformers import T5EncoderModel, T5Tokenizer 22 | from diffusers import ( 23 | CogVideoXDDIMScheduler, 24 | CogVideoXDPMScheduler, 25 | AutoencoderKLCogVideoX 26 | ) 27 | from diffusers.utils import export_to_video, load_video 28 | from controlnet_aux import HEDdetector, CannyDetector 29 | 30 | from controlnet_pipeline import ControlnetCogVideoXPipeline 31 | from cogvideo_transformer import CustomCogVideoXTransformer3DModel 32 | from cogvideo_controlnet import CogVideoXControlnet 33 | 34 | 35 | def init_controlnet_processor(controlnet_type): 36 | if controlnet_type in ['canny', 'lineart']: 37 | return controlnet_mapping[controlnet_type]() 38 | return controlnet_mapping[controlnet_type].from_pretrained('lllyasviel/Annotators').to(device='cuda') 39 | 40 | 41 | controlnet_mapping = { 42 | 'hed': HEDdetector, 43 | 'canny': CannyDetector, 44 | } 45 | 46 | 47 | @torch.no_grad() 48 | def generate_video( 49 | prompt: str, 50 | video_path: str, 51 | base_model_path: str, 52 | controlnet_model_path: str, 53 | controlnet_type: str, 54 | controlnet_weights: float = 1.0, 55 | controlnet_guidance_start: float = 0.0, 56 | controlnet_guidance_end: float = 1.0, 57 | lora_path: str = None, 58 | lora_rank: int = 128, 59 | output_path: str = "./output.mp4", 60 | num_inference_steps: int = 50, 61 | guidance_scale: float = 6.0, 62 | num_videos_per_prompt: int = 1, 63 | dtype: torch.dtype = torch.bfloat16, 64 | seed: int = 42, 65 | ): 66 | """ 67 | Generates a video based on the given prompt and saves it to the specified path. 68 | 69 | Parameters: 70 | - prompt (str): The description of the video to be generated. 71 | - video_path (str): The video for controlnet processing. 72 | - base_model_path (str): The path of the pre-trained model to be used. 73 | - controlnet_model_path (str): The path of the pre-trained conrolnet model to be used. 74 | - controlnet_type (str): Type of controlnet model (e.g. canny, hed). 75 | - controlnet_weights (float): Strenght of controlnet 76 | - controlnet_guidance_start (float): The stage when the controlnet starts to be applied 77 | - controlnet_guidance_end (float): The stage when the controlnet end to be applied 78 | - lora_path (str): The path of the LoRA weights to be used. 79 | - lora_rank (int): The rank of the LoRA weights. 80 | - output_path (str): The path where the generated video will be saved. 81 | - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality. 82 | - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt. 83 | - num_videos_per_prompt (int): Number of videos to generate per prompt. 84 | - dtype (torch.dtype): The data type for computation (default is torch.bfloat16). 85 | - seed (int): The seed for reproducibility. 86 | """ 87 | 88 | # 0. Load selected controlnet processor 89 | controlnet_processor = init_controlnet_processor(controlnet_type) 90 | # 1. Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16). 91 | tokenizer = T5Tokenizer.from_pretrained( 92 | base_model_path, subfolder="tokenizer" 93 | ) 94 | text_encoder = T5EncoderModel.from_pretrained( 95 | base_model_path, subfolder="text_encoder" 96 | ) 97 | transformer = CustomCogVideoXTransformer3DModel.from_pretrained( 98 | base_model_path, subfolder="transformer" 99 | ) 100 | vae = AutoencoderKLCogVideoX.from_pretrained( 101 | base_model_path, subfolder="vae" 102 | ) 103 | scheduler = CogVideoXDDIMScheduler.from_pretrained( 104 | base_model_path, subfolder="scheduler" 105 | ) 106 | controlnet = CogVideoXControlnet.from_pretrained( 107 | controlnet_model_path 108 | ) 109 | 110 | pipe = ControlnetCogVideoXPipeline( 111 | tokenizer=tokenizer, 112 | text_encoder=text_encoder, 113 | transformer=transformer, 114 | vae=vae, 115 | controlnet=controlnet, 116 | scheduler=scheduler, 117 | ) 118 | video = load_video(video_path)[:49] 119 | controlnet_frames = [controlnet_processor(x) for x in video] 120 | # If you're using with lora, add this code 121 | if lora_path: 122 | pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1") 123 | pipe.fuse_lora(lora_scale=1 / lora_rank) 124 | 125 | # 2. Set Scheduler. 126 | # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`. 127 | # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B. 128 | # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V. 129 | 130 | # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") 131 | pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") 132 | 133 | # 3. Enable CPU offload for the model. 134 | # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference 135 | # and enable to("cuda") 136 | 137 | # pipe.to("cuda") 138 | pipe = pipe.to(dtype=dtype) 139 | pipe.enable_sequential_cpu_offload() 140 | 141 | pipe.vae.enable_slicing() 142 | pipe.vae.enable_tiling() 143 | 144 | # 4. Generate the video frames based on the prompt. 145 | # `num_frames` is the Number of frames to generate. 146 | # This is the default value for 6 seconds video and 8 fps and will plus 1 frame for the first frame and 49 frames. 147 | video_generate = pipe( 148 | prompt=prompt, 149 | controlnet_frames=controlnet_frames, # The path of the image to be used as the background of the video 150 | num_videos_per_prompt=num_videos_per_prompt, # Number of videos to generate per prompt 151 | num_inference_steps=num_inference_steps, # Number of inference steps 152 | num_frames=49, # Number of frames to generate，changed to 49 for diffusers version `0.30.3` and after. 153 | use_dynamic_cfg=True, # This id used for DPM Sechduler, for DDIM scheduler, it should be False 154 | guidance_scale=guidance_scale, 155 | generator=torch.Generator().manual_seed(seed), # Set the seed for reproducibility 156 | controlnet_weights=controlnet_weights, 157 | controlnet_guidance_start=controlnet_guidance_start, 158 | controlnet_guidance_end=controlnet_guidance_end, 159 | ).frames[0] 160 | 161 | # 5. Export the generated frames to a video file. fps must be 8 for original video. 162 | export_to_video(video_generate, output_path, fps=8) 163 | 164 | 165 | if __name__ == "__main__": 166 | parser = argparse.ArgumentParser(description="Generate a video from a text prompt using CogVideoX") 167 | parser.add_argument("--prompt", type=str, required=True, help="The description of the video to be generated") 168 | parser.add_argument( 169 | "--video_path", 170 | type=str, 171 | required=True, 172 | help="The path of the video for controlnet processing.", 173 | ) 174 | parser.add_argument( 175 | "--base_model_path", type=str, default="THUDM/CogVideoX-5b", help="The path of the pre-trained model to be used" 176 | ) 177 | parser.add_argument( 178 | "--controlnet_model_path", type=str, default="TheDenk/cogvideox-5b-controlnet-hed-v1", help="The path of the controlnet pre-trained model to be used" 179 | ) 180 | parser.add_argument("--controlnet_type", type=str, default='hed', help="Type of controlnet model (e.g. canny, hed)") 181 | parser.add_argument("--controlnet_weights", type=float, default=0.8, help="Strenght of controlnet") 182 | parser.add_argument("--controlnet_guidance_start", type=float, default=0.0, help="The stage when the controlnet starts to be applied") 183 | parser.add_argument("--controlnet_guidance_end", type=float, default=0.5, help="The stage when the controlnet end to be applied") 184 | parser.add_argument("--lora_path", type=str, default=None, help="The path of the LoRA weights to be used") 185 | parser.add_argument("--lora_rank", type=int, default=128, help="The rank of the LoRA weights") 186 | parser.add_argument( 187 | "--output_path", type=str, default="./output.mp4", help="The path where the generated video will be saved" 188 | ) 189 | parser.add_argument("--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance") 190 | parser.add_argument( 191 | "--num_inference_steps", type=int, default=50, help="Number of steps for the inference process" 192 | ) 193 | parser.add_argument("--num_videos_per_prompt", type=int, default=1, help="Number of videos to generate per prompt") 194 | parser.add_argument( 195 | "--dtype", type=str, default="bfloat16", help="The data type for computation (e.g., 'float16' or 'bfloat16')" 196 | ) 197 | parser.add_argument("--seed", type=int, default=42, help="The seed for reproducibility") 198 | 199 | args = parser.parse_args() 200 | dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 201 | generate_video( 202 | prompt=args.prompt, 203 | video_path=args.video_path, 204 | base_model_path=args.base_model_path, 205 | controlnet_model_path=args.controlnet_model_path, 206 | controlnet_type=args.controlnet_type, 207 | controlnet_weights=args.controlnet_weights, 208 | controlnet_guidance_start=args.controlnet_guidance_start, 209 | controlnet_guidance_end=args.controlnet_guidance_end, 210 | lora_path=args.lora_path, 211 | lora_rank=args.lora_rank, 212 | output_path=args.output_path, 213 | num_inference_steps=args.num_inference_steps, 214 | guidance_scale=args.guidance_scale, 215 | num_videos_per_prompt=args.num_videos_per_prompt, 216 | dtype=dtype, 217 | seed=args.seed, 218 | ) 219 | -------------------------------------------------------------------------------- /inference/gradio_web_demo.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import threading 4 | import time 5 | 6 | import gradio as gr 7 | import torch 8 | from transformers import T5EncoderModel, T5Tokenizer 9 | from diffusers.utils import export_to_video, load_video 10 | from diffusers import ( 11 | CogVideoXDDIMScheduler, 12 | CogVideoXDPMScheduler, 13 | AutoencoderKLCogVideoX 14 | ) 15 | from datetime import datetime, timedelta 16 | import moviepy.editor as mp 17 | from controlnet_aux import HEDdetector, CannyDetector 18 | 19 | from controlnet_pipeline import ControlnetCogVideoXPipeline 20 | from cogvideo_transformer import CustomCogVideoXTransformer3DModel 21 | from cogvideo_controlnet import CogVideoXControlnet 22 | 23 | 24 | os.makedirs("./output", exist_ok=True) 25 | os.makedirs("./gradio_tmp", exist_ok=True) 26 | 27 | controlnet_mapping = { 28 | 'hed': HEDdetector, 29 | 'canny': CannyDetector, 30 | } 31 | 32 | 33 | def init_controlnet_processor(controlnet_type): 34 | if controlnet_type in ['canny', 'lineart']: 35 | return controlnet_mapping[controlnet_type]() 36 | return controlnet_mapping[controlnet_type].from_pretrained('lllyasviel/Annotators').to(device='cuda') 37 | 38 | 39 | def save_video(tensor): 40 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") 41 | video_path = f"./output/{timestamp}.mp4" 42 | os.makedirs(os.path.dirname(video_path), exist_ok=True) 43 | export_to_video(tensor, video_path) 44 | return video_path 45 | 46 | 47 | def convert_to_gif(video_path): 48 | clip = mp.VideoFileClip(video_path) 49 | clip = clip.set_fps(8) 50 | clip = clip.resize(height=240) 51 | gif_path = video_path.replace(".mp4", ".gif") 52 | clip.write_gif(gif_path, fps=8) 53 | return gif_path 54 | 55 | 56 | def delete_old_files(): 57 | while True: 58 | now = datetime.now() 59 | cutoff = now - timedelta(minutes=10) 60 | directories = ["./output", "./gradio_tmp"] 61 | 62 | for directory in directories: 63 | for filename in os.listdir(directory): 64 | file_path = os.path.join(directory, filename) 65 | if os.path.isfile(file_path): 66 | file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path)) 67 | if file_mtime < cutoff: 68 | os.remove(file_path) 69 | time.sleep(600) 70 | 71 | 72 | threading.Thread(target=delete_old_files, daemon=True).start() 73 | 74 | def main(args): 75 | # 0. Load selected controlnet processor 76 | controlnet_processor = init_controlnet_processor(args.controlnet_type) 77 | # 1. Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16). 78 | tokenizer = T5Tokenizer.from_pretrained( 79 | args.base_model_path, subfolder="tokenizer" 80 | ) 81 | text_encoder = T5EncoderModel.from_pretrained( 82 | args.base_model_path, subfolder="text_encoder" 83 | ) 84 | transformer = CustomCogVideoXTransformer3DModel.from_pretrained( 85 | args.base_model_path, subfolder="transformer" 86 | ) 87 | vae = AutoencoderKLCogVideoX.from_pretrained( 88 | args.base_model_path, subfolder="vae" 89 | ) 90 | scheduler = CogVideoXDDIMScheduler.from_pretrained( 91 | args.base_model_path, subfolder="scheduler" 92 | ) 93 | controlnet = CogVideoXControlnet.from_pretrained( 94 | args.controlnet_model_path 95 | ) 96 | 97 | pipe = ControlnetCogVideoXPipeline( 98 | tokenizer=tokenizer, 99 | text_encoder=text_encoder, 100 | transformer=transformer, 101 | vae=vae, 102 | controlnet=controlnet, 103 | scheduler=scheduler, 104 | ) 105 | 106 | if args.lora_path: 107 | pipe.load_lora_weights(args.lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1") 108 | pipe.fuse_lora(lora_scale=1 / args.lora_rank) 109 | 110 | # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") 111 | pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") 112 | 113 | dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 114 | # pipe.to("cuda") 115 | pipe = pipe.to(dtype=dtype) 116 | pipe.enable_sequential_cpu_offload() 117 | 118 | pipe.vae.enable_slicing() 119 | pipe.vae.enable_tiling() 120 | 121 | 122 | def infer(prompt: str, controlnet_frames: list, num_inference_steps: int, guidance_scale: float, seed: int, progress=gr.Progress(track_tqdm=True)): 123 | torch.cuda.empty_cache() 124 | video = pipe( 125 | prompt=prompt, 126 | controlnet_frames=controlnet_frames, 127 | num_videos_per_prompt=1, 128 | num_inference_steps=num_inference_steps, 129 | num_frames=49, 130 | guidance_scale=guidance_scale, 131 | generator=torch.Generator().manual_seed(seed), 132 | ).frames[0] 133 | 134 | return video 135 | 136 | with gr.Blocks() as demo: 137 | gr.Markdown(""" 138 |