fix broken num_images_per_prompt

pipeline.py

@@ -324,13 +324,9 @@ def get_weighted_text_embeddings(
             uncond_embeddings *= uncond_weights.unsqueeze(-1)
             uncond_embeddings *= previous_mean / uncond_embeddings.mean(axis=[-2, -1])
 
-    # For classifier free guidance, we need to do two forward passes.
-    # Here we concatenate the unconditional and text embeddings into a single batch
-    # to avoid doing two forward passes
     if uncond_prompt is not None:
-        text_embeddings = torch.concat([uncond_embeddings, text_embeddings])
-
-    return text_embeddings
+        return text_embeddings, uncond_embeddings
+    return text_embeddings, None
 
 
 def preprocess_image(image):
@@ -598,13 +594,22 @@ class StableDiffusionLongPromptWeightingPipeline(DiffusionPipeline):
             else:
                 uncond_tokens = negative_prompt
 
-        text_embeddings = get_weighted_text_embeddings(
+        text_embeddings, uncond_embeddings  = get_weighted_text_embeddings(
             pipe=self,
             prompt=prompt,
             uncond_prompt=uncond_tokens if do_classifier_free_guidance else None,
             max_embeddings_multiples=max_embeddings_multiples,
             **kwargs
         )
+        bs_embed, seq_len, _ = text_embeddings.shape
+        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
+        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        if do_classifier_free_guidance:
+            bs_embed, seq_len, _ = uncond_embeddings.shape
+            uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)
+            uncond_embeddings = uncond_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
+            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
 
         # set timesteps
         self.scheduler.set_timesteps(num_inference_steps)