Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -971,7 +971,7 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         prompt_embeds, ram_encoder_hidden_states = self._encode_prompt(
             prompt,
             device,
-            num_images_per_prompt * num_particles,
+            num_images_per_prompt,
             do_classifier_free_guidance,
             negative_prompt,
             prompt_embeds=prompt_embeds,
@@ -984,7 +984,7 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
             image=image,
             width=width,
             height=height,
-            batch_size=batch_size * num_images_per_prompt * num_particles,
+            batch_size=batch_size * num_images_per_prompt,
             num_images_per_prompt=num_images_per_prompt,
             device=device,
             dtype=controlnet.dtype,
@@ -999,7 +999,7 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         # 6. Prepare latent variables
         num_channels_latents = self.unet.config.in_channels
         latents = self.prepare_latents(
-            batch_size * num_images_per_prompt * num_particles,
+            batch_size * num_images_per_prompt,
             num_channels_latents,
             height,
             width,
@@ -1024,6 +1024,9 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
         if use_KDS:
+            latents = latents.repeat_interleave(num_particles, dim=0)
+            image = image.repeat_interleave(num_particles, dim=0)
+            prompt_embeds = prompt_embeds.repeat_interleave(num_particles, dim=0)
             latents.requires_grad_(True)
 
         # 8. Denoising loop