bugfixes

Author: clarencechen

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_diffedit.py

@@ -1242,7 +1242,7 @@ def invert(
                         callback(i, t, latents)
 
         assert len(inverted_latents) == len(timesteps)
-        latents = torch.stack(list(reversed(inverted_latents)), 0)
+        latents = torch.cat(list(reversed(inverted_latents)))
 
         # 8. Post-processing
         image = self.decode_latents(latents.detach())
@@ -1478,7 +1478,7 @@ def __call__(
         )
 
         # 7. Preprocess inverted latents
-        inverted_latents_shape = (len(timesteps), batch_size, num_channels_latents, *vae_latent_size)
+        inverted_latents_shape = (len(timesteps) * batch_size, num_channels_latents, *vae_latent_size)
         if inverted_latents is None:
             raise ValueError(
                 "`inverted_latents` input cannot be undefined. Use `invert()` to compute `inverted_latents`."
@@ -1489,6 +1489,10 @@ def __call__(
             )
         if isinstance(inverted_latents, np.ndarray):
             inverted_latents = torch.from_numpy(inverted_latents)
+        inverted_latents = torch.cat(
+            [inverted_latents.reshape(-1, batch_size, num_channels_latents, *vae_latent_size)] * num_images_per_prompt,
+            1,
+        )
         inverted_latents = inverted_latents.to(device=device, dtype=prompt_embeds.dtype)
 
         # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline