Update enhance.py

enhance.py

@@ -76,6 +76,9 @@ class LTXEnhanceAttnProcessor2_0:
         attention_mask = None,
         **kwargs
     ) -> torch.Tensor:
+        # Store original dtype first
+        orig_dtype = hidden_states.dtype
+    
         # The shape could be [batch_size, sequence_length, channels] or [batch_size, sequence_length, num_heads, head_dim]
         # We need to handle both cases
         if hidden_states.ndim == 4:
@@ -95,7 +98,11 @@ class LTXEnhanceAttnProcessor2_0:
         query = attn.to_q(hidden_states)
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
-
+    
+        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)
+    
         if attn.upcast_attention:
             query = query.float()
             key = key.float()
@@ -125,7 +132,7 @@ class LTXEnhanceAttnProcessor2_0:
         )
     
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, inner_dim)
-        hidden_states = hidden_states.to(orig_dtype)  # Ensure we're back to original dtype
+        hidden_states = hidden_states.to(orig_dtype)  # Now orig_dtype is defined
     
         if is_enhance_enabled() and enhance_scores is not None:
             hidden_states = hidden_states * enhance_scores