diff --git a/cogvideo_controlnet.py b/cogvideo_controlnet.py index 04334e9..b360d0c 100644 --- a/cogvideo_controlnet.py +++ b/cogvideo_controlnet.py @@ -14,6 +14,8 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config class CogVideoXControlnet(ModelMixin, ConfigMixin, PeftAdapterMixin): + _supports_gradient_checkpointing = True + @register_to_config def __init__( self, @@ -42,6 +44,7 @@ class CogVideoXControlnet(ModelMixin, ConfigMixin, PeftAdapterMixin): temporal_interpolation_scale: float = 1.0, use_rotary_positional_embeddings: bool = False, use_learned_positional_embeddings: bool = False, + out_proj_dim = None, ): super().__init__() inner_dim = num_attention_heads * attention_head_dim @@ -109,7 +112,16 @@ class CogVideoXControlnet(ModelMixin, ConfigMixin, PeftAdapterMixin): ] ) + self.out_projectors = None + if out_proj_dim is not None: + self.out_projectors = nn.ModuleList( + [nn.Linear(inner_dim, out_proj_dim) for _ in range(num_layers)] + ) + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value def compress_time(self, x, num_frames): x = rearrange(x, '(b f) c h w -> b f c h w', f=num_frames) @@ -197,7 +209,11 @@ class CogVideoXControlnet(ModelMixin, ConfigMixin, PeftAdapterMixin): temb=emb, image_rotary_emb=image_rotary_emb, ) - controlnet_hidden_states += (hidden_states,) + + if self.out_projectors is not None: + controlnet_hidden_states += (self.out_projectors[i](hidden_states),) + else: + controlnet_hidden_states += (hidden_states,) if not return_dict: return (controlnet_hidden_states,) diff --git a/cogvideox_fun/pipeline_cogvideox_control.py b/cogvideox_fun/pipeline_cogvideox_control.py index 545e084..85687fe 100644 --- a/cogvideox_fun/pipeline_cogvideox_control.py +++ b/cogvideox_fun/pipeline_cogvideox_control.py @@ -528,6 +528,7 @@ class CogVideoX_Fun_Pipeline_Control(VideoSysPipeline): context_stride: Optional[int] = None, context_overlap: Optional[int] = None, freenoise: Optional[bool] = True, + tora: Optional[dict] = None, ) -> Union[CogVideoX_Fun_PipelineOutput, Tuple]: """ Function invoked when calling the pipeline for generation. @@ -720,7 +721,13 @@ class CogVideoX_Fun_Pipeline_Control(VideoSysPipeline): if self.transformer.config.use_rotary_positional_embeddings else None ) - + if tora is not None and do_classifier_free_guidance: + video_flow_features = tora["video_flow_features"].repeat(1, 2, 1, 1, 1).contiguous() + + if tora is not None: + for module in self.transformer.fuser_list: + for param in module.parameters(): + param.data = param.data.to(device) with self.progress_bar(total=num_inference_steps) as progress_bar: # for DPM-solver++ @@ -910,6 +917,8 @@ class CogVideoX_Fun_Pipeline_Control(VideoSysPipeline): image_rotary_emb=image_rotary_emb, return_dict=False, control_latents=current_control_latents, + video_flow_features=video_flow_features if (tora is not None and tora["start_percent"] <= current_step_percentage <= tora["end_percent"]) else None, + )[0] noise_pred = noise_pred.float() diff --git a/cogvideox_fun/pipeline_cogvideox_inpaint.py b/cogvideox_fun/pipeline_cogvideox_inpaint.py index 459e845..f2a8680 100644 --- a/cogvideox_fun/pipeline_cogvideox_inpaint.py +++ b/cogvideox_fun/pipeline_cogvideox_inpaint.py @@ -349,7 +349,7 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): noise[:, place_idx:place_idx + delta, :, :, :] = noise[:, list_idx, :, :, :] # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(video_latents, noise, timestep) + latents = noise if is_strength_max else self.scheduler.add_noise(video_latents.to(noise), noise, timestep) # if pure noise then scale the initial latents by the Scheduler's init sigma latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents latents = latents.to(device) @@ -610,6 +610,7 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): context_stride: Optional[int] = None, context_overlap: Optional[int] = None, freenoise: Optional[bool] = True, + tora: Optional[dict] = None, ) -> Union[CogVideoX_Fun_PipelineOutput, Tuple]: """ Function invoked when calling the pipeline for generation. @@ -889,6 +890,13 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): if self.transformer.config.use_rotary_positional_embeddings else None ) + if tora is not None and do_classifier_free_guidance: + video_flow_features = tora["video_flow_features"].repeat(1, 2, 1, 1, 1).contiguous() + + if tora is not None: + for module in self.transformer.fuser_list: + for param in module.parameters(): + param.data = param.data.to(device) # 8. Denoising loop num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) @@ -980,7 +988,6 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): pbar.update(1) # ========================================== elif use_context_schedule: - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) @@ -993,6 +1000,8 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): counter = torch.zeros_like(latent_model_input) noise_pred = torch.zeros_like(latent_model_input) + current_step_percentage = i / num_inference_steps + image_rotary_emb = ( self._prepare_rotary_positional_embeddings(height, width, context_frames, device) if self.transformer.config.use_rotary_positional_embeddings @@ -1003,6 +1012,13 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): partial_latent_model_input = latent_model_input[:, c, :, :, :] partial_inpaint_latents = inpaint_latents[:, c, :, :, :] partial_inpaint_latents[:, 0, :, :, :] = inpaint_latents[:, 0, :, :, :] + if (tora is not None and tora["start_percent"] <= current_step_percentage <= tora["end_percent"]): + if do_classifier_free_guidance: + partial_video_flow_features = tora["video_flow_features"][:, c, :, :, :].repeat(1, 2, 1, 1, 1).contiguous() + else: + partial_video_flow_features = tora["video_flow_features"][:, c, :, :, :] + else: + partial_video_flow_features = None # predict noise model_output noise_pred[:, c, :, :, :] += self.transformer( @@ -1012,6 +1028,7 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): image_rotary_emb=image_rotary_emb, return_dict=False, inpaint_latents=partial_inpaint_latents, + video_flow_features=partial_video_flow_features )[0] counter[:, c, :, :, :] += 1 @@ -1061,6 +1078,8 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latent_model_input.shape[0]) + current_step_percentage = i / num_inference_steps + # predict noise model_output noise_pred = self.transformer( hidden_states=latent_model_input, @@ -1069,6 +1088,8 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline): image_rotary_emb=image_rotary_emb, return_dict=False, inpaint_latents=inpaint_latents, + video_flow_features=video_flow_features if (tora is not None and tora["start_percent"] <= current_step_percentage <= tora["end_percent"]) else None, + )[0] noise_pred = noise_pred.float() diff --git a/cogvideox_fun/transformer_3d.py b/cogvideox_fun/transformer_3d.py index 8a607b4..2b57923 100644 --- a/cogvideox_fun/transformer_3d.py +++ b/cogvideox_fun/transformer_3d.py @@ -34,6 +34,7 @@ from diffusers.models.normalization import AdaLayerNorm, CogVideoXLayerNormZero logger = logging.get_logger(__name__) # pylint: disable=invalid-name +from einops import rearrange try: from sageattention import sageattn SAGEATTN_IS_AVAVILABLE = True @@ -42,6 +43,23 @@ except: logger.info("sageattn not found, using sdpa") SAGEATTN_IS_AVAVILABLE = False +def fft(tensor): + tensor_fft = torch.fft.fft2(tensor) + tensor_fft_shifted = torch.fft.fftshift(tensor_fft) + B, C, H, W = tensor.size() + radius = min(H, W) // 5 + + Y, X = torch.meshgrid(torch.arange(H), torch.arange(W)) + center_x, center_y = W // 2, H // 2 + mask = (X - center_x) ** 2 + (Y - center_y) ** 2 <= radius ** 2 + low_freq_mask = mask.unsqueeze(0).unsqueeze(0).to(tensor.device) + high_freq_mask = ~low_freq_mask + + low_freq_fft = tensor_fft_shifted * low_freq_mask + high_freq_fft = tensor_fft_shifted * high_freq_mask + + return low_freq_fft, high_freq_fft + class CogVideoXAttnProcessor2_0: r""" Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on @@ -315,6 +333,11 @@ class CogVideoXBlock(nn.Module): encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + video_flow_feature: Optional[torch.Tensor] = None, + fuser=None, + fastercache_counter=0, + fastercache_start_step=15, + fastercache_device="cuda:0", ) -> torch.Tensor: text_seq_length = encoder_hidden_states.size(1) @@ -322,7 +345,39 @@ class CogVideoXBlock(nn.Module): norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( hidden_states, encoder_hidden_states, temb ) - + # Tora Motion-guidance Fuser + if video_flow_feature is not None: + H, W = video_flow_feature.shape[-2:] + T = norm_hidden_states.shape[1] // H // W + h = rearrange(norm_hidden_states, "B (T H W) C -> (B T) C H W", H=H, W=W) + h = fuser(h, video_flow_feature.to(h), T=T) + norm_hidden_states = rearrange(h, "(B T) C H W -> B (T H W) C", T=T) + del h, fuser + #fastercache + B = norm_hidden_states.shape[0] + if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B: + attn_hidden_states = ( + self.cached_hidden_states[1][:B] + + (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) + * 0.3 + ).to(norm_hidden_states.device, non_blocking=True) + attn_encoder_hidden_states = ( + self.cached_encoder_hidden_states[1][:B] + + (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B]) + * 0.3 + ).to(norm_hidden_states.device, non_blocking=True) + else: + attn_hidden_states, attn_encoder_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + if fastercache_counter == fastercache_start_step: + self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)] + self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)] + elif fastercache_counter > fastercache_start_step: + self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device)) + self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device)) # attention attn_hidden_states, attn_encoder_hidden_states = self.attn1( hidden_states=norm_hidden_states, @@ -497,6 +552,15 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): self.gradient_checkpointing = False + self.fuser_list = None + + self.use_fastercache = False + self.fastercache_counter = 0 + self.fastercache_start_step = 15 + self.fastercache_lf_step = 40 + self.fastercache_hf_step = 30 + self.fastercache_device = "cuda" + def _set_gradient_checkpointing(self, module, value=False): self.gradient_checkpointing = value @@ -609,6 +673,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): inpaint_latents: Optional[torch.Tensor] = None, control_latents: Optional[torch.Tensor] = None, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + video_flow_features: Optional[torch.Tensor] = None, return_dict: bool = True, ): batch_size, num_frames, channels, height, width = hidden_states.shape @@ -649,50 +714,101 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): encoder_hidden_states = hidden_states[:, :text_seq_length] hidden_states = hidden_states[:, text_seq_length:] - # 4. Transformer blocks - for i, block in enumerate(self.transformer_blocks): - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - encoder_hidden_states, - emb, - image_rotary_emb, - **ckpt_kwargs, + if self.use_fastercache: + self.fastercache_counter+=1 + if self.fastercache_counter >= self.fastercache_start_step + 3 and self.fastercache_counter % 5 !=0: + # 4. Transformer blocks + for i, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = block( + hidden_states=hidden_states[:1], + encoder_hidden_states=encoder_hidden_states[:1], + temb=emb[:1], + image_rotary_emb=image_rotary_emb, + video_flow_feature=video_flow_features[i][:1] if video_flow_features is not None else None, + fuser = self.fuser_list[i] if self.fuser_list is not None else None, + fastercache_counter = self.fastercache_counter, + fastercache_device = self.fastercache_device ) + + if not self.config.use_rotary_positional_embeddings: + # CogVideoX-2B + hidden_states = self.norm_final(hidden_states) else: + # CogVideoX-5B + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + hidden_states = self.norm_final(hidden_states) + hidden_states = hidden_states[:, text_seq_length:] + + # 5. Final block + hidden_states = self.norm_out(hidden_states, temb=emb[:1]) + hidden_states = self.proj_out(hidden_states) + + # 6. Unpatchify + p = self.config.patch_size + output = hidden_states.reshape(1, num_frames, height // p, width // p, channels, p, p) + output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) + + (bb, tt, cc, hh, ww) = output.shape + cond = rearrange(output, "B T C H W -> (B T) C H W", B=bb, C=cc, T=tt, H=hh, W=ww) + lf_c, hf_c = fft(cond.float()) + #lf_step = 40 + #hf_step = 30 + if self.fastercache_counter <= self.fastercache_lf_step: + self.delta_lf = self.delta_lf * 1.1 + if self.fastercache_counter >= self.fastercache_hf_step: + self.delta_hf = self.delta_hf * 1.1 + + new_hf_uc = self.delta_hf + hf_c + new_lf_uc = self.delta_lf + lf_c + + combine_uc = new_lf_uc + new_hf_uc + combined_fft = torch.fft.ifftshift(combine_uc) + recovered_uncond = torch.fft.ifft2(combined_fft).real + recovered_uncond = rearrange(recovered_uncond.to(output.dtype), "(B T) C H W -> B T C H W", B=bb, C=cc, T=tt, H=hh, W=ww) + output = torch.cat([output, recovered_uncond]) + else: + # 4. Transformer blocks + for i, block in enumerate(self.transformer_blocks): hidden_states, encoder_hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=emb, image_rotary_emb=image_rotary_emb, + video_flow_feature=video_flow_features[i] if video_flow_features is not None else None, + fuser = self.fuser_list[i] if self.fuser_list is not None else None, + fastercache_counter = self.fastercache_counter, + fastercache_device = self.fastercache_device ) - if not self.config.use_rotary_positional_embeddings: - # CogVideoX-2B - hidden_states = self.norm_final(hidden_states) - else: - # CogVideoX-5B - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - hidden_states = self.norm_final(hidden_states) - hidden_states = hidden_states[:, text_seq_length:] + if not self.config.use_rotary_positional_embeddings: + # CogVideoX-2B + hidden_states = self.norm_final(hidden_states) + else: + # CogVideoX-5B + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + hidden_states = self.norm_final(hidden_states) + hidden_states = hidden_states[:, text_seq_length:] - # 5. Final block - hidden_states = self.norm_out(hidden_states, temb=emb) - hidden_states = self.proj_out(hidden_states) + # 5. Final block + hidden_states = self.norm_out(hidden_states, temb=emb) + hidden_states = self.proj_out(hidden_states) - # 6. Unpatchify - p = self.config.patch_size - output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, channels, p, p) - output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) + # 6. Unpatchify + p = self.config.patch_size + output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, channels, p, p) + output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) + + if self.fastercache_counter >= self.fastercache_start_step + 1: + (bb, tt, cc, hh, ww) = output.shape + cond = rearrange(output[0:1].float(), "B T C H W -> (B T) C H W", B=bb//2, C=cc, T=tt, H=hh, W=ww) + uncond = rearrange(output[1:2].float(), "B T C H W -> (B T) C H W", B=bb//2, C=cc, T=tt, H=hh, W=ww) + + lf_c, hf_c = fft(cond) + lf_uc, hf_uc = fft(uncond) + + self.delta_lf = lf_uc - lf_c + self.delta_hf = hf_uc - hf_c + if not return_dict: return (output,) diff --git a/custom_cogvideox_transformer_3d.py b/custom_cogvideox_transformer_3d.py index aa6a3fb..dc33eec 100644 --- a/custom_cogvideox_transformer_3d.py +++ b/custom_cogvideox_transformer_3d.py @@ -20,9 +20,10 @@ from torch import nn import torch.nn.functional as F import numpy as np +from einops import rearrange from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.utils import is_torch_version, logging +from diffusers.utils import logging from diffusers.utils.torch_utils import maybe_allow_in_graph from diffusers.models.attention import Attention, FeedForward from diffusers.models.attention_processor import AttentionProcessor @@ -36,11 +37,28 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name try: from sageattention import sageattn - SAGEATTN_IS_AVAVILABLE = True + SAGEATTN_IS_AVAILABLE = True logger.info("Using sageattn") except: logger.info("sageattn not found, using sdpa") - SAGEATTN_IS_AVAVILABLE = False + SAGEATTN_IS_AVAILABLE = False + +def fft(tensor): + tensor_fft = torch.fft.fft2(tensor) + tensor_fft_shifted = torch.fft.fftshift(tensor_fft) + B, C, H, W = tensor.size() + radius = min(H, W) // 5 + + Y, X = torch.meshgrid(torch.arange(H), torch.arange(W)) + center_x, center_y = W // 2, H // 2 + mask = (X - center_x) ** 2 + (Y - center_y) ** 2 <= radius ** 2 + low_freq_mask = mask.unsqueeze(0).unsqueeze(0).to(tensor.device) + high_freq_mask = ~low_freq_mask + + low_freq_fft = tensor_fft_shifted * low_freq_mask + high_freq_fft = tensor_fft_shifted * high_freq_mask + + return low_freq_fft, high_freq_fft class CogVideoXAttnProcessor2_0: r""" @@ -51,7 +69,7 @@ class CogVideoXAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - + @torch.compiler.disable() def __call__( self, attn: Attention, @@ -96,7 +114,7 @@ class CogVideoXAttnProcessor2_0: if not attn.is_cross_attention: key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb) - if SAGEATTN_IS_AVAVILABLE: + if SAGEATTN_IS_AVAILABLE: hidden_states = sageattn(query, key, value, is_causal=False) else: hidden_states = F.scaled_dot_product_attention( @@ -125,7 +143,7 @@ class FusedCogVideoXAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - + @torch.compiler.disable() def __call__( self, attn: Attention, @@ -170,7 +188,7 @@ class FusedCogVideoXAttnProcessor2_0: if not attn.is_cross_attention: key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb) - if SAGEATTN_IS_AVAVILABLE: + if SAGEATTN_IS_AVAILABLE: hidden_states = sageattn(query, key, value, is_causal=False) else: hidden_states = F.scaled_dot_product_attention( @@ -191,6 +209,7 @@ class FusedCogVideoXAttnProcessor2_0: @maybe_allow_in_graph class CogVideoXBlock(nn.Module): + r""" Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model. @@ -269,27 +288,60 @@ class CogVideoXBlock(nn.Module): inner_dim=ff_inner_dim, bias=ff_bias, ) - + self.cached_hidden_states = [] + self.cached_encoder_hidden_states = [] + def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + video_flow_feature: Optional[torch.Tensor] = None, + fuser=None, + fastercache_counter=0, + fastercache_start_step=15, + fastercache_device="cuda:0", ) -> torch.Tensor: text_seq_length = encoder_hidden_states.size(1) - # norm & modulate norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( hidden_states, encoder_hidden_states, temb ) - - # attention - attn_hidden_states, attn_encoder_hidden_states = self.attn1( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - ) + # Tora Motion-guidance Fuser + if video_flow_feature is not None: + H, W = video_flow_feature.shape[-2:] + T = norm_hidden_states.shape[1] // H // W + h = rearrange(norm_hidden_states, "B (T H W) C -> (B T) C H W", H=H, W=W) + h = fuser(h, video_flow_feature.to(h), T=T) + norm_hidden_states = rearrange(h, "(B T) C H W -> B (T H W) C", T=T) + del h, fuser + + #fastercache + B = norm_hidden_states.shape[0] + if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B: + attn_hidden_states = ( + self.cached_hidden_states[1][:B] + + (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) + * 0.3 + ).to(norm_hidden_states.device, non_blocking=True) + attn_encoder_hidden_states = ( + self.cached_encoder_hidden_states[1][:B] + + (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B]) + * 0.3 + ).to(norm_hidden_states.device, non_blocking=True) + else: + attn_hidden_states, attn_encoder_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + if fastercache_counter == fastercache_start_step: + self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)] + self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)] + elif fastercache_counter > fastercache_start_step: + self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device)) + self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device)) hidden_states = hidden_states + gate_msa * attn_hidden_states encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states @@ -458,6 +510,14 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): self.gradient_checkpointing = False + self.fuser_list = None + self.use_fastercache = False + self.fastercache_counter = 0 + self.fastercache_start_step = 15 + self.fastercache_lf_step = 40 + self.fastercache_hf_step = 30 + self.fastercache_device = "cuda" + def _set_gradient_checkpointing(self, module, value=False): self.gradient_checkpointing = value @@ -570,6 +630,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, controlnet_states: torch.Tensor = None, controlnet_weights: Optional[Union[float, int, list, np.ndarray, torch.FloatTensor]] = 1.0, + video_flow_features: Optional[torch.Tensor] = None, return_dict: bool = True, ): batch_size, num_frames, channels, height, width = hidden_states.shape @@ -591,32 +652,82 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): text_seq_length = encoder_hidden_states.shape[1] encoder_hidden_states = hidden_states[:, :text_seq_length] hidden_states = hidden_states[:, text_seq_length:] + if self.use_fastercache: + self.fastercache_counter+=1 + if self.fastercache_counter >= self.fastercache_start_step + 3 and self.fastercache_counter % 5 !=0: + # 3. Transformer blocks + for i, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = block( + hidden_states=hidden_states[:1], + encoder_hidden_states=encoder_hidden_states[:1], + temb=emb[:1], + image_rotary_emb=image_rotary_emb, + video_flow_feature=video_flow_features[i][:1] if video_flow_features is not None else None, + fuser = self.fuser_list[i] if self.fuser_list is not None else None, + fastercache_counter = self.fastercache_counter, + fastercache_device = self.fastercache_device + ) - # 3. Transformer blocks - for i, block in enumerate(self.transformer_blocks): - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - encoder_hidden_states, - emb, - image_rotary_emb, - **ckpt_kwargs, - ) + if (controlnet_states is not None) and (i < len(controlnet_states)): + controlnet_states_block = controlnet_states[i] + controlnet_block_weight = 1.0 + if isinstance(controlnet_weights, (list, np.ndarray)) or torch.is_tensor(controlnet_weights): + controlnet_block_weight = controlnet_weights[i] + elif isinstance(controlnet_weights, (float, int)): + controlnet_block_weight = controlnet_weights + + hidden_states = hidden_states + controlnet_states_block * controlnet_block_weight + + if not self.config.use_rotary_positional_embeddings: + # CogVideoX-2B + hidden_states = self.norm_final(hidden_states) else: + # CogVideoX-5B + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + hidden_states = self.norm_final(hidden_states) + hidden_states = hidden_states[:, text_seq_length:] + + # 4. Final block + hidden_states = self.norm_out(hidden_states, temb=emb[:1]) + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + # Note: we use `-1` instead of `channels`: + # - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels) + # - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels) + p = self.config.patch_size + output = hidden_states.reshape(1, num_frames, height // p, width // p, -1, p, p) + output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) + + (bb, tt, cc, hh, ww) = output.shape + cond = rearrange(output, "B T C H W -> (B T) C H W", B=bb, C=cc, T=tt, H=hh, W=ww) + lf_c, hf_c = fft(cond.float()) + #lf_step = 40 + #hf_step = 30 + if self.fastercache_counter <= self.fastercache_lf_step: + self.delta_lf = self.delta_lf * 1.1 + if self.fastercache_counter >= self.fastercache_hf_step: + self.delta_hf = self.delta_hf * 1.1 + + new_hf_uc = self.delta_hf + hf_c + new_lf_uc = self.delta_lf + lf_c + + combine_uc = new_lf_uc + new_hf_uc + combined_fft = torch.fft.ifftshift(combine_uc) + recovered_uncond = torch.fft.ifft2(combined_fft).real + recovered_uncond = rearrange(recovered_uncond.to(output.dtype), "(B T) C H W -> B T C H W", B=bb, C=cc, T=tt, H=hh, W=ww) + output = torch.cat([output, recovered_uncond]) + else: + for i, block in enumerate(self.transformer_blocks): hidden_states, encoder_hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=emb, image_rotary_emb=image_rotary_emb, + video_flow_feature=video_flow_features[i] if video_flow_features is not None else None, + fuser = self.fuser_list[i] if self.fuser_list is not None else None, + fastercache_counter = self.fastercache_counter, + fastercache_device = self.fastercache_device ) if (controlnet_states is not None) and (i < len(controlnet_states)): @@ -628,28 +739,40 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): controlnet_block_weight = controlnet_weights hidden_states = hidden_states + controlnet_states_block * controlnet_block_weight + + if not self.config.use_rotary_positional_embeddings: + # CogVideoX-2B + hidden_states = self.norm_final(hidden_states) + else: + # CogVideoX-5B + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + hidden_states = self.norm_final(hidden_states) + hidden_states = hidden_states[:, text_seq_length:] - if not self.config.use_rotary_positional_embeddings: - # CogVideoX-2B - hidden_states = self.norm_final(hidden_states) - else: - # CogVideoX-5B - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - hidden_states = self.norm_final(hidden_states) - hidden_states = hidden_states[:, text_seq_length:] + # 4. Final block + hidden_states = self.norm_out(hidden_states, temb=emb) + hidden_states = self.proj_out(hidden_states) - # 4. Final block - hidden_states = self.norm_out(hidden_states, temb=emb) - hidden_states = self.proj_out(hidden_states) + # 5. Unpatchify + # Note: we use `-1` instead of `channels`: + # - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels) + # - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels) + p = self.config.patch_size + output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p) + output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) - # 5. Unpatchify - # Note: we use `-1` instead of `channels`: - # - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels) - # - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels) - p = self.config.patch_size - output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p) - output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) + if self.fastercache_counter >= self.fastercache_start_step + 1: + (bb, tt, cc, hh, ww) = output.shape + cond = rearrange(output[0:1].float(), "B T C H W -> (B T) C H W", B=bb//2, C=cc, T=tt, H=hh, W=ww) + uncond = rearrange(output[1:2].float(), "B T C H W -> (B T) C H W", B=bb//2, C=cc, T=tt, H=hh, W=ww) + + lf_c, hf_c = fft(cond) + lf_uc, hf_uc = fft(uncond) + + self.delta_lf = lf_uc - lf_c + self.delta_hf = hf_uc - hf_c if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) + \ No newline at end of file diff --git a/pipeline_cogvideox.py b/pipeline_cogvideox.py index 64208f0..d787f33 100644 --- a/pipeline_cogvideox.py +++ b/pipeline_cogvideox.py @@ -164,6 +164,8 @@ class CogVideoXPipeline(VideoSysPipeline): if pab_config is not None: set_pab_manager(pab_config) + self.input_with_padding = True + def prepare_latents( self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, timesteps, denoise_strength, num_inference_steps, latents=None, freenoise=True, context_size=None, context_overlap=None @@ -381,13 +383,13 @@ class CogVideoXPipeline(VideoSysPipeline): prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, device = torch.device("cuda"), - scheduler_name: str = "DPM", context_schedule: Optional[str] = None, context_frames: Optional[int] = None, context_stride: Optional[int] = None, context_overlap: Optional[int] = None, freenoise: Optional[bool] = True, controlnet: Optional[dict] = None, + tora: Optional[dict] = None, ): """ @@ -501,16 +503,33 @@ class CogVideoXPipeline(VideoSysPipeline): # 5.5. if image_cond_latents is not None: - padding_shape = ( + if image_cond_latents.shape[1] > 1: + logger.info("More than one image conditioning frame received, interpolating") + padding_shape = ( batch_size, - (latents.shape[1] - 1), + (latents.shape[1] - 2), self.vae.config.latent_channels, height // self.vae_scale_factor_spatial, width // self.vae_scale_factor_spatial, - ) - latent_padding = torch.zeros(padding_shape, device=device, dtype=self.vae.dtype) - image_cond_latents = torch.cat([image_cond_latents, latent_padding], dim=1) - + ) + latent_padding = torch.zeros(padding_shape, device=device, dtype=self.vae.dtype) + + image_cond_latents = torch.cat([image_cond_latents[:, 0, :, :, :].unsqueeze(1), latent_padding, image_cond_latents[:, -1, :, :, :].unsqueeze(1)], dim=1) + logger.info(f"image cond latents shape: {image_cond_latents.shape}") + else: + logger.info("Only one image conditioning frame received, img2vid") + if self.input_with_padding: + padding_shape = ( + batch_size, + (latents.shape[1] - 1), + self.vae.config.latent_channels, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + latent_padding = torch.zeros(padding_shape, device=device, dtype=self.vae.dtype) + image_cond_latents = torch.cat([image_cond_latents, latent_padding], dim=1) + else: + image_cond_latents = image_cond_latents.repeat(1, latents.shape[1], 1, 1, 1) # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) @@ -524,15 +543,15 @@ class CogVideoXPipeline(VideoSysPipeline): # masks if self.original_mask is not None: mask = self.original_mask.to(device) - print("self.original_mask: ", self.original_mask.shape) + logger.info(f"self.original_mask: {self.original_mask.shape}") mask = F.interpolate(self.original_mask.unsqueeze(1), size=(latents.shape[-2], latents.shape[-1]), mode='bilinear', align_corners=False) if mask.shape[0] != latents.shape[1]: mask = mask.unsqueeze(1).repeat(1, latents.shape[1], 16, 1, 1) else: mask = mask.unsqueeze(0).repeat(1, 1, 16, 1, 1) - print("latents: ", latents.shape) - print("mask: ", mask.shape) + logger.info(f"latents: {latents.shape}") + logger.info(f"mask: {mask.shape}") # 7. Denoising loop num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) @@ -544,11 +563,11 @@ class CogVideoXPipeline(VideoSysPipeline): t_tile_overlap = context_overlap t_tile_weights = self._gaussian_weights(t_tile_length=t_tile_length, t_batch_size=1).to(latents.device).to(self.vae.dtype) use_temporal_tiling = True - print("Temporal tiling enabled") + logger.info("Temporal tiling enabled") elif context_schedule is not None: if image_cond_latents is not None: raise NotImplementedError("Context schedule not currently supported with image conditioning") - print(f"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap") + logger.info(f"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap") use_temporal_tiling = False use_context_schedule = True from .cogvideox_fun.context import get_context_scheduler @@ -557,15 +576,17 @@ class CogVideoXPipeline(VideoSysPipeline): else: use_temporal_tiling = False use_context_schedule = False - print("Temporal tiling and context schedule disabled") + logger.info("Temporal tiling and context schedule disabled") # 7. Create rotary embeds if required image_rotary_emb = ( self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) if self.transformer.config.use_rotary_positional_embeddings else None ) - # 9. Controlnet + if tora is not None and do_classifier_free_guidance: + video_flow_features = tora["video_flow_features"].repeat(1, 2, 1, 1, 1).contiguous() + # 9. Controlnet if controlnet is not None: self.controlnet = controlnet["control_model"].to(device) if self.transformer.dtype == torch.float8_e4m3fn: @@ -584,13 +605,18 @@ class CogVideoXPipeline(VideoSysPipeline): control_frames = controlnet["control_frames"].to(device).to(self.controlnet.dtype).contiguous() control_frames = torch.cat([control_frames] * 2) if do_classifier_free_guidance else control_frames control_weights = controlnet["control_weights"] - print("Controlnet enabled with weights: ", control_weights) + logger.info(f"Controlnet enabled with weights: {control_weights}") control_start = controlnet["control_start"] control_end = controlnet["control_end"] else: controlnet_states = None control_weights= None + if tora is not None: + for module in self.transformer.fuser_list: + for param in module.parameters(): + param.data = param.data.to(device) + # 10. Denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: old_pred_original_sample = None # for DPM-solver++ @@ -764,6 +790,13 @@ class CogVideoXPipeline(VideoSysPipeline): else: for c in context_queue: partial_latent_model_input = latent_model_input[:, c, :, :, :] + if (tora is not None and tora["start_percent"] <= current_step_percentage <= tora["end_percent"]): + if do_classifier_free_guidance: + partial_video_flow_features = tora["video_flow_features"][:, c, :, :, :].repeat(1, 2, 1, 1, 1).contiguous() + else: + partial_video_flow_features = tora["video_flow_features"][:, c, :, :, :] + else: + partial_video_flow_features = None # predict noise model_output noise_pred[:, c, :, :, :] += self.transformer( @@ -771,6 +804,7 @@ class CogVideoXPipeline(VideoSysPipeline): encoder_hidden_states=prompt_embeds, timestep=timestep, image_rotary_emb=image_rotary_emb, + video_flow_features=partial_video_flow_features, return_dict=False )[0] @@ -829,7 +863,8 @@ class CogVideoXPipeline(VideoSysPipeline): if isinstance(controlnet_states, (tuple, list)): controlnet_states = [x.to(dtype=self.vae.dtype) for x in controlnet_states] else: - controlnet_states = controlnet_states.to(dtype=self.vae.dtype) + controlnet_states = controlnet_states.to(dtype=self.vae.dtype) + # predict noise model_output noise_pred = self.transformer( @@ -840,6 +875,8 @@ class CogVideoXPipeline(VideoSysPipeline): return_dict=False, controlnet_states=controlnet_states, controlnet_weights=control_weights, + video_flow_features=video_flow_features if (tora is not None and tora["start_percent"] <= current_step_percentage <= tora["end_percent"]) else None, + )[0] noise_pred = noise_pred.float() @@ -887,4 +924,4 @@ class CogVideoXPipeline(VideoSysPipeline): # Offload all models self.maybe_free_model_hooks() - return latents + return latents \ No newline at end of file