test latent fix

pipeline.py

@@ -617,7 +617,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             extra_step_kwargs["generator"] = generator
         return extra_step_kwargs
 
-    def prepare_latents(self, image, timestep, batch_size, height, width, dtype, device, generator, latents=None):
+    def prepare_latents(self, image, timestep, batch_size, height, width, dtype, device, generator,upscale=False, latents=None):
         if image is None:
             shape = (
                 batch_size,
@@ -637,9 +637,11 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
                     raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
                 latents = latents.to(device)
 
+                if upscale:
+                        return latents, None, None
             # scale the initial noise by the standard deviation required by the scheduler
-            latents = latents * self.scheduler.init_noise_sigma
-            return latents, None, None
+                latents = latents * self.scheduler.init_noise_sigma
+                return latents, None, None
         else:
             init_latent_dist = self.vae.encode(image).latent_dist
             init_latents = init_latent_dist.sample(generator=generator)
@@ -673,6 +675,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         eta: float = 0.0,
         generator: Optional[torch.Generator] = None,
         latents: Optional[torch.FloatTensor] = None,
+        upscale: bool = False,
         return_latents: bool = False,
         max_embeddings_multiples: Optional[int] = 3,
         output_type: Optional[str] = "pil",
@@ -1023,6 +1026,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             generator=generator,
             latents=latents,
             return_latents=True,
+            upscale=False,
             max_embeddings_multiples=max_embeddings_multiples,
             output_type=output_type,
             return_dict=return_dict,
@@ -1046,6 +1050,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             generator=generator,
             latents=latents,
             return_latents=False,
+            upscale=True,
             max_embeddings_multiples=max_embeddings_multiples,
             output_type=output_type,
             return_dict=return_dict,