Update

1 files changed, 8 insertions, 5 deletions

diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index aa3dbc6..aa446ec 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -386,7 +386,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
 
     def decode_latents(self, latents):
         latents = 1 / self.vae.config.scaling_factor * latents
-        image = self.vae.decode(latents.to(dtype=self.vae.dtype)).sample
+        image = self.vae.decode(latents.to(dtype=self.vae.dtype), return_dict=False)[0]
         image = (image / 2 + 0.5).clamp(0, 1)
         # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
         image = image.cpu().permute(0, 2, 3, 1).float().numpy()
@@ -545,7 +545,8 @@ class VlpnStableDiffusion(DiffusionPipeline):
                     t,
                     encoder_hidden_states=prompt_embeds,
                     cross_attention_kwargs=cross_attention_kwargs,
-                ).sample
+                    return_dict=False,
+                )[0]
 
                 # perform guidance
                 if do_classifier_free_guidance:
@@ -567,7 +568,8 @@ class VlpnStableDiffusion(DiffusionPipeline):
                         )
                         uncond_emb, _ = prompt_embeds.chunk(2)
                         # forward and give guidance
-                        degraded_pred = self.unet(degraded_latents, t, encoder_hidden_states=uncond_emb).sample
+                        degraded_pred = self.unet(
+                            degraded_latents, t, encoder_hidden_states=uncond_emb, return_dict=False)[0]
                         noise_pred += sag_scale * (noise_pred_uncond - degraded_pred)
                     else:
                         # DDIM-like prediction of x0
@@ -579,11 +581,12 @@ class VlpnStableDiffusion(DiffusionPipeline):
                             pred_x0, cond_attn, t, self.pred_epsilon(latents, noise_pred, t)
                         )
                         # forward and give guidance
-                        degraded_pred = self.unet(degraded_latents, t, encoder_hidden_states=prompt_embeds).sample
+                        degraded_pred = self.unet(
+                            degraded_latents, t, encoder_hidden_states=prompt_embeds, return_dict=False)[0]
                         noise_pred += sag_scale * (noise_pred - degraded_pred)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
 
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):