Update enhance.py

enhance.py

@@ -73,7 +73,7 @@ class LTXEnhanceAttnProcessor2_0:
         enhance_scores = mean_scores.mean() * (num_frames + 4.0)
         enhance_scores = enhance_scores.clamp(min=1)
         return enhance_scores
-
+    
     def __call__(
         self,
         attn: Attention,
@@ -92,12 +92,13 @@ class LTXEnhanceAttnProcessor2_0:
         num_heads = attn.heads
         head_dim = inner_dim // num_heads
     
+        # Get query, key, value projections
         query = attn.to_q(hidden_states)
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
     
-        # Reshape query, key, value to match expected dimensions
-        query = query.view(batch_size, -1, num_heads, head_dim).transpose(1, 2)
+        # Reshape projections
+        query = query.view(batch_size, sequence_length, num_heads, head_dim).transpose(1, 2)
         key = key.view(batch_size, -1, num_heads, head_dim).transpose(1, 2)
         value = value.view(batch_size, -1, num_heads, head_dim).transpose(1, 2)
     
@@ -105,6 +106,7 @@ class LTXEnhanceAttnProcessor2_0:
             query = query.float()
             key = key.float()
     
+        # Process attention
         enhance_scores = None
         if is_enhance_enabled():
             try:
@@ -118,6 +120,12 @@ class LTXEnhanceAttnProcessor2_0:
             except ValueError as e:
                 print(f"Warning: Could not calculate enhance scores: {e}")
     
+        # Make sure attention_mask has correct shape
+        if attention_mask is not None:
+            attention_mask = attention_mask.view(batch_size, 1, 1, attention_mask.shape[-1])
+            attention_mask = attention_mask.expand(-1, num_heads, -1, -1)
+    
+        # Compute attention with correct shapes
         hidden_states = torch.nn.functional.scaled_dot_product_attention(
             query, key, value,
             attn_mask=attention_mask,
@@ -125,6 +133,7 @@ class LTXEnhanceAttnProcessor2_0:
             is_causal=False
         )
     
+        # Reshape output
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, inner_dim)
         hidden_states = hidden_states.to(query.dtype)
     
@@ -132,11 +141,12 @@ class LTXEnhanceAttnProcessor2_0:
         if is_enhance_enabled() and enhance_scores is not None:
             hidden_states = hidden_states * enhance_scores
     
+        # Output projection
         hidden_states = attn.to_out[0](hidden_states)
         hidden_states = attn.to_out[1](hidden_states)
     
         return hidden_states
-
+        
 def inject_enhance_for_ltx(model: nn.Module) -> None:
     """
     Inject enhance score for LTX model.