Update teacache.py

teacache.py

@@ -5,135 +5,142 @@ from typing import Optional, Dict, Union, Any
 from functools import wraps
 
 class TeaCacheConfig:
-    """Configuration for TeaCache acceleration"""
-    def __init__(
-        self,
-        rel_l1_thresh: float = 0.15,
-        enable: bool = True
-    ):
-        self.rel_l1_thresh = rel_l1_thresh
-        self.enable = enable
-        self._reset_state()
-    
-    def _reset_state(self):
-        """Reset internal state"""
-        self.cnt = 0
-        self.accumulated_rel_l1_distance = 0
-        self.previous_modulated_input = None
-        self.previous_residual = None
+   """Configuration for TeaCache acceleration"""
+   def __init__(
+       self,
+       rel_l1_thresh: float = 0.15,
+       enable: bool = True
+   ):
+       self.rel_l1_thresh = rel_l1_thresh
+       self.enable = enable
+       self._reset_state()
+   
+   def _reset_state(self):
+       """Reset internal state"""
+       self.cnt = 0
+       self.accumulated_rel_l1_distance = 0
+       self.previous_modulated_input = None
+       self.previous_residual = None
 
 def create_teacache_forward(original_forward):
-    """Factory function to create a TeaCache-enabled forward pass"""
-    @wraps(original_forward)
-    def teacache_forward(
-        self,
-        hidden_states: torch.Tensor,
-        timestep: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        pooled_projections: Optional[torch.Tensor] = None,
-        guidance: Optional[torch.Tensor] = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = True,
-    ):
-        # Skip TeaCache if not enabled
-        if not hasattr(self, 'teacache_config') or not self.teacache_config.enable:
-            return original_forward(
-                self,
-                hidden_states=hidden_states,
-                timestep=timestep,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_attention_mask=encoder_attention_mask,
-                pooled_projections=pooled_projections,
-                guidance=guidance,
-                attention_kwargs=attention_kwargs,
-                return_dict=return_dict
-            )
+   """Factory function to create a TeaCache-enabled forward pass"""
+   @wraps(original_forward)
+   def teacache_forward(
+       self,
+       hidden_states: torch.Tensor,
+       timestep: torch.Tensor,
+       encoder_hidden_states: Optional[torch.Tensor] = None,
+       encoder_attention_mask: Optional[torch.Tensor] = None,
+       pooled_projections: Optional[torch.Tensor] = None,
+       guidance: Optional[torch.Tensor] = None,
+       attention_kwargs: Optional[Dict[str, Any]] = None,
+       return_dict: bool = True,
+   ):
+       # Skip TeaCache if not enabled
+       if not hasattr(self, 'teacache_config') or not self.teacache_config.enable:
+           return original_forward(
+               self,
+               hidden_states=hidden_states,
+               timestep=timestep,
+               encoder_hidden_states=encoder_hidden_states,
+               encoder_attention_mask=encoder_attention_mask,
+               pooled_projections=pooled_projections,
+               guidance=guidance,
+               attention_kwargs=attention_kwargs,
+               return_dict=return_dict
+           )
 
-        config = self.teacache_config
-        
-        # Initial embeddings
-        # Changed from time_proj to time_in to match HunyuanVideo implementation
-        t_emb = self.time_in(timestep)
+       config = self.teacache_config
 
-        if pooled_projections is not None:
-            t_emb = t_emb + self.vector_in(pooled_projections)
+       # Prepare modulation vectors similar to HunyuanVideo implementation
+       if pooled_projections is not None:
+           vec = self.vector_in(pooled_projections)
 
-        if guidance is not None:
-            t_emb = t_emb + self.guidance_in(guidance)
+       if guidance is not None:
+           if vec is None:
+               vec = self.guidance_in(guidance)
+           else:
+               vec = vec + self.guidance_in(guidance)
 
-        # TeaCache optimization logic
-        inp = hidden_states.clone()
-        normed_inp = self.transformer_blocks[0].norm1(inp)
-        modulated_inp = self.transformer_blocks[0].attn1.to_q(normed_inp)
+       # TeaCache optimization logic
+       inp = hidden_states.clone()
+       if hasattr(self.double_blocks[0], 'img_norm1'):
+           # HunyuanVideo specific modulation
+           img_mod1_shift, img_mod1_scale, _, _, _, _ = self.double_blocks[0].img_mod(vec).chunk(6, dim=-1)
+           normed_inp = self.double_blocks[0].img_norm1(inp)
+           modulated_inp = normed_inp * (1 + img_mod1_scale) + img_mod1_shift
+       else:
+           # Fallback modulation
+           normed_inp = self.transformer_blocks[0].norm1(inp)
+           modulated_inp = normed_inp
 
-        # Determine if we should calculate or use cache
-        should_calc = True
-        if config.cnt == 0 or config.cnt == self.num_inference_steps - 1:
-            should_calc = True
-            config.accumulated_rel_l1_distance = 0
-        elif config.previous_modulated_input is not None:
-            coefficients = [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, 
-                          -3.14987800e+00, 9.61237896e-02]
-            rescale_func = np.poly1d(coefficients)
-            
-            rel_l1 = ((modulated_inp - config.previous_modulated_input).abs().mean() / 
-                     config.previous_modulated_input.abs().mean()).cpu().item()
-            config.accumulated_rel_l1_distance += rescale_func(rel_l1)
-            
-            should_calc = config.accumulated_rel_l1_distance >= config.rel_l1_thresh
-            if should_calc:
-                config.accumulated_rel_l1_distance = 0
+       # Determine if we should calculate or use cache
+       should_calc = True
+       if config.cnt == 0 or config.cnt == self.num_inference_steps - 1:
+           should_calc = True
+           config.accumulated_rel_l1_distance = 0
+       elif config.previous_modulated_input is not None:
+           coefficients = [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, 
+                         -3.14987800e+00, 9.61237896e-02]
+           rescale_func = np.poly1d(coefficients)
+           
+           rel_l1 = ((modulated_inp - config.previous_modulated_input).abs().mean() / 
+                    config.previous_modulated_input.abs().mean()).cpu().item()
+           config.accumulated_rel_l1_distance += rescale_func(rel_l1)
+           
+           should_calc = config.accumulated_rel_l1_distance >= config.rel_l1_thresh
+           if should_calc:
+               config.accumulated_rel_l1_distance = 0
 
-        config.previous_modulated_input = modulated_inp
-        config.cnt += 1
-        if config.cnt >= self.num_inference_steps:
-            config.cnt = 0
+       config.previous_modulated_input = modulated_inp
+       config.cnt += 1
+       if config.cnt >= self.num_inference_steps:
+           config.cnt = 0
 
-        # Use cache or calculate new result
-        if not should_calc and config.previous_residual is not None:
-            hidden_states += config.previous_residual
-        else:
-            ori_hidden_states = hidden_states.clone()
-            
-            # Process through transformer blocks
-            for block in self.transformer_blocks:
-                hidden_states = block(
-                    hidden_states=hidden_states,
-                    temb=t_emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=encoder_attention_mask,
-                    **attention_kwargs if attention_kwargs else {}
-                )
+       # Use cache or calculate new result
+       if not should_calc and config.previous_residual is not None:
+           hidden_states += config.previous_residual
+       else:
+           ori_hidden_states = hidden_states.clone()
+           
+           # Use original forward pass
+           out = original_forward(
+               self,
+               hidden_states=hidden_states,
+               timestep=timestep,
+               encoder_hidden_states=encoder_hidden_states,
+               encoder_attention_mask=encoder_attention_mask,
+               pooled_projections=pooled_projections,
+               guidance=guidance,
+               attention_kwargs=attention_kwargs,
+               return_dict=True
+           )
+           hidden_states = out["sample"]
 
-            # Store residual for future use
-            config.previous_residual = hidden_states - ori_hidden_states
+           # Store residual for future use
+           config.previous_residual = hidden_states - ori_hidden_states
 
-        # Final layer normalization and projection
-        hidden_states = self.norm_out(hidden_states)
-        output = self.proj_out(hidden_states)
+       if not return_dict:
+           return (hidden_states,)
+           
+       return {"sample": hidden_states}
 
-        if not return_dict:
-            return (output,)
-            
-        return {"sample": output}
-
-    return teacache_forward
+   return teacache_forward
 
 def enable_teacache(model: Any, num_inference_steps: int, rel_l1_thresh: float = 0.15):
-    """Enable TeaCache acceleration for a model"""
-    if not hasattr(model, '_original_forward'):
-        model._original_forward = model.forward
-    
-    model.teacache_config = TeaCacheConfig(rel_l1_thresh=rel_l1_thresh)
-    model.num_inference_steps = num_inference_steps
-    model.forward = create_teacache_forward(model._original_forward).__get__(model)
+   """Enable TeaCache acceleration for a model"""
+   if not hasattr(model, '_original_forward'):
+       model._original_forward = model.forward
+   
+   model.teacache_config = TeaCacheConfig(rel_l1_thresh=rel_l1_thresh)
+   model.num_inference_steps = num_inference_steps
+   model.forward = create_teacache_forward(model._original_forward).__get__(model)
 
 def disable_teacache(model: Any):
-    """Disable TeaCache acceleration for a model"""
-    if hasattr(model, '_original_forward'):
-        model.forward = model._original_forward
-        del model._original_forward
-    
-    if hasattr(model, 'teacache_config'):
-        del model.teacache_config
+   """Disable TeaCache acceleration for a model"""
+   if hasattr(model, '_original_forward'):
+       model.forward = model._original_forward
+       del model._original_forward
+   
+   if hasattr(model, 'teacache_config'):
+       del model.teacache_config