From 2d97544dc70b4869ab9b5b85a35d7a09856d0c54 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Sun, 12 Jun 2022 19:56:13 +0000 Subject: [PATCH] add more tests schedulers --- .../modeling_latent_diffusion.py | 2 +- src/diffusers/pipelines/pipeline_ddim.py | 6 +- .../pipelines/pipeline_latent_diffusion.py | 2 +- .../schedulers/classifier_free_guidance.py | 4 +- src/diffusers/schedulers/ddim.py | 8 +- src/diffusers/schedulers/gaussian_ddpm.py | 6 +- src/diffusers/schedulers/glide_ddim.py | 6 +- tests/test_modeling_utils.py | 5 +- tests/test_scheduler.py | 277 ++++++++++++++++-- 9 files changed, 274 insertions(+), 42 deletions(-) diff --git a/src/diffusers/pipelines/old/latent_diffusion/modeling_latent_diffusion.py b/src/diffusers/pipelines/old/latent_diffusion/modeling_latent_diffusion.py index 33c73ff25d..9dd778e51d 100644 --- a/src/diffusers/pipelines/old/latent_diffusion/modeling_latent_diffusion.py +++ b/src/diffusers/pipelines/old/latent_diffusion/modeling_latent_diffusion.py @@ -33,7 +33,7 @@ class LatentDiffusion(DiffusionPipeline): text_input = self.tokenizer(prompt, padding="max_length", max_length=77, return_tensors='pt').to(torch_device) text_embedding = self.bert(text_input.input_ids)[0] - num_trained_timesteps = self.noise_scheduler.num_timesteps + num_trained_timesteps = self.noise_scheduler.timesteps inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps) image = self.noise_scheduler.sample_noise( diff --git a/src/diffusers/pipelines/pipeline_ddim.py b/src/diffusers/pipelines/pipeline_ddim.py index 4f7a14178a..530945238b 100644 --- a/src/diffusers/pipelines/pipeline_ddim.py +++ b/src/diffusers/pipelines/pipeline_ddim.py @@ -17,7 +17,7 @@ import torch import tqdm -from .. import DiffusionPipeline +from ..pipeline_utils import DiffusionPipeline class DDIM(DiffusionPipeline): @@ -30,7 +30,7 @@ class DDIM(DiffusionPipeline): if torch_device is None: torch_device = "cuda" if torch.cuda.is_available() else "cpu" - num_trained_timesteps = self.noise_scheduler.num_timesteps + num_trained_timesteps = self.noise_scheduler.timesteps inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps) self.unet.to(torch_device) @@ -64,7 +64,7 @@ class DDIM(DiffusionPipeline): variance = 0 if eta > 0: noise = self.noise_scheduler.sample_noise(image.shape, device=image.device, generator=generator) - variance = self.noise_scheduler.get_variance(t).sqrt() * eta * noise + variance = self.noise_scheduler.get_variance(t, num_inference_steps).sqrt() * eta * noise # 4. set current image to prev_image: x_t -> x_t-1 image = pred_prev_image + variance diff --git a/src/diffusers/pipelines/pipeline_latent_diffusion.py b/src/diffusers/pipelines/pipeline_latent_diffusion.py index b723313972..2d57f88968 100644 --- a/src/diffusers/pipelines/pipeline_latent_diffusion.py +++ b/src/diffusers/pipelines/pipeline_latent_diffusion.py @@ -883,7 +883,7 @@ class LatentDiffusion(DiffusionPipeline): text_input = self.tokenizer(prompt, padding="max_length", max_length=77, return_tensors='pt').to(torch_device) text_embedding = self.bert(text_input.input_ids)[0] - num_trained_timesteps = self.noise_scheduler.num_timesteps + num_trained_timesteps = self.noise_scheduler.timesteps inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps) image = self.noise_scheduler.sample_noise( diff --git a/src/diffusers/schedulers/classifier_free_guidance.py b/src/diffusers/schedulers/classifier_free_guidance.py index 2cd8152144..ec4356557b 100644 --- a/src/diffusers/schedulers/classifier_free_guidance.py +++ b/src/diffusers/schedulers/classifier_free_guidance.py @@ -61,7 +61,7 @@ class ClassifierFreeGuidanceScheduler(nn.Module, ConfigMixin): timesteps=timesteps, beta_schedule=beta_schedule, ) - self.num_timesteps = int(timesteps) + self.timesteps = int(timesteps) if beta_schedule == "squaredcos_cap_v2": # GLIDE cosine schedule @@ -94,4 +94,4 @@ class ClassifierFreeGuidanceScheduler(nn.Module, ConfigMixin): return torch.randn(shape, generator=generator).to(device) def __len__(self): - return self.num_timesteps + return self.timesteps diff --git a/src/diffusers/schedulers/ddim.py b/src/diffusers/schedulers/ddim.py index abb26dc57a..42b8f0d029 100644 --- a/src/diffusers/schedulers/ddim.py +++ b/src/diffusers/schedulers/ddim.py @@ -42,7 +42,7 @@ class DDIMScheduler(nn.Module, ConfigMixin): beta_end=beta_end, beta_schedule=beta_schedule, ) - self.num_timesteps = int(timesteps) + self.timesteps = int(timesteps) self.clip_image = clip_predicted_image if beta_schedule == "linear": @@ -90,7 +90,7 @@ class DDIMScheduler(nn.Module, ConfigMixin): def get_orig_t(self, t, num_inference_steps): if t < 0: return -1 - return self.num_timesteps // num_inference_steps * t + return self.timesteps // num_inference_steps * t def get_variance(self, t, num_inference_steps): orig_t = self.get_orig_t(t, num_inference_steps) @@ -105,7 +105,7 @@ class DDIMScheduler(nn.Module, ConfigMixin): return variance - def step(self, residual, image, t, num_inference_steps, eta, output_pred_x_0=False): + def step(self, residual, image, t, num_inference_steps, eta): # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf # Ideally, read DDIM paper in-detail understanding @@ -152,4 +152,4 @@ class DDIMScheduler(nn.Module, ConfigMixin): return torch.randn(shape, generator=generator).to(device) def __len__(self): - return self.num_timesteps + return self.timesteps diff --git a/src/diffusers/schedulers/gaussian_ddpm.py b/src/diffusers/schedulers/gaussian_ddpm.py index a6540438ac..c3e1b1fad1 100644 --- a/src/diffusers/schedulers/gaussian_ddpm.py +++ b/src/diffusers/schedulers/gaussian_ddpm.py @@ -44,7 +44,7 @@ class GaussianDDPMScheduler(nn.Module, ConfigMixin): variance_type=variance_type, clip_predicted_image=clip_predicted_image, ) - self.num_timesteps = int(timesteps) + self.timesteps = int(timesteps) self.clip_image = clip_predicted_image self.variance_type = variance_type @@ -107,7 +107,7 @@ class GaussianDDPMScheduler(nn.Module, ConfigMixin): return variance - def step(self, residual, image, t, output_pred_x_0=False): + def step(self, residual, image, t): # 1. compute alphas, betas alpha_prod_t = self.get_alpha_prod(t) alpha_prod_t_prev = self.get_alpha_prod(t - 1) @@ -138,4 +138,4 @@ class GaussianDDPMScheduler(nn.Module, ConfigMixin): return torch.randn(shape, generator=generator).to(device) def __len__(self): - return self.num_timesteps + return self.timesteps diff --git a/src/diffusers/schedulers/glide_ddim.py b/src/diffusers/schedulers/glide_ddim.py index 8b5d86bd3c..ade4e01ed9 100644 --- a/src/diffusers/schedulers/glide_ddim.py +++ b/src/diffusers/schedulers/glide_ddim.py @@ -32,12 +32,12 @@ class GlideDDIMScheduler(nn.Module, ConfigMixin): timesteps=timesteps, beta_schedule=beta_schedule, ) - self.num_timesteps = int(timesteps) + self.timesteps = int(timesteps) if beta_schedule == "linear": # Linear schedule from Ho et al, extended to work for any number of # diffusion steps. - scale = 1000 / self.num_timesteps + scale = 1000 / self.timesteps beta_start = scale * 0.0001 beta_end = scale * 0.02 betas = linear_beta_schedule(timesteps, beta_start=beta_start, beta_end=beta_end) @@ -88,4 +88,4 @@ class GlideDDIMScheduler(nn.Module, ConfigMixin): return torch.randn(shape, generator=generator).to(device) def __len__(self): - return self.num_timesteps + return self.timesteps diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 2e4301ddd1..40ed0b5da4 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -75,16 +75,18 @@ class ModelTesterMixin(unittest.TestCase): sizes = (32, 32) noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor([10]) + time_step = torch.tensor([10]).to(torch_device) return (noise, time_step) def test_from_pretrained_save_pretrained(self): model = UNetModel(ch=32, ch_mult=(1, 2), num_res_blocks=2, attn_resolutions=(16,), resolution=32) + model.to(torch_device) with tempfile.TemporaryDirectory() as tmpdirname: model.save_pretrained(tmpdirname) new_model = UNetModel.from_pretrained(tmpdirname) + new_model.to(torch_device) dummy_input = self.dummy_input @@ -95,6 +97,7 @@ class ModelTesterMixin(unittest.TestCase): def test_from_pretrained_hub(self): model = UNetModel.from_pretrained("fusing/ddpm_dummy") + model.to(torch_device) image = model(*self.dummy_input) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index bcef600896..20943aa525 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -26,8 +26,8 @@ torch.backends.cuda.matmul.allow_tf32 = False class SchedulerCommonTest(unittest.TestCase): - - scheduler_class = None + scheduler_classes = () + forward_default_kwargs = () @property def dummy_image(self): @@ -38,42 +38,271 @@ class SchedulerCommonTest(unittest.TestCase): image = np.random.rand(batch_size, num_channels, height, width) - return image + return torch.tensor(image) + + @property + def dummy_image_deter(self): + batch_size = 4 + num_channels = 3 + height = 8 + width = 8 + + num_elems = batch_size * num_channels * height * width + image = np.arange(num_elems) + image = image.reshape(num_channels, height, width, batch_size) + image = image / num_elems + image = image.transpose(3, 0, 1, 2) + + return torch.tensor(image) def get_scheduler_config(self): raise NotImplementedError def dummy_model(self): - def model(image, residual, t, *args): - return (image + residual) * t / (t + 1) + def model(image, t, *args): + return image * t / (t + 1) return model + def check_over_configs(self, time_step=0, **config): + kwargs = dict(self.forward_default_kwargs) + + for scheduler_class in self.scheduler_classes: + scheduler_class = self.scheduler_classes[0] + image = self.dummy_image + residual = 0.1 * image + + scheduler_config = self.get_scheduler_config(**config) + scheduler = scheduler_class(**scheduler_config) + + with tempfile.TemporaryDirectory() as tmpdirname: + scheduler.save_config(tmpdirname) + new_scheduler = scheduler_class.from_config(tmpdirname) + + output = scheduler.step(residual, image, time_step, **kwargs) + new_output = new_scheduler.step(residual, image, time_step, **kwargs) + + assert (output - new_output).abs().sum() < 1e-5, "Scheduler outputs are not identical" + + def check_over_forward(self, time_step=0, **forward_kwargs): + kwargs = dict(self.forward_default_kwargs) + kwargs.update(forward_kwargs) + + for scheduler_class in self.scheduler_classes: + scheduler_class = self.scheduler_classes[0] + image = self.dummy_image + residual = 0.1 * image + + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + with tempfile.TemporaryDirectory() as tmpdirname: + scheduler.save_config(tmpdirname) + new_scheduler = scheduler_class.from_config(tmpdirname) + + output = scheduler.step(residual, image, time_step, **kwargs) + new_output = new_scheduler.step(residual, image, time_step, **kwargs) + + assert (output - new_output).abs().sum() < 1e-5, "Scheduler outputs are not identical" + def test_from_pretrained_save_pretrained(self): - image = self.dummy_image - residual = 0.1 * image + kwargs = dict(self.forward_default_kwargs) + for scheduler_class in self.scheduler_classes: + image = self.dummy_image + residual = 0.1 * image + + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + with tempfile.TemporaryDirectory() as tmpdirname: + scheduler.save_config(tmpdirname) + new_scheduler = scheduler_class.from_config(tmpdirname) + + output = scheduler.step(residual, image, 1, **kwargs) + new_output = new_scheduler.step(residual, image, 1, **kwargs) + + assert (output - new_output).abs().sum() < 1e-5, "Scheduler outputs are not identical" + + def test_step_shape(self): + kwargs = dict(self.forward_default_kwargs) + + for scheduler_class in self.scheduler_classes: + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + image = self.dummy_image + residual = 0.1 * image + + output_0 = scheduler.step(residual, image, 0, **kwargs) + output_1 = scheduler.step(residual, image, 1, **kwargs) + + self.assertEqual(output_0.shape, image.shape) + self.assertEqual(output_0.shape, output_1.shape) + + +class DDPMSchedulerTest(SchedulerCommonTest): + scheduler_classes = (GaussianDDPMScheduler,) + + def get_scheduler_config(self, **kwargs): + config = { + "timesteps": 1000, + "beta_start": 0.0001, + "beta_end": 0.02, + "beta_schedule": "linear", + "variance_type": "fixed_small", + "clip_predicted_image": True + } + + config.update(**kwargs) + return config + + def test_timesteps(self): + for timesteps in [1, 5, 100, 1000]: + self.check_over_configs(timesteps=timesteps) + + def test_betas(self): + for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1], [0.002, 0.02, 0.2, 2]): + self.check_over_configs(beta_start=beta_start, beta_end=beta_end) + + def test_schedules(self): + for schedule in ["linear", "squaredcos_cap_v2"]: + self.check_over_configs(beta_schedule=schedule) + + def test_variance_type(self): + for variance in ["fixed_small", "fixed_large", "other"]: + self.check_over_configs(variance_type=variance) + + def test_clip_image(self): + for clip_predicted_image in [True, False]: + self.check_over_configs(clip_predicted_image=clip_predicted_image) + + def test_time_indices(self): + for t in [0, 500, 999]: + self.check_over_forward(time_step=t) + + def test_variance(self): + scheduler_class = self.scheduler_classes[0] scheduler_config = self.get_scheduler_config() - scheduler = self.scheduler_class(scheduler_config()) + scheduler = scheduler_class(**scheduler_config) - with tempfile.TemporaryDirectory() as tmpdirname: - scheduler.save_pretrained(tmpdirname) - new_scheduler = self.scheduler_class.from_config(tmpdirname) + assert (scheduler.get_variance(0) - 0.0).abs().sum() < 1e-5 + assert (scheduler.get_variance(487) - 0.00979).abs().sum() < 1e-5 + assert (scheduler.get_variance(999) - 0.02).abs().sum() < 1e-5 - output = scheduler(residual, image, 1) - new_output = new_scheduler(residual, image, 1) - - import ipdb; ipdb.set_trace() - - def test_step(self): + def test_full_loop_no_noise(self): + scheduler_class = self.scheduler_classes[0] scheduler_config = self.get_scheduler_config() - scheduler = self.scheduler_class(scheduler_config()) + scheduler = scheduler_class(**scheduler_config) - image = self.dummy_image - residual = 0.1 * image + num_trained_timesteps = len(scheduler) - output_0 = scheduler(residual, image, 0) - output_1 = scheduler(residual, image, 1) + model = self.dummy_model() + image = self.dummy_image_deter - self.assertEqual(output_0.shape, image.shape) - self.assertEqual(output_0.shape, output_1.shape) + for t in reversed(range(num_trained_timesteps)): + # 1. predict noise residual + residual = model(image, t) + + # 2. predict previous mean of image x_t-1 + pred_prev_image = scheduler.step(residual, image, t) + + if t > 0: + noise = self.dummy_image_deter + variance = scheduler.get_variance(t).sqrt() * noise + + image = pred_prev_image + variance + + result_sum = image.abs().sum() + result_mean = image.abs().mean() + + assert result_sum.item() - 732.9947 < 1e-3 + assert result_mean.item() - 0.9544 < 1e-3 + + +class DDIMSchedulerTest(SchedulerCommonTest): + scheduler_classes = (DDIMScheduler,) + forward_default_kwargs = (("num_inference_steps", 50), ("eta", 0.0)) + + def get_scheduler_config(self, **kwargs): + config = { + "timesteps": 1000, + "beta_start": 0.0001, + "beta_end": 0.02, + "beta_schedule": "linear", + "clip_predicted_image": True + } + + config.update(**kwargs) + return config + + def test_timesteps(self): + for timesteps in [1, 5, 100, 1000]: + self.check_over_configs(timesteps=timesteps) + + def test_betas(self): + for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1], [0.002, 0.02, 0.2, 2]): + self.check_over_configs(beta_start=beta_start, beta_end=beta_end) + + def test_schedules(self): + for schedule in ["linear", "squaredcos_cap_v2"]: + self.check_over_configs(beta_schedule=schedule) + + def test_clip_image(self): + for clip_predicted_image in [True, False]: + self.check_over_configs(clip_predicted_image=clip_predicted_image) + + def test_time_indices(self): + for t in [1, 10, 49]: + self.check_over_forward(time_step=t) + + def test_inference_steps(self): + for t, num_inference_steps in zip([1, 10, 50], [10, 50, 500]): + self.check_over_forward(time_step=t, num_inference_steps=num_inference_steps) + + def test_eta(self): + for t, eta in zip([1, 10, 49], [0.0, 0.5, 1.0]): + self.check_over_forward(time_step=t, eta=eta) + + def test_variance(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + assert (scheduler.get_variance(0, 50) - 0.0).abs().sum() < 1e-5 + assert (scheduler.get_variance(21, 50) - 0.14771).abs().sum() < 1e-5 + assert (scheduler.get_variance(49, 50) - 0.32460).abs().sum() < 1e-5 + assert (scheduler.get_variance(0, 1000) - 0.0).abs().sum() < 1e-5 + assert (scheduler.get_variance(487, 1000) - 0.00979).abs().sum() < 1e-5 + assert (scheduler.get_variance(999, 1000) - 0.02).abs().sum() < 1e-5 + + def test_full_loop_no_noise(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + num_inference_steps, eta = 10, 0.1 + num_trained_timesteps = len(scheduler) + + inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps) + + model = self.dummy_model() + image = self.dummy_image_deter + + for t in reversed(range(num_inference_steps)): + residual = model(image, inference_step_times[t]) + + pred_prev_image = scheduler.step(residual, image, t, num_inference_steps, eta) + + variance = 0 + if eta > 0: + noise = self.dummy_image_deter + variance = scheduler.get_variance(t, num_inference_steps).sqrt() * eta * noise + + image = pred_prev_image + variance + + result_sum = image.abs().sum() + result_mean = image.abs().mean() + + assert result_sum.item() - 270.6214 < 1e-3 + assert result_mean.item() - 0.3524 < 1e-3