From 27266abc9ff8cda37a892e6ddfdd6a5caab94e66 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 13 Jun 2022 10:39:53 +0200 Subject: [PATCH] rename schedulers --- README.md | 4 ++-- src/diffusers/__init__.py | 6 ++++-- src/diffusers/pipelines/old/ddim/run_ddpm.py | 4 ++-- src/diffusers/pipelines/old/ddpm/run_ddpm.py | 4 ++-- src/diffusers/schedulers/__init__.py | 9 +++++---- src/diffusers/schedulers/{ddim.py => scheduling_ddim.py} | 0 .../schedulers/{gaussian_ddpm.py => scheduling_ddpm.py} | 2 +- tests/test_modeling_utils.py | 6 +++--- tests/test_scheduler.py | 4 ++-- 9 files changed, 21 insertions(+), 18 deletions(-) rename src/diffusers/schedulers/{ddim.py => scheduling_ddim.py} (100%) rename src/diffusers/schedulers/{gaussian_ddpm.py => scheduling_ddpm.py} (98%) diff --git a/README.md b/README.md index 51832cca6f..10d3655b21 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ Both models and schedulers should be load- and saveable from the Hub. ```python import torch -from diffusers import UNetModel, GaussianDDPMScheduler +from diffusers import UNetModel, DDPMScheduler import PIL import numpy as np import tqdm @@ -44,7 +44,7 @@ generator = torch.manual_seed(0) torch_device = "cuda" if torch.cuda.is_available() else "cpu" # 1. Load models -noise_scheduler = GaussianDDPMScheduler.from_config("fusing/ddpm-lsun-church", tensor_format="pt") +noise_scheduler = DDPMScheduler.from_config("fusing/ddpm-lsun-church", tensor_format="pt") unet = UNetModel.from_pretrained("fusing/ddpm-lsun-church").to(torch_device) # 2. Sample gaussian noise diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 9a6132a433..0f7d134b91 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -10,8 +10,10 @@ from .models.unet_glide import GLIDESuperResUNetModel, GLIDETextToImageUNetModel from .models.unet_ldm import UNetLDMModel from .pipeline_utils import DiffusionPipeline from .pipelines import DDIM, DDPM, GLIDE, LatentDiffusion + from .schedulers import SchedulerMixin +from .schedulers.scheduling_ddim import DDIMScheduler +from .schedulers.scheduling_ddpm import DDPMScheduler + from .schedulers.classifier_free_guidance import ClassifierFreeGuidanceScheduler -from .schedulers.ddim import DDIMScheduler -from .schedulers.gaussian_ddpm import GaussianDDPMScheduler from .schedulers.glide_ddim import GlideDDIMScheduler diff --git a/src/diffusers/pipelines/old/ddim/run_ddpm.py b/src/diffusers/pipelines/old/ddim/run_ddpm.py index 88de931381..64752be002 100755 --- a/src/diffusers/pipelines/old/ddim/run_ddpm.py +++ b/src/diffusers/pipelines/old/ddim/run_ddpm.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 import torch -from diffusers import GaussianDDPMScheduler, UNetModel +from diffusers import DDPMScheduler, UNetModel model = UNetModel(dim=64, dim_mults=(1, 2, 4, 8)) -diffusion = GaussianDDPMScheduler(model, image_size=128, timesteps=1000, loss_type="l1") # number of steps # L1 or L2 +diffusion = DDPMScheduler(model, image_size=128, timesteps=1000, loss_type="l1") # number of steps # L1 or L2 training_images = torch.randn(8, 3, 128, 128) # your images need to be normalized from a range of -1 to +1 loss = diffusion(training_images) diff --git a/src/diffusers/pipelines/old/ddpm/run_ddpm.py b/src/diffusers/pipelines/old/ddpm/run_ddpm.py index 88de931381..64752be002 100755 --- a/src/diffusers/pipelines/old/ddpm/run_ddpm.py +++ b/src/diffusers/pipelines/old/ddpm/run_ddpm.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 import torch -from diffusers import GaussianDDPMScheduler, UNetModel +from diffusers import DDPMScheduler, UNetModel model = UNetModel(dim=64, dim_mults=(1, 2, 4, 8)) -diffusion = GaussianDDPMScheduler(model, image_size=128, timesteps=1000, loss_type="l1") # number of steps # L1 or L2 +diffusion = DDPMScheduler(model, image_size=128, timesteps=1000, loss_type="l1") # number of steps # L1 or L2 training_images = torch.randn(8, 3, 128, 128) # your images need to be normalized from a range of -1 to +1 loss = diffusion(training_images) diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 6eb952234f..8cc7284bf7 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -16,8 +16,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .classifier_free_guidance import ClassifierFreeGuidanceScheduler -from .ddim import DDIMScheduler -from .gaussian_ddpm import GaussianDDPMScheduler -from .glide_ddim import GlideDDIMScheduler +from .scheduling_ddim import DDIMScheduler +from .scheduling_ddpm import DDPMScheduler from .schedulers_utils import SchedulerMixin + +from .classifier_free_guidance import ClassifierFreeGuidanceScheduler +from .glide_ddim import GlideDDIMScheduler diff --git a/src/diffusers/schedulers/ddim.py b/src/diffusers/schedulers/scheduling_ddim.py similarity index 100% rename from src/diffusers/schedulers/ddim.py rename to src/diffusers/schedulers/scheduling_ddim.py diff --git a/src/diffusers/schedulers/gaussian_ddpm.py b/src/diffusers/schedulers/scheduling_ddpm.py similarity index 98% rename from src/diffusers/schedulers/gaussian_ddpm.py rename to src/diffusers/schedulers/scheduling_ddpm.py index f64dcaf32a..5d5549cf0e 100644 --- a/src/diffusers/schedulers/gaussian_ddpm.py +++ b/src/diffusers/schedulers/scheduling_ddpm.py @@ -19,7 +19,7 @@ from ..configuration_utils import ConfigMixin from .schedulers_utils import SchedulerMixin, betas_for_alpha_bar, linear_beta_schedule -class GaussianDDPMScheduler(SchedulerMixin, ConfigMixin): +class DDPMScheduler(SchedulerMixin, ConfigMixin): def __init__( self, timesteps=1000, diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 6b84a728d1..9382736d01 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -19,7 +19,7 @@ import unittest import torch -from diffusers import DDIM, DDPM, DDIMScheduler, GaussianDDPMScheduler, LatentDiffusion, UNetModel +from diffusers import DDIM, DDPM, DDIMScheduler, DDPMScheduler, LatentDiffusion, UNetModel from diffusers.configuration_utils import ConfigMixin from diffusers.pipeline_utils import DiffusionPipeline from diffusers.testing_utils import floats_tensor, slow, torch_device @@ -107,7 +107,7 @@ class PipelineTesterMixin(unittest.TestCase): def test_from_pretrained_save_pretrained(self): # 1. Load models model = UNetModel(ch=32, ch_mult=(1, 2), num_res_blocks=2, attn_resolutions=(16,), resolution=32) - schedular = GaussianDDPMScheduler(timesteps=10) + schedular = DDPMScheduler(timesteps=10) ddpm = DDPM(model, schedular) @@ -147,7 +147,7 @@ class PipelineTesterMixin(unittest.TestCase): model_id = "fusing/ddpm-cifar10" unet = UNetModel.from_pretrained(model_id) - noise_scheduler = GaussianDDPMScheduler.from_config(model_id) + noise_scheduler = DDPMScheduler.from_config(model_id) noise_scheduler = noise_scheduler.set_format("pt") ddpm = DDPM(unet=unet, noise_scheduler=noise_scheduler) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 4b4540a286..13d4592372 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -20,7 +20,7 @@ import unittest import numpy as np import torch -from diffusers import DDIMScheduler, GaussianDDPMScheduler +from diffusers import DDIMScheduler, DDPMScheduler torch.backends.cuda.matmul.allow_tf32 = False @@ -163,7 +163,7 @@ class SchedulerCommonTest(unittest.TestCase): class DDPMSchedulerTest(SchedulerCommonTest): - scheduler_classes = (GaussianDDPMScheduler,) + scheduler_classes = (DDPMScheduler,) def get_scheduler_config(self, **kwargs): config = {