Upload matryoshka.py

matryoshka.py

@@ -1631,10 +1631,10 @@ class MatryoshkaFusedAttnProcessor1_0_or_2_0:
             split_size = kv.shape[-1] // 2
             key, value = torch.split(kv, split_size, dim=-1)
 
-        if self_attention_output is None:
-            query = query.permute(0, 2, 1)
-        key = key.permute(0, 2, 1)
-        value = value.permute(0, 2, 1)
+        # if self_attention_output is None:
+        #     query = query.permute(0, 2, 1)
+        # key = key.permute(0, 2, 1)
+        # value = value.permute(0, 2, 1)
 
         if attn.norm_q is not None:
             query = attn.norm_q(query)
@@ -1665,7 +1665,7 @@ class MatryoshkaFusedAttnProcessor1_0_or_2_0:
         )
 
         hidden_states = hidden_states.to(query.dtype)
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, height * width, channel)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
 
         if self_attention_output is not None:
             hidden_states = hidden_states + self_attention_output