diff options
Diffstat (limited to 'pipelines')
-rw-r--r-- | pipelines/stable_diffusion/vlpn_stable_diffusion.py | 13 |
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): | |||
386 | 386 | ||
387 | def decode_latents(self, latents): | 387 | def decode_latents(self, latents): |
388 | latents = 1 / self.vae.config.scaling_factor * latents | 388 | latents = 1 / self.vae.config.scaling_factor * latents |
389 | image = self.vae.decode(latents.to(dtype=self.vae.dtype)).sample | 389 | image = self.vae.decode(latents.to(dtype=self.vae.dtype), return_dict=False)[0] |
390 | image = (image / 2 + 0.5).clamp(0, 1) | 390 | image = (image / 2 + 0.5).clamp(0, 1) |
391 | # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 | 391 | # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 |
392 | image = image.cpu().permute(0, 2, 3, 1).float().numpy() | 392 | image = image.cpu().permute(0, 2, 3, 1).float().numpy() |
@@ -545,7 +545,8 @@ class VlpnStableDiffusion(DiffusionPipeline): | |||
545 | t, | 545 | t, |
546 | encoder_hidden_states=prompt_embeds, | 546 | encoder_hidden_states=prompt_embeds, |
547 | cross_attention_kwargs=cross_attention_kwargs, | 547 | cross_attention_kwargs=cross_attention_kwargs, |
548 | ).sample | 548 | return_dict=False, |
549 | )[0] | ||
549 | 550 | ||
550 | # perform guidance | 551 | # perform guidance |
551 | if do_classifier_free_guidance: | 552 | if do_classifier_free_guidance: |
@@ -567,7 +568,8 @@ class VlpnStableDiffusion(DiffusionPipeline): | |||
567 | ) | 568 | ) |
568 | uncond_emb, _ = prompt_embeds.chunk(2) | 569 | uncond_emb, _ = prompt_embeds.chunk(2) |
569 | # forward and give guidance | 570 | # forward and give guidance |
570 | degraded_pred = self.unet(degraded_latents, t, encoder_hidden_states=uncond_emb).sample | 571 | degraded_pred = self.unet( |
572 | degraded_latents, t, encoder_hidden_states=uncond_emb, return_dict=False)[0] | ||
571 | noise_pred += sag_scale * (noise_pred_uncond - degraded_pred) | 573 | noise_pred += sag_scale * (noise_pred_uncond - degraded_pred) |
572 | else: | 574 | else: |
573 | # DDIM-like prediction of x0 | 575 | # DDIM-like prediction of x0 |
@@ -579,11 +581,12 @@ class VlpnStableDiffusion(DiffusionPipeline): | |||
579 | pred_x0, cond_attn, t, self.pred_epsilon(latents, noise_pred, t) | 581 | pred_x0, cond_attn, t, self.pred_epsilon(latents, noise_pred, t) |
580 | ) | 582 | ) |
581 | # forward and give guidance | 583 | # forward and give guidance |
582 | degraded_pred = self.unet(degraded_latents, t, encoder_hidden_states=prompt_embeds).sample | 584 | degraded_pred = self.unet( |
585 | degraded_latents, t, encoder_hidden_states=prompt_embeds, return_dict=False)[0] | ||
583 | noise_pred += sag_scale * (noise_pred - degraded_pred) | 586 | noise_pred += sag_scale * (noise_pred - degraded_pred) |
584 | 587 | ||
585 | # compute the previous noisy sample x_t -> x_t-1 | 588 | # compute the previous noisy sample x_t -> x_t-1 |
586 | latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample | 589 | latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] |
587 | 590 | ||
588 | # call the callback, if provided | 591 | # call the callback, if provided |
589 | if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): | 592 | if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): |