Upload diffusion.py with huggingface_hub

diffusion.py

@@ -0,0 +1,322 @@
+"""E3Diff Gaussian Diffusion - exact copy from original with fixed imports."""
+
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+from inspect import isfunction
+from functools import partial
+import numpy as np
+
+
+def _warmup_beta(linear_start, linear_end, n_timestep, warmup_frac):
+    betas = linear_end * np.ones(n_timestep, dtype=np.float64)
+    warmup_time = int(n_timestep * warmup_frac)
+    betas[:warmup_time] = np.linspace(
+        linear_start, linear_end, warmup_time, dtype=np.float64)
+    return betas
+
+
+def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+    if schedule == 'quad':
+        betas = np.linspace(linear_start ** 0.5, linear_end ** 0.5,
+                            n_timestep, dtype=np.float64) ** 2
+    elif schedule == 'linear':
+        betas = np.linspace(linear_start, linear_end,
+                            n_timestep, dtype=np.float64)
+    elif schedule == 'warmup10':
+        betas = _warmup_beta(linear_start, linear_end, n_timestep, 0.1)
+    elif schedule == 'warmup50':
+        betas = _warmup_beta(linear_start, linear_end, n_timestep, 0.5)
+    elif schedule == 'const':
+        betas = linear_end * np.ones(n_timestep, dtype=np.float64)
+    elif schedule == 'jsd':
+        betas = 1. / np.linspace(n_timestep, 1, n_timestep, dtype=np.float64)
+    elif schedule == "cosine":
+        timesteps = (
+            torch.arange(n_timestep + 1, dtype=torch.float64) /
+            n_timestep + cosine_s
+        )
+        alphas = timesteps / (1 + cosine_s) * math.pi / 2
+        alphas = torch.cos(alphas).pow(2)
+        alphas = alphas / alphas[0]
+        betas = 1 - alphas[1:] / alphas[:-1]
+        betas = betas.clamp(max=0.999)
+    else:
+        raise NotImplementedError(schedule)
+    return betas
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+class GaussianDiffusion(nn.Module):
+    def __init__(
+        self,
+        denoise_fn,
+        image_size,
+        channels=3,
+        loss_type='l1',
+        conditional=True,
+        schedule_opt=None,
+        xT_noise_r=0.1,
+        seed=1,
+        opt=None
+    ):
+        super().__init__()
+        self.lq_noiselevel_val = schedule_opt["lq_noiselevel"]
+        self.opt = opt
+        self.channels = channels
+        self.image_size = image_size
+        self.denoise_fn = denoise_fn
+        self.loss_type = loss_type
+        self.conditional = conditional
+        self.ddim = schedule_opt['ddim']
+        self.xT_noise_r = xT_noise_r
+        self.seed = seed
+
+    def set_loss(self, device):
+        if self.loss_type == 'l1':
+            self.loss_func = nn.L1Loss(reduction='sum').to(device)
+        elif self.loss_type == 'l2':
+            self.loss_func = nn.MSELoss(reduction='sum').to(device)
+        else:
+            raise NotImplementedError()
+
+    def set_new_noise_schedule(self, schedule_opt, device, num_train_timesteps=1000):
+        self.ddim = schedule_opt['ddim']
+        self.num_train_timesteps = num_train_timesteps
+        to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
+
+        betas = make_beta_schedule(
+            schedule=schedule_opt['schedule'],
+            n_timestep=num_train_timesteps,
+            linear_start=schedule_opt['linear_start'],
+            linear_end=schedule_opt['linear_end'])
+        betas = betas.detach().cpu().numpy() if isinstance(
+            betas, torch.Tensor) else betas
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+        self.sqrt_alphas_cumprod_prev = np.sqrt(
+            np.append(1., alphas_cumprod))
+
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        self.register_buffer('posterior_variance', to_torch(posterior_variance))
+        self.register_buffer('posterior_log_variance_clipped', to_torch(
+            np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer('posterior_mean_coef1', to_torch(
+            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+        self.register_buffer('posterior_mean_coef2', to_torch(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
+
+        self.schedule_type = schedule_opt['schedule']
+        if self.ddim > 0:
+            self.ddim_num_steps = schedule_opt['n_timestep']
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return self.sqrt_recip_alphas_cumprod[t] * x_t - \
+            self.sqrt_recipm1_alphas_cumprod[t] * noise
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = self.posterior_mean_coef1[t] * \
+            x_start + self.posterior_mean_coef2[t] * x_t
+        posterior_log_variance_clipped = self.posterior_log_variance_clipped[t]
+        return posterior_mean, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, clip_denoised: bool, condition_x=None):
+        batch_size = x.shape[0]
+        noise_level = torch.FloatTensor(
+            [self.sqrt_alphas_cumprod_prev[t+1]]).repeat(batch_size, 1).to(x.device)
+        if condition_x is not None:
+            x_recon = self.predict_start_from_noise(
+                x, t=t, noise=self.denoise_fn(torch.cat([condition_x, x], dim=1), noise_level, t))
+        else:
+            x_recon = self.predict_start_from_noise(
+                x, t=t, noise=self.denoise_fn(x, noise_level))
+
+        if clip_denoised:
+            x_recon.clamp_(-1., 1.)
+
+        model_mean, posterior_log_variance = self.q_posterior(
+            x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_log_variance, x_recon
+
+    def ddim_sample(self, condition_x, img_or_shape, device, seed=1, img_s1=None):
+        if self.schedule_type == 'linear':
+            self.ddim_sampling_eta = 0.8
+            simple_var = False
+            threshold_x = False
+        elif self.schedule_type == 'cosine':
+            self.ddim_sampling_eta = 0.8
+            simple_var = False
+            threshold_x = False
+
+        batch, total_timesteps, sampling_timesteps, eta = \
+            img_or_shape[0], self.num_train_timesteps, \
+            self.ddim_num_steps, self.ddim_sampling_eta
+
+        noisy_img_s1 = None
+
+        if simple_var:
+            eta = 1
+        ts = torch.linspace(total_timesteps, 0, (sampling_timesteps + 1)).to(device).to(torch.long)
+
+        x = torch.randn(img_or_shape).to(device)
+        batch_size = x.shape[0]
+        imgs = [x]
+        img_onestep = [condition_x[:, :self.channels, ...]]
+        
+        tbar = range(1, sampling_timesteps + 1)
+        for i in tbar:
+            cur_t = ts[i - 1] - 1
+            prev_t = ts[i] - 1
+            noise_level = torch.FloatTensor(
+                [self.sqrt_alphas_cumprod_prev[cur_t]]).repeat(batch_size, 1).to(x.device)
+
+            alpha_prod_t = self.alphas_cumprod[cur_t]
+            alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else 1
+            beta_prod_t = 1 - alpha_prod_t
+
+            # pred noise
+            model_output = self.denoise_fn(torch.cat([condition_x, x], dim=1), noise_level)
+
+            sigma_2 = eta * (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+            noise = torch.randn_like(x)
+
+            pred_original_sample = (x - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+
+            if threshold_x:
+                pred_original_sample = self._threshold_sample(pred_original_sample)
+            else:
+                pred_original_sample = pred_original_sample.clamp(-1, 1)
+
+            pred_sample_direction = (1 - alpha_prod_t_prev - sigma_2) ** (0.5) * model_output
+
+            if simple_var:
+                third_term = (1 - alpha_prod_t / alpha_prod_t_prev) ** 0.5 * noise
+            else:
+                third_term = sigma_2 ** 0.5 * noise
+                
+            x = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction + third_term
+            imgs.append(x)
+            img_onestep.append(pred_original_sample)
+
+        imgs = torch.concat(imgs, dim=0)
+        img_onestep = torch.concat(img_onestep, dim=0)
+
+        return imgs, img_onestep
+
+    @torch.no_grad()
+    def p_sample(self, x, t, clip_denoised=True, condition_x=None):
+        model_mean, model_log_variance, x_recon = self.p_mean_variance(
+            x=x, t=t, clip_denoised=clip_denoised, condition_x=condition_x)
+        noise = torch.randn_like(x) if t > 0 else torch.zeros_like(x)
+        return model_mean + noise * (0.5 * model_log_variance).exp(), x_recon
+
+    @torch.no_grad()
+    def p_sample_loop(self, x_in, continous=False, seed=1, img_s1=None):
+        device = self.betas.device
+        sample_inter = 1
+        
+        if not self.conditional:
+            shape = x_in
+            img = torch.randn(shape, device=device)
+            ret_img = img
+            if not self.ddim:
+                for i in reversed(range(0, self.num_timesteps)):
+                    img, x_recon = self.p_sample(img, i)
+                    if i % sample_inter == 0:
+                        ret_img = torch.cat([ret_img, img], dim=0)
+            else:
+                for i in range(0, len(self.ddim_timesteps)):
+                    ddim_t = self.ddim_timesteps[i]
+                    img = self.ddim_sample(img, ddim_t)
+                    if i % sample_inter == 0:
+                        ret_img = torch.cat([ret_img, img], dim=0)
+        else:
+            x = x_in
+            shape = (x.shape[0], self.channels, x.shape[-2], x.shape[-1])
+
+            if self.xT_noise_r > 0:
+                img0 = torch.randn(shape, device=device)
+                x_start = x_in[:, 0:1, ...]
+                continuous_sqrt_alpha_cumprod = torch.FloatTensor(
+                    np.random.uniform(
+                        self.sqrt_alphas_cumprod_prev[self.num_timesteps-1],
+                        self.sqrt_alphas_cumprod_prev[self.num_timesteps],
+                        size=x_start.shape[0]
+                    )).to(x_start.device)
+                continuous_sqrt_alpha_cumprod = continuous_sqrt_alpha_cumprod.view(x_start.shape[0], -1)
+
+                noise = default(x_start, lambda: torch.randn_like(x_start))
+                img = self.q_sample(
+                    x_start=x_start, continuous_sqrt_alpha_cumprod=continuous_sqrt_alpha_cumprod.view(-1, 1, 1, 1), noise=noise)
+                img = self.xT_noise_r * img + (1 - self.xT_noise_r) * img0
+            else:
+                img = torch.randn(shape, device=device)
+
+            ret_img = x
+            img_onestep = x
+
+            if self.opt['stage'] != 2:
+                if not self.ddim:
+                    for i in reversed(range(0, self.num_timesteps)):
+                        img, x_recon = self.p_sample(img, i, condition_x=x)
+                        if i % sample_inter == 0:
+                            ret_img = torch.cat([ret_img[:, :self.channels, ...], img], dim=0)
+                        if i % sample_inter == 0 or i == self.num_timesteps - 1:
+                            img_onestep = torch.cat([img_onestep[:, :self.channels, ...], x_recon], dim=0)
+                else:
+                    ret_img, img_onestep = self.ddim_sample(condition_x=x, img_or_shape=shape, device=device, seed=seed, img_s1=img_s1)
+
+                if continous:
+                    return ret_img, img_onestep
+                else:
+                    return ret_img[-x_in.shape[0]:], img_onestep
+            else:
+                self.ddim_num_steps = self.opt['ddim_steps']
+                ret_img, img_onestep = self.ddim_sample(condition_x=x, img_or_shape=shape, device=device, seed=seed, img_s1=img_s1)
+
+                if continous:
+                    return ret_img, img_onestep
+                else:
+                    return ret_img[-x_in.shape[0]:], img_onestep
+
+    @torch.no_grad()
+    def sample(self, batch_size=1, continous=False):
+        image_size = self.image_size
+        channels = self.channels
+        return self.p_sample_loop((batch_size, channels, image_size, image_size), continous)
+
+    @torch.no_grad()
+    def super_resolution(self, x_in, continous=False, seed=1, img_s1=None):
+        return self.p_sample_loop(x_in, continous, seed=seed, img_s1=img_s1)
+
+    def q_sample(self, x_start, continuous_sqrt_alpha_cumprod, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (
+            continuous_sqrt_alpha_cumprod * x_start +
+            (1 - continuous_sqrt_alpha_cumprod ** 2).sqrt() * noise
+        )