From 49a37b054ea7c1cdd8c0d7c44f3809ab8bee0693 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Thu, 6 Oct 2022 17:15:22 +0200 Subject: Update --- .../stable_diffusion/vlpn_stable_diffusion.py | 35 ++++------------------ 1 file changed, 5 insertions(+), 30 deletions(-) (limited to 'pipelines/stable_diffusion') diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py index 8fbe5f9..a198cf6 100644 --- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py +++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py @@ -216,7 +216,6 @@ class VlpnStableDiffusion(DiffusionPipeline): offset = self.scheduler.config.get("steps_offset", 0) init_timestep = num_inference_steps + offset - ensure_sigma = not isinstance(latents, PIL.Image.Image) # get the initial random noise unless the user supplied it @@ -246,13 +245,8 @@ class VlpnStableDiffusion(DiffusionPipeline): init_timestep = int(num_inference_steps * strength) + offset init_timestep = min(init_timestep, num_inference_steps) - if isinstance(self.scheduler, LMSDiscreteScheduler): - timesteps = torch.tensor( - [num_inference_steps - init_timestep] * batch_size, dtype=torch.long, device=self.device - ) - else: - timesteps = self.scheduler.timesteps[-init_timestep] - timesteps = torch.tensor([timesteps] * batch_size, dtype=torch.long, device=self.device) + timesteps = self.scheduler.timesteps[-init_timestep] + timesteps = torch.tensor([timesteps] * batch_size, device=self.device) # add noise to latents using the timesteps noise = torch.randn(latents.shape, generator=generator, device=self.device) @@ -263,13 +257,6 @@ class VlpnStableDiffusion(DiffusionPipeline): if latents.device != self.device: raise ValueError(f"Unexpected latents device, got {latents.device}, expected {self.device}") - # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas - if ensure_sigma: - if isinstance(self.scheduler, LMSDiscreteScheduler): - latents = latents * self.scheduler.sigmas[0] - elif isinstance(self.scheduler, EulerAScheduler): - latents = latents * self.scheduler.sigmas[0] - t_start = max(num_inference_steps - init_timestep + offset, 0) # Some schedulers like PNDM have timesteps as arrays @@ -290,19 +277,13 @@ class VlpnStableDiffusion(DiffusionPipeline): extra_step_kwargs["generator"] = generator for i, t in enumerate(self.progress_bar(timesteps_tensor)): - t_index = t_start + i - # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - if isinstance(self.scheduler, LMSDiscreteScheduler): - sigma = self.scheduler.sigmas[t_index] - # the model input needs to be scaled to match the continuous ODE formulation in K-LMS - latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) noise_pred = None if isinstance(self.scheduler, EulerAScheduler): - sigma = self.scheduler.sigmas[t].reshape(1) + sigma = t.reshape(1) sigma_in = torch.cat([sigma] * latent_model_input.shape[0]) noise_pred = CFGDenoiserForward(self.unet, latent_model_input, sigma_in, text_embeddings, guidance_scale, quantize=True, DSsigmas=self.scheduler.DSsigmas) @@ -316,13 +297,7 @@ class VlpnStableDiffusion(DiffusionPipeline): noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - if isinstance(self.scheduler, LMSDiscreteScheduler): - latents = self.scheduler.step(noise_pred, t_index, latents, **extra_step_kwargs).prev_sample - elif isinstance(self.scheduler, EulerAScheduler): - latents = self.scheduler.step(noise_pred, t_index, t_index + 1, - latents, **extra_step_kwargs).prev_sample - else: - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # scale and decode the image latents with vae latents = 1 / 0.18215 * latents -- cgit v1.2.3-70-g09d2