Merge master

README.md

@@ -34,6 +34,7 @@ import numpy as np
 
 generator = torch.Generator()
 generator = generator.manual_seed(6694729458485568)
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 1. Load models
 scheduler = GaussianDDPMScheduler.from_config("fusing/ddpm-lsun-church")
@@ -45,10 +46,10 @@ image = scheduler.sample_noise((1, model.in_channels, model.resolution, model.re
 # 3. Denoise                                                                                                                                           
 for t in reversed(range(len(scheduler))):
     # i) define coefficients for time step t
-    clip_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
-    clip_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
+    clipped_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
+    clipped_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
     image_coeff = (1 - scheduler.get_alpha_prod(t - 1)) * torch.sqrt(scheduler.get_alpha(t)) / (1 - scheduler.get_alpha_prod(t))
-    clip_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
+    clipped_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
 
     # ii) predict noise residual
     with torch.no_grad():
@@ -56,9 +57,9 @@ for t in reversed(range(len(scheduler))):
 
     # 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 = clipped_image_coeff * image - clipped_noise_coeff * noise_residual
     pred_mean = torch.clamp(pred_mean, -1, 1)
-    prev_image = clip_coeff * pred_mean + image_coeff * image
+    prev_image = clipped_coeff * pred_mean + image_coeff * image
 
     # iv) sample variance
     prev_variance = scheduler.sample_variance(t, prev_image.shape, device=torch_device, generator=generator)
@@ -83,12 +84,12 @@ image_pil.save("test.png")
 Example:
 
 ```python
-from modeling_ddpm import DDPM
+from diffusers import DiffusionPipeline
 import PIL.Image
 import numpy as np
 
 # load model and scheduler
-ddpm = DDPM.from_pretrained("fusing/ddpm-lsun-bedroom-pipe")
+ddpm = DiffusionPipeline.from_pretrained("fusing/ddpm-lsun-bedroom")
 
 # run pipeline in inference (sample random noise and denoise)
 image = ddpm()

examples/sample_loop.py

-#!/usr/bin/env python3
-from diffusers import UNetModel, GaussianDDPMScheduler
-import torch
-import torch.nn.functional as F
-import numpy as np
-import PIL.Image
-import tqdm
-
-#torch_device = "cuda"
-#
-#unet = UNetModel.from_pretrained("/home/patrick/ddpm-lsun-church")
-#unet.to(torch_device)
-#
-#TIME_STEPS = 10
-#
-#scheduler = GaussianDDPMScheduler.from_config("/home/patrick/ddpm-lsun-church", timesteps=TIME_STEPS)
-#
-#diffusion_config = {
-#    "beta_start": 0.0001,
-#    "beta_end": 0.02,
-#    "num_diffusion_timesteps": TIME_STEPS,
-#}
-#
-# 2. Do one denoising step with model
-#batch_size, num_channels, height, width = 1, 3, 256, 256
-#
-#torch.manual_seed(0)
-#noise_image = torch.randn(batch_size, num_channels, height, width, device="cuda")
-#
-#
-# Helper
-#def noise_like(shape, device, repeat=False):
-#    def repeat_noise():
-#        return torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-#
-#    def noise():
-#        return torch.randn(shape, device=device)
-#
-#    return repeat_noise() if repeat else noise()
-#
-#
-#betas = np.linspace(diffusion_config["beta_start"], diffusion_config["beta_end"], diffusion_config["num_diffusion_timesteps"], dtype=np.float64)
-#betas = torch.tensor(betas, device=torch_device)
-#alphas = 1.0 - betas
-#
-#alphas_cumprod = torch.cumprod(alphas, axis=0)
-#alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
-#
-#posterior_mean_coef1 = betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
-#posterior_mean_coef2 = (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod)
-#
-#posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
-#posterior_log_variance_clipped = torch.log(posterior_variance.clamp(min=1e-20))
-#
-#
-#sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod)
-#sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod - 1)
-#
-#
-#noise_coeff = (1 - alphas) / torch.sqrt(1 - alphas_cumprod)
-#coeff = 1 / torch.sqrt(alphas)
-
-
-def real_fn():
-    # Compare the following to Algorithm 2 Sampling of paper: https://arxiv.org/pdf/2006.11239.pdf
-    # 1: x_t ~ N(0,1)
-    x_t = noise_image
-    # 2: for t = T, ...., 1 do
-    for i in reversed(range(TIME_STEPS)):
-        t = torch.tensor([i]).to(torch_device)
-        # 3: z ~ N(0, 1)
-        noise = noise_like(x_t.shape, torch_device)
-
-        # 4:  √1αtxt − √1−αt1−α¯tθ(xt, t) + σtz
-        # ------------------------- MODEL ------------------------------------#
-        with torch.no_grad():
-            pred_noise = unet(x_t, t)  # pred epsilon_theta
-
-    #    pred_x = sqrt_recip_alphas_cumprod[t] * x_t - sqrt_recipm1_alphas_cumprod[t] * pred_noise
-    #    pred_x.clamp_(-1.0, 1.0)
-        # pred mean
-    #    posterior_mean = posterior_mean_coef1[t] * pred_x + posterior_mean_coef2[t] * x_t
-        # --------------------------------------------------------------------#
-
-        posterior_mean = coeff[t] * (x_t - noise_coeff[t] * pred_noise)
-
-        # ------------------------- Variance Scheduler -----------------------#
-        # pred variance
-        posterior_log_variance = posterior_log_variance_clipped[t]
-
-        b, *_, device = *x_t.shape, x_t.device
-        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x_t.shape) - 1)))
-        posterior_variance = nonzero_mask * (0.5 * posterior_log_variance).exp()
-        # --------------------------------------------------------------------#
-
-        x_t_1 = (posterior_mean + posterior_variance * noise).to(torch.float32)
-        x_t = x_t_1
-
-        print(x_t.abs().sum())
-
-
-def post_process_to_image(x_t):
-    image = x_t.cpu().permute(0, 2, 3, 1)
-    image = (image + 1.0) * 127.5
-    image = image.numpy().astype(np.uint8)
-
-    return PIL.Image.fromarray(image[0])
-
-
-from pytorch_diffusion import Diffusion
-
-#diffusion = Diffusion.from_pretrained("lsun_church")
-#samples = diffusion.denoise(1)
-#
-#image = post_process_to_image(samples)
-#image.save("check.png")
-#import ipdb; ipdb.set_trace()
-
-
-device = "cuda"
-scheduler = GaussianDDPMScheduler.from_config("/home/patrick/ddpm-lsun-church", timesteps=10)
-
-import ipdb; ipdb.set_trace()
-
-model = UNetModel.from_pretrained("/home/patrick/ddpm-lsun-church").to(device)
-
-
-torch.manual_seed(0)
-next_image = scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=device)
-
-for t in tqdm.tqdm(reversed(range(len(scheduler))), total=len(scheduler)):
-    # define coefficients for time step t
-    clip_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
-    clip_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
-    image_coeff = (1 - scheduler.get_alpha_prod(t - 1)) * torch.sqrt(scheduler.get_alpha(t)) / (1 - scheduler.get_alpha_prod(t))
-    clip_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
-
-    # predict noise residual
-    with torch.no_grad():
-        noise_residual = model(next_image, t)
-
-    # compute prev image from noise
-    pred_mean = clip_image_coeff * next_image - clip_noise_coeff * noise_residual
-    pred_mean = torch.clamp(pred_mean, -1, 1)
-    image = clip_coeff * pred_mean + image_coeff * next_image
-
-    # sample variance
-    variance = scheduler.sample_variance(t, image.shape, device=device)
-
-    # sample previous image
-    sampled_image = image + variance
-
-    next_image = sampled_image
-
-
-image = post_process_to_image(next_image)
-image.save("example_new.png")

models/vision/ddpm/example.py

 #!/usr/bin/env python3
-import tempfile
-import sys
-
+import os
+import pathlib
 from modeling_ddpm import DDPM
+import PIL.Image
+import numpy as np
 
-model_id = sys.argv[1]
+model_ids = ["ddpm-lsun-cat", "ddpm-lsun-cat-ema", "ddpm-lsun-church-ema", "ddpm-lsun-church", "ddpm-lsun-bedroom", "ddpm-lsun-bedroom-ema", "ddpm-cifar10-ema", "ddpm-cifar10", "ddpm-celeba-hq", "ddpm-celeba-hq-ema"]
 
-ddpm = DDPM.from_pretrained(model_id)
-image = ddpm()
+for model_id in model_ids:
+    path = os.path.join("/home/patrick/images/hf", model_id)
+    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
 
-import PIL.Image
-import numpy as np
-image_processed = image.cpu().permute(0, 2, 3, 1)
-image_processed = (image_processed + 1.0) * 127.5
-image_processed = image_processed.numpy().astype(np.uint8)
-image_pil = PIL.Image.fromarray(image_processed[0])
-image_pil.save("test.png")
+    ddpm = DDPM.from_pretrained("fusing/" + model_id)
+    image = ddpm(batch_size=4)
+
+    image_processed = image.cpu().permute(0, 2, 3, 1)
+    image_processed = (image_processed + 1.0) * 127.5
+    image_processed = image_processed.numpy().astype(np.uint8)
 
-import ipdb; ipdb.set_trace()
+    for i in range(image_processed.shape[0]):
+        image_pil = PIL.Image.fromarray(image_processed[i])
+        image_pil.save(os.path.join(path, f"image_{i}.png"))

models/vision/ddpm/modeling_ddpm.py

@@ -23,22 +23,23 @@ class DDPM(DiffusionPipeline):
 
     modeling_file = "modeling_ddpm.py"
 
-    def __init__(self, unet, noise_scheduler, vqvae):
+    def __init__(self, unet, 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"
+    def __call__(self, batch_size=1, generator=None, torch_device=None):
+        if torch_device is 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)
+        image = self.noise_scheduler.sample_noise((batch_size, 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)
+            clipped_image_coeff = 1 / torch.sqrt(self.noise_scheduler.get_alpha_prod(t))
+            clipped_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))
+            clipped_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():
@@ -46,9 +47,9 @@ class DDPM(DiffusionPipeline):
 
             # 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 = clipped_image_coeff * image - clipped_noise_coeff * noise_residual
             pred_mean = torch.clamp(pred_mean, -1, 1)
-            prev_image = clip_coeff * pred_mean + image_coeff * image
+            prev_image = clipped_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)

src/diffusers/configuration_utils.py

@@ -89,6 +89,7 @@ class ConfigMixin:
 
         self.to_json_file(output_config_file)
         logger.info(f"ConfigMixinuration saved in {output_config_file}")
+    
 
     @classmethod
     def get_config_dict(
@@ -182,35 +183,42 @@ class ConfigMixin:
             logger.info(f"loading configuration file {config_file}")
         else:
             logger.info(f"loading configuration file {config_file} from cache at {resolved_config_file}")
+        
+        return config_dict
 
+    @classmethod
+    def extract_init_dict(cls, config_dict, **kwargs):
         expected_keys = set(dict(inspect.signature(cls.__init__).parameters).keys())
         expected_keys.remove("self")
-
+        init_dict = {}
         for key in expected_keys:
             if key in kwargs:
                 # overwrite key
-                config_dict[key] = kwargs.pop(key)
+                init_dict[key] = kwargs.pop(key)
+            elif key in config_dict:
+                # use value from config dict
+                init_dict[key] = config_dict.pop(key)
 
-        passed_keys = set(config_dict.keys())
-
-        unused_kwargs = kwargs
-        for key in passed_keys - expected_keys:
-            unused_kwargs[key] = config_dict.pop(key)
 
+        unused_kwargs = config_dict.update(kwargs)
+        
+        passed_keys = set(init_dict.keys())
         if len(expected_keys - passed_keys) > 0:
             logger.warn(
                 f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values."
             )
 
-        return config_dict, unused_kwargs
+        return init_dict, unused_kwargs
 
     @classmethod
     def from_config(cls, pretrained_model_name_or_path: Union[str, os.PathLike], return_unused_kwargs=False, **kwargs):
-        config_dict, unused_kwargs = cls.get_config_dict(
+        config_dict = cls.get_config_dict(
             pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs
         )
 
-        model = cls(**config_dict)
+        init_dict, unused_kwargs = cls.extract_init_dict(config_dict, **kwargs)
+
+        model = cls(**init_dict)
 
         if return_unused_kwargs:
             return model, unused_kwargs

src/diffusers/dynamic_modules_utils.py

+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. team.
+#
+# 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.
+"""Utilities to dynamically load objects from the Hub."""
+
+import importlib
+import os
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+from huggingface_hub import HfFolder, model_info
+
+from transformers.utils import (
+    HF_MODULES_CACHE,
+    TRANSFORMERS_DYNAMIC_MODULE_NAME,
+    cached_path,
+    hf_bucket_url,
+    is_offline_mode,
+    logging,
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def init_hf_modules():
+    """
+    Creates the cache directory for modules with an init, and adds it to the Python path.
+    """
+    # This function has already been executed if HF_MODULES_CACHE already is in the Python path.
+    if HF_MODULES_CACHE in sys.path:
+        return
+
+    sys.path.append(HF_MODULES_CACHE)
+    os.makedirs(HF_MODULES_CACHE, exist_ok=True)
+    init_path = Path(HF_MODULES_CACHE) / "__init__.py"
+    if not init_path.exists():
+        init_path.touch()
+
+
+def create_dynamic_module(name: Union[str, os.PathLike]):
+    """
+    Creates a dynamic module in the cache directory for modules.
+    """
+    init_hf_modules()
+    dynamic_module_path = Path(HF_MODULES_CACHE) / name
+    # If the parent module does not exist yet, recursively create it.
+    if not dynamic_module_path.parent.exists():
+        create_dynamic_module(dynamic_module_path.parent)
+    os.makedirs(dynamic_module_path, exist_ok=True)
+    init_path = dynamic_module_path / "__init__.py"
+    if not init_path.exists():
+        init_path.touch()
+
+
+def get_relative_imports(module_file):
+    """
+    Get the list of modules that are relatively imported in a module file.
+
+    Args:
+        module_file (`str` or `os.PathLike`): The module file to inspect.
+    """
+    with open(module_file, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    # Imports of the form `import .xxx`
+    relative_imports = re.findall("^\s*import\s+\.(\S+)\s*$", content, flags=re.MULTILINE)
+    # Imports of the form `from .xxx import yyy`
+    relative_imports += re.findall("^\s*from\s+\.(\S+)\s+import", content, flags=re.MULTILINE)
+    # Unique-ify
+    return list(set(relative_imports))
+
+
+def get_relative_import_files(module_file):
+    """
+    Get the list of all files that are needed for a given module. Note that this function recurses through the relative
+    imports (if a imports b and b imports c, it will return module files for b and c).
+
+    Args:
+        module_file (`str` or `os.PathLike`): The module file to inspect.
+    """
+    no_change = False
+    files_to_check = [module_file]
+    all_relative_imports = []
+
+    # Let's recurse through all relative imports
+    while not no_change:
+        new_imports = []
+        for f in files_to_check:
+            new_imports.extend(get_relative_imports(f))
+
+        module_path = Path(module_file).parent
+        new_import_files = [str(module_path / m) for m in new_imports]
+        new_import_files = [f for f in new_import_files if f not in all_relative_imports]
+        files_to_check = [f"{f}.py" for f in new_import_files]
+
+        no_change = len(new_import_files) == 0
+        all_relative_imports.extend(files_to_check)
+
+    return all_relative_imports
+
+
+def check_imports(filename):
+    """
+    Check if the current Python environment contains all the libraries that are imported in a file.
+    """
+    with open(filename, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    # Imports of the form `import xxx`
+    imports = re.findall("^\s*import\s+(\S+)\s*$", content, flags=re.MULTILINE)
+    # Imports of the form `from xxx import yyy`
+    imports += re.findall("^\s*from\s+(\S+)\s+import", content, flags=re.MULTILINE)
+    # Only keep the top-level module
+    imports = [imp.split(".")[0] for imp in imports if not imp.startswith(".")]
+
+    # Unique-ify and test we got them all
+    imports = list(set(imports))
+    missing_packages = []
+    for imp in imports:
+        try:
+            importlib.import_module(imp)
+        except ImportError:
+            missing_packages.append(imp)
+
+    if len(missing_packages) > 0:
+        raise ImportError(
+            "This modeling file requires the following packages that were not found in your environment: "
+            f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`"
+        )
+
+    return get_relative_imports(filename)
+
+
+def get_class_in_module(class_name, module_path):
+    """
+    Import a module on the cache directory for modules and extract a class from it.
+    """
+    module_path = module_path.replace(os.path.sep, ".")
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
+def get_cached_module_file(
+    pretrained_model_name_or_path: Union[str, os.PathLike],
+    module_file: str,
+    cache_dir: Optional[Union[str, os.PathLike]] = None,
+    force_download: bool = False,
+    resume_download: bool = False,
+    proxies: Optional[Dict[str, str]] = None,
+    use_auth_token: Optional[Union[bool, str]] = None,
+    revision: Optional[str] = None,
+    local_files_only: bool = False,
+):
+    """
+    Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached
+    Transformers module.
+
+    Args:
+        pretrained_model_name_or_path (`str` or `os.PathLike`):
+            This can be either:
+
+            - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+              huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced
+              under a user or organization name, like `dbmdz/bert-base-german-cased`.
+            - a path to a *directory* containing a configuration file saved using the
+              [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+        module_file (`str`):
+            The name of the module file containing the class to look for.
+        cache_dir (`str` or `os.PathLike`, *optional*):
+            Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+            cache should not be used.
+        force_download (`bool`, *optional*, defaults to `False`):
+            Whether or not to force to (re-)download the configuration files and override the cached versions if they
+            exist.
+        resume_download (`bool`, *optional*, defaults to `False`):
+            Whether or not to delete incompletely received file. Attempts to resume the download if such a file exists.
+        proxies (`Dict[str, str]`, *optional*):
+            A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+            'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+        use_auth_token (`str` or *bool*, *optional*):
+            The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+            when running `transformers-cli login` (stored in `~/.huggingface`).
+        revision (`str`, *optional*, defaults to `"main"`):
+            The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+            git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+            identifier allowed by git.
+        local_files_only (`bool`, *optional*, defaults to `False`):
+            If `True`, will only try to load the tokenizer configuration from local files.
+
+    <Tip>
+
+    Passing `use_auth_token=True` is required when you want to use a private model.
+
+    </Tip>
+
+    Returns:
+        `str`: The path to the module inside the cache.
+    """
+    # Download and cache module_file from the repo `pretrained_model_name_or_path` of grab it if it's a local file.
+    pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+    module_file_or_url = os.path.join(pretrained_model_name_or_path, module_file)
+    submodule = "local"
+
+    try:
+        # Load from URL or cache if already cached
+        resolved_module_file = cached_path(
+            module_file_or_url,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            resume_download=resume_download,
+            local_files_only=local_files_only,
+            use_auth_token=use_auth_token,
+        )
+
+    except EnvironmentError:
+        logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.")
+        raise
+
+    # Check we have all the requirements in our environment
+    modules_needed = check_imports(resolved_module_file)
+
+    # Now we move the module inside our cached dynamic modules.
+    full_submodule = TRANSFORMERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
+    create_dynamic_module(full_submodule)
+    submodule_path = Path(HF_MODULES_CACHE) / full_submodule
+    # We always copy local files (we could hash the file to see if there was a change, and give them the name of
+    # that hash, to only copy when there is a modification but it seems overkill for now).
+    # The only reason we do the copy is to avoid putting too many folders in sys.path.
+    shutil.copy(resolved_module_file, submodule_path / module_file)
+    for module_needed in modules_needed:
+        module_needed = f"{module_needed}.py"
+        shutil.copy(os.path.join(pretrained_model_name_or_path, module_needed), submodule_path / module_needed)
+    return os.path.join(full_submodule, module_file)
+
+
+def get_class_from_dynamic_module(
+    pretrained_model_name_or_path: Union[str, os.PathLike],
+    module_file: str,
+    class_name: str,
+    cache_dir: Optional[Union[str, os.PathLike]] = None,
+    force_download: bool = False,
+    resume_download: bool = False,
+    proxies: Optional[Dict[str, str]] = None,
+    use_auth_token: Optional[Union[bool, str]] = None,
+    revision: Optional[str] = None,
+    local_files_only: bool = False,
+    **kwargs,
+):
+    """
+    Extracts a class from a module file, present in the local folder or repository of a model.
+
+    <Tip warning={true}>
+
+    Calling this function will execute the code in the module file found locally or downloaded from the Hub. It should
+    therefore only be called on trusted repos.
+
+    </Tip>
+
+    Args:
+        pretrained_model_name_or_path (`str` or `os.PathLike`):
+            This can be either:
+
+            - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+              huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced
+              under a user or organization name, like `dbmdz/bert-base-german-cased`.
+            - a path to a *directory* containing a configuration file saved using the
+              [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+        module_file (`str`):
+            The name of the module file containing the class to look for.
+        class_name (`str`):
+            The name of the class to import in the module.
+        cache_dir (`str` or `os.PathLike`, *optional*):
+            Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+            cache should not be used.
+        force_download (`bool`, *optional*, defaults to `False`):
+            Whether or not to force to (re-)download the configuration files and override the cached versions if they
+            exist.
+        resume_download (`bool`, *optional*, defaults to `False`):
+            Whether or not to delete incompletely received file. Attempts to resume the download if such a file exists.
+        proxies (`Dict[str, str]`, *optional*):
+            A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+            'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+        use_auth_token (`str` or `bool`, *optional*):
+            The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+            when running `transformers-cli login` (stored in `~/.huggingface`).
+        revision (`str`, *optional*, defaults to `"main"`):
+            The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+            git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+            identifier allowed by git.
+        local_files_only (`bool`, *optional*, defaults to `False`):
+            If `True`, will only try to load the tokenizer configuration from local files.
+
+    <Tip>
+
+    Passing `use_auth_token=True` is required when you want to use a private model.
+
+    </Tip>
+
+    Returns:
+        `type`: The class, dynamically imported from the module.
+
+    Examples:
+
+    ```python
+    # Download module `modeling.py` from huggingface.co and cache then extract the class `MyBertModel` from this
+    # module.
+    cls = get_class_from_dynamic_module("sgugger/my-bert-model", "modeling.py", "MyBertModel")
+    ```"""
+    # And lastly we get the class inside our newly created module
+    final_module = get_cached_module_file(
+        pretrained_model_name_or_path,
+        module_file,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        resume_download=resume_download,
+        proxies=proxies,
+        use_auth_token=use_auth_token,
+        revision=revision,
+        local_files_only=local_files_only,
+    )
+    return get_class_in_module(class_name, final_module.replace(".py", ""))

src/diffusers/models/unet_glide.py

@@ -8,6 +8,7 @@ import torch.nn.functional as F
 from ..configuration_utils import ConfigMixin
 from ..modeling_utils import ModelMixin
 
+
 def convert_module_to_f16(l):
     """
     Convert primitive modules to float16.

src/diffusers/pipeline_utils.py

@@ -23,6 +23,7 @@ from huggingface_hub import snapshot_download
 from transformers.utils import logging
 
 from .configuration_utils import ConfigMixin
+from .dynamic_modules_utils import get_class_from_dynamic_module
 
 
 INDEX_FILE = "diffusion_model.pt"
@@ -54,19 +55,23 @@ class DiffusionPipeline(ConfigMixin):
             class_name = module.__class__.__name__
 
             register_dict = {name: (library, class_name)}
-            register_dict["_module"] = self.__module__
+            
 
             # save model index config
             self.register(**register_dict)
 
             # set models
             setattr(self, name, module)
+        
+        register_dict = {"_module" : self.__module__.split(".")[-1] + ".py"}
+        self.register(**register_dict)
 
     def save_pretrained(self, save_directory: Union[str, os.PathLike]):
         self.save_config(save_directory)
 
         model_index_dict = self._dict_to_save
         model_index_dict.pop("_class_name")
+        model_index_dict.pop("_module")
 
         for name, (library_name, class_name) in self._dict_to_save.items():
             importable_classes = LOADABLE_CLASSES[library_name]
@@ -90,17 +95,27 @@ 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
-        cached_folder = snapshot_download(pretrained_model_name_or_path)
-        config_dict, pipeline_kwargs = cls.get_config_dict(cached_folder)
-
-        module = pipeline_kwargs["_module"]
-        # TODO(Suraj) - make from hub import work
-        # Make `ddpm = DiffusionPipeline.from_pretrained("fusing/ddpm-lsun-bedroom-pipe")` work
-        # Add Sylvains code from transformers
+        if not os.path.isdir(pretrained_model_name_or_path):
+            cached_folder = snapshot_download(pretrained_model_name_or_path)
+        else:
+            cached_folder = pretrained_model_name_or_path
+
+        config_dict = cls.get_config_dict(cached_folder)
+        
+        module = config_dict["_module"]
+        class_name_ = config_dict["_class_name"]
+        
+        if class_name_ == cls.__name__:
+            pipeline_class = cls
+        else:
+            pipeline_class = get_class_from_dynamic_module(cached_folder, module, class_name_, cached_folder)
+        
+
+        init_dict, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
 
         init_kwargs = {}
 
-        for name, (library_name, class_name) in config_dict.items():
+        for name, (library_name, class_name) in init_dict.items():
             importable_classes = LOADABLE_CLASSES[library_name]
 
             if library_name == module:
@@ -118,9 +133,12 @@ class DiffusionPipeline(ConfigMixin):
 
             load_method = getattr(class_obj, load_method_name)
 
-            loaded_sub_model = load_method(os.path.join(cached_folder, 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(...)
 
-        model = cls(**init_kwargs)
+        model = pipeline_class(**init_kwargs)
         return model

src/diffusers/schedulers/gaussian_ddpm.py

@@ -108,7 +108,7 @@ class GaussianDDPMScheduler(nn.Module, ConfigMixin):
 
     def sample_variance(self, time_step, shape, device, generator=None):
         variance = self.log_variance[time_step]
-        nonzero_mask = torch.tensor([1 - (time_step == 0)], device=device).float()[None, :].repeat(shape[0], 1)
+        nonzero_mask = torch.tensor([1 - (time_step == 0)], device=device).float()[None, :]
 
         noise = self.sample_noise(shape, device=device, generator=generator)
 

tests/test_modeling_utils.py

@@ -22,6 +22,8 @@ from distutils.util import strtobool
 import torch
 
 from diffusers import GaussianDDPMScheduler, UNetModel
+from diffusers.pipeline_utils import DiffusionPipeline
+from models.vision.ddpm.modeling_ddpm import DDPM
 
 
 global_rng = random.Random()
@@ -76,7 +78,7 @@ def floats_tensor(shape, scale=1.0, rng=None, name=None):
 class ModelTesterMixin(unittest.TestCase):
     @property
     def dummy_input(self):
-        batch_size = 1
+        batch_size = 4
         num_channels = 3
         sizes = (32, 32)
 
@@ -124,10 +126,10 @@ class SamplerTesterMixin(unittest.TestCase):
         # 3. Denoise
         for t in reversed(range(len(scheduler))):
             # i) define coefficients for time step t
-            clip_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
-            clip_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
+            clipped_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
+            clipped_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
             image_coeff = (1 - scheduler.get_alpha_prod(t - 1)) * torch.sqrt(scheduler.get_alpha(t)) / (1 - scheduler.get_alpha_prod(t))
-            clip_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
+            clipped_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
 
             # ii) predict noise residual
             with torch.no_grad():
@@ -135,9 +137,9 @@ class SamplerTesterMixin(unittest.TestCase):
 
             # 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 = clipped_image_coeff * image - clipped_noise_coeff * noise_residual
             pred_mean = torch.clamp(pred_mean, -1, 1)
-            prev_image = clip_coeff * pred_mean + image_coeff * image
+            prev_image = clipped_coeff * pred_mean + image_coeff * image
 
             # iv) sample variance
             prev_variance = scheduler.sample_variance(t, prev_image.shape, device=torch_device, generator=generator)
@@ -174,10 +176,10 @@ class SamplerTesterMixin(unittest.TestCase):
         # 3. Denoise
         for t in reversed(range(len(scheduler))):
             # i) define coefficients for time step t
-            clip_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
-            clip_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
+            clipped_image_coeff = 1 / torch.sqrt(scheduler.get_alpha_prod(t))
+            clipped_noise_coeff = torch.sqrt(1 / scheduler.get_alpha_prod(t) - 1)
             image_coeff = (1 - scheduler.get_alpha_prod(t - 1)) * torch.sqrt(scheduler.get_alpha(t)) / (1 - scheduler.get_alpha_prod(t))
-            clip_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
+            clipped_coeff = torch.sqrt(scheduler.get_alpha_prod(t - 1)) * scheduler.get_beta(t) / (1 - scheduler.get_alpha_prod(t))
 
             # ii) predict noise residual
             with torch.no_grad():
@@ -185,9 +187,9 @@ class SamplerTesterMixin(unittest.TestCase):
 
             # 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 = clipped_image_coeff * image - clipped_noise_coeff * noise_residual
             pred_mean = torch.clamp(pred_mean, -1, 1)
-            prev_image = clip_coeff * pred_mean + image_coeff * image
+            prev_image = clipped_coeff * pred_mean + image_coeff * image
 
             # iv) sample variance
             prev_variance = scheduler.sample_variance(t, prev_image.shape, device=torch_device, generator=generator)
@@ -199,3 +201,46 @@ class SamplerTesterMixin(unittest.TestCase):
         assert image.shape == (1, 3, 256, 256)
         image_slice = image[0, -1, -3:, -3:].cpu()
         assert (image_slice - torch.tensor([[0.1746, 0.5125, -0.7920], [-0.5734, -0.2910, -0.1984], [0.4090, -0.7740, -0.3941]])).abs().sum() < 1e-3
+
+
+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)
+
+        ddpm = DDPM(model, schedular)
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            ddpm.save_pretrained(tmpdirname)
+            new_ddpm = DDPM.from_pretrained(tmpdirname)
+        
+        generator = torch.Generator()
+        generator = generator.manual_seed(669472945848556)
+
+        image = ddpm(generator=generator)
+        generator = generator.manual_seed(669472945848556)
+        new_image = new_ddpm(generator=generator)
+
+        assert (image - new_image).abs().sum() < 1e-5, "Models don't give the same forward pass"
+    
+
+    @slow
+    def test_from_pretrained_hub(self):
+        model_path = "fusing/ddpm-cifar10"
+
+        ddpm = DDPM.from_pretrained(model_path)
+        ddpm_from_hub = DiffusionPipeline.from_pretrained(model_path)
+
+        ddpm.noise_scheduler.num_timesteps = 10
+        ddpm_from_hub.noise_scheduler.num_timesteps = 10
+
+
+        generator = torch.Generator(device=torch_device)
+        generator = generator.manual_seed(669472945848556)
+
+        image = ddpm(generator=generator)
+        generator = generator.manual_seed(669472945848556)
+        new_image = ddpm_from_hub(generator=generator)
+
+        assert (image - new_image).abs().sum() < 1e-5, "Models don't give the same forward pass"