Update pipeline.py

pipeline.py

@@ -42,6 +42,7 @@ if is_torch_xla_available():
 else:
     XLA_AVAILABLE = False
 
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 # Constants for shift calculation
 BASE_SEQ_LEN = 256
@@ -188,6 +189,7 @@ class FluxWithCFGPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
         prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+        pooled_prompt_embeds = prompt_embeds[0]
 
         # Use pooled output of CLIPTextModel
         prompt_embeds = prompt_embeds.pooler_output
@@ -196,8 +198,11 @@ class FluxWithCFGPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
         prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+        
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
 
-        return prompt_embeds
+        return prompt_embeds, pooled_prompt_embeds
 
     def encode_prompt(
         self,