Merge master

Author: anton-l

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,20 +46,20 @@ 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():
         noise_residual = model(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 = 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

@@ -1,20 +1,23 @@
 #!/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/example.py

@@ -23,32 +23,33 @@ 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():
                 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 = 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)

models/vision/ddpm/modeling_ddpm.py

@@ -89,6 +89,7 @@ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool
 
         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 @@ def get_config_dict(
             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