'update'

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,)

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()
 

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()
 

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,)

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)
+    

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