Merge branch 'main' of https://github.com/huggingface/diffusers into main

models/vision/glide/modeling_glide.py

+# Copyright 2022 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 diffusers import DiffusionPipeline
+from diffusers import UNetGLIDEModel
+
+import tqdm
+import torch
+
+
+class GLIDE(DiffusionPipeline):
+    def __init__(self, unet: UNetGLIDEModel, noise_scheduler):
+        super().__init__()
+        self.register_modules(unet=unet, noise_scheduler=noise_scheduler)
+
+    def __call__(self, generator=None, torch_device=None):
+        torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.unet.to(torch_device)
+        # 1. Sample gaussian noise
+        image = self.noise_scheduler.sample_noise((1, self.unet.in_channels, self.unet.resolution, self.unet.resolution), device=torch_device, generator=generator)
+        for t in tqdm.tqdm(reversed(range(len(self.noise_scheduler))), total=len(self.noise_scheduler)):
+            # i) define coefficients for time step t
+            clip_image_coeff = 1 / torch.sqrt(self.noise_scheduler.get_alpha_prod(t))
+            clip_noise_coeff = torch.sqrt(1 / self.noise_scheduler.get_alpha_prod(t) - 1)
+            image_coeff = (1 - self.noise_scheduler.get_alpha_prod(t - 1)) * torch.sqrt(self.noise_scheduler.get_alpha(t)) / (1 - self.noise_scheduler.get_alpha_prod(t))
+            clip_coeff = torch.sqrt(self.noise_scheduler.get_alpha_prod(t - 1)) * self.noise_scheduler.get_beta(t) / (1 - self.noise_scheduler.get_alpha_prod(t))
+
+            # ii) predict noise residual
+            with torch.no_grad():
+                noise_residual = self.unet(image, t)
+
+            # iii) compute predicted image from residual
+            # See 2nd formula at https://github.com/hojonathanho/diffusion/issues/5#issue-896554416 for comparison
+            pred_mean = clip_image_coeff * image - clip_noise_coeff * noise_residual
+            pred_mean = torch.clamp(pred_mean, -1, 1)
+            prev_image = clip_coeff * pred_mean + image_coeff * image
+
+            # iv) sample variance
+            prev_variance = self.noise_scheduler.sample_variance(t, prev_image.shape, device=torch_device, generator=generator)
+
+            # v) sample  x_{t-1} ~ N(prev_image, prev_variance)
+            sampled_prev_image = prev_image + prev_variance
+            image = sampled_prev_image
+
+        return image

models/vision/glide/run_glide.py

+import torch
+from .modeling_glide import GLIDE
+from diffusers import UNetGLIDEModel, GaussianDDPMScheduler
+
+generator = torch.Generator()
+generator = generator.manual_seed(0)
+
+# 1. Load models
+scheduler = GaussianDDPMScheduler.from_config("fusing/glide-base")
+model = UNetGLIDEModel.from_pretrained("fusing/glide-base")
+
+pipeline = GLIDE(model, scheduler)
+
+img = pipeline(generator)
+
+print(img)

src/diffusers/__init__.py

@@ -6,5 +6,6 @@ __version__ = "0.0.1"
 
 from .modeling_utils import ModelMixin
 from .models.unet import UNetModel
+from .models.unet_glide import UNetGLIDEModel
 from .pipeline_utils import DiffusionPipeline
 from .schedulers.gaussian_ddpm import GaussianDDPMScheduler

src/diffusers/models/__init__.py

@@ -17,3 +17,4 @@
 # limitations under the License.
 
 from .unet import UNetModel
+from .unet_glide import UNetGLIDEModel

src/diffusers/pipeline_utils.py

@@ -92,11 +92,17 @@ class DiffusionPipeline(ConfigMixin):
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
         # use snapshot download here to get it working from from_pretrained
+        if not os.path.isdir(pretrained_model_name_or_path):
             cached_folder = snapshot_download(pretrained_model_name_or_path)
-        _, config_dict = cls.get_config_dict(cached_folder)
+        else:
+            cached_folder = pretrained_model_name_or_path
 
-        module = config_dict.pop("_module", None)
-        class_name_ = config_dict.pop("_class_name")
+        config_dict, pipeline_kwargs = cls.get_config_dict(cached_folder)
+
+        module = pipeline_kwargs.pop("_module", None)
+        # TODO(Suraj) - make from hub import work
+        # Make `ddpm = DiffusionPipeline.from_pretrained("fusing/ddpm-lsun-bedroom-pipe")` work
+        # Add Sylvains code from transformers
 
         init_kwargs = {}
 
@@ -118,7 +124,10 @@ class DiffusionPipeline(ConfigMixin):
 
             load_method = getattr(class_obj, load_method_name)
 
+            if os.path.isdir(os.path.join(cached_folder, name)):
                 loaded_sub_model = load_method(os.path.join(cached_folder, name))
+            else:
+                loaded_sub_model = load_method(cached_folder)
 
             init_kwargs[name] = loaded_sub_model  # UNet(...), # DiffusionSchedule(...)
 

src/diffusers/schedulers/gaussian_ddpm.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
+import math
 from torch import nn
 
 from ..configuration_utils import ConfigMixin
@@ -24,6 +25,26 @@ def linear_beta_schedule(timesteps, beta_start, beta_end):
     return torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
 
 
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float64)
+
+
 class GaussianDDPMScheduler(nn.Module, ConfigMixin):
 
     config_name = SAMPLING_CONFIG_NAME
@@ -48,6 +69,12 @@ class GaussianDDPMScheduler(nn.Module, ConfigMixin):
 
         if beta_schedule == "linear":
             betas = linear_beta_schedule(timesteps, beta_start=beta_start, beta_end=beta_end)
+        elif beta_schedule == "squaredcos_cap_v2":
+            # GLIDE cosine schedule
+            betas = betas_for_alpha_bar(
+                timesteps,
+                lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+            )
         else:
             raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")