Update pipeline.py

pipeline.py

@@ -308,6 +308,9 @@ class FluxWithCFGPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile
         negative_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
         negative_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
 
+        negative_prompt_embeds = torch.transpose(negative_prompt_embeds, 0, 1)
+        negative_pooled_prompt_embeds = torch.transpose(negative_pooled_prompt_embeds, 0, 1)
+
         return prompt_embeds, pooled_prompt_embeds, text_ids, negative_prompt_embeds, negative_pooled_prompt_embeds
 
     def check_inputs(
@@ -579,8 +582,8 @@ class FluxWithCFGPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile
         )
 
         if self.do_classifier_free_guidance:
-            prompt_embeds = torch.stack([negative_prompt_embeds, prompt_embeds], dim=0)
-            pooled_prompt_embeds = torch.stack([negative_prompt_embeds, pooled_prompt_embeds], dim=0)
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            pooled_prompt_embeds = torch.cat([negative_prompt_embeds, pooled_prompt_embeds], dim=0)
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels // 4