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

models/vision/ddim/example.py

+#!/usr/bin/env python3
+import os
+import pathlib
+from modeling_ddim import DDIM
+import PIL.Image
+import numpy as np
+
+model_ids = ["ddim-celeba-hq", "ddim-lsun-church", "ddim-lsun-bedroom"]
+
+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)
+
+    ddpm = DDIM.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)
+
+    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/ddim/modeling_ddim.py

@@ -19,12 +19,6 @@ import tqdm
 import torch
 
 
-def compute_alpha(beta, t):
-    beta = torch.cat([torch.zeros(1).to(beta.device), beta], dim=0)
-    a = (1 - beta).cumprod(dim=0).index_select(0, t + 1).view(-1, 1, 1, 1)
-    return a
-
-
 class DDIM(DiffusionPipeline):
 
     def __init__(self, unet, noise_scheduler):
@@ -32,7 +26,7 @@ class DDIM(DiffusionPipeline):
         self.register_modules(unet=unet, noise_scheduler=noise_scheduler)
 
     def __call__(self, batch_size=1, generator=None, torch_device=None, eta=0.0, num_inference_steps=50):
-        # eta is η in paper
+        # eta corresponds to η in paper and should be between [0, 1]
         if torch_device is None:
             torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -59,15 +53,20 @@ class DDIM(DiffusionPipeline):
             coeff_1 = (alpha_prod_t_prev - alpha_prod_t).sqrt() * alpha_prod_t_prev_rsqrt * beta_prod_t_prev_sqrt / beta_prod_t_sqrt * eta
             coeff_2 = ((1 - alpha_prod_t_prev) - coeff_1 ** 2).sqrt()
 
+            # model forward
             with torch.no_grad():
                 noise_residual = self.unet(image, train_step)
 
-            print(train_step)
-
-            pred_mean = (image - noise_residual * beta_prod_t_sqrt) * alpha_prod_t_rsqrt
-            xt_next = alpha_prod_t_prev.sqrt() * pred_mean + coeff_1 * torch.randn_like(image) + coeff_2 * noise_residual
-#            xt_next = 1 / alpha_prod_t_rsqrt * pred_mean + coeff_1 * torch.randn_like(image) + coeff_2 * noise_residual
-            # eta
-            image = xt_next
+            # predict mean of prev image
+            pred_mean = alpha_prod_t_rsqrt * (image - beta_prod_t_sqrt * noise_residual)
+            pred_mean = torch.clamp(pred_mean, -1, 1)
+            pred_mean = (1 / alpha_prod_t_prev_rsqrt) * pred_mean + coeff_2 * noise_residual
+
+            # if eta > 0.0 add noise. Note eta = 1.0 essentially corresponds to DDPM
+            if eta > 0.0:
+                noise = self.noise_scheduler.sample_noise(image.shape, device=image.device, generator=generator)
+                image = pred_mean + coeff_1 * noise
+            else:
+                image = pred_mean
 
         return image

src/diffusers/__init__.py

@@ -9,6 +9,7 @@ from .models.clip_text_transformer import CLIPTextModel
 from .models.unet import UNetModel
 from .models.unet_glide import GLIDETextToImageUNetModel, GLIDESuperResUNetModel
 from .models.unet_ldm import UNetLDMModel
+from .models.vqvae import VQModel
 from .pipeline_utils import DiffusionPipeline
 from .schedulers.classifier_free_guidance import ClassifierFreeGuidanceScheduler
 from .schedulers.gaussian_ddpm import GaussianDDPMScheduler

src/diffusers/models/__init__.py

@@ -20,3 +20,4 @@ from .clip_text_transformer import CLIPTextModel
 from .unet import UNetModel
 from .unet_glide import GLIDETextToImageUNetModel, GLIDESuperResUNetModel
 from .unet_ldm import UNetLDMModel
+from .vqvae import VQModel
\ No newline at end of file

src/diffusers/models/unet_ldm.py

@@ -6,7 +6,12 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from einops import repeat, rearrange
+
+try:
+    from einops import repeat, rearrange
+except:
+    print("Einops is not installed")
+    pass
 
 from ..configuration_utils import ConfigMixin
 from ..modeling_utils import ModelMixin