From 3691a1e94f4a6b90f333404d99df077cca3b646b Mon Sep 17 00:00:00 2001 From: anton-l Date: Thu, 11 Aug 2022 14:41:08 +0200 Subject: [PATCH 1/8] test LMS with LDM --- src/diffusers/__init__.py | 11 +- src/diffusers/pipelines/__init__.py | 1 + .../pipelines/lms_discrete/__init__.py | 5 + .../lms_discrete/pipeline_lms_discrete.py | 92 +++++++++++++ src/diffusers/schedulers/__init__.py | 1 + .../schedulers/scheduling_lms_discrete.py | 126 ++++++++++++++++++ tests/test_modeling_utils.py | 23 ++++ 7 files changed, 258 insertions(+), 1 deletion(-) create mode 100644 src/diffusers/pipelines/lms_discrete/__init__.py create mode 100644 src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py create mode 100644 src/diffusers/schedulers/scheduling_lms_discrete.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index f8313509eed8..f5b448f76b65 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -18,11 +18,20 @@ get_scheduler, ) from .pipeline_utils import DiffusionPipeline -from .pipelines import DDIMPipeline, DDPMPipeline, KarrasVePipeline, LDMPipeline, PNDMPipeline, ScoreSdeVePipeline +from .pipelines import ( + DDIMPipeline, + DDPMPipeline, + KarrasVePipeline, + LDMPipeline, + LmsTextToImagePipeline, + PNDMPipeline, + ScoreSdeVePipeline, +) from .schedulers import ( DDIMScheduler, DDPMScheduler, KarrasVeScheduler, + LmsDiscreteScheduler, PNDMScheduler, SchedulerMixin, ScoreSdeVeScheduler, diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index c1b2068a3e23..d5c83f1b9329 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -2,6 +2,7 @@ from .ddim import DDIMPipeline from .ddpm import DDPMPipeline from .latent_diffusion_uncond import LDMPipeline +from .lms_discrete import LmsTextToImagePipeline from .pndm import PNDMPipeline from .score_sde_ve import ScoreSdeVePipeline from .stochatic_karras_ve import KarrasVePipeline diff --git a/src/diffusers/pipelines/lms_discrete/__init__.py b/src/diffusers/pipelines/lms_discrete/__init__.py new file mode 100644 index 000000000000..be9017af2d14 --- /dev/null +++ b/src/diffusers/pipelines/lms_discrete/__init__.py @@ -0,0 +1,5 @@ +from ...utils import is_transformers_available + + +if is_transformers_available(): + from .pipeline_lms_discrete import LDMBertModel, LmsTextToImagePipeline diff --git a/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py b/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py new file mode 100644 index 000000000000..60bbe1207807 --- /dev/null +++ b/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py @@ -0,0 +1,92 @@ +import torch +import torch.utils.checkpoint + +from tqdm.auto import tqdm + +from ...models import AutoencoderKL, UNet2DConditionModel +from ...pipeline_utils import DiffusionPipeline +from ...schedulers import LmsDiscreteScheduler +from ..latent_diffusion import LDMBertModel + + +class LmsTextToImagePipeline(DiffusionPipeline): + scheduler: LmsDiscreteScheduler + unet: UNet2DConditionModel + vae: AutoencoderKL + text_encoder: LDMBertModel + + def __init__(self, unet, scheduler, vae, text_encoder, tokenizer): + super().__init__() + scheduler = scheduler.set_format("pt") + self.register_modules(unet=unet, scheduler=scheduler, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer) + + @torch.no_grad() + def __call__( + self, + prompt, + batch_size=1, + generator=None, + torch_device=None, + guidance_scale=1.0, + num_inference_steps=20, + output_type="pil", + ): + if torch_device is None: + torch_device = "cuda" if torch.cuda.is_available() else "cpu" + batch_size = len(prompt) + + self.unet.to(torch_device) + self.vae.to(torch_device) + self.text_encoder.to(torch_device) + + # get unconditional embeddings for classifier free guidance + if guidance_scale != 1.0: + uncond_input = self.tokenizer([""] * batch_size, padding="max_length", max_length=77, return_tensors="pt") + uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(torch_device))[0] + + # get prompt text embeddings + text_input = self.tokenizer(prompt, padding="max_length", max_length=77, return_tensors="pt") + text_embeddings = self.text_encoder(text_input.input_ids.to(torch_device))[0] + + latents = torch.randn( + (batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size), + generator=generator, + ) + latents = latents.to(torch_device) + + self.scheduler.set_timesteps(num_inference_steps) + derivatives = [] + for t in tqdm(self.scheduler.timesteps): + if guidance_scale == 1.0: + # guidance_scale of 1 means no guidance + latents_input = latents + context = text_embeddings + else: + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + latents_input = torch.cat([latents] * 2) + context = torch.cat([uncond_embeddings, text_embeddings]) + + # predict the noise residual + sigma = self.scheduler.sigmas[t] + scaled_input = latents_input / (sigma**2 + 1) ** 0.5 + noise_pred = self.unet(scaled_input, t, encoder_hidden_states=context)["sample"] + # perform guidance + if guidance_scale != 1.0: + noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, derivatives=derivatives)["prev_sample"] + + # scale and decode the image latents with vae + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents) + + image = (image / 2 + 0.5).clamp(0, 1) + image = image.cpu().permute(0, 2, 3, 1).numpy() + if output_type == "pil": + image = self.numpy_to_pil(image) + + return {"sample": image} diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 42a536aae203..4c010b55ff82 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -19,6 +19,7 @@ from .scheduling_ddim import DDIMScheduler from .scheduling_ddpm import DDPMScheduler from .scheduling_karras_ve import KarrasVeScheduler +from .scheduling_lms_discrete import LmsDiscreteScheduler from .scheduling_pndm import PNDMScheduler from .scheduling_sde_ve import ScoreSdeVeScheduler from .scheduling_sde_vp import ScoreSdeVpScheduler diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py new file mode 100644 index 000000000000..b7c71a921544 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -0,0 +1,126 @@ +# Copyright 2022 Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Union + +import numpy as np +import torch + +from scipy import integrate + +from ..configuration_utils import ConfigMixin, register_to_config +from .scheduling_utils import SchedulerMixin + + +class LmsDiscreteScheduler(SchedulerMixin, ConfigMixin): + @register_to_config + def __init__( + self, + num_train_timesteps=1000, + beta_start=0.0001, + beta_end=0.02, + beta_schedule="linear", + trained_betas=None, + timestep_values=None, + tensor_format="pt", + ): + + if beta_schedule == "linear": + self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32) + elif beta_schedule == "scaled_linear": + # this schedule is very specific to the latent diffusion model. + self.betas = np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2 + else: + raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") + + self.alphas = 1.0 - self.betas + self.alphas_cumprod = np.cumprod(self.alphas, axis=0) + + self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 + + # setable values + self.num_inference_steps = None + self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy() + + self.tensor_format = tensor_format + self.set_format(tensor_format=tensor_format) + + def get_lms_coefficient(self, order, t, current_order): + """ + Compute a linear multistep coefficient + """ + + def lms_derivative(tau): + prod = 1.0 + for k in range(order): + if current_order == k: + continue + prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k]) + return prod + + integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0] + + return integrated_coeff + + def set_timesteps(self, num_inference_steps): + self.num_inference_steps = num_inference_steps + self.timesteps = np.linspace(num_inference_steps - 1, 0, num_inference_steps, dtype=int) + + train_timesteps = np.linspace(self.num_train_timesteps - 1, 0, num_inference_steps, dtype=float) + low_idx = np.floor(train_timesteps).astype(int) + high_idx = np.ceil(train_timesteps).astype(int) + frac = np.mod(train_timesteps, 1.0) + sigmas = np.array(self.sigmas) + sigmas = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx] + self.sigmas = np.concatenate([sigmas, [0.]]) + + self.set_format(tensor_format=self.tensor_format) + + def step( + self, + model_output: Union[torch.FloatTensor, np.ndarray], + timestep: int, + sample: Union[torch.FloatTensor, np.ndarray], + order: int = 4, + derivatives: List[Union[torch.FloatTensor, np.ndarray]] = None, + ): + sigma = self.sigmas[timestep] + + # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise + pred_original_sample = sample - sigma * model_output + + # 2. Convert to an ODE derivative + derivative = (sample - pred_original_sample) / sigma + derivatives.append(derivative) + if len(derivatives) > order: + derivatives.pop(0) + + # 3. Compute linear multistep coefficients + order = min(timestep + 1, order) + lms_coeffs = [self.get_lms_coefficient(order, timestep, curr_order) for curr_order in range(order)] + + # 4. Compute previous sample based on the derivatives path + prev_sample = sample + sum(coeff * derivative for coeff, derivative in zip(lms_coeffs, derivatives)) + + return {"prev_sample": prev_sample, "derivatives": derivatives} + + def add_noise(self, original_samples, noise, timesteps): + alpha_prod = self.alphas_cumprod[timesteps] + alpha_prod = self.match_shape(alpha_prod, original_samples) + + noisy_samples = (alpha_prod**0.5) * original_samples + ((1 - alpha_prod) ** 0.5) * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 072109e84ce8..691149c38aa1 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -33,6 +33,8 @@ KarrasVeScheduler, LDMPipeline, LDMTextToImagePipeline, + LmsDiscreteScheduler, + LmsTextToImagePipeline, PNDMPipeline, PNDMScheduler, ScoreSdeVePipeline, @@ -927,3 +929,24 @@ def test_karras_ve_pipeline(self): assert image.shape == (1, 256, 256, 3) expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + + # @slow + def test_lms_text2img_pipeline(self): + model_id = "CompVis/ldm-text2im-large-256" + ldm = LDMTextToImagePipeline.from_pretrained(model_id) + scheduler = LmsDiscreteScheduler.from_config(model_id, subfolder="scheduler") + lms_pipe = LmsTextToImagePipeline( + unet=ldm.unet, scheduler=scheduler, vae=ldm.vqvae, text_encoder=ldm.bert, tokenizer=ldm.tokenizer + ) + + prompt = "A painting of a squirrel eating a burger" + generator = torch.manual_seed(0) + image = lms_pipe( + [prompt], generator=generator, guidance_scale=6.0, num_inference_steps=20, output_type="numpy" + )["sample"] + + image_slice = image[0, -3:, -3:, -1] + + assert image.shape == (1, 256, 256, 3) + expected_slice = np.array([0.9256, 0.9340, 0.8933, 0.9361, 0.9113, 0.8727, 0.9122, 0.8745, 0.8099]) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 From f5017b73636c857320fdc955f5615ee1f70e838a Mon Sep 17 00:00:00 2001 From: anton-l Date: Mon, 15 Aug 2022 11:31:12 +0200 Subject: [PATCH 2/8] test LMS with LDM --- .../pipelines/lms_discrete/pipeline_lms_discrete.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py b/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py index 60bbe1207807..0c79929156fa 100644 --- a/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py +++ b/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py @@ -36,7 +36,7 @@ def __call__( batch_size = len(prompt) self.unet.to(torch_device) - self.vae.to(torch_device) + #self.vae.to(torch_device) self.text_encoder.to(torch_device) # get unconditional embeddings for classifier free guidance @@ -51,8 +51,9 @@ def __call__( latents = torch.randn( (batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size), generator=generator, - ) + ) * self.scheduler.sigmas[0] latents = latents.to(torch_device) + print(latents.sum()) self.scheduler.set_timesteps(num_inference_steps) derivatives = [] From 8be31148257a55c3e6c0e50b250d7253ea756f47 Mon Sep 17 00:00:00 2001 From: anton-l Date: Mon, 15 Aug 2022 16:19:32 +0200 Subject: [PATCH 3/8] Interchangeable sigma and timestep. Added dummy objects --- src/diffusers/__init__.py | 19 ++-- src/diffusers/pipelines/__init__.py | 1 - .../pipelines/lms_discrete/__init__.py | 5 - .../lms_discrete/pipeline_lms_discrete.py | 93 ------------------- .../pipeline_stable_diffusion.py | 17 +++- src/diffusers/schedulers/__init__.py | 6 +- .../schedulers/scheduling_lms_discrete.py | 42 ++++++--- src/diffusers/utils/__init__.py | 19 ++++ src/diffusers/utils/dummy_scipy_objects.py | 24 +++++ tests/test_modeling_utils.py | 31 +++---- 10 files changed, 110 insertions(+), 147 deletions(-) delete mode 100644 src/diffusers/pipelines/lms_discrete/__init__.py delete mode 100644 src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py create mode 100644 src/diffusers/utils/dummy_scipy_objects.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 715312b4d90e..2c302633cbcd 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -1,7 +1,7 @@ # flake8: noqa # There's no way to ignore "F401 '...' imported but unused" warnings in this # module, but to preserve other warnings. So, don't check this module at all. -from .utils import is_inflect_available, is_transformers_available, is_unidecode_available +from .utils import is_inflect_available, is_scipy_available, is_transformers_available, is_unidecode_available __version__ = "0.1.3" @@ -18,29 +18,26 @@ get_scheduler, ) from .pipeline_utils import DiffusionPipeline -from .pipelines import ( - DDIMPipeline, - DDPMPipeline, - KarrasVePipeline, - LDMPipeline, - LmsTextToImagePipeline, - PNDMPipeline, - ScoreSdeVePipeline, -) +from .pipelines import DDIMPipeline, DDPMPipeline, KarrasVePipeline, LDMPipeline, PNDMPipeline, ScoreSdeVePipeline from .schedulers import ( DDIMScheduler, DDPMScheduler, KarrasVeScheduler, - LmsDiscreteScheduler, PNDMScheduler, SchedulerMixin, ScoreSdeVeScheduler, ) + + +if is_scipy_available(): + from .schedulers import LmsDiscreteScheduler + from .training_utils import EMAModel if is_transformers_available(): from .pipelines import LDMTextToImagePipeline, StableDiffusionPipeline + else: from .utils.dummy_transformers_objects import * diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 4fc7a63104a2..19dafb477f9d 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -2,7 +2,6 @@ from .ddim import DDIMPipeline from .ddpm import DDPMPipeline from .latent_diffusion_uncond import LDMPipeline -from .lms_discrete import LmsTextToImagePipeline from .pndm import PNDMPipeline from .score_sde_ve import ScoreSdeVePipeline from .stochatic_karras_ve import KarrasVePipeline diff --git a/src/diffusers/pipelines/lms_discrete/__init__.py b/src/diffusers/pipelines/lms_discrete/__init__.py deleted file mode 100644 index be9017af2d14..000000000000 --- a/src/diffusers/pipelines/lms_discrete/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from ...utils import is_transformers_available - - -if is_transformers_available(): - from .pipeline_lms_discrete import LDMBertModel, LmsTextToImagePipeline diff --git a/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py b/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py deleted file mode 100644 index 0c79929156fa..000000000000 --- a/src/diffusers/pipelines/lms_discrete/pipeline_lms_discrete.py +++ /dev/null @@ -1,93 +0,0 @@ -import torch -import torch.utils.checkpoint - -from tqdm.auto import tqdm - -from ...models import AutoencoderKL, UNet2DConditionModel -from ...pipeline_utils import DiffusionPipeline -from ...schedulers import LmsDiscreteScheduler -from ..latent_diffusion import LDMBertModel - - -class LmsTextToImagePipeline(DiffusionPipeline): - scheduler: LmsDiscreteScheduler - unet: UNet2DConditionModel - vae: AutoencoderKL - text_encoder: LDMBertModel - - def __init__(self, unet, scheduler, vae, text_encoder, tokenizer): - super().__init__() - scheduler = scheduler.set_format("pt") - self.register_modules(unet=unet, scheduler=scheduler, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer) - - @torch.no_grad() - def __call__( - self, - prompt, - batch_size=1, - generator=None, - torch_device=None, - guidance_scale=1.0, - num_inference_steps=20, - output_type="pil", - ): - if torch_device is None: - torch_device = "cuda" if torch.cuda.is_available() else "cpu" - batch_size = len(prompt) - - self.unet.to(torch_device) - #self.vae.to(torch_device) - self.text_encoder.to(torch_device) - - # get unconditional embeddings for classifier free guidance - if guidance_scale != 1.0: - uncond_input = self.tokenizer([""] * batch_size, padding="max_length", max_length=77, return_tensors="pt") - uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(torch_device))[0] - - # get prompt text embeddings - text_input = self.tokenizer(prompt, padding="max_length", max_length=77, return_tensors="pt") - text_embeddings = self.text_encoder(text_input.input_ids.to(torch_device))[0] - - latents = torch.randn( - (batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size), - generator=generator, - ) * self.scheduler.sigmas[0] - latents = latents.to(torch_device) - print(latents.sum()) - - self.scheduler.set_timesteps(num_inference_steps) - derivatives = [] - for t in tqdm(self.scheduler.timesteps): - if guidance_scale == 1.0: - # guidance_scale of 1 means no guidance - latents_input = latents - context = text_embeddings - else: - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - latents_input = torch.cat([latents] * 2) - context = torch.cat([uncond_embeddings, text_embeddings]) - - # predict the noise residual - sigma = self.scheduler.sigmas[t] - scaled_input = latents_input / (sigma**2 + 1) ** 0.5 - noise_pred = self.unet(scaled_input, t, encoder_hidden_states=context)["sample"] - # perform guidance - if guidance_scale != 1.0: - noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, derivatives=derivatives)["prev_sample"] - - # scale and decode the image latents with vae - latents = 1 / 0.18215 * latents - image = self.vae.decode(latents) - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - return {"sample": image} diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 9db646af09b6..7fc2fa9307d3 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -8,7 +8,7 @@ from ...models import AutoencoderKL, UNet2DConditionModel from ...pipeline_utils import DiffusionPipeline -from ...schedulers import DDIMScheduler, PNDMScheduler +from ...schedulers import DDIMScheduler, LmsDiscreteScheduler, PNDMScheduler class StableDiffusionPipeline(DiffusionPipeline): @@ -18,7 +18,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler], + scheduler: Union[DDIMScheduler, PNDMScheduler, LmsDiscreteScheduler], ): super().__init__() scheduler = scheduler.set_format("pt") @@ -92,13 +92,20 @@ def __call__( extra_kwargs["eta"] = eta self.scheduler.set_timesteps(num_inference_steps) + if isinstance(self.scheduler, LmsDiscreteScheduler): + latents = latents * self.scheduler.sigmas[0] - for t in tqdm(self.scheduler.timesteps): + for i, t in tqdm(enumerate(self.scheduler.timesteps)): # 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[i] + latent_model_input = latent_model_input / (sigma**2 + 1) ** 0.5 # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] + noise_pred = self.unet( + latent_model_input, self.scheduler.sigma_to_timestep(sigma), encoder_hidden_states=text_embeddings + )["sample"] # perform guidance if do_classifier_free_guidance: @@ -106,7 +113,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs)["prev_sample"] + latents = self.scheduler.step(noise_pred, i, latents, **extra_kwargs)["prev_sample"] # scale and decode the image latents with vae latents = 1 / 0.18215 * latents diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 4c010b55ff82..dcc47ee3c5bf 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -16,11 +16,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from ..utils import is_scipy_available from .scheduling_ddim import DDIMScheduler from .scheduling_ddpm import DDPMScheduler from .scheduling_karras_ve import KarrasVeScheduler -from .scheduling_lms_discrete import LmsDiscreteScheduler from .scheduling_pndm import PNDMScheduler from .scheduling_sde_ve import ScoreSdeVeScheduler from .scheduling_sde_vp import ScoreSdeVpScheduler from .scheduling_utils import SchedulerMixin + + +if is_scipy_available(): + from .scheduling_lms_discrete import LmsDiscreteScheduler diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index b7c71a921544..c8a1e4d6c056 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -35,6 +35,11 @@ def __init__( timestep_values=None, tensor_format="pt", ): + """ + Linear Multistep Scheduler for discrete beta schedules. + Based on the original k-diffusion implementation by Katherine Crowson: + https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 + """ if beta_schedule == "linear": self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32) @@ -52,6 +57,7 @@ def __init__( # setable values self.num_inference_steps = None self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy() + self.derivatives = [] self.tensor_format = tensor_format self.set_format(tensor_format=tensor_format) @@ -73,17 +79,28 @@ def lms_derivative(tau): return integrated_coeff + def sigma_to_timestep(self, sigma): + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + dists = torch.abs(sigma - sigmas) + low_idx, high_idx = np.argsort(dists)[:2] + low, high = sigmas[low_idx], sigmas[high_idx] + w = (low - sigma) / (low - high) + w = w.clamp(0, 1) + timestep = (1 - w) * low_idx + w * high_idx + return timestep + def set_timesteps(self, num_inference_steps): self.num_inference_steps = num_inference_steps - self.timesteps = np.linspace(num_inference_steps - 1, 0, num_inference_steps, dtype=int) + self.timesteps = np.linspace(self.num_train_timesteps - 1, 0, num_inference_steps, dtype=float) - train_timesteps = np.linspace(self.num_train_timesteps - 1, 0, num_inference_steps, dtype=float) - low_idx = np.floor(train_timesteps).astype(int) - high_idx = np.ceil(train_timesteps).astype(int) - frac = np.mod(train_timesteps, 1.0) - sigmas = np.array(self.sigmas) + low_idx = np.floor(self.timesteps).astype(int) + high_idx = np.ceil(self.timesteps).astype(int) + frac = np.mod(self.timesteps, 1.0) + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) sigmas = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx] - self.sigmas = np.concatenate([sigmas, [0.]]) + self.sigmas = np.concatenate([sigmas, [0.0]]) + + self.derivatives = [] self.set_format(tensor_format=self.tensor_format) @@ -93,7 +110,6 @@ def step( timestep: int, sample: Union[torch.FloatTensor, np.ndarray], order: int = 4, - derivatives: List[Union[torch.FloatTensor, np.ndarray]] = None, ): sigma = self.sigmas[timestep] @@ -102,18 +118,18 @@ def step( # 2. Convert to an ODE derivative derivative = (sample - pred_original_sample) / sigma - derivatives.append(derivative) - if len(derivatives) > order: - derivatives.pop(0) + self.derivatives.append(derivative) + if len(self.derivatives) > order: + self.derivatives.pop(0) # 3. Compute linear multistep coefficients order = min(timestep + 1, order) lms_coeffs = [self.get_lms_coefficient(order, timestep, curr_order) for curr_order in range(order)] # 4. Compute previous sample based on the derivatives path - prev_sample = sample + sum(coeff * derivative for coeff, derivative in zip(lms_coeffs, derivatives)) + prev_sample = sample + sum(coeff * derivative for coeff, derivative in zip(lms_coeffs, self.derivatives)) - return {"prev_sample": prev_sample, "derivatives": derivatives} + return {"prev_sample": prev_sample} def add_noise(self, original_samples, noise, timesteps): alpha_prod = self.alphas_cumprod[timesteps] diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index c063cfaeb1fd..efbd89e27f4b 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -69,6 +69,14 @@ _modelcards_available = False +_scipy_available = importlib.util.find_spec("scipy") is not None +try: + _scipy_version = importlib_metadata.version("scipy") + logger.debug(f"Successfully imported transformers version {_scipy_version}") +except importlib_metadata.PackageNotFoundError: + _scipy_available = False + + def is_transformers_available(): return _transformers_available @@ -85,6 +93,10 @@ def is_modelcards_available(): return _modelcards_available +def is_scipy_available(): + return _scipy_available + + class RepositoryNotFoundError(HTTPError): """ Raised when trying to access a hf.co URL with an invalid repository name, or with a private repo name the user does @@ -118,11 +130,18 @@ class RevisionNotFoundError(HTTPError): """ +SCIPY_IMPORT_ERROR = """ +{0} requires the scipy library but it was not found in your environment. You can install it with pip: `pip install +scipy` +""" + + BACKENDS_MAPPING = OrderedDict( [ ("transformers", (is_transformers_available, TRANSFORMERS_IMPORT_ERROR)), ("unidecode", (is_unidecode_available, UNIDECODE_IMPORT_ERROR)), ("inflect", (is_inflect_available, INFLECT_IMPORT_ERROR)), + ("scipy", (is_scipy_available, SCIPY_IMPORT_ERROR)), ] ) diff --git a/src/diffusers/utils/dummy_scipy_objects.py b/src/diffusers/utils/dummy_scipy_objects.py new file mode 100644 index 000000000000..34680c2ab073 --- /dev/null +++ b/src/diffusers/utils/dummy_scipy_objects.py @@ -0,0 +1,24 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +# flake8: noqa +from ..utils import DummyObject, requires_backends + + +class LmsDiscreteScheduler(metaclass=DummyObject): + _backends = ["scipy"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["scipy"]) + + +class LDMTextToImagePipeline(metaclass=DummyObject): + _backends = ["scipy"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["scipy"]) + + +class StableDiffusionPipeline(metaclass=DummyObject): + _backends = ["scipy"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["scipy"]) diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 345e335c43ab..855e5f66c496 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -34,11 +34,11 @@ LDMPipeline, LDMTextToImagePipeline, LmsDiscreteScheduler, - LmsTextToImagePipeline, PNDMPipeline, PNDMScheduler, ScoreSdeVePipeline, ScoreSdeVeScheduler, + StableDiffusionPipeline, UNet2DModel, VQModel, ) @@ -47,8 +47,6 @@ from diffusers.testing_utils import floats_tensor, slow, torch_device from diffusers.training_utils import EMAModel -from ..src.diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipeline - torch.backends.cuda.matmul.allow_tf32 = False @@ -966,23 +964,20 @@ def test_karras_ve_pipeline(self): expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - # @slow - def test_lms_text2img_pipeline(self): - model_id = "CompVis/ldm-text2im-large-256" - ldm = LDMTextToImagePipeline.from_pretrained(model_id) - scheduler = LmsDiscreteScheduler.from_config(model_id, subfolder="scheduler") - lms_pipe = LmsTextToImagePipeline( - unet=ldm.unet, scheduler=scheduler, vae=ldm.vqvae, text_encoder=ldm.bert, tokenizer=ldm.tokenizer - ) + @slow + def test_lms_stable_diffusion_pipeline(self): + model_id = "CompVis/stable-diffusion-v1-1-diffusers" + pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True) + scheduler = LmsDiscreteScheduler.from_config(model_id, subfolder="scheduler", use_auth_token=True) + pipe.scheduler = scheduler - prompt = "A painting of a squirrel eating a burger" + prompt = "a photograph of an astronaut riding a horse" generator = torch.manual_seed(0) - image = lms_pipe( - [prompt], generator=generator, guidance_scale=6.0, num_inference_steps=20, output_type="numpy" - )["sample"] + image = pipe([prompt], generator=generator, guidance_scale=7.5, num_inference_steps=50, output_type="numpy")[ + "sample" + ] image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 256, 256, 3) - expected_slice = np.array([0.9256, 0.9340, 0.8933, 0.9361, 0.9113, 0.8727, 0.9122, 0.8745, 0.8099]) + assert image.shape == (1, 512, 512, 3) + expected_slice = np.array([0.0926, 0.0537, 0.0521, 0.0740, 0.0619, 0.0390, 0.0363, 0.0316, 0.0603]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 From eb1aee4ef5cc759e192aa8bb73160be84d5352c7 Mon Sep 17 00:00:00 2001 From: anton-l Date: Tue, 16 Aug 2022 10:47:47 +0200 Subject: [PATCH 4/8] Debug --- .../pipeline_stable_diffusion.py | 5 ++++- .../schedulers/scheduling_lms_discrete.py | 16 ++++++---------- tests/test_modeling_utils.py | 13 ++++++++++++- 3 files changed, 22 insertions(+), 12 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 7fc2fa9307d3..d0f8b23278e3 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -96,15 +96,18 @@ def __call__( latents = latents * self.scheduler.sigmas[0] for i, t in tqdm(enumerate(self.scheduler.timesteps)): + print() # 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[i] + print("c_in", (1 / (sigma**2 + 1) ** 0.5).item()) + print("sigma_to_t", t) latent_model_input = latent_model_input / (sigma**2 + 1) ** 0.5 # predict the noise residual noise_pred = self.unet( - latent_model_input, self.scheduler.sigma_to_timestep(sigma), encoder_hidden_states=text_embeddings + latent_model_input, t, encoder_hidden_states=text_embeddings )["sample"] # perform guidance diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index c8a1e4d6c056..9fa0f3f51f59 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -79,16 +79,6 @@ def lms_derivative(tau): return integrated_coeff - def sigma_to_timestep(self, sigma): - sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) - dists = torch.abs(sigma - sigmas) - low_idx, high_idx = np.argsort(dists)[:2] - low, high = sigmas[low_idx], sigmas[high_idx] - w = (low - sigma) / (low - high) - w = w.clamp(0, 1) - timestep = (1 - w) * low_idx + w * high_idx - return timestep - def set_timesteps(self, num_inference_steps): self.num_inference_steps = num_inference_steps self.timesteps = np.linspace(self.num_train_timesteps - 1, 0, num_inference_steps, dtype=float) @@ -112,12 +102,18 @@ def step( order: int = 4, ): sigma = self.sigmas[timestep] + print("timestep", timestep) + print("sigma", sigma.item()) # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise + print("sample", sample.sum().item()) + print("model_output", model_output.sum().item()) pred_original_sample = sample - sigma * model_output + print("pred_original_sample", pred_original_sample.sum().item()) # 2. Convert to an ODE derivative derivative = (sample - pred_original_sample) / sigma + print("derivative", derivative.sum().item()) self.derivatives.append(derivative) if len(self.derivatives) > order: self.derivatives.pop(0) diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 855e5f66c496..b53afedc2912 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -964,8 +964,10 @@ def test_karras_ve_pipeline(self): expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - @slow + #@slow def test_lms_stable_diffusion_pipeline(self): + from PIL import Image + model_id = "CompVis/stable-diffusion-v1-1-diffusers" pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True) scheduler = LmsDiscreteScheduler.from_config(model_id, subfolder="scheduler", use_auth_token=True) @@ -977,6 +979,15 @@ def test_lms_stable_diffusion_pipeline(self): "sample" ] + # def make_grid(images, rows, cols): + # w, h = images[0].size + # grid = Image.new('RGB', size=(cols * w, rows * h)) + # for i, image in enumerate(images): + # grid.paste(image, box=(i % cols * w, i // cols * h)) + # return grid + # + # image = make_grid(image, 2, 2) + image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 512, 512, 3) expected_slice = np.array([0.0926, 0.0537, 0.0521, 0.0740, 0.0619, 0.0390, 0.0363, 0.0316, 0.0603]) From ab5016867c5de00fe152f058a491102899a84973 Mon Sep 17 00:00:00 2001 From: anton-l Date: Tue, 16 Aug 2022 11:39:13 +0200 Subject: [PATCH 5/8] cuda generator --- tests/test_modeling_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index ee2102adf074..5a82e03e6346 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -988,7 +988,8 @@ def test_lms_stable_diffusion_pipeline(self): pipe.scheduler = scheduler prompt = "a photograph of an astronaut riding a horse" - generator = torch.manual_seed(0) + generator = torch.Generator(device='cuda') + generator.manual_seed(0) image = pipe([prompt], generator=generator, guidance_scale=7.5, num_inference_steps=50, output_type="numpy")[ "sample" ] From 5fe7895dda520bd2640116f2c850496cf0c6dfe0 Mon Sep 17 00:00:00 2001 From: anton-l Date: Tue, 16 Aug 2022 14:36:17 +0200 Subject: [PATCH 6/8] Fix derivatives --- .../pipeline_stable_diffusion.py | 9 ++------- .../schedulers/scheduling_lms_discrete.py | 10 +++------- tests/test_modeling_utils.py | 20 ++++--------------- 3 files changed, 9 insertions(+), 30 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 7b0ae2b05c4b..1652140702da 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -110,19 +110,14 @@ def __call__( latents = latents * self.scheduler.sigmas[0] for i, t in tqdm(enumerate(self.scheduler.timesteps)): - print() # 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[i] - print("c_in", (1 / (sigma**2 + 1) ** 0.5).item()) - print("sigma_to_t", t) - latent_model_input = latent_model_input / (sigma**2 + 1) ** 0.5 + latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) # predict the noise residual - noise_pred = self.unet( - latent_model_input, t, encoder_hidden_states=text_embeddings - )["sample"] + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] # perform guidance if do_classifier_free_guidance: diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 9fa0f3f51f59..99f7376cecd7 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -102,18 +102,12 @@ def step( order: int = 4, ): sigma = self.sigmas[timestep] - print("timestep", timestep) - print("sigma", sigma.item()) # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - print("sample", sample.sum().item()) - print("model_output", model_output.sum().item()) pred_original_sample = sample - sigma * model_output - print("pred_original_sample", pred_original_sample.sum().item()) # 2. Convert to an ODE derivative derivative = (sample - pred_original_sample) / sigma - print("derivative", derivative.sum().item()) self.derivatives.append(derivative) if len(self.derivatives) > order: self.derivatives.pop(0) @@ -123,7 +117,9 @@ def step( lms_coeffs = [self.get_lms_coefficient(order, timestep, curr_order) for curr_order in range(order)] # 4. Compute previous sample based on the derivatives path - prev_sample = sample + sum(coeff * derivative for coeff, derivative in zip(lms_coeffs, self.derivatives)) + prev_sample = sample + sum( + coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(self.derivatives)) + ) return {"prev_sample": prev_sample} diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 5a82e03e6346..696bac20042c 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -978,32 +978,20 @@ def test_karras_ve_pipeline(self): expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - #@slow + # @slow def test_lms_stable_diffusion_pipeline(self): - from PIL import Image - model_id = "CompVis/stable-diffusion-v1-1-diffusers" pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True) scheduler = LmsDiscreteScheduler.from_config(model_id, subfolder="scheduler", use_auth_token=True) pipe.scheduler = scheduler prompt = "a photograph of an astronaut riding a horse" - generator = torch.Generator(device='cuda') - generator.manual_seed(0) - image = pipe([prompt], generator=generator, guidance_scale=7.5, num_inference_steps=50, output_type="numpy")[ + generator = torch.Generator(device=torch_device).manual_seed(0) + image = pipe([prompt], generator=generator, guidance_scale=7.5, num_inference_steps=10, output_type="numpy")[ "sample" ] - # def make_grid(images, rows, cols): - # w, h = images[0].size - # grid = Image.new('RGB', size=(cols * w, rows * h)) - # for i, image in enumerate(images): - # grid.paste(image, box=(i % cols * w, i // cols * h)) - # return grid - # - # image = make_grid(image, 2, 2) - image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 512, 512, 3) - expected_slice = np.array([0.0926, 0.0537, 0.0521, 0.0740, 0.0619, 0.0390, 0.0363, 0.0316, 0.0603]) + expected_slice = np.array([0.9077, 0.9254, 0.9181, 0.9227, 0.9213, 0.9367, 0.9399, 0.9406, 0.9024]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 From 1be007191eb4c10b1cdaaa89a5237b605d175ab1 Mon Sep 17 00:00:00 2001 From: anton-l Date: Tue, 16 Aug 2022 15:58:28 +0200 Subject: [PATCH 7/8] Update tests --- .../stable_diffusion/pipeline_stable_diffusion.py | 5 ++++- tests/test_modeling_utils.py | 7 ++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 1652140702da..51234c4bb8b1 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -125,7 +125,10 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs)["prev_sample"] + if isinstance(self.scheduler, LmsDiscreteScheduler): + latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs)["prev_sample"] + else: + 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 diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 4f7c2ed6edeb..6abf67398dde 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -842,7 +842,7 @@ def test_ldm_text2img_fast(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @slow - @unittest.skipIf(torch_device == "cpu", "Stable diffusion is suppused to run on GPU") + @unittest.skipIf(torch_device == "cpu", "Stable diffusion is supposed to run on GPU") def test_stable_diffusion(self): # make sure here that pndm scheduler skips prk sd_pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-1-diffusers") @@ -863,7 +863,7 @@ def test_stable_diffusion(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @slow - @unittest.skipIf(torch_device == "cpu", "Stable diffusion is suppused to run on GPU") + @unittest.skipIf(torch_device == "cpu", "Stable diffusion is supposed to run on GPU") def test_stable_diffusion_fast_ddim(self): sd_pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-1-diffusers") @@ -979,7 +979,8 @@ def test_karras_ve_pipeline(self): expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - # @slow + @slow + @unittest.skipIf(torch_device == "cpu", "Stable diffusion is supposed to run on GPU") def test_lms_stable_diffusion_pipeline(self): model_id = "CompVis/stable-diffusion-v1-1-diffusers" pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True) From 4aa07f5cb8cfcca78cb58e61a533ec2d5284eefb Mon Sep 17 00:00:00 2001 From: anton-l Date: Tue, 16 Aug 2022 16:42:27 +0200 Subject: [PATCH 8/8] Rename Lms->LMS --- src/diffusers/__init__.py | 2 +- .../stable_diffusion/pipeline_stable_diffusion.py | 10 +++++----- src/diffusers/schedulers/__init__.py | 2 +- src/diffusers/schedulers/scheduling_lms_discrete.py | 2 +- tests/test_modeling_utils.py | 4 ++-- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2c302633cbcd..6ef61806115c 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -30,7 +30,7 @@ if is_scipy_available(): - from .schedulers import LmsDiscreteScheduler + from .schedulers import LMSDiscreteScheduler from .training_utils import EMAModel diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 51234c4bb8b1..769cf9c342b5 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -8,7 +8,7 @@ from ...models import AutoencoderKL, UNet2DConditionModel from ...pipeline_utils import DiffusionPipeline -from ...schedulers import DDIMScheduler, LmsDiscreteScheduler, PNDMScheduler +from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler class StableDiffusionPipeline(DiffusionPipeline): @@ -18,7 +18,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LmsDiscreteScheduler], + scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], ): super().__init__() scheduler = scheduler.set_format("pt") @@ -106,13 +106,13 @@ def __call__( extra_step_kwargs["eta"] = eta self.scheduler.set_timesteps(num_inference_steps) - if isinstance(self.scheduler, LmsDiscreteScheduler): + if isinstance(self.scheduler, LMSDiscreteScheduler): latents = latents * self.scheduler.sigmas[0] for i, t in tqdm(enumerate(self.scheduler.timesteps)): # 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): + if isinstance(self.scheduler, LMSDiscreteScheduler): sigma = self.scheduler.sigmas[i] latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) @@ -125,7 +125,7 @@ def __call__( 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): + if isinstance(self.scheduler, LMSDiscreteScheduler): latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs)["prev_sample"] else: latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)["prev_sample"] diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index dcc47ee3c5bf..6360dbe0c625 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -27,4 +27,4 @@ if is_scipy_available(): - from .scheduling_lms_discrete import LmsDiscreteScheduler + from .scheduling_lms_discrete import LMSDiscreteScheduler diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 99f7376cecd7..f8867ed802e4 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -23,7 +23,7 @@ from .scheduling_utils import SchedulerMixin -class LmsDiscreteScheduler(SchedulerMixin, ConfigMixin): +class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin): @register_to_config def __init__( self, diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 6abf67398dde..5873761af4a7 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -33,7 +33,7 @@ KarrasVeScheduler, LDMPipeline, LDMTextToImagePipeline, - LmsDiscreteScheduler, + LMSDiscreteScheduler, PNDMPipeline, PNDMScheduler, ScoreSdeVePipeline, @@ -984,7 +984,7 @@ def test_karras_ve_pipeline(self): def test_lms_stable_diffusion_pipeline(self): model_id = "CompVis/stable-diffusion-v1-1-diffusers" pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True) - scheduler = LmsDiscreteScheduler.from_config(model_id, subfolder="scheduler", use_auth_token=True) + scheduler = LMSDiscreteScheduler.from_config(model_id, subfolder="scheduler", use_auth_token=True) pipe.scheduler = scheduler prompt = "a photograph of an astronaut riding a horse"