Corrected einops removal

Author: guiyrt

src/diffusers/models/attention_processor.py

@@ -5155,11 +5155,11 @@ def __call__(
             ip_value = self.to_v_ip(norm_ip_hidden_states)
 
             # Reshape
-            img_query = img_query.view(batch_size, head_dim, attn.heads, -1).transpose(1,2)            
-            img_key = img_key.view(batch_size, head_dim, attn.heads, -1).transpose(1,2)
-            img_value = img_value.view(batch_size, head_dim, attn.heads, -1).transpose(1,2)
-            ip_key = ip_key.view(batch_size, head_dim, attn.heads, -1).transpose(1,2)
-            ip_value = ip_value.view(batch_size, head_dim, attn.heads, -1).transpose(1,2)
+            img_query = img_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            img_key = img_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            img_value = img_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
 
             # Norm
             img_query = self.norm_q(img_query)
@@ -5171,7 +5171,7 @@ def __call__(
             img_value = torch.cat([img_value, ip_value], dim=2)
 
             ip_hidden_states = F.scaled_dot_product_attention(img_query, img_key, img_value, dropout_p=0.0, is_causal=False)
-            ip_hidden_states = ip_hidden_states.transpose(1,2).view(batch_size, head_dim, -1)
+            ip_hidden_states = ip_hidden_states.transpose(1, 2).view(batch_size, -1, attn.heads * head_dim)
             ip_hidden_states = ip_hidden_states.to(img_query.dtype)
 
             hidden_states = hidden_states + ip_hidden_states * self.scale

src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -994,8 +994,8 @@ def __call__(
                     timestep=timestep,
                     encoder_hidden_states=prompt_embeds,
                     pooled_projections=pooled_prompt_embeds,
-                    return_dict=False,
                     joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
                 )[0]
 
                 # perform guidance