[
  {
    "path": "EMA.py",
    "content": "# ---------------------------------------------------------------\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n#\n# This work is licensed under the NVIDIA Source Code License\n# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.\n# ---------------------------------------------------------------\n\n'''\nCodes adapted from https://github.com/NVlabs/LSGM/blob/main/util/ema.py\n'''\nimport warnings\n\nimport torch\nfrom torch.optim import Optimizer\n\n\nclass EMA(Optimizer):\n    def __init__(self, opt, ema_decay):\n        self.ema_decay = ema_decay\n        self.apply_ema = self.ema_decay > 0.\n        self.optimizer = opt\n        self.state = opt.state\n        self.param_groups = opt.param_groups\n\n    def step(self, *args, **kwargs):\n        retval = self.optimizer.step(*args, **kwargs)\n\n        # stop here if we are not applying EMA\n        if not self.apply_ema:\n            return retval\n\n        ema, params = {}, {}\n        for group in self.optimizer.param_groups:\n            for i, p in enumerate(group['params']):\n                if p.grad is None:\n                    continue\n                state = self.optimizer.state[p]\n\n                # State initialization\n                if 'ema' not in state:\n                    state['ema'] = p.data.clone()\n\n                if p.shape not in params:\n                    params[p.shape] = {'idx': 0, 'data': []}\n                    ema[p.shape] = []\n\n                params[p.shape]['data'].append(p.data)\n                ema[p.shape].append(state['ema'])\n\n            for i in params:\n                params[i]['data'] = torch.stack(params[i]['data'], dim=0)\n                ema[i] = torch.stack(ema[i], dim=0)\n                ema[i].mul_(self.ema_decay).add_(params[i]['data'], alpha=1. - self.ema_decay)\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n                idx = params[p.shape]['idx']\n                self.optimizer.state[p]['ema'] = ema[p.shape][idx, :]\n                params[p.shape]['idx'] += 1\n\n        return retval\n\n    def load_state_dict(self, state_dict):\n        super(EMA, self).load_state_dict(state_dict)\n        # load_state_dict loads the data to self.state and self.param_groups. We need to pass this data to\n        # the underlying optimizer too.\n        self.optimizer.state = self.state\n        self.optimizer.param_groups = self.param_groups\n\n    def swap_parameters_with_ema(self, store_params_in_ema):\n        \"\"\" This function swaps parameters with their ema values. It records original parameters in the ema\n        parameters, if store_params_in_ema is true.\"\"\"\n\n        # stop here if we are not applying EMA\n        if not self.apply_ema:\n            warnings.warn('swap_parameters_with_ema was called when there is no EMA weights.')\n            return\n\n        for group in self.optimizer.param_groups:\n            for i, p in enumerate(group['params']):\n                if not p.requires_grad:\n                    continue\n                ema = self.optimizer.state[p]['ema']\n                if store_params_in_ema:\n                    tmp = p.data.detach()\n                    p.data = ema.detach()\n                    self.optimizer.state[p]['ema'] = tmp\n                else:\n                    p.data = ema.detach()\n"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2022 Sangyun Lee\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.md",
    "content": "# blur-diffusion\nThis is the codebase for our paper Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image Synthesis.\n\n![Teaser image](./images/main_fig.jpg)\n\n> **Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image Synthesis**<br>\n> Sangyun Lee<sup>1</sup>, Hyungjin Chung<sup>2</sup>, Jaehyeon Kim<sup>3</sup>, ‪Jong Chul Ye<sup>2</sup>\n\n> <sup>1</sup>Soongsil University, <sup>2</sup>KAIST, <sup>3</sup>Kakao Enterprise<br>\n\n> Paper: https://arxiv.org/abs/2207.11192<br>\n\n> **Abstract:** *Recently, diffusion models have shown remarkable results in image synthesis by gradually removing noise and amplifying signals. \nAlthough the simple generative process surprisingly works well, is this the best way to generate image data? For instance, despite the fact that human perception is more sensitive to the low-frequencies of an image, diffusion models themselves do not consider any relative importance of each frequency component. Therefore, to incorporate the inductive bias for image data, we propose a novel generative process that synthesizes images in a coarse-to-fine manner. First, we generalize the standard diffusion models by enabling diffusion in a rotated coordinate system with different velocities for each component of the vector. We further propose a blur diffusion as a special case, where each frequency component of an image is diffused at different speeds. Specifically, the proposed blur diffusion consists of a forward process that blurs an image and adds noise gradually, after which  a corresponding reverse process deblurs an image and removes noise progressively. Experiments show that proposed model outperforms the previous method in FID on LSUN bedroom and church datasets.*\n\n## Train\n```bash train.sh```\n## Visualization\n```bash eval_x0hat.sh```\n## Dataset\nImage files are required to compute FID during training.\n"
  },
  {
    "path": "blur_diffusion.py",
    "content": "import torch\nimport torch.nn.functional as F\nimport time\nimport utils\nimport os\nimport numpy as np\n\nfrom utils import clear\n\ndef gaussian_kernel_1d(kernel_size, sigma):\n    assert sigma > 0.00001\n    x_cord = torch.arange(kernel_size)\n    mean = (kernel_size - 1) / 2.\n\n    # pdf of 1d gaussian not considering normalization constant\n    gaussian_kernel = torch.exp(-0.5 * ((x_cord - mean)/sigma)**2)\n    gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)\n\n    return gaussian_kernel\n\n\ndef betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):\n    \"\"\"\n    Ported from https://github.com/openai/guided-diffusion\n    Utility function for cosine noise schedule\n    \"\"\"\n    betas = []\n    for i in range(num_diffusion_timesteps):\n        t1 = i / num_diffusion_timesteps\n        t2 = (i + 1) / num_diffusion_timesteps\n        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))\n    return np.array(betas)\n\nclass ExpSchedule:\n    def __init__(self, N, offset = 1e-4):\n        self.N = N\n        def f(i):\n            #return 1 - np.sin((i / N + 1) * np.pi / 2) + offset\n            return np.exp(5*i/N - 5) + offset\n        idxs = np.arange(N+1)\n        self.alphas_bar = 1 - f(idxs) / f(idxs[-1])\n        self.alphas_bar_left_shifted = 1 - f(idxs-1) / f(idxs[-1])\n        self.alphas = self.alphas_bar / self.alphas_bar_left_shifted\n        self.betas = 1 - self.alphas\n    def get_betas(self):\n        return self.betas\n        \n\n\n# Increase blur strength as i increases\nclass ForwardBlurIncreasing:\n    def __init__(self, N, beta_min, beta_max, sig, sig_min, sig_max, D_diag, blur=None, noise_schedule='linear',\n                 channel=3, resolution=32, pad=None, device='cuda:0', f_type = 'linear'):\n        # N: total number of discretizations used\n        self.device = device\n        self.N = N\n        self.beta_min = beta_min\n        self.beta_max = beta_max\n        self.sig = sig\n        self.sig_min = sig_min # \\sigma_1\n        self.sig_max = sig_max # \\sigma_N\n        \n        self.D_diag = D_diag\n        # total dimension of image\n        self.dim = self.D_diag.shape[0]\n        # \"deblur\" class\n        self.blur = blur\n\n        self.resolution = resolution\n        self.channel = channel\n        self.pad = pad\n\n        self.noise_schedule = noise_schedule\n        if noise_schedule == 'linear':\n            self.betas = torch.linspace(beta_min, beta_max, N, device=device)\n        elif noise_schedule == 'cosine':\n            self.betas = torch.tensor(betas_for_alpha_bar(\n                N,\n                lambda t: np.cos((t + 0.008) / 1.008 * np.pi / 2) ** 2,\n            ), device=device, dtype=torch.float)\n        elif noise_schedule == 'exp':\n            Exp = ExpSchedule(N)\n            betas = torch.tensor(Exp.get_betas())\n            self.betas = torch.tensor(betas, device=device)\n        self.betas = F.pad(self.betas, (1, 0), value=0.0)\n        self.alphas = 1 - self.betas\n        self.sqrt_alphas = torch.sqrt(self.alphas)\n        self.alphas_bar = torch.cumprod(self.alphas, dim=0)\n        self.sqrt_alphas_bar = torch.cumprod(self.sqrt_alphas, dim=0)\n        self.alphas_bar_prev = F.pad(self.alphas_bar, [1, 0], value=1)[:N]\n\n        # def f(i):\n        #     a = (self.sig_max**2 - self.sig_min**2) / (self.sig**2 * (self.N**2 - 1))\n        #     b = (self.sig_min / self.sig) ** 2 - a\n        #     return a * i**2 + b\n        def f_linear(i):\n            f_N = (self.sig_max / self.sig) ** 2\n            f_1 = (self.sig_min / self.sig) ** 2\n            return (f_N - f_1)/ (N-1) * (i - 1) + f_1\n        \n        idxs = torch.tensor(np.arange(N+1), device=device, dtype=torch.float)\n        \n        def f_log(i):\n            log = lambda x: torch.log(x + 1e-6) / (10*np.log(N))\n            f_N = (self.sig_max / self.sig) ** 2\n            f_1 = (self.sig_min / self.sig) ** 2\n            a = (f_N - f_1) / log(torch.tensor(N))\n            b = f_1\n            return a*log(i) + b\n        def f_quadratic(i):\n            f_N = (sig_max / sig) ** 2\n            f_1 = (sig_min / sig) ** 2\n            a = (f_N - f_1) / (N**2 - 1)\n            b = f_1 - (f_N - f_1) / (N**2 - 1)\n            return a*i**2 + b\n        def cubic(i):\n            f_N = (sig_max / sig) ** 2\n            f_1 = (sig_min / sig) ** 2\n            a = (f_N-f_1) / N**3\n            b = f_1\n            return a*i**3 + b\n        def quartic(i):\n            f_N = (sig_max / sig) ** 2\n            f_1 = (sig_min / sig) ** 2\n            a = (f_N-f_1) / N**4\n            b = f_1\n            return a*i**4 + b\n        def triangular(i):\n            less_than_N_2 = (i < N/2).type(torch.int)\n            return f_linear(i) * less_than_N_2 + f_linear(N-i) * (1-less_than_N_2 )\n\n        if f_type == 'linear':\n            f = f_linear\n        elif f_type == 'log':\n            f = f_log\n        elif f_type == 'quadratic':\n            f = f_quadratic\n        elif f_type == 'cubic':\n            f = cubic\n        elif f_type == 'quartic':\n            f = quartic\n        elif f_type == 'triangular':\n            f = triangular\n        else:\n            raise NotImplementedError\n        self.fs = f(idxs)\n        print(\"fs: \", self.fs)\n        self.fs_cum = torch.cumsum(self.fs, dim=0)\n        idxs_long = torch.tensor(np.arange(N+1), device=device, dtype=torch.long)\n        def B(i):\n            p = (2 * f(i)).unsqueeze(-1).repeat(1, self.D_diag.shape[0])\n            print(\"p: \", p.shape)\n            D = self.D_diag.repeat(p.shape[0], 1)\n            print(\"D: \", D.shape)\n            return self.alphas[i].unsqueeze(-1) * self.D_diag ** p\n        self.Bs = B(idxs_long)\n        \n\n        \n        self.Bs_bar = F.pad(torch.cumprod(self.Bs[1:], dim=0), (0,0,1, 0), value=0)\n\n        \n\n       \n        self.one_minus_Bs_bar = 1 - self.Bs_bar\n        self.one_minus_Bs_bar_sqrt = torch.sqrt(self.one_minus_Bs_bar)\n        self.Bs_sqrt = torch.sqrt(self.Bs)\n        self.Bs_squared = self.Bs ** 2\n        self.Bs_bar_sqrt = torch.sqrt(self.Bs_bar)\n        \n        \n        # fname = f'res{resolution}_N{N}_sig{sig}_sigmin_{sig_min}_sigmax_{sig_max}_b_mat_sq-nc{channel}_{noise_schedule}_inc_nopad_rotated.npz'\n        # if os.path.isfile(fname):\n        #     with np.load(fname, allow_pickle=True) as data:\n        #         print(\"data\", list(data.keys()))\n        #         data = data['arr'].astype(np.float32)\n        #         b_mat_sq = torch.tensor(data, device=device)\n        #     self.b_mat_sq = b_mat_sq\n        # else:\n        #     tic = time.time()\n        #     self.b_mat_sq = torch.zeros([N + 1, self.dim]).to(device)\n        # self.b_mat = torch.sqrt(self.b_mat_sq)\n\n    def _beta_i(self, i):\n        return self.betas[i]\n    \n    def _alpha_i(self, i):\n        return self.alphas[i]\n\n    def _alphas_bar_i(self, i):\n        return self.alphas_bar[i]\n\n    def _Bhat_i(self, i):\n        return self.b_mat[i, :]\n\n    # Forward blur\n    def _Bhat_sq_i(self, i):\n        assert i != 0\n        return self.b_mat_sq[i, :]\n\n\n    \n    # Forward blur\n    def get_mean(self, x0, i):\n        mat = self.Bs_bar_sqrt[i]\n        mean = self.blur.U(mat * self.blur.Ut(x0))\n        return mean\n\n    # Forward blur\n    def get_std(self,i, noise):\n        mat = self.one_minus_Bs_bar_sqrt[i]\n        std = self.blur.U(mat * self.blur.Ut(noise))\n        return std\n\n    # Forward blur\n    def get_var(self, x0, i, noise=None):\n        raise NotImplementedError\n        Bhat_i = self._Bhat_sq_i(i)\n        noise = noise if noise is not None else torch.randn_like(x0)\n        std = self.blur.U(Bhat_i * self.blur.Ut(noise))\n        return std\n\n    def W(self, x, i):\n        mat = self.Bs_sqrt[i]\n        batch = x.shape[0]\n        blurred = self.blur.U(mat * self.blur.Ut(x)).view(batch,\n                                          self.channel,\n                                          self.resolution,\n                                          self.resolution)\n        return blurred\n        \n\n    def W_inv(self, x, i):\n        mat = self.Bs_squared[i]\n        batch = x.shape[0]\n        blurred = self.blur.U(mat * self.blur.Ut(x)).view(batch,\n                                          self.channel,\n                                          self.resolution,\n                                          self.resolution)\n        return blurred\n    def U_I_minus_B_Ut(self, x, i): # U(I-B)UT\n        batch = x.shape[0]\n        mat = 1 - self.Bs[i]\n        result = self.blur.U(mat * self.blur.Ut(x)).view(batch,\n                                          self.channel,\n                                          self.resolution,\n                                          self.resolution)\n        return result\n    \n    def U_I_minus_B_sqrt_Ut(self, x, i): # U(sqrt(I-B))UT\n        batch = x.shape[0]\n        mat = torch.sqrt(1 - self.Bs[i])\n        result = self.blur.U(mat * self.blur.Ut(x)).view(batch,\n                                          self.channel,\n                                          self.resolution,\n                                          self.resolution)\n        return result\n    \n    # Forward blur\n    def get_x_i(self, x0, i, return_eps = False):\n        assert 0 not in i\n        if len(x0.shape) == 4:\n            batch = x0.shape[0]\n        # reparam trick\n        mean = self.get_mean(x0, i)\n        noise = torch.randn_like(x0)\n        std = self.get_std(i, noise = noise)\n        # crop\n        if len(x0.shape) == 4:\n            img = (mean + std).view(batch,\n                                                self.channel,\n                                                self.resolution,\n                                                self.resolution)\n        else:\n            img = (mean + std).view(self.channel,\n                                                self.resolution,\n                                                self.resolution)\n        \n        # cropped = cropped + torch.randn_like(cropped) * torch.sqrt(self.betas[i-1])\n        if return_eps:\n            return img, noise\n        return img\n\n    def get_x_N(self, x0_shape, N):\n        \"\"\"\n        prior sampling\n        \"\"\"\n        if len(x0_shape) == 4:\n            batch = x0_shape[0]\n        # pad\n        x0 = torch.zeros(x0_shape, device=self.device)\n        noise = torch.randn_like(x0)\n        std = self.get_std(N, noise)\n        # crop\n        if len(x0_shape) == 4:\n            img = std.view(batch,\n                                   self.channel,\n                                   self.resolution,\n                                   self.resolution)\n        else:\n            img = std.view(self.channel,\n                                   self.resolution,\n                                   self.resolution)\n     \n        return img\n\n    \n    def get_x0_from_eps(self, xi, eps, i):\n        batch = xi.shape[0]\n        std = self.get_std(i, noise = eps).view(batch,\n                                   self.channel,\n                                   self.resolution,\n                                   self.resolution)\n        mean = xi - std\n        return mean / torch.sqrt(self.alphas_bar[i]).view(-1, 1, 1, 1)\n        mat = self.Bs_bar_sqrt[i] ** (-1)\n        x0 = self.blur.U(mat * self.blur.Ut(mean)).view(batch,\n                                   self.channel,\n                                   self.resolution,\n                                   self.resolution)\n        return x0\n    \n    def get_score_gt(self, xi, x0, i):\n        # Return the ground-truth score of x_i\n        assert torch.all(i>0)\n        batch = x0.shape[0]\n        \n        \n        mean = self.get_mean(x0, i).view(batch,\n                                                 self.channel,\n                                                 self.resolution,\n                                                 self.resolution)\n\n        diff = xi - mean\n        mat = (self.one_minus_Bs_bar[i])**(-1)\n        score = -self.blur.U(mat * self.blur.Ut(diff)).view(batch,\n                                                            self.channel,\n                                                            self.resolution,\n                                                            self.resolution)                                 \n        return score\n    def get_score_gt2(self, xi, x0, i):\n        # Return the ground-truth score of x_i\n        batch = x0.shape[0]\n        x0_pad = utils.pad(x0, self.pad)\n        # a_iW^ix_0\n        mean_pad = self.get_mean(x0_pad, i).view(batch,\n                                                 self.channel,\n                                                 self.resolution + self.pad * 2,\n                                                 self.resolution + self.pad * 2)\n\n        mean = utils.crop(mean_pad, self.pad)\n        diff = xi - mean\n        diff = utils.pad(diff, self.pad)\n        score = self.blur.U(self._Bhat_i(i) ** (-2) * self.blur.Ut(diff)).view(batch,\n                                                                               self.channel,\n                                                                               self.resolution + self.pad * 2,\n                                                                               self.resolution + self.pad * 2)\n        score = -utils.crop(score, self.pad) \n        return score\n    def sanity(self, x_0, i):\n        batch = x_0.shape[0]\n        tol = 1e-2\n        xi, eps = self.get_x_i(x_0, i, return_eps = True)\n        score1 = self.get_score_from_eps(eps, i)\n        score2 = self.get_score_gt(xi, x_0, i)\n        # utils.tensor_imsave(score1[0], \"./\", \"score1.jpg\")\n        # utils.tensor_imsave(score2[0], \"./\", \"score2.jpg\")\n        # rms = lambda x: torch.sqrt(torch.mean(x**2))\n        # print(f\"rms(score1) = {rms(score1)}, rms(score2) = {rms(score2)}\")\n        \n        # noise = torch.randn_like(x_0)\n        # noise_pad = utils.pad(noise, self.pad)\n        \n        # std = self.blur.U(self._Bhat_i(i) * self.blur.Ut(noise_pad)).view(batch,\n        #                                          self.channel,\n        #                                          self.resolution + self.pad * 2,\n        #                                          self.resolution + self.pad * 2)\n        # std = utils.crop(std, self.pad)\n        # std = utils.pad(std, self.pad)\n        # # crop -> pad makes error\n        # score1 = self.blur.U(self._Bhat_i(i) ** (-2) * self.blur.Ut(std)).view(batch,\n        #                                                                        self.channel,\n        #                                                                        self.resolution + self.pad * 2,\n        #                                                                        self.resolution + self.pad * 2)\n        # score1 = -utils.crop(score1, self.pad)\n        \n        # score2 = self.blur.U(self._Bhat_i(i) ** (-1) * self.blur.Ut(noise_pad)).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2)\n        # score2 = -utils.crop(score2, self.pad)\n        MAE = torch.mean(torch.abs(score1 - score2))\n        assert MAE < tol, f'MAE = {MAE}'\n        print(f\"MAE = {MAE}\")\n        \n        # score3 = self.get_score_gt2(xi, x_0, i)\n        # print(f\"score3 = {score3[:5]}\")\n        # print(f\"score2 = {score2[:5]}\")\n        \n        # MAE = torch.mean(torch.abs(score3 - score2))\n        # print(f\"MAE = {MAE}, fraction = {MAE/torch.mean(torch.abs(score3))}, norm: {torch.mean(torch.abs(score3))}\")\n        # assert MAE <1e-5, f'MAE = {MAE}'\n        # MAE = torch.mean(torch.abs(self.b_mat_sq[1:] - self.one_minus_Bs_bar[1:]))\n        # assert MAE <1e-5, f'MAE = {MAE}'\n    def get_score_from_eps(self, eps, i):\n        # Return the score of x_i\n        batch = eps.shape[0]\n        mat = self.one_minus_Bs_bar_sqrt[i]**(-1)\n        score = -self.blur.U(mat * self.blur.Ut(eps)).view(batch, self.channel, self.resolution, self.resolution)\n        \n        return score\n    def get_score_from_std(self, std, i):\n        # Return the score of x_i\n        batch = std.shape[0]\n        mat = self.one_minus_Bs_bar[i] ** (-1)\n        score = -self.blur.U(mat * self.blur.Ut(std)).view(batch, self.channel, self.resolution, self.resolution)\n        return score\n    def get_loss_i_exact(self, model, x0, xi, i):\n        # Calculate the DSM objective (exact)\n        batch = x0.shape[0]\n        pred = model(xi, i)\n        score = self.get_score_gt(x0, xi, i)\n        loss = torch.mean((pred - score) ** 2)\n        return loss\n\n    def get_loss_i_simple(self, model, x0, xi, i):\n        raise NotImplementedError\n        # Calculate the DSM objective (simple)\n        batch = x0.shape[0]\n        pred = model(xi, i)\n        # UB**2U^Ts\n        pred = utils.pad(pred, self.pad)\n        pred = self.blur.U(self._Bhat_i(i) ** 2 * self.blur.Ut(pred)).view(batch,\n                                                                           self.channel,\n                                                                           self.resolution + self.pad * 2,\n                                                                           self.resolution + self.pad * 2)\n        pred = utils.crop(pred, self.pad)\n        # ai[UDUT]x0\n        x0_pad = utils.pad(x0, self.pad)\n        mean = self.get_mean(x0_pad, i).view(batch,\n                                             self.channel,\n                                             self.resolution + self.pad * 2,\n                                             self.resolution + self.pad * 2)\n        mean = utils.crop(mean, self.pad)\n        loss = torch.mean((pred + xi - mean) ** 2)\n        return loss\n    def get_loss_i_eps_simple(self, model, x_i, i, eps):\n        pred = model(x_i, i)\n        loss = torch.mean((pred - eps) ** 2)\n        return loss\n    def get_loss_i_std_matching(self, model, x_i, i, eps):\n        pred = model(x_i, i)\n        batch = x_i.shape[0]\n        std = self.get_std(i, noise = eps).view(batch,\n                                   self.channel,\n                                   self.resolution,\n                                   self.resolution)\n        loss = torch.mean((pred - std) ** 2)\n        return loss\n\nclass H_functions:\n    \"\"\"\n    Ported from https://github.com/bahjat-kawar/ddrm\n    A class replacing the SVD of a matrix H, perhaps efficiently.\n    All input vectors are of shape (Batch, ...).\n    All output vectors are of shape (Batch, DataDimension).\n    \"\"\"\n\n    def V(self, vec):\n        \"\"\"\n        Multiplies the input vector by V\n        \"\"\"\n        raise NotImplementedError()\n\n    def Vt(self, vec):\n        \"\"\"\n        Multiplies the input vector by V transposed\n        \"\"\"\n        raise NotImplementedError()\n\n    def U(self, vec):\n        \"\"\"\n        Multiplies the input vector by U\n        \"\"\"\n        raise NotImplementedError()\n\n    def Ut(self, vec):\n        \"\"\"\n        Multiplies the input vector by U transposed\n        \"\"\"\n        raise NotImplementedError()\n\n    def singulars(self):\n        \"\"\"\n        Returns a vector containing the singular values. The shape of the vector should be the same as the smaller dimension (like U)\n        \"\"\"\n        raise NotImplementedError()\n\n    def add_zeros(self, vec):\n        \"\"\"\n        Adds trailing zeros to turn a vector from the small dimension (U) to the big dimension (V)\n        \"\"\"\n        raise NotImplementedError()\n\n    def H(self, vec):\n        \"\"\"\n        Multiplies the input vector by H\n        \"\"\"\n        temp = self.Vt(vec)\n        singulars = self.singulars()\n        return self.U(singulars * temp[:, :singulars.shape[0]])\n\n    def Ht(self, vec):\n        \"\"\"\n        Multiplies the input vector by H transposed\n        \"\"\"\n        temp = self.Ut(vec)\n        singulars = self.singulars()\n        return self.V(self.add_zeros(singulars * temp[:, :singulars.shape[0]]))\n\n    def H_pinv(self, vec):\n        \"\"\"\n        Multiplies the input vector by the pseudo inverse of H\n        \"\"\"\n        temp = self.Ut(vec)\n        singulars = self.singulars()\n        temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] / singulars\n        return self.V(self.add_zeros(temp))\n\n\nclass Deblurring(H_functions):\n    def mat_by_img(self, M, v):\n        return torch.matmul(M, v.reshape(v.shape[0] * self.channels, self.img_dim,\n                                         self.img_dim)).reshape(v.shape[0], self.channels, M.shape[0], self.img_dim)\n\n    def img_by_mat(self, v, M):\n        return torch.matmul(v.reshape(v.shape[0] * self.channels, self.img_dim,\n                                      self.img_dim), M).reshape(v.shape[0], self.channels, self.img_dim, M.shape[1])\n\n    def __init__(self, kernel, channels, img_dim, device):\n        self.img_dim = img_dim\n        self.channels = channels\n        # build 1D conv matrix\n        H_small = torch.zeros(img_dim, img_dim, device=device)\n        for i in range(img_dim):\n            for j in range(i - kernel.shape[0] // 2, i + kernel.shape[0] // 2):\n                if j < 0 or j >= img_dim: continue\n                H_small[i, j] = kernel[j - i + kernel.shape[0] // 2]\n        # torch.cholesky(H_small) # raise error if H_small is not positive definite\n        self.H_small = H_small # positive definite matrix\n        # get the evd of the 1D conv\n        self.U_small, self.singulars_small, _ = torch.svd(H_small, some=False)\n        self.V_small = self.U_small # V should equal U since H is symmetric\n        \n        ZERO = 3e-2\n        # truncate\n        self.singulars_small[self.singulars_small < ZERO] = ZERO\n        # calculate the singular values of the big matrix using Kronecker product\n        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1),\n                                       self.singulars_small.reshape(1, img_dim)).reshape(img_dim ** 2)\n        self._singulars[self._singulars > 1] = 1\n        print(\"contains zero?\", torch.any(self._singulars == 0))\n        # raise NotImplementedError\n        # sort the big matrix singulars and save the permutation\n        self._singulars, self._perm = self._singulars.sort(descending=True)  # , stable=True)\n\n    def V(self, vec):\n        # invert the permutation\n        temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)\n        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels)\n        temp = temp.permute(0, 2, 1)\n        # multiply the image by V from the left and by V^T from the right\n        out = self.mat_by_img(self.V_small, temp)\n        out = self.img_by_mat(out, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1)\n        return out\n\n    def Vt(self, vec):\n        # multiply the image by V^T from the left and by V from the right\n        temp = self.mat_by_img(self.V_small.transpose(0, 1), vec.clone())\n        temp = self.img_by_mat(temp, self.V_small).reshape(vec.shape[0], self.channels, -1)\n        # permute the entries according to the singular values\n        temp = temp[:, :, self._perm].permute(0, 2, 1)\n        return temp.reshape(vec.shape[0], -1)\n\n    def U(self, vec):\n        # invert the permutation\n        temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)\n        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels)\n        temp = temp.permute(0, 2, 1)\n        # multiply the image by U from the left and by U^T from the right\n        out = self.mat_by_img(self.U_small, temp)\n        out = self.img_by_mat(out, self.U_small.transpose(0, 1)).reshape(vec.shape[0], -1)\n        return out\n\n    def Ut(self, vec):\n        # multiply the image by U^T from the left and by U from the right\n        temp = self.mat_by_img(self.U_small.transpose(0, 1), vec.clone())\n        temp = self.img_by_mat(temp, self.U_small).reshape(vec.shape[0], self.channels, -1)\n        # permute the entries according to the singular values\n        temp = temp[:, :, self._perm].permute(0, 2, 1)\n        return temp.reshape(vec.shape[0], -1)\n\n    def conv1d_col_matmul(self, x):\n        return torch.matmul(self.H_small, x)\n\n    def conv1d_row_matmul(self, x):\n        return torch.matmul(x, self.H_small)\n\n    def conv2d_sep_matmul(self, x):\n        return torch.matmul(torch.matmul(self.H_small, x), self.H_small)\n\n    def singulars(self):\n        if self.channels == 3:\n            return self._singulars.repeat(1, 3).reshape(-1)\n        else:\n            return self._singulars.reshape(-1)\n\n    def update_singulars(self, new_singulars):\n        self._singulars = new_singulars\n\n    def add_zeros(self, vec):\n        return vec.clone().reshape(vec.shape[0], -1)"
  },
  {
    "path": "eval_x0hat.py",
    "content": "import torch\nfrom PIL import Image\nfrom collections import defaultdict\nimport torch.nn.functional as TF\nimport torchvision.datasets as dsets\nfrom torchvision import transforms\nimport numpy as np\n\nimport torch.optim as optim\nimport torchvision\nimport matplotlib.pyplot as plt\nimport utils\nfrom guided_diffusion.unet import UNetModel\nimport math\nfrom tensorboardX import SummaryWriter\nimport os\nimport json\nfrom collections import namedtuple\nimport argparse\nfrom torchvision.utils import save_image\nfrom tqdm import tqdm\nfrom blur_diffusion import Deblurring, ForwardBlurIncreasing, gaussian_kernel_1d\nfrom utils import normalize_np, clear\nfrom EMA import EMA\nfrom torch.nn import DataParallel\n\nparser = argparse.ArgumentParser(description='Configs')\nparser.add_argument('--gpu', type=str, help='gpu num')\nparser.add_argument('--name', type=str, help='Saving directory name')\nparser.add_argument('--ckpt', default='', type=str, help='UNet checkpoint')\n\nparser.add_argument('--bsize', default=16, type=int, help='batchsize')\nparser.add_argument('--N', default=1000, type=int, help='Max diffusion timesteps')\nparser.add_argument('--sig', default=0.3, type=float, help='sigma value for blur kernel')\nparser.add_argument('--sig_min', default=0.3, type=float, help='sigma value for blur kernel')\nparser.add_argument('--sig_max', default=1.5, type=float, help='sigma value for blur kernel')\nparser.add_argument('--noise_schedule', default='linear', type=str, help='Type of noise schedule to use')\nparser.add_argument('--betamin', default=0.0001, type=float, help='beta (min). get_score(1) can diverge if this is too low.')\nparser.add_argument('--betamax', default=0.02, type=float, help='beta (max)')\nparser.add_argument('--res',  type=int, help='resolution')\nparser.add_argument('--nc',  type=int, help='channels')\nparser.add_argument('--loss_type', type=str, default = 'sm_simple', choices=['sm_simple', 'eps_simple', 'sm_exact'])\nparser.add_argument('--f_type', type=str, default = 'linear', choices=['linear', 'log', 'quadratic', 'quartic'])\nparser.add_argument('--dropout', default=0, type=float, help='dropout')\nparser.add_argument('--freq_feat', action='store_true', help = \"concat Utx_i\")\nopt = parser.parse_args()\n\nksize = opt.res * 2 - 1\npad = 0\nbsize = opt.bsize\nbeta_min = opt.betamin\nbeta_max = opt.betamax\ndevice = torch.device(f'cuda:{opt.gpu}')\n# define forward blur\nkernel = gaussian_kernel_1d(ksize, opt.sig)\nblur = Deblurring(kernel, opt.nc, opt.res, device=device)\nprint(\"blur.U_small.shape:\", blur.U_small.shape)\nD_diag = blur.singulars()\nfb = ForwardBlurIncreasing(N=opt.N, beta_min=beta_min, beta_max=beta_max, sig=opt.sig, sig_max = opt.sig_max, sig_min = opt.sig_min, D_diag=D_diag,\n                    blur=blur, channel=opt.nc, device=device, noise_schedule=opt.noise_schedule, resolution=opt.res, pad=pad, f_type=opt.f_type)\niter = opt.ckpt.split('/')[-1].split('.')[0]\ndir = os.path.join('experiments', opt.name, f'inferences-{iter}')\nif not os.path.exists(dir):\n    os.mkdir(dir)\nmodel = UNetModel(opt.res, opt.nc, 128, opt.nc, blur = blur, dropout=opt.dropout, freq_feat = opt.freq_feat)\nmodel.load_state_dict(torch.load(opt.ckpt))\nmodel.to(device)\nmodel.eval()\nprint(\"input_nc\", opt.nc, \"resolution\", opt.res)\n\nnum_samples = bsize\nid = 0\nfor _ in tqdm(range(num_samples // bsize)):\n    with torch.no_grad():\n        i = np.array([opt.N - 1] * bsize)\n        i = torch.from_numpy(i).to(device)\n        pred = fb.get_x_N([bsize, opt.nc, opt.res, opt.res], i)\n        x0_list = []\n        for i in reversed(range(1, opt.N)):\n            i = np.array([i] * bsize)\n            i = torch.from_numpy(i).to(device)\n            if opt.loss_type == \"sm_simple\":\n                s = model(pred, i)\n            elif opt.loss_type == \"eps_simple\":\n                eps = model(pred, i)\n                s = fb.get_score_from_eps(eps, i)\n                x0_hat = fb.get_x0_from_eps(pred, eps, i)\n\n                if i[0] % 100 == 0 or i[0] == 1:\n                    x0_list.append(x0_hat)\n            elif opt.loss_type == \"sm_exact\":\n                s = model(pred, i)\n            else:\n                raise NotImplementedError\n            s = fb.U_I_minus_B_Ut(s, i)\n            rms = lambda x: torch.sqrt(torch.mean(x ** 2))\n            print(f\"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}\")\n            hf = pred - fb.W(pred, i)\n            # Anderson theorem\n            pred1 = pred + hf # unsharpening mask filtering\n            pred2 = pred1 + s  # # denoising\n            if i[0] > 2:\n                pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise\n            else:\n                pred = pred2\n            print(f\"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}\" )\n        grid = torch.cat(x0_list, dim=3) # grid_sample.shape: (bsize, channel, H, W * 12)\n        # (batch_size, channel, H, W * 12) -> (channel, H * bsize, W * 12)\n        grid = grid.permute(1, 0, 2, 3).contiguous().view(grid.shape[1], -1, grid.shape[3])\n        # (bsize, channel, H, W) -> (channel, H, W * bsize)\n        save_image(grid, os.path.join(dir, f'xhat.png'))\n"
  },
  {
    "path": "eval_x0hat.sh",
    "content": "#!/bin/bash\nexport OMP_NUM_THREADS=1\nexport CUDA_VISIBLE_DEVICES=1\n\npython eval_x0hat.py \\\n--name test \\\n--noise_schedule linear \\\n--gpu 0 \\\n--bsize 4 \\\n--sig 0.4 \\\n--sig_min 0 \\\n--sig_max 0.15 \\\n--f_type quartic \\\n--res 64 \\\n--nc 3 \\\n--loss_type eps_simple \\\n--ckpt /home/ubuntu/code/sangyoon/forward-blur-new/experiments/lsun-simplified-64-sigmax0.15-quartic/model_450001.ckpt\n"
  },
  {
    "path": "fid.py",
    "content": "from pytorch_fid.fid_score import compute_statistics_of_path, calculate_frechet_distance\nfrom pytorch_fid.inception import InceptionV3\nimport os\nclass FID(object):\n    def __init__(self, real_dir, device, bsize = 128):\n        # https://github.com/mseitzer/pytorch-fid\n        # If real_dir contains a .npz file, then we'll just use that.\n        # Otherwise it's assumed to be a directory.\n        # In that case, statistics will be computed, and .npz file will be saved.\n        if os.path.exists(os.path.join(real_dir, \"inception_statistics.npz\")):\n            real_dir = os.path.join(real_dir, \"inception_statistics.npz\")\n        self.dims = 2048\n        self.bsize = bsize\n        self.device = device\n        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims]\n        self.model = InceptionV3([block_idx]).to(device)\n        m, s = compute_statistics_of_path(real_dir, self.model, bsize, self.dims, device)\n        self.mu_real = m\n        self.sig_real = s\n    def __call__(self, fake_dir):\n        m, s = compute_statistics_of_path(fake_dir, self.model, self.bsize, self.dims, self.device)\n        fid_value = calculate_frechet_distance(m, s, self.mu_real, self.sig_real)\n        return fid_value"
  },
  {
    "path": "guided_diffusion/__init__.py",
    "content": "\"\"\"\nCodebase for \"Improved Denoising Diffusion Probabilistic Models\".\n\"\"\"\n"
  },
  {
    "path": "guided_diffusion/dist_util.py",
    "content": "\"\"\"\nHelpers for distributed training.\n\"\"\"\n\nimport io\nimport os\nimport socket\n\nimport blobfile as bf\nfrom mpi4py import MPI\nimport torch as th\nimport torch.distributed as dist\n\n# Change this to reflect your cluster layout.\n# The GPU for a given rank is (rank % GPUS_PER_NODE).\nGPUS_PER_NODE = 8\n\nSETUP_RETRY_COUNT = 3\n\n\ndef setup_dist():\n    \"\"\"\n    Setup a distributed process group.\n    \"\"\"\n    if dist.is_initialized():\n        return\n    os.environ[\"CUDA_VISIBLE_DEVICES\"] = f\"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}\"\n\n    comm = MPI.COMM_WORLD\n    backend = \"gloo\" if not th.cuda.is_available() else \"nccl\" # backend = \"nccl\"\n    if backend == \"gloo\":\n        hostname = \"localhost\"\n    else:\n        hostname = socket.gethostbyname(socket.getfqdn())\n    # print(hostname)\n    os.environ[\"MASTER_ADDR\"] = comm.bcast(hostname, root=0)\n    os.environ[\"RANK\"] = str(comm.rank)\n    os.environ[\"WORLD_SIZE\"] = str(comm.size)\n\n    port = comm.bcast(_find_free_port(), root=0)\n    os.environ[\"MASTER_PORT\"] = str(port)\n    \n    # print(port)\n    # print(os.environ[\"MASTER_ADDR\"], os.environ[\"RANK\"], os.environ[\"WORLD_SIZE\"])\n    \n    print(r\"dist.init_process_group(backend=backend, init_method=\\\"env://\\\") -- start\")\n    dist.init_process_group(backend=backend, init_method=\"env://\")\n    print(r\"dist.init_process_group(backend=backend, init_method=\\\"env://\\\") -- end\")\n\n\ndef dev():\n    \"\"\"\n    Get the device to use for torch.distributed.\n    \"\"\"\n    if th.cuda.is_available():\n        return th.device(f\"cuda\")\n    return th.device(\"cpu\")\n\n\ndef load_state_dict(path, **kwargs):\n    \"\"\"\n    Load a PyTorch file without redundant fetches across MPI ranks.\n    \"\"\"\n    chunk_size = 2 ** 30  # MPI has a relatively small size limit\n    if MPI.COMM_WORLD.Get_rank() == 0:\n        with bf.BlobFile(path, \"rb\") as f:\n            data = f.read()\n        num_chunks = len(data) // chunk_size\n        if len(data) % chunk_size:\n            num_chunks += 1\n        MPI.COMM_WORLD.bcast(num_chunks)\n        for i in range(0, len(data), chunk_size):\n            MPI.COMM_WORLD.bcast(data[i : i + chunk_size])\n    else:\n        num_chunks = MPI.COMM_WORLD.bcast(None)\n        data = bytes()\n        for _ in range(num_chunks):\n            data += MPI.COMM_WORLD.bcast(None)\n\n    return th.load(io.BytesIO(data), **kwargs)\n\n\ndef sync_params(params):\n    \"\"\"\n    Synchronize a sequence of Tensors across ranks from rank 0.\n    \"\"\"\n    for p in params:\n        with th.no_grad():\n            dist.broadcast(p, 0)\n\n\ndef _find_free_port():\n    try:\n        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n        s.bind((\"\", 0))\n        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)\n        return s.getsockname()[1]\n    finally:\n        s.close()\n"
  },
  {
    "path": "guided_diffusion/fp16_util.py",
    "content": "\"\"\"\nHelpers to train with 16-bit precision.\n\"\"\"\n\nimport numpy as np\nimport torch as th\nimport torch.nn as nn\nfrom torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors\n\nfrom . import logger\n\nINITIAL_LOG_LOSS_SCALE = 20.0\n\n\ndef convert_module_to_f16(l):\n    \"\"\"\n    Convert primitive modules to float16.\n    \"\"\"\n    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):\n        l.weight.data = l.weight.data.half()\n        if l.bias is not None:\n            l.bias.data = l.bias.data.half()\n\n\ndef convert_module_to_f32(l):\n    \"\"\"\n    Convert primitive modules to float32, undoing convert_module_to_f16().\n    \"\"\"\n    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):\n        l.weight.data = l.weight.data.float()\n        if l.bias is not None:\n            l.bias.data = l.bias.data.float()\n\n\ndef make_master_params(param_groups_and_shapes):\n    \"\"\"\n    Copy model parameters into a (differently-shaped) list of full-precision\n    parameters.\n    \"\"\"\n    master_params = []\n    for param_group, shape in param_groups_and_shapes:\n        master_param = nn.Parameter(\n            _flatten_dense_tensors(\n                [param.detach().float() for (_, param) in param_group]\n            ).view(shape)\n        )\n        master_param.requires_grad = True\n        master_params.append(master_param)\n    return master_params\n\n\ndef model_grads_to_master_grads(param_groups_and_shapes, master_params):\n    \"\"\"\n    Copy the gradients from the model parameters into the master parameters\n    from make_master_params().\n    \"\"\"\n    for master_param, (param_group, shape) in zip(\n        master_params, param_groups_and_shapes\n    ):\n        master_param.grad = _flatten_dense_tensors(\n            [param_grad_or_zeros(param) for (_, param) in param_group]\n        ).view(shape)\n\n\ndef master_params_to_model_params(param_groups_and_shapes, master_params):\n    \"\"\"\n    Copy the master parameter data back into the model parameters.\n    \"\"\"\n    # Without copying to a list, if a generator is passed, this will\n    # silently not copy any parameters.\n    for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes):\n        for (_, param), unflat_master_param in zip(\n            param_group, unflatten_master_params(param_group, master_param.view(-1))\n        ):\n            param.detach().copy_(unflat_master_param)\n\n\ndef unflatten_master_params(param_group, master_param):\n    return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group])\n\n\ndef get_param_groups_and_shapes(named_model_params):\n    named_model_params = list(named_model_params)\n    scalar_vector_named_params = (\n        [(n, p) for (n, p) in named_model_params if p.ndim <= 1],\n        (-1),\n    )\n    matrix_named_params = (\n        [(n, p) for (n, p) in named_model_params if p.ndim > 1],\n        (1, -1),\n    )\n    return [scalar_vector_named_params, matrix_named_params]\n\n\ndef master_params_to_state_dict(\n    model, param_groups_and_shapes, master_params, use_fp16\n):\n    if use_fp16:\n        state_dict = model.state_dict()\n        for master_param, (param_group, _) in zip(\n            master_params, param_groups_and_shapes\n        ):\n            for (name, _), unflat_master_param in zip(\n                param_group, unflatten_master_params(param_group, master_param.view(-1))\n            ):\n                assert name in state_dict\n                state_dict[name] = unflat_master_param\n    else:\n        state_dict = model.state_dict()\n        for i, (name, _value) in enumerate(model.named_parameters()):\n            assert name in state_dict\n            state_dict[name] = master_params[i]\n    return state_dict\n\n\ndef state_dict_to_master_params(model, state_dict, use_fp16):\n    if use_fp16:\n        named_model_params = [\n            (name, state_dict[name]) for name, _ in model.named_parameters()\n        ]\n        param_groups_and_shapes = get_param_groups_and_shapes(named_model_params)\n        master_params = make_master_params(param_groups_and_shapes)\n    else:\n        master_params = [state_dict[name] for name, _ in model.named_parameters()]\n    return master_params\n\n\ndef zero_master_grads(master_params):\n    for param in master_params:\n        param.grad = None\n\n\ndef zero_grad(model_params):\n    for param in model_params:\n        # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group\n        if param.grad is not None:\n            param.grad.detach_()\n            param.grad.zero_()\n\n\ndef param_grad_or_zeros(param):\n    if param.grad is not None:\n        return param.grad.data.detach()\n    else:\n        return th.zeros_like(param)\n\n\nclass MixedPrecisionTrainer:\n    def __init__(\n        self,\n        *,\n        model,\n        use_fp16=False,\n        fp16_scale_growth=1e-3,\n        initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE,\n    ):\n        self.model = model\n        self.use_fp16 = use_fp16\n        self.fp16_scale_growth = fp16_scale_growth\n\n        self.model_params = list(self.model.parameters())\n        self.master_params = self.model_params\n        self.param_groups_and_shapes = None\n        self.lg_loss_scale = initial_lg_loss_scale\n\n        if self.use_fp16:\n            self.param_groups_and_shapes = get_param_groups_and_shapes(\n                self.model.named_parameters()\n            )\n            self.master_params = make_master_params(self.param_groups_and_shapes)\n            self.model.convert_to_fp16()\n\n    def zero_grad(self):\n        zero_grad(self.model_params)\n\n    def backward(self, loss: th.Tensor):\n        if self.use_fp16:\n            loss_scale = 2 ** self.lg_loss_scale\n            (loss * loss_scale).backward()\n        else:\n            loss.backward()\n\n    def optimize(self, opt: th.optim.Optimizer):\n        if self.use_fp16:\n            return self._optimize_fp16(opt)\n        else:\n            return self._optimize_normal(opt)\n\n    def _optimize_fp16(self, opt: th.optim.Optimizer):\n        logger.logkv_mean(\"lg_loss_scale\", self.lg_loss_scale)\n        model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params)\n        grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale)\n        if check_overflow(grad_norm):\n            self.lg_loss_scale -= 1\n            logger.log(f\"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}\")\n            zero_master_grads(self.master_params)\n            return False\n\n        logger.logkv_mean(\"grad_norm\", grad_norm)\n        logger.logkv_mean(\"param_norm\", param_norm)\n\n        self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale))\n        opt.step()\n        zero_master_grads(self.master_params)\n        master_params_to_model_params(self.param_groups_and_shapes, self.master_params)\n        self.lg_loss_scale += self.fp16_scale_growth\n        return True\n\n    def _optimize_normal(self, opt: th.optim.Optimizer):\n        grad_norm, param_norm = self._compute_norms()\n        logger.logkv_mean(\"grad_norm\", grad_norm)\n        logger.logkv_mean(\"param_norm\", param_norm)\n        opt.step()\n        return True\n\n    def _compute_norms(self, grad_scale=1.0):\n        grad_norm = 0.0\n        param_norm = 0.0\n        for p in self.master_params:\n            with th.no_grad():\n                param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2\n                if p.grad is not None:\n                    grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2\n        return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm)\n\n    def master_params_to_state_dict(self, master_params):\n        return master_params_to_state_dict(\n            self.model, self.param_groups_and_shapes, master_params, self.use_fp16\n        )\n\n    def state_dict_to_master_params(self, state_dict):\n        return state_dict_to_master_params(self.model, state_dict, self.use_fp16)\n\n\ndef check_overflow(value):\n    return (value == float(\"inf\")) or (value == -float(\"inf\")) or (value != value)\n"
  },
  {
    "path": "guided_diffusion/gaussian_diffusion.py",
    "content": "\"\"\"\nThis code started out as a PyTorch port of Ho et al's diffusion models:\nhttps://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py\n\nDocstrings have been added, as well as DDIM sampling and a new collection of beta schedules.\n\"\"\"\n\nimport enum\nimport math\n\nimport numpy as np\nimport torch as th\nfrom torchvision import transforms\nfrom .nn import mean_flat\nfrom .losses import normal_kl, discretized_gaussian_log_likelihood\nimport os\nfrom PIL import Image\ndef get_named_beta_schedule(schedule_name, num_diffusion_timesteps):\n    \"\"\"\n    Get a pre-defined beta schedule for the given name.\n\n    The beta schedule library consists of beta schedules which remain similar\n    in the limit of num_diffusion_timesteps.\n    Beta schedules may be added, but should not be removed or changed once\n    they are committed to maintain backwards compatibility.\n    \"\"\"\n    if schedule_name == \"linear\":\n        # Linear schedule from Ho et al, extended to work for any number of\n        # diffusion steps.\n        scale = 1000 / num_diffusion_timesteps\n        beta_start = scale * 0.0001\n        beta_end = scale * 0.02\n        return np.linspace(\n            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64\n        )\n    elif schedule_name == \"cosine\":\n        return betas_for_alpha_bar(\n            num_diffusion_timesteps,\n            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,\n        )\n    else:\n        raise NotImplementedError(f\"unknown beta schedule: {schedule_name}\")\n\n\ndef betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):\n    \"\"\"\n    Create a beta schedule that discretizes the given alpha_t_bar function,\n    which defines the cumulative product of (1-beta) over time from t = [0,1].\n\n    :param num_diffusion_timesteps: the number of betas to produce.\n    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and\n                      produces the cumulative product of (1-beta) up to that\n                      part of the diffusion process.\n    :param max_beta: the maximum beta to use; use values lower than 1 to\n                     prevent singularities.\n    \"\"\"\n    betas = []\n    for i in range(num_diffusion_timesteps):\n        t1 = i / num_diffusion_timesteps\n        t2 = (i + 1) / num_diffusion_timesteps\n        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))\n    return np.array(betas)\n\n\nclass ModelMeanType(enum.Enum):\n    \"\"\"\n    Which type of output the model predicts.\n    \"\"\"\n\n    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}\n    START_X = enum.auto()  # the model predicts x_0\n    EPSILON = enum.auto()  # the model predicts epsilon\n\n\nclass ModelVarType(enum.Enum):\n    \"\"\"\n    What is used as the model's output variance.\n\n    The LEARNED_RANGE option has been added to allow the model to predict\n    values between FIXED_SMALL and FIXED_LARGE, making its job easier.\n    \"\"\"\n\n    LEARNED = enum.auto()\n    FIXED_SMALL = enum.auto()\n    FIXED_LARGE = enum.auto()\n    LEARNED_RANGE = enum.auto()\n\n\nclass LossType(enum.Enum):\n    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)\n    RESCALED_MSE = (\n        enum.auto()\n    )  # use raw MSE loss (with RESCALED_KL when learning variances)\n    KL = enum.auto()  # use the variational lower-bound\n    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB\n\n    def is_vb(self):\n        return self == LossType.KL or self == LossType.RESCALED_KL\n\n\nclass GaussianDiffusion:\n    \"\"\"\n    Utilities for training and sampling diffusion models.\n\n    Ported directly from here, and then adapted over time to further experimentation.\n    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42\n\n    :param betas: a 1-D numpy array of betas for each diffusion timestep,\n                  starting at T and going to 1.\n    :param model_mean_type: a ModelMeanType determining what the model outputs.\n    :param model_var_type: a ModelVarType determining how variance is output.\n    :param loss_type: a LossType determining the loss function to use.\n    :param rescale_timesteps: if True, pass floating point timesteps into the\n                              model so that they are always scaled like in the\n                              original paper (0 to 1000).\n    \"\"\"\n\n    def __init__(\n        self,\n        *,\n        betas,\n        model_mean_type,\n        model_var_type,\n        loss_type,\n        rescale_timesteps=False,\n    ):\n        self.model_mean_type = model_mean_type\n        self.model_var_type = model_var_type\n        self.loss_type = loss_type\n        self.rescale_timesteps = rescale_timesteps\n\n        # Use float64 for accuracy.\n        betas = np.array(betas, dtype=np.float64)\n        self.betas = betas\n        assert len(betas.shape) == 1, \"betas must be 1-D\"\n        assert (betas > 0).all() and (betas <= 1).all()\n\n        self.num_timesteps = int(betas.shape[0])\n        \n        alphas = 1.0 - betas\n        self.alphas_cumprod = np.cumprod(alphas, axis=0)\n        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])\n        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)\n        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)\n\n        # calculations for diffusion q(x_t | x_{t-1}) and others\n        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)\n        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)\n        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)\n        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)\n        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)\n\n        # calculations for posterior q(x_{t-1} | x_t, x_0)\n        self.posterior_variance = (\n            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)\n        )\n        # log calculation clipped because the posterior variance is 0 at the\n        # beginning of the diffusion chain.\n        self.posterior_log_variance_clipped = np.log(\n            np.append(self.posterior_variance[1], self.posterior_variance[1:])\n        )\n        self.posterior_mean_coef1 = (\n            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)\n        )\n        self.posterior_mean_coef2 = (\n            (1.0 - self.alphas_cumprod_prev)\n            * np.sqrt(alphas)\n            / (1.0 - self.alphas_cumprod)\n        )\n\n    def q_mean_variance(self, x_start, t):\n        \"\"\"\n        Get the distribution q(x_t | x_0).\n\n        :param x_start: the [N x C x ...] tensor of noiseless inputs.\n        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.\n        :return: A tuple (mean, variance, log_variance), all of x_start's shape.\n        \"\"\"\n        mean = (\n            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start\n        )\n        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)\n        log_variance = _extract_into_tensor(\n            self.log_one_minus_alphas_cumprod, t, x_start.shape\n        )\n        return mean, variance, log_variance\n\n    def q_sample(self, x_start, t, noise=None):\n        \"\"\"\n        Diffuse the data for a given number of diffusion steps.\n\n        In other words, sample from q(x_t | x_0).\n\n        :param x_start: the initial data batch.\n        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.\n        :param noise: if specified, the split-out normal noise.\n        :return: A noisy version of x_start.\n        \"\"\"\n        raise NotImplementedError\n        if noise is None:\n            noise = th.randn_like(x_start)\n        assert noise.shape == x_start.shape\n        return (\n            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start\n            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)\n            * noise\n        )\n    def q_sample_blur(self, x_start, t, noise=None):\n        \"\"\"\n        Diffuse the data for a given number of diffusion steps.\n\n        In other words, sample from q(x_t | x_0).\n\n        :param x_start: the initial data batch.\n        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.\n        :param noise: if specified, the split-out normal noise.\n        :return: A noisy version of x_start.\n        \"\"\"\n        \n        assert th.all(t == t[0])\n        assert float((5*t[0]+1)/self.num_timesteps) >0 , f\"t[0]: {t[0]}\"\n        out = transforms.functional.gaussian_blur(x_start, 15, float((10*t[0]+1)/self.num_timesteps))\n\n        return out\n\n    def q_posterior_mean_variance(self, x_start, x_t, t):\n        \"\"\"\n        Compute the mean and variance of the diffusion posterior:\n\n            q(x_{t-1} | x_t, x_0)\n\n        \"\"\"\n        assert x_start.shape == x_t.shape\n        posterior_mean = (\n            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start\n            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t\n        )\n        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)\n        posterior_log_variance_clipped = _extract_into_tensor(\n            self.posterior_log_variance_clipped, t, x_t.shape\n        )\n        assert (\n            posterior_mean.shape[0]\n            == posterior_variance.shape[0]\n            == posterior_log_variance_clipped.shape[0]\n            == x_start.shape[0]\n        )\n        return posterior_mean, posterior_variance, posterior_log_variance_clipped\n\n    def p_mean_variance(\n        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None\n    ):\n        \"\"\"\n        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of\n        the initial x, x_0.\n\n        :param model: the model, which takes a signal and a batch of timesteps\n                      as input.\n        :param x: the [N x C x ...] tensor at time t.\n        :param t: a 1-D Tensor of timesteps.\n        :param clip_denoised: if True, clip the denoised signal into [-1, 1].\n        :param denoised_fn: if not None, a function which applies to the\n            x_start prediction before it is used to sample. Applies before\n            clip_denoised.\n        :param model_kwargs: if not None, a dict of extra keyword arguments to\n            pass to the model. This can be used for conditioning.\n        :return: a dict with the following keys:\n                 - 'mean': the model mean output.\n                 - 'variance': the model variance output.\n                 - 'log_variance': the log of 'variance'.\n                 - 'pred_xstart': the prediction for x_0.\n        \"\"\"\n        if model_kwargs is None:\n            model_kwargs = {}\n\n        B, C = x.shape[:2]\n        assert t.shape == (B,)\n        model_output = model(x, self._scale_timesteps(t), **model_kwargs)\n\n        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:\n            assert model_output.shape == (B, C * 2, *x.shape[2:])\n            model_output, model_var_values = th.split(model_output, C, dim=1)\n            if self.model_var_type == ModelVarType.LEARNED:\n                model_log_variance = model_var_values\n                model_variance = th.exp(model_log_variance)\n            else:\n                min_log = _extract_into_tensor(\n                    self.posterior_log_variance_clipped, t, x.shape\n                )\n                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)\n                # The model_var_values is [-1, 1] for [min_var, max_var].\n                frac = (model_var_values + 1) / 2\n                model_log_variance = frac * max_log + (1 - frac) * min_log\n                model_variance = th.exp(model_log_variance)\n        else:\n            model_variance, model_log_variance = {\n                # for fixedlarge, we set the initial (log-)variance like so\n                # to get a better decoder log likelihood.\n                ModelVarType.FIXED_LARGE: (\n                    np.append(self.posterior_variance[1], self.betas[1:]),\n                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),\n                ),\n                ModelVarType.FIXED_SMALL: (\n                    self.posterior_variance,\n                    self.posterior_log_variance_clipped,\n                ),\n            }[self.model_var_type]\n            model_variance = _extract_into_tensor(model_variance, t, x.shape)\n            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)\n\n        def process_xstart(x):\n            if denoised_fn is not None:\n                x = denoised_fn(x)\n            if clip_denoised:\n                return x.clamp(-1, 1)\n            return x\n\n        if self.model_mean_type == ModelMeanType.PREVIOUS_X:\n            pred_xstart = process_xstart(\n                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)\n            )\n            model_mean = model_output\n        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:\n            if self.model_mean_type == ModelMeanType.START_X:\n                pred_xstart = process_xstart(model_output)\n            else:\n                pred_xstart = process_xstart(\n                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)\n                )\n            model_mean, _, _ = self.q_posterior_mean_variance(\n                x_start=pred_xstart, x_t=x, t=t\n            )\n        else:\n            raise NotImplementedError(self.model_mean_type)\n\n        assert (\n            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape\n        )\n        return {\n            \"mean\": model_mean,\n            \"variance\": model_variance,\n            \"log_variance\": model_log_variance,\n            \"pred_xstart\": pred_xstart,\n        }\n\n    def _predict_xstart_from_eps(self, x_t, t, eps):\n        assert x_t.shape == eps.shape\n        return (\n            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t\n            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps\n        )\n\n    def _predict_xstart_from_xprev(self, x_t, t, xprev):\n        assert x_t.shape == xprev.shape\n        return (  # (xprev - coef2*x_t) / coef1\n            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev\n            - _extract_into_tensor(\n                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape\n            )\n            * x_t\n        )\n\n    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):\n        return (\n            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t\n            - pred_xstart\n        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)\n\n    def _scale_timesteps(self, t):\n        if self.rescale_timesteps:\n            return t.float() * (1000.0 / self.num_timesteps)\n        return t\n\n    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):\n        \"\"\"\n        Compute the mean for the previous step, given a function cond_fn that\n        computes the gradient of a conditional log probability with respect to\n        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to\n        condition on y.\n\n        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).\n        \"\"\"\n        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)\n        new_mean = (\n            p_mean_var[\"mean\"].float() + p_mean_var[\"variance\"] * gradient.float()\n        )\n        return new_mean\n\n    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):\n        \"\"\"\n        Compute what the p_mean_variance output would have been, should the\n        model's score function be conditioned by cond_fn.\n\n        See condition_mean() for details on cond_fn.\n\n        Unlike condition_mean(), this instead uses the conditioning strategy\n        from Song et al (2020).\n        \"\"\"\n        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)\n\n        eps = self._predict_eps_from_xstart(x, t, p_mean_var[\"pred_xstart\"])\n        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(\n            x, self._scale_timesteps(t), **model_kwargs\n        )\n\n        out = p_mean_var.copy()\n        out[\"pred_xstart\"] = self._predict_xstart_from_eps(x, t, eps)\n        out[\"mean\"], _, _ = self.q_posterior_mean_variance(\n            x_start=out[\"pred_xstart\"], x_t=x, t=t\n        )\n        return out\n    def p_sample(\n        self,\n        model,\n        x,\n        t,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n    ):\n        raise NotImplementedError\n        \"\"\"\n        Sample x_{t-1} from the model at the given timestep.\n        :param model: the model to sample from.\n        :param x: the current tensor at x_{t-1}.\n        :param t: the value of t, starting at 0 for the first diffusion step.\n        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].\n        :param denoised_fn: if not None, a function which applies to the\n            x_start prediction before it is used to sample.\n        :param cond_fn: if not None, this is a gradient function that acts\n                        similarly to the model.\n        :param model_kwargs: if not None, a dict of extra keyword arguments to\n            pass to the model. This can be used for conditioning.\n        :return: a dict containing the following keys:\n                 - 'sample': a random sample from the model.\n                 - 'pred_xstart': a prediction of x_0.\n        \"\"\"\n        out = self.p_mean_variance(\n            model,\n            x,\n            t,\n            clip_denoised=clip_denoised,\n            denoised_fn=denoised_fn,\n            model_kwargs=model_kwargs,\n        )\n        noise = th.randn_like(x)\n        nonzero_mask = (\n            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))\n        )  # no noise when t == 0\n        if cond_fn is not None:\n            out[\"mean\"] = self.condition_mean(\n                cond_fn, out, x, t, model_kwargs=model_kwargs\n            )\n        sample = out[\"mean\"] + nonzero_mask * th.exp(0.5 * out[\"log_variance\"]) * noise\n        return {\"sample\": sample, \"pred_xstart\": out[\"pred_xstart\"]}\n\n    def p_sample_blur(\n        self,\n        model,\n        x,\n        t,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n    ):\n        \"\"\"\n        Sample x_{t-1} from the model at the given timestep.\n\n        :param model: the model to sample from.\n        :param x: the current tensor at x_{t-1}.\n        :param t: the value of t, starting at 0 for the first diffusion step.\n        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].\n        :param denoised_fn: if not None, a function which applies to the\n            x_start prediction before it is used to sample.\n        :param cond_fn: if not None, this is a gradient function that acts\n                        similarly to the model.\n        :param model_kwargs: if not None, a dict of extra keyword arguments to\n            pass to the model. This can be used for conditioning.\n        :return: a dict containing the following keys:\n                 - 'sample': a random sample from the model.\n                 - 'pred_xstart': a prediction of x_0.\n        \"\"\"\n        # out = self.p_mean_variance(\n        #     model,\n        #     x,\n        #     t,\n        #     clip_denoised=clip_denoised,\n        #     denoised_fn=denoised_fn,\n        #     model_kwargs=model_kwargs,\n        # )\n        # noise = th.randn_like(x)\n        # nonzero_mask = (\n        #     (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))\n        # )  # no noise when t == 0\n        if cond_fn is not None:\n            raise NotImplementedError\n            out[\"mean\"] = self.condition_mean(\n                cond_fn, out, x, t, model_kwargs=model_kwargs\n            )\n        model_output = model(x, self._scale_timesteps(t), **model_kwargs)\n        # sample = out[\"mean\"] + nonzero_mask * th.exp(0.5 * out[\"log_variance\"]) * noise\n        return {\"sample\": model_output}\n\n    def p_sample_loop(\n        self,\n        model,\n        shape,\n        noise=None,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n        device=None,\n        progress=False,\n    ):\n        \"\"\"\n        Generate samples from the model.\n\n        :param model: the model module.\n        :param shape: the shape of the samples, (N, C, H, W).\n        :param noise: if specified, the noise from the encoder to sample.\n                      Should be of the same shape as `shape`.\n        :param clip_denoised: if True, clip x_start predictions to [-1, 1].\n        :param denoised_fn: if not None, a function which applies to the\n            x_start prediction before it is used to sample.\n        :param cond_fn: if not None, this is a gradient function that acts\n                        similarly to the model.\n        :param model_kwargs: if not None, a dict of extra keyword arguments to\n            pass to the model. This can be used for conditioning.\n        :param device: if specified, the device to create the samples on.\n                       If not specified, use a model parameter's device.\n        :param progress: if True, show a tqdm progress bar.\n        :return: a non-differentiable batch of samples.\n        \"\"\"\n        final = None\n        i=0\n        for sample in self.p_sample_loop_progressive(\n            model,\n            shape,\n            noise=noise,\n            clip_denoised=clip_denoised,\n            denoised_fn=denoised_fn,\n            cond_fn=cond_fn,\n            model_kwargs=model_kwargs,\n            device=device,\n            progress=progress,\n        ):\n            print(\"p_sample_loop\", i)\n            i+=1\n            tensor_imsave(sample[\"sample\"][0], \"./debug\", f\"{self.num_timesteps-i}.png\", prt=True)\n            \n            final = sample\n        return final[\"sample\"]\n\n    def p_sample_loop_progressive(\n        self,\n        model,\n        shape,\n        noise=None,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n        device=None,\n        progress=False,\n    ):\n        \"\"\"\n        Generate samples from the model and yield intermediate samples from\n        each timestep of diffusion.\n\n        Arguments are the same as p_sample_loop().\n        Returns a generator over dicts, where each dict is the return value of\n        p_sample().\n        \"\"\"\n        if device is None:\n            device = next(model.parameters()).device\n        assert isinstance(shape, (tuple, list))\n        if noise is not None:\n            img = noise\n        else:\n            img = th.randn(*shape, device=device)\n            T = th.tensor(1000, device=device).repeat(shape[0])\n            print(\"T:\", T)\n            img = self.q_sample_blur(img, T)\n        indices = list(range(self.num_timesteps))[::-1]\n\n        if progress:\n            # Lazy import so that we don't depend on tqdm.\n            from tqdm.auto import tqdm\n\n            indices = tqdm(indices)\n\n        for i in indices:\n            t = th.tensor([i] * shape[0], device=device)\n            with th.no_grad():\n                out = self.p_sample_blur(\n                    model,\n                    img,\n                    t,\n                    clip_denoised=clip_denoised,\n                    denoised_fn=denoised_fn,\n                    cond_fn=cond_fn,\n                    model_kwargs=model_kwargs,\n                )\n                yield out\n                img = out[\"sample\"]\n\n    def ddim_sample(\n        self,\n        model,\n        x,\n        t,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n        eta=0.0,\n    ):\n        \"\"\"\n        Sample x_{t-1} from the model using DDIM.\n\n        Same usage as p_sample().\n        \"\"\"\n        out = self.p_mean_variance(\n            model,\n            x,\n            t,\n            clip_denoised=clip_denoised,\n            denoised_fn=denoised_fn,\n            model_kwargs=model_kwargs,\n        )\n        if cond_fn is not None:\n            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)\n\n        # Usually our model outputs epsilon, but we re-derive it\n        # in case we used x_start or x_prev prediction.\n        eps = self._predict_eps_from_xstart(x, t, out[\"pred_xstart\"])\n\n        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)\n        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)\n        sigma = (\n            eta\n            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))\n            * th.sqrt(1 - alpha_bar / alpha_bar_prev)\n        )\n        # Equation 12.\n        noise = th.randn_like(x)\n        mean_pred = (\n            out[\"pred_xstart\"] * th.sqrt(alpha_bar_prev)\n            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps\n        )\n        nonzero_mask = (\n            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))\n        )  # no noise when t == 0\n        sample = mean_pred + nonzero_mask * sigma * noise\n        return {\"sample\": sample, \"pred_xstart\": out[\"pred_xstart\"]}\n\n    def ddim_reverse_sample(\n        self,\n        model,\n        x,\n        t,\n        clip_denoised=True,\n        denoised_fn=None,\n        model_kwargs=None,\n        eta=0.0,\n    ):\n        \"\"\"\n        Sample x_{t+1} from the model using DDIM reverse ODE.\n        \"\"\"\n        assert eta == 0.0, \"Reverse ODE only for deterministic path\"\n        out = self.p_mean_variance(\n            model,\n            x,\n            t,\n            clip_denoised=clip_denoised,\n            denoised_fn=denoised_fn,\n            model_kwargs=model_kwargs,\n        )\n        # Usually our model outputs epsilon, but we re-derive it\n        # in case we used x_start or x_prev prediction.\n        eps = (\n            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x\n            - out[\"pred_xstart\"]\n        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)\n        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)\n\n        # Equation 12. reversed\n        mean_pred = (\n            out[\"pred_xstart\"] * th.sqrt(alpha_bar_next)\n            + th.sqrt(1 - alpha_bar_next) * eps\n        )\n\n        return {\"sample\": mean_pred, \"pred_xstart\": out[\"pred_xstart\"]}\n\n    def ddim_sample_loop(\n        self,\n        model,\n        shape,\n        noise=None,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n        device=None,\n        progress=False,\n        eta=0.0,\n    ):\n        \"\"\"\n        Generate samples from the model using DDIM.\n\n        Same usage as p_sample_loop().\n        \"\"\"\n        final = None\n        for sample in self.ddim_sample_loop_progressive(\n            model,\n            shape,\n            noise=noise,\n            clip_denoised=clip_denoised,\n            denoised_fn=denoised_fn,\n            cond_fn=cond_fn,\n            model_kwargs=model_kwargs,\n            device=device,\n            progress=progress,\n            eta=eta,\n        ):\n            final = sample\n        return final[\"sample\"]\n\n    def ddim_sample_loop_progressive(\n        self,\n        model,\n        shape,\n        noise=None,\n        clip_denoised=True,\n        denoised_fn=None,\n        cond_fn=None,\n        model_kwargs=None,\n        device=None,\n        progress=False,\n        eta=0.0,\n    ):\n        \"\"\"\n        Use DDIM to sample from the model and yield intermediate samples from\n        each timestep of DDIM.\n\n        Same usage as p_sample_loop_progressive().\n        \"\"\"\n        if device is None:\n            device = next(model.parameters()).device\n        assert isinstance(shape, (tuple, list))\n        if noise is not None:\n            img = noise\n        else:\n            img = th.randn(*shape, device=device)\n        indices = list(range(self.num_timesteps))[::-1]\n\n        if progress:\n            # Lazy import so that we don't depend on tqdm.\n            from tqdm.auto import tqdm\n\n            indices = tqdm(indices)\n\n        for i in indices:\n            t = th.tensor([i] * shape[0], device=device)\n            with th.no_grad():\n                out = self.ddim_sample(\n                    model,\n                    img,\n                    t,\n                    clip_denoised=clip_denoised,\n                    denoised_fn=denoised_fn,\n                    cond_fn=cond_fn,\n                    model_kwargs=model_kwargs,\n                    eta=eta,\n                )\n                yield out\n                img = out[\"sample\"]\n\n    def _vb_terms_bpd(\n        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None\n    ):\n        \"\"\"\n        Get a term for the variational lower-bound.\n\n        The resulting units are bits (rather than nats, as one might expect).\n        This allows for comparison to other papers.\n\n        :return: a dict with the following keys:\n                 - 'output': a shape [N] tensor of NLLs or KLs.\n                 - 'pred_xstart': the x_0 predictions.\n        \"\"\"\n        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(\n            x_start=x_start, x_t=x_t, t=t\n        )\n        out = self.p_mean_variance(\n            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs\n        )\n        kl = normal_kl(\n            true_mean, true_log_variance_clipped, out[\"mean\"], out[\"log_variance\"]\n        )\n        kl = mean_flat(kl) / np.log(2.0)\n\n        decoder_nll = -discretized_gaussian_log_likelihood(\n            x_start, means=out[\"mean\"], log_scales=0.5 * out[\"log_variance\"]\n        )\n        assert decoder_nll.shape == x_start.shape\n        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)\n\n        # At the first timestep return the decoder NLL,\n        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))\n        output = th.where((t == 0), decoder_nll, kl)\n        return {\"output\": output, \"pred_xstart\": out[\"pred_xstart\"]}\n\n    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):\n        \"\"\"\n        Compute training losses for a single timestep.\n\n        :param model: the model to evaluate loss on.\n        :param x_start: the [N x C x ...] tensor of inputs.\n        :param t: a batch of timestep indices.\n        :param model_kwargs: if not None, a dict of extra keyword arguments to\n            pass to the model. This can be used for conditioning.\n        :param noise: if specified, the specific Gaussian noise to try to remove.\n        :return: a dict with the key \"loss\" containing a tensor of shape [N].\n                 Some mean or variance settings may also have other keys.\n        \"\"\"\n        if model_kwargs is None:\n            model_kwargs = {}\n        if noise is None:\n            noise = th.randn_like(x_start)\n        # x_t = self.q_sample(x_start, t, noise=noise)\n        \n        x_t = self.q_sample_blur(x_start, t)\n        x_t_prev = self.q_sample_blur(x_start, t-1) # x_{t-1}\n        terms = {}\n\n        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:\n            raise NotImplementedError\n            terms[\"loss\"] = self._vb_terms_bpd(\n                model=model,\n                x_start=x_start,\n                x_t=x_t,\n                t=t,\n                clip_denoised=False,\n                model_kwargs=model_kwargs,\n            )[\"output\"]\n            if self.loss_type == LossType.RESCALED_KL:\n                terms[\"loss\"] *= self.num_timesteps\n        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:\n            model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)\n\n            if self.model_var_type in [\n                ModelVarType.LEARNED,\n                ModelVarType.LEARNED_RANGE,\n            ]:\n                raise NotImplementedError\n                B, C = x_t.shape[:2]\n                assert model_output.shape == (B, C * 2, *x_t.shape[2:])\n                model_output, model_var_values = th.split(model_output, C, dim=1)\n                # Learn the variance using the variational bound, but don't let\n                # it affect our mean prediction.\n                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)\n                terms[\"vb\"] = self._vb_terms_bpd(\n                    model=lambda *args, r=frozen_out: r,\n                    x_start=x_start,\n                    x_t=x_t,\n                    t=t,\n                    clip_denoised=False,\n                )[\"output\"]\n                if self.loss_type == LossType.RESCALED_MSE:\n                    # Divide by 1000 for equivalence with initial implementation.\n                    # Without a factor of 1/1000, the VB term hurts the MSE term.\n                    terms[\"vb\"] *= self.num_timesteps / 1000.0\n\n            # target = {\n            #     ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(\n            #         x_start=x_start, x_t=x_t, t=t\n            #     )[0],\n            #     ModelMeanType.START_X: x_start,\n            #     ModelMeanType.EPSILON: noise,\n            # }[self.model_mean_type]\n            target = x_t_prev\n            assert model_output.shape == target.shape == x_start.shape\n            terms[\"mse\"] = mean_flat((target - model_output) ** 2)\n            if \"vb\" in terms:\n                terms[\"loss\"] = terms[\"mse\"] + terms[\"vb\"]\n            else:\n                terms[\"loss\"] = terms[\"mse\"]\n        else:\n            raise NotImplementedError(self.loss_type)\n\n        return terms\n\n    def _prior_bpd(self, x_start):\n        \"\"\"\n        Get the prior KL term for the variational lower-bound, measured in\n        bits-per-dim.\n\n        This term can't be optimized, as it only depends on the encoder.\n\n        :param x_start: the [N x C x ...] tensor of inputs.\n        :return: a batch of [N] KL values (in bits), one per batch element.\n        \"\"\"\n        batch_size = x_start.shape[0]\n        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)\n        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)\n        kl_prior = normal_kl(\n            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0\n        )\n        return mean_flat(kl_prior) / np.log(2.0)\n\n    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):\n        \"\"\"\n        Compute the entire variational lower-bound, measured in bits-per-dim,\n        as well as other related quantities.\n\n        :param model: the model to evaluate loss on.\n        :param x_start: the [N x C x ...] tensor of inputs.\n        :param clip_denoised: if True, clip denoised samples.\n        :param model_kwargs: if not None, a dict of extra keyword arguments to\n            pass to the model. This can be used for conditioning.\n\n        :return: a dict containing the following keys:\n                 - total_bpd: the total variational lower-bound, per batch element.\n                 - prior_bpd: the prior term in the lower-bound.\n                 - vb: an [N x T] tensor of terms in the lower-bound.\n                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.\n                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.\n        \"\"\"\n        device = x_start.device\n        batch_size = x_start.shape[0]\n\n \n        xstart_mse = []\n        mse = []\n        for t in list(range(self.num_timesteps))[::-1]:\n            t_batch = th.tensor([t] * batch_size, device=device)\n\n            x_t = self.q_sample_blur(x_start=x_start, t=t_batch)\n            x_t_prev = self.q_sample_blur(x_start=x_start, t=t_batch-1)\n            # Calculate VLB term at the current timestep\n            with th.no_grad():\n                model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)\n                mse.append(mean_flat((x_t_prev - model_output)**2))\n\n        mse = th.stack(mse, dim=1)\n\n        prior_bpd = self._prior_bpd(x_start)\n        \n        return {'mse' : mse}\n\n\n\n\n        # device = x_start.device\n        # batch_size = x_start.shape[0]\n\n        # vb = []\n        # xstart_mse = []\n        # mse = []\n        # for t in list(range(self.num_timesteps))[::-1]:\n        #     t_batch = th.tensor([t] * batch_size, device=device)\n        #     noise = th.randn_like(x_start)\n        #     x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)\n        #     # Calculate VLB term at the current timestep\n        #     with th.no_grad():\n        #         out = self._vb_terms_bpd(\n        #             model,\n        #             x_start=x_start,\n        #             x_t=x_t,\n        #             t=t_batch,\n        #             clip_denoised=clip_denoised,\n        #             model_kwargs=model_kwargs,\n        #         )\n        #     vb.append(out[\"output\"])\n        #     xstart_mse.append(mean_flat((out[\"pred_xstart\"] - x_start) ** 2))\n        #     eps = self._predict_eps_from_xstart(x_t, t_batch, out[\"pred_xstart\"])\n        #     mse.append(mean_flat((eps - noise) ** 2))\n\n        # vb = th.stack(vb, dim=1)\n        # xstart_mse = th.stack(xstart_mse, dim=1)\n        # mse = th.stack(mse, dim=1)\n\n        # prior_bpd = self._prior_bpd(x_start)\n        # total_bpd = vb.sum(dim=1) + prior_bpd\n        # return {\n        #     \"total_bpd\": total_bpd,\n        #     \"prior_bpd\": prior_bpd,\n        #     \"vb\": vb,\n        #     \"xstart_mse\": xstart_mse,\n        #     \"mse\": mse,\n        # }\n\n\ndef _extract_into_tensor(arr, timesteps, broadcast_shape):\n    \"\"\"\n    Extract values from a 1-D numpy array for a batch of indices.\n\n    :param arr: the 1-D numpy array.\n    :param timesteps: a tensor of indices into the array to extract.\n    :param broadcast_shape: a larger shape of K dimensions with the batch\n                            dimension equal to the length of timesteps.\n    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.\n    \"\"\"\n    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()\n    while len(res.shape) < len(broadcast_shape):\n        res = res[..., None]\n    return res.expand(broadcast_shape)\n\ndef tensor2pil(t):\n    img = transforms.functional.to_pil_image(denorm(t))\n    return img\ndef denorm(t):\n    return t*0.5 + 0.5\n\ndef tensor_imsave(t, path, fname, denormalization=True, prt=False):\n    # Save tensor as .png file\n    check_folder(path)\n    if denormalization: t = denorm(t)\n    t = th.clamp(input=t, min=0, max=1)\n    img = transforms.functional.to_pil_image(t.detach().cpu())\n    img.save(os.path.join(path,fname))\n    if prt:\n        print(f\"Saved to {os.path.join(path,fname)}\")\ndef pil_loader(path):\n    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\n    with open(path, 'rb') as f:##Automatic closing using with. rb = read binary\n        img = Image.open(f)\n        return img.convert('RGB')\ndef check_folder(log_dir):\n    if not os.path.exists(log_dir):\n        os.makedirs(log_dir)\n    return log_dir"
  },
  {
    "path": "guided_diffusion/image_datasets.py",
    "content": "import math\nimport random\n\nfrom PIL import Image\nimport blobfile as bf\nfrom mpi4py import MPI\nimport numpy as np\nfrom torch.utils.data import DataLoader, Dataset\n\n\ndef load_data(\n    *,\n    data_dir,\n    batch_size,\n    image_size,\n    class_cond=False,\n    deterministic=False,\n    random_crop=False,\n    random_flip=True,\n):\n    \"\"\"\n    For a dataset, create a generator over (images, kwargs) pairs.\n\n    Each images is an NCHW float tensor, and the kwargs dict contains zero or\n    more keys, each of which map to a batched Tensor of their own.\n    The kwargs dict can be used for class labels, in which case the key is \"y\"\n    and the values are integer tensors of class labels.\n\n    :param data_dir: a dataset directory.\n    :param batch_size: the batch size of each returned pair.\n    :param image_size: the size to which images are resized.\n    :param class_cond: if True, include a \"y\" key in returned dicts for class\n                       label. If classes are not available and this is true, an\n                       exception will be raised.\n    :param deterministic: if True, yield results in a deterministic order.\n    :param random_crop: if True, randomly crop the images for augmentation.\n    :param random_flip: if True, randomly flip the images for augmentation.\n    \"\"\"\n    if not data_dir:\n        raise ValueError(\"unspecified data directory\")\n    all_files = _list_image_files_recursively(data_dir)\n    classes = None\n    if class_cond:\n        # Assume classes are the first part of the filename,\n        # before an underscore.\n        class_names = [bf.basename(path).split(\"_\")[0] for path in all_files]\n        sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}\n        classes = [sorted_classes[x] for x in class_names]\n    dataset = ImageDataset(\n        image_size,\n        all_files,\n        classes=classes,\n        shard=MPI.COMM_WORLD.Get_rank(),\n        num_shards=MPI.COMM_WORLD.Get_size(),\n        random_crop=random_crop,\n        random_flip=random_flip,\n    )\n    if deterministic:\n        loader = DataLoader(\n            dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True\n        )\n    else:\n        loader = DataLoader(\n            dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True\n        )\n    while True:\n        yield from loader\n\n\ndef _list_image_files_recursively(data_dir):\n    results = []\n    for entry in sorted(bf.listdir(data_dir)):\n        full_path = bf.join(data_dir, entry)\n        ext = entry.split(\".\")[-1]\n        if \".\" in entry and ext.lower() in [\"jpg\", \"jpeg\", \"png\", \"gif\"]:\n            results.append(full_path)\n        elif bf.isdir(full_path):\n            results.extend(_list_image_files_recursively(full_path))\n    return results\n\n\nclass ImageDataset(Dataset):\n    def __init__(\n        self,\n        resolution,\n        image_paths,\n        classes=None,\n        shard=0,\n        num_shards=1,\n        random_crop=False,\n        random_flip=True,\n    ):\n        super().__init__()\n        self.resolution = resolution\n        self.local_images = image_paths[shard:][::num_shards]\n        self.local_classes = None if classes is None else classes[shard:][::num_shards]\n        self.random_crop = random_crop\n        self.random_flip = random_flip\n\n    def __len__(self):\n        return len(self.local_images)\n\n    def __getitem__(self, idx):\n        path = self.local_images[idx]\n        with bf.BlobFile(path, \"rb\") as f:\n            pil_image = Image.open(f)\n            pil_image.load()\n        pil_image = pil_image.convert(\"RGB\")\n\n        if self.random_crop:\n            arr = random_crop_arr(pil_image, self.resolution)\n        else:\n            arr = center_crop_arr(pil_image, self.resolution)\n\n        if self.random_flip and random.random() < 0.5:\n            arr = arr[:, ::-1]\n\n        arr = arr.astype(np.float32) / 127.5 - 1\n\n        out_dict = {}\n        if self.local_classes is not None:\n            out_dict[\"y\"] = np.array(self.local_classes[idx], dtype=np.int64)\n        return np.transpose(arr, [2, 0, 1]), out_dict\n\n\ndef center_crop_arr(pil_image, image_size):\n    # We are not on a new enough PIL to support the `reducing_gap`\n    # argument, which uses BOX downsampling at powers of two first.\n    # Thus, we do it by hand to improve downsample quality.\n    while min(*pil_image.size) >= 2 * image_size:\n        pil_image = pil_image.resize(\n            tuple(x // 2 for x in pil_image.size), resample=Image.BOX\n        )\n\n    scale = image_size / min(*pil_image.size)\n    pil_image = pil_image.resize(\n        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC\n    )\n\n    arr = np.array(pil_image)\n    crop_y = (arr.shape[0] - image_size) // 2\n    crop_x = (arr.shape[1] - image_size) // 2\n    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]\n\n\ndef random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):\n    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)\n    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)\n    smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)\n\n    # We are not on a new enough PIL to support the `reducing_gap`\n    # argument, which uses BOX downsampling at powers of two first.\n    # Thus, we do it by hand to improve downsample quality.\n    while min(*pil_image.size) >= 2 * smaller_dim_size:\n        pil_image = pil_image.resize(\n            tuple(x // 2 for x in pil_image.size), resample=Image.BOX\n        )\n\n    scale = smaller_dim_size / min(*pil_image.size)\n    pil_image = pil_image.resize(\n        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC\n    )\n\n    arr = np.array(pil_image)\n    crop_y = random.randrange(arr.shape[0] - image_size + 1)\n    crop_x = random.randrange(arr.shape[1] - image_size + 1)\n    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]\n"
  },
  {
    "path": "guided_diffusion/logger.py",
    "content": "\"\"\"\nLogger copied from OpenAI baselines to avoid extra RL-based dependencies:\nhttps://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py\n\"\"\"\n\nimport os\nimport sys\nimport shutil\nimport os.path as osp\nimport json\nimport time\nimport datetime\nimport tempfile\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import contextmanager\n\nDEBUG = 10\nINFO = 20\nWARN = 30\nERROR = 40\n\nDISABLED = 50\n\n\nclass KVWriter(object):\n    def writekvs(self, kvs):\n        raise NotImplementedError\n\n\nclass SeqWriter(object):\n    def writeseq(self, seq):\n        raise NotImplementedError\n\n\nclass HumanOutputFormat(KVWriter, SeqWriter):\n    def __init__(self, filename_or_file):\n        if isinstance(filename_or_file, str):\n            self.file = open(filename_or_file, \"wt\")\n            self.own_file = True\n        else:\n            assert hasattr(filename_or_file, \"read\"), (\n                \"expected file or str, got %s\" % filename_or_file\n            )\n            self.file = filename_or_file\n            self.own_file = False\n\n    def writekvs(self, kvs):\n        # Create strings for printing\n        key2str = {}\n        for (key, val) in sorted(kvs.items()):\n            if hasattr(val, \"__float__\"):\n                valstr = \"%-8.3g\" % val\n            else:\n                valstr = str(val)\n            key2str[self._truncate(key)] = self._truncate(valstr)\n\n        # Find max widths\n        if len(key2str) == 0:\n            print(\"WARNING: tried to write empty key-value dict\")\n            return\n        else:\n            keywidth = max(map(len, key2str.keys()))\n            valwidth = max(map(len, key2str.values()))\n\n        # Write out the data\n        dashes = \"-\" * (keywidth + valwidth + 7)\n        lines = [dashes]\n        for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()):\n            lines.append(\n                \"| %s%s | %s%s |\"\n                % (key, \" \" * (keywidth - len(key)), val, \" \" * (valwidth - len(val)))\n            )\n        lines.append(dashes)\n        self.file.write(\"\\n\".join(lines) + \"\\n\")\n\n        # Flush the output to the file\n        self.file.flush()\n\n    def _truncate(self, s):\n        maxlen = 30\n        return s[: maxlen - 3] + \"...\" if len(s) > maxlen else s\n\n    def writeseq(self, seq):\n        seq = list(seq)\n        for (i, elem) in enumerate(seq):\n            self.file.write(elem)\n            if i < len(seq) - 1:  # add space unless this is the last one\n                self.file.write(\" \")\n        self.file.write(\"\\n\")\n        self.file.flush()\n\n    def close(self):\n        if self.own_file:\n            self.file.close()\n\n\nclass JSONOutputFormat(KVWriter):\n    def __init__(self, filename):\n        self.file = open(filename, \"wt\")\n\n    def writekvs(self, kvs):\n        for k, v in sorted(kvs.items()):\n            if hasattr(v, \"dtype\"):\n                kvs[k] = float(v)\n        self.file.write(json.dumps(kvs) + \"\\n\")\n        self.file.flush()\n\n    def close(self):\n        self.file.close()\n\n\nclass CSVOutputFormat(KVWriter):\n    def __init__(self, filename):\n        self.file = open(filename, \"w+t\")\n        self.keys = []\n        self.sep = \",\"\n\n    def writekvs(self, kvs):\n        # Add our current row to the history\n        extra_keys = list(kvs.keys() - self.keys)\n        extra_keys.sort()\n        if extra_keys:\n            self.keys.extend(extra_keys)\n            self.file.seek(0)\n            lines = self.file.readlines()\n            self.file.seek(0)\n            for (i, k) in enumerate(self.keys):\n                if i > 0:\n                    self.file.write(\",\")\n                self.file.write(k)\n            self.file.write(\"\\n\")\n            for line in lines[1:]:\n                self.file.write(line[:-1])\n                self.file.write(self.sep * len(extra_keys))\n                self.file.write(\"\\n\")\n        for (i, k) in enumerate(self.keys):\n            if i > 0:\n                self.file.write(\",\")\n            v = kvs.get(k)\n            if v is not None:\n                self.file.write(str(v))\n        self.file.write(\"\\n\")\n        self.file.flush()\n\n    def close(self):\n        self.file.close()\n\n\nclass TensorBoardOutputFormat(KVWriter):\n    \"\"\"\n    Dumps key/value pairs into TensorBoard's numeric format.\n    \"\"\"\n\n    def __init__(self, dir):\n        os.makedirs(dir, exist_ok=True)\n        self.dir = dir\n        self.step = 1\n        prefix = \"events\"\n        path = osp.join(osp.abspath(dir), prefix)\n        import tensorflow as tf\n        from tensorflow.python import pywrap_tensorflow\n        from tensorflow.core.util import event_pb2\n        from tensorflow.python.util import compat\n\n        self.tf = tf\n        self.event_pb2 = event_pb2\n        self.pywrap_tensorflow = pywrap_tensorflow\n        self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path))\n\n    def writekvs(self, kvs):\n        def summary_val(k, v):\n            kwargs = {\"tag\": k, \"simple_value\": float(v)}\n            return self.tf.Summary.Value(**kwargs)\n\n        summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()])\n        event = self.event_pb2.Event(wall_time=time.time(), summary=summary)\n        event.step = (\n            self.step\n        )  # is there any reason why you'd want to specify the step?\n        self.writer.WriteEvent(event)\n        self.writer.Flush()\n        self.step += 1\n\n    def close(self):\n        if self.writer:\n            self.writer.Close()\n            self.writer = None\n\n\ndef make_output_format(format, ev_dir, log_suffix=\"\"):\n    os.makedirs(ev_dir, exist_ok=True)\n    if format == \"stdout\":\n        return HumanOutputFormat(sys.stdout)\n    elif format == \"log\":\n        return HumanOutputFormat(osp.join(ev_dir, \"log%s.txt\" % log_suffix))\n    elif format == \"json\":\n        return JSONOutputFormat(osp.join(ev_dir, \"progress%s.json\" % log_suffix))\n    elif format == \"csv\":\n        return CSVOutputFormat(osp.join(ev_dir, \"progress%s.csv\" % log_suffix))\n    elif format == \"tensorboard\":\n        return TensorBoardOutputFormat(osp.join(ev_dir, \"tb%s\" % log_suffix))\n    else:\n        raise ValueError(\"Unknown format specified: %s\" % (format,))\n\n\n# ================================================================\n# API\n# ================================================================\n\n\ndef logkv(key, val):\n    \"\"\"\n    Log a value of some diagnostic\n    Call this once for each diagnostic quantity, each iteration\n    If called many times, last value will be used.\n    \"\"\"\n    get_current().logkv(key, val)\n\n\ndef logkv_mean(key, val):\n    \"\"\"\n    The same as logkv(), but if called many times, values averaged.\n    \"\"\"\n    get_current().logkv_mean(key, val)\n\n\ndef logkvs(d):\n    \"\"\"\n    Log a dictionary of key-value pairs\n    \"\"\"\n    for (k, v) in d.items():\n        logkv(k, v)\n\n\ndef dumpkvs():\n    \"\"\"\n    Write all of the diagnostics from the current iteration\n    \"\"\"\n    return get_current().dumpkvs()\n\n\ndef getkvs():\n    return get_current().name2val\n\n\ndef log(*args, level=INFO):\n    \"\"\"\n    Write the sequence of args, with no separators, to the console and output files (if you've configured an output file).\n    \"\"\"\n    get_current().log(*args, level=level)\n\n\ndef debug(*args):\n    log(*args, level=DEBUG)\n\n\ndef info(*args):\n    log(*args, level=INFO)\n\n\ndef warn(*args):\n    log(*args, level=WARN)\n\n\ndef error(*args):\n    log(*args, level=ERROR)\n\n\ndef set_level(level):\n    \"\"\"\n    Set logging threshold on current logger.\n    \"\"\"\n    get_current().set_level(level)\n\n\ndef set_comm(comm):\n    get_current().set_comm(comm)\n\n\ndef get_dir():\n    \"\"\"\n    Get directory that log files are being written to.\n    will be None if there is no output directory (i.e., if you didn't call start)\n    \"\"\"\n    return get_current().get_dir()\n\n\nrecord_tabular = logkv\ndump_tabular = dumpkvs\n\n\n@contextmanager\ndef profile_kv(scopename):\n    logkey = \"wait_\" + scopename\n    tstart = time.time()\n    try:\n        yield\n    finally:\n        get_current().name2val[logkey] += time.time() - tstart\n\n\ndef profile(n):\n    \"\"\"\n    Usage:\n    @profile(\"my_func\")\n    def my_func(): code\n    \"\"\"\n\n    def decorator_with_name(func):\n        def func_wrapper(*args, **kwargs):\n            with profile_kv(n):\n                return func(*args, **kwargs)\n\n        return func_wrapper\n\n    return decorator_with_name\n\n\n# ================================================================\n# Backend\n# ================================================================\n\n\ndef get_current():\n    if Logger.CURRENT is None:\n        _configure_default_logger()\n\n    return Logger.CURRENT\n\n\nclass Logger(object):\n    DEFAULT = None  # A logger with no output files. (See right below class definition)\n    # So that you can still log to the terminal without setting up any output files\n    CURRENT = None  # Current logger being used by the free functions above\n\n    def __init__(self, dir, output_formats, comm=None):\n        self.name2val = defaultdict(float)  # values this iteration\n        self.name2cnt = defaultdict(int)\n        self.level = INFO\n        self.dir = dir\n        self.output_formats = output_formats\n        self.comm = comm\n\n    # Logging API, forwarded\n    # ----------------------------------------\n    def logkv(self, key, val):\n        self.name2val[key] = val\n\n    def logkv_mean(self, key, val):\n        oldval, cnt = self.name2val[key], self.name2cnt[key]\n        self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1)\n        self.name2cnt[key] = cnt + 1\n\n    def dumpkvs(self):\n        if self.comm is None:\n            d = self.name2val\n        else:\n            d = mpi_weighted_mean(\n                self.comm,\n                {\n                    name: (val, self.name2cnt.get(name, 1))\n                    for (name, val) in self.name2val.items()\n                },\n            )\n            if self.comm.rank != 0:\n                d[\"dummy\"] = 1  # so we don't get a warning about empty dict\n        out = d.copy()  # Return the dict for unit testing purposes\n        for fmt in self.output_formats:\n            if isinstance(fmt, KVWriter):\n                fmt.writekvs(d)\n        self.name2val.clear()\n        self.name2cnt.clear()\n        return out\n\n    def log(self, *args, level=INFO):\n        if self.level <= level:\n            self._do_log(args)\n\n    # Configuration\n    # ----------------------------------------\n    def set_level(self, level):\n        self.level = level\n\n    def set_comm(self, comm):\n        self.comm = comm\n\n    def get_dir(self):\n        return self.dir\n\n    def close(self):\n        for fmt in self.output_formats:\n            fmt.close()\n\n    # Misc\n    # ----------------------------------------\n    def _do_log(self, args):\n        for fmt in self.output_formats:\n            if isinstance(fmt, SeqWriter):\n                fmt.writeseq(map(str, args))\n\n\ndef get_rank_without_mpi_import():\n    # check environment variables here instead of importing mpi4py\n    # to avoid calling MPI_Init() when this module is imported\n    for varname in [\"PMI_RANK\", \"OMPI_COMM_WORLD_RANK\"]:\n        if varname in os.environ:\n            return int(os.environ[varname])\n    return 0\n\n\ndef mpi_weighted_mean(comm, local_name2valcount):\n    \"\"\"\n    Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110\n    Perform a weighted average over dicts that are each on a different node\n    Input: local_name2valcount: dict mapping key -> (value, count)\n    Returns: key -> mean\n    \"\"\"\n    all_name2valcount = comm.gather(local_name2valcount)\n    if comm.rank == 0:\n        name2sum = defaultdict(float)\n        name2count = defaultdict(float)\n        for n2vc in all_name2valcount:\n            for (name, (val, count)) in n2vc.items():\n                try:\n                    val = float(val)\n                except ValueError:\n                    if comm.rank == 0:\n                        warnings.warn(\n                            \"WARNING: tried to compute mean on non-float {}={}\".format(\n                                name, val\n                            )\n                        )\n                else:\n                    name2sum[name] += val * count\n                    name2count[name] += count\n        return {name: name2sum[name] / name2count[name] for name in name2sum}\n    else:\n        return {}\n\n\ndef configure(dir=None, format_strs=None, comm=None, log_suffix=\"\"):\n    \"\"\"\n    If comm is provided, average all numerical stats across that comm\n    \"\"\"\n    dir = \"./cifar10-forwardblur-fromblurrednoise\"\n    if dir is None:\n        dir = os.getenv(\"OPENAI_LOGDIR\")\n    if dir is None:\n        dir = osp.join(\n            tempfile.gettempdir(),\n            datetime.datetime.now().strftime(\"openai-%Y-%m-%d-%H-%M-%S-%f\"),\n        )\n    assert isinstance(dir, str)\n    dir = os.path.expanduser(dir)\n    os.makedirs(os.path.expanduser(dir), exist_ok=True)\n\n    rank = get_rank_without_mpi_import()\n    if rank > 0:\n        log_suffix = log_suffix + \"-rank%03i\" % rank\n\n    if format_strs is None:\n        if rank == 0:\n            format_strs = os.getenv(\"OPENAI_LOG_FORMAT\", \"stdout,log,csv\").split(\",\")\n        else:\n            format_strs = os.getenv(\"OPENAI_LOG_FORMAT_MPI\", \"log\").split(\",\")\n    format_strs = filter(None, format_strs)\n    output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs]\n\n    Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm)\n    if output_formats:\n        log(\"Logging to %s\" % dir)\n\n\ndef _configure_default_logger():\n    configure()\n    Logger.DEFAULT = Logger.CURRENT\n\n\ndef reset():\n    if Logger.CURRENT is not Logger.DEFAULT:\n        Logger.CURRENT.close()\n        Logger.CURRENT = Logger.DEFAULT\n        log(\"Reset logger\")\n\n\n@contextmanager\ndef scoped_configure(dir=None, format_strs=None, comm=None):\n    prevlogger = Logger.CURRENT\n    configure(dir=dir, format_strs=format_strs, comm=comm)\n    try:\n        yield\n    finally:\n        Logger.CURRENT.close()\n        Logger.CURRENT = prevlogger\n\n"
  },
  {
    "path": "guided_diffusion/losses.py",
    "content": "\"\"\"\nHelpers for various likelihood-based losses. These are ported from the original\nHo et al. diffusion models codebase:\nhttps://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py\n\"\"\"\n\nimport numpy as np\n\nimport torch as th\n\n\ndef normal_kl(mean1, logvar1, mean2, logvar2):\n    \"\"\"\n    Compute the KL divergence between two gaussians.\n\n    Shapes are automatically broadcasted, so batches can be compared to\n    scalars, among other use cases.\n    \"\"\"\n    tensor = None\n    for obj in (mean1, logvar1, mean2, logvar2):\n        if isinstance(obj, th.Tensor):\n            tensor = obj\n            break\n    assert tensor is not None, \"at least one argument must be a Tensor\"\n\n    # Force variances to be Tensors. Broadcasting helps convert scalars to\n    # Tensors, but it does not work for th.exp().\n    logvar1, logvar2 = [\n        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)\n        for x in (logvar1, logvar2)\n    ]\n\n    return 0.5 * (\n        -1.0\n        + logvar2\n        - logvar1\n        + th.exp(logvar1 - logvar2)\n        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)\n    )\n\n\ndef approx_standard_normal_cdf(x):\n    \"\"\"\n    A fast approximation of the cumulative distribution function of the\n    standard normal.\n    \"\"\"\n    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))\n\n\ndef discretized_gaussian_log_likelihood(x, *, means, log_scales):\n    \"\"\"\n    Compute the log-likelihood of a Gaussian distribution discretizing to a\n    given image.\n\n    :param x: the target images. It is assumed that this was uint8 values,\n              rescaled to the range [-1, 1].\n    :param means: the Gaussian mean Tensor.\n    :param log_scales: the Gaussian log stddev Tensor.\n    :return: a tensor like x of log probabilities (in nats).\n    \"\"\"\n    assert x.shape == means.shape == log_scales.shape\n    centered_x = x - means\n    inv_stdv = th.exp(-log_scales)\n    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)\n    cdf_plus = approx_standard_normal_cdf(plus_in)\n    min_in = inv_stdv * (centered_x - 1.0 / 255.0)\n    cdf_min = approx_standard_normal_cdf(min_in)\n    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))\n    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))\n    cdf_delta = cdf_plus - cdf_min\n    log_probs = th.where(\n        x < -0.999,\n        log_cdf_plus,\n        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),\n    )\n    assert log_probs.shape == x.shape\n    return log_probs\n"
  },
  {
    "path": "guided_diffusion/nn.py",
    "content": "\"\"\"\nVarious utilities for neural networks.\n\"\"\"\n\nimport math\n\nimport torch as th\nimport torch.nn as nn\n\n\n# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.\nclass SiLU(nn.Module):\n    def forward(self, x):\n        return x * th.sigmoid(x)\n\n\nclass GroupNorm32(nn.GroupNorm):\n    def forward(self, x):\n        return super().forward(x.float()).type(x.dtype)\n\n\ndef conv_nd(dims, *args, **kwargs):\n    \"\"\"\n    Create a 1D, 2D, or 3D convolution module.\n    \"\"\"\n    if dims == 1:\n        return nn.Conv1d(*args, **kwargs)\n    elif dims == 2:\n        return nn.Conv2d(*args, **kwargs)\n    elif dims == 3:\n        return nn.Conv3d(*args, **kwargs)\n    raise ValueError(f\"unsupported dimensions: {dims}\")\n\n\ndef linear(*args, **kwargs):\n    \"\"\"\n    Create a linear module.\n    \"\"\"\n    return nn.Linear(*args, **kwargs)\n\n\ndef avg_pool_nd(dims, *args, **kwargs):\n    \"\"\"\n    Create a 1D, 2D, or 3D average pooling module.\n    \"\"\"\n    if dims == 1:\n        return nn.AvgPool1d(*args, **kwargs)\n    elif dims == 2:\n        return nn.AvgPool2d(*args, **kwargs)\n    elif dims == 3:\n        return nn.AvgPool3d(*args, **kwargs)\n    raise ValueError(f\"unsupported dimensions: {dims}\")\n\n\ndef update_ema(target_params, source_params, rate=0.99):\n    \"\"\"\n    Update target parameters to be closer to those of source parameters using\n    an exponential moving average.\n\n    :param target_params: the target parameter sequence.\n    :param source_params: the source parameter sequence.\n    :param rate: the EMA rate (closer to 1 means slower).\n    \"\"\"\n    for targ, src in zip(target_params, source_params):\n        targ.detach().mul_(rate).add_(src, alpha=1 - rate)\n\n\ndef zero_module(module):\n    \"\"\"\n    Zero out the parameters of a module and return it.\n    \"\"\"\n    for p in module.parameters():\n        p.detach().zero_()\n    return module\n\n\ndef scale_module(module, scale):\n    \"\"\"\n    Scale the parameters of a module and return it.\n    \"\"\"\n    for p in module.parameters():\n        p.detach().mul_(scale)\n    return module\n\n\ndef mean_flat(tensor):\n    \"\"\"\n    Take the mean over all non-batch dimensions.\n    \"\"\"\n    return tensor.mean(dim=list(range(1, len(tensor.shape))))\n\n\ndef normalization(channels):\n    \"\"\"\n    Make a standard normalization layer.\n\n    :param channels: number of input channels.\n    :return: an nn.Module for normalization.\n    \"\"\"\n    return GroupNorm32(32, channels)\n\n\ndef timestep_embedding(timesteps, dim, max_period=10000):\n    \"\"\"\n    Create sinusoidal timestep embeddings.\n\n    :param timesteps: a 1-D Tensor of N indices, one per batch element.\n                      These may be fractional.\n    :param dim: the dimension of the output.\n    :param max_period: controls the minimum frequency of the embeddings.\n    :return: an [N x dim] Tensor of positional embeddings.\n    \"\"\"\n    half = dim // 2\n    freqs = th.exp(\n        -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half\n    ).to(device=timesteps.device)\n    args = timesteps[:, None].float() * freqs[None]\n    embedding = th.cat([th.cos(args), th.sin(args)], dim=-1)\n    if dim % 2:\n        embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1)\n    return embedding\n\n\ndef checkpoint(func, inputs, params, flag):\n    \"\"\"\n    Evaluate a function without caching intermediate activations, allowing for\n    reduced memory at the expense of extra compute in the backward pass.\n\n    :param func: the function to evaluate.\n    :param inputs: the argument sequence to pass to `func`.\n    :param params: a sequence of parameters `func` depends on but does not\n                   explicitly take as arguments.\n    :param flag: if False, disable gradient checkpointing.\n    \"\"\"\n    if flag:\n        args = tuple(inputs) + tuple(params)\n        return CheckpointFunction.apply(func, len(inputs), *args)\n    else:\n        return func(*inputs)\n\n\nclass CheckpointFunction(th.autograd.Function):\n    @staticmethod\n    def forward(ctx, run_function, length, *args):\n        ctx.run_function = run_function\n        ctx.input_tensors = list(args[:length])\n        ctx.input_params = list(args[length:])\n        with th.no_grad():\n            output_tensors = ctx.run_function(*ctx.input_tensors)\n        return output_tensors\n\n    @staticmethod\n    def backward(ctx, *output_grads):\n        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]\n        with th.enable_grad():\n            # Fixes a bug where the first op in run_function modifies the\n            # Tensor storage in place, which is not allowed for detach()'d\n            # Tensors.\n            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]\n            output_tensors = ctx.run_function(*shallow_copies)\n        input_grads = th.autograd.grad(\n            output_tensors,\n            ctx.input_tensors + ctx.input_params,\n            output_grads,\n            allow_unused=True,\n        )\n        del ctx.input_tensors\n        del ctx.input_params\n        del output_tensors\n        return (None, None) + input_grads\n"
  },
  {
    "path": "guided_diffusion/resample.py",
    "content": "from abc import ABC, abstractmethod\n\nimport numpy as np\nimport torch as th\nimport torch.distributed as dist\n\n\ndef create_named_schedule_sampler(name, diffusion):\n    \"\"\"\n    Create a ScheduleSampler from a library of pre-defined samplers.\n\n    :param name: the name of the sampler.\n    :param diffusion: the diffusion object to sample for.\n    \"\"\"\n    if name == \"uniform\":\n        return UniformSampler(diffusion)\n    elif name == \"loss-second-moment\":\n        return LossSecondMomentResampler(diffusion)\n    else:\n        raise NotImplementedError(f\"unknown schedule sampler: {name}\")\n\n\nclass ScheduleSampler(ABC):\n    \"\"\"\n    A distribution over timesteps in the diffusion process, intended to reduce\n    variance of the objective.\n\n    By default, samplers perform unbiased importance sampling, in which the\n    objective's mean is unchanged.\n    However, subclasses may override sample() to change how the resampled\n    terms are reweighted, allowing for actual changes in the objective.\n    \"\"\"\n\n    @abstractmethod\n    def weights(self):\n        \"\"\"\n        Get a numpy array of weights, one per diffusion step.\n\n        The weights needn't be normalized, but must be positive.\n        \"\"\"\n\n    def sample(self, batch_size, device):\n        \"\"\"\n        Importance-sample timesteps for a batch.\n\n        :param batch_size: the number of timesteps.\n        :param device: the torch device to save to.\n        :return: a tuple (timesteps, weights):\n                 - timesteps: a tensor of timestep indices.\n                 - weights: a tensor of weights to scale the resulting losses.\n        \"\"\"\n        w = self.weights()\n        p = w / np.sum(w) # loss별 weight를 normalize\n        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)[0]\n        if indices_np == 0: indices_np += 1\n        assert indices_np > 0, f\"indices_np : {indices_np}\"\n        indices_np = np.repeat(indices_np, batch_size)\n        indices = th.from_numpy(indices_np).long().to(device)\n        weights_np = 1 / (len(p) * p[indices_np])\n        weights = th.from_numpy(weights_np).float().to(device)\n        return indices, weights\n\n\nclass UniformSampler(ScheduleSampler):\n    def __init__(self, diffusion):\n        self.diffusion = diffusion\n        self._weights = np.ones([diffusion.num_timesteps])\n\n    def weights(self):\n        return self._weights\n\n\nclass LossAwareSampler(ScheduleSampler):\n    def update_with_local_losses(self, local_ts, local_losses):\n        \"\"\"\n        Update the reweighting using losses from a model.\n\n        Call this method from each rank with a batch of timesteps and the\n        corresponding losses for each of those timesteps.\n        This method will perform synchronization to make sure all of the ranks\n        maintain the exact same reweighting.\n\n        :param local_ts: an integer Tensor of timesteps.\n        :param local_losses: a 1D Tensor of losses.\n        \"\"\"\n        batch_sizes = [\n            th.tensor([0], dtype=th.int32, device=local_ts.device)\n            for _ in range(dist.get_world_size())\n        ]\n        dist.all_gather(\n            batch_sizes,\n            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),\n        )\n\n        # Pad all_gather batches to be the maximum batch size.\n        batch_sizes = [x.item() for x in batch_sizes]\n        max_bs = max(batch_sizes)\n\n        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]\n        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]\n        dist.all_gather(timestep_batches, local_ts)\n        dist.all_gather(loss_batches, local_losses)\n        timesteps = [\n            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]\n        ]\n        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]\n        self.update_with_all_losses(timesteps, losses)\n\n    @abstractmethod\n    def update_with_all_losses(self, ts, losses):\n        \"\"\"\n        Update the reweighting using losses from a model.\n\n        Sub-classes should override this method to update the reweighting\n        using losses from the model.\n\n        This method directly updates the reweighting without synchronizing\n        between workers. It is called by update_with_local_losses from all\n        ranks with identical arguments. Thus, it should have deterministic\n        behavior to maintain state across workers.\n\n        :param ts: a list of int timesteps.\n        :param losses: a list of float losses, one per timestep.\n        \"\"\"\n\n\nclass LossSecondMomentResampler(LossAwareSampler):\n    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):\n        raise NotImplementedError\n        self.diffusion = diffusion\n        self.history_per_term = history_per_term\n        self.uniform_prob = uniform_prob\n        self._loss_history = np.zeros(\n            [diffusion.num_timesteps, history_per_term], dtype=np.float64\n        )\n        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)\n\n    def weights(self):\n        if not self._warmed_up():\n            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)\n        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))\n        weights /= np.sum(weights)\n        weights *= 1 - self.uniform_prob\n        weights += self.uniform_prob / len(weights)\n        return weights\n\n    def update_with_all_losses(self, ts, losses):\n        for t, loss in zip(ts, losses):\n            if self._loss_counts[t] == self.history_per_term:\n                # Shift out the oldest loss term.\n                self._loss_history[t, :-1] = self._loss_history[t, 1:]\n                self._loss_history[t, -1] = loss\n            else:\n                self._loss_history[t, self._loss_counts[t]] = loss\n                self._loss_counts[t] += 1\n\n    def _warmed_up(self):\n        return (self._loss_counts == self.history_per_term).all()\n"
  },
  {
    "path": "guided_diffusion/respace.py",
    "content": "import numpy as np\nimport torch as th\n\nfrom .gaussian_diffusion import GaussianDiffusion\n\n\ndef space_timesteps(num_timesteps, section_counts):\n    \"\"\"\n    Create a list of timesteps to use from an original diffusion process,\n    given the number of timesteps we want to take from equally-sized portions\n    of the original process.\n\n    For example, if there's 300 timesteps and the section counts are [10,15,20]\n    then the first 100 timesteps are strided to be 10 timesteps, the second 100\n    are strided to be 15 timesteps, and the final 100 are strided to be 20.\n\n    If the stride is a string starting with \"ddim\", then the fixed striding\n    from the DDIM paper is used, and only one section is allowed.\n\n    :param num_timesteps: the number of diffusion steps in the original\n                          process to divide up.\n    :param section_counts: either a list of numbers, or a string containing\n                           comma-separated numbers, indicating the step count\n                           per section. As a special case, use \"ddimN\" where N\n                           is a number of steps to use the striding from the\n                           DDIM paper.\n    :return: a set of diffusion steps from the original process to use.\n    \"\"\"\n    if isinstance(section_counts, str):\n        if section_counts.startswith(\"ddim\"):\n            desired_count = int(section_counts[len(\"ddim\") :])\n            for i in range(1, num_timesteps):\n                if len(range(0, num_timesteps, i)) == desired_count:\n                    return set(range(0, num_timesteps, i))\n            raise ValueError(\n                f\"cannot create exactly {num_timesteps} steps with an integer stride\"\n            )\n        section_counts = [int(x) for x in section_counts.split(\",\")]\n    size_per = num_timesteps // len(section_counts)\n    extra = num_timesteps % len(section_counts)\n    start_idx = 0\n    all_steps = []\n    for i, section_count in enumerate(section_counts):\n        size = size_per + (1 if i < extra else 0)\n        if size < section_count:\n            raise ValueError(\n                f\"cannot divide section of {size} steps into {section_count}\"\n            )\n        if section_count <= 1:\n            frac_stride = 1\n        else:\n            frac_stride = (size - 1) / (section_count - 1)\n        cur_idx = 0.0\n        taken_steps = []\n        for _ in range(section_count):\n            taken_steps.append(start_idx + round(cur_idx))\n            cur_idx += frac_stride\n        all_steps += taken_steps\n        start_idx += size\n    return set(all_steps)\n\n\nclass SpacedDiffusion(GaussianDiffusion):\n    \"\"\"\n    A diffusion process which can skip steps in a base diffusion process.\n\n    :param use_timesteps: a collection (sequence or set) of timesteps from the\n                          original diffusion process to retain.\n    :param kwargs: the kwargs to create the base diffusion process.\n    \"\"\"\n\n    def __init__(self, use_timesteps, **kwargs):\n        self.use_timesteps = set(use_timesteps)\n        self.timestep_map = []\n        self.original_num_steps = len(kwargs[\"betas\"])\n\n        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa\n        last_alpha_cumprod = 1.0\n        new_betas = []\n        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):\n            if i in self.use_timesteps:\n                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)\n                last_alpha_cumprod = alpha_cumprod\n                self.timestep_map.append(i)\n        kwargs[\"betas\"] = np.array(new_betas)\n        super().__init__(**kwargs)\n\n    def p_mean_variance(\n        self, model, *args, **kwargs\n    ):  # pylint: disable=signature-differs\n        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)\n\n    def training_losses(\n        self, model, *args, **kwargs\n    ):  # pylint: disable=signature-differs\n        return super().training_losses(self._wrap_model(model), *args, **kwargs)\n\n    def condition_mean(self, cond_fn, *args, **kwargs):\n        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)\n\n    def condition_score(self, cond_fn, *args, **kwargs):\n        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)\n\n    def _wrap_model(self, model):\n        if isinstance(model, _WrappedModel):\n            return model\n        return _WrappedModel(\n            model, self.timestep_map, self.rescale_timesteps, self.original_num_steps\n        )\n\n    def _scale_timesteps(self, t):\n        # Scaling is done by the wrapped model.\n        return t\n\n\nclass _WrappedModel:\n    def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):\n        self.model = model\n        self.timestep_map = timestep_map\n        self.rescale_timesteps = rescale_timesteps\n        self.original_num_steps = original_num_steps\n\n    def __call__(self, x, ts, **kwargs):\n        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)\n        new_ts = map_tensor[ts]\n        if self.rescale_timesteps:\n            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)\n        return self.model(x, new_ts, **kwargs)\n"
  },
  {
    "path": "guided_diffusion/script_util.py",
    "content": "import argparse\nimport inspect\n\nfrom . import gaussian_diffusion as gd\nfrom .respace import SpacedDiffusion, space_timesteps\nfrom .unet import SuperResModel, UNetModel, EncoderUNetModel\n\nNUM_CLASSES = 1000\n\n\ndef diffusion_defaults():\n    \"\"\"\n    Defaults for image and classifier training.\n    \"\"\"\n    return dict(\n        learn_sigma=False,\n        diffusion_steps=1000,\n        noise_schedule=\"linear\",\n        timestep_respacing=\"\",\n        use_kl=False,\n        predict_xstart=False,\n        rescale_timesteps=False,\n        rescale_learned_sigmas=False,\n    )\n\n\ndef classifier_defaults():\n    \"\"\"\n    Defaults for classifier models.\n    \"\"\"\n    return dict(\n        image_size=64,\n        classifier_use_fp16=False,\n        classifier_width=128,\n        classifier_depth=2,\n        classifier_attention_resolutions=\"32,16,8\",  # 16\n        classifier_use_scale_shift_norm=True,  # False\n        classifier_resblock_updown=True,  # False\n        classifier_pool=\"attention\",\n    )\n\n\ndef model_and_diffusion_defaults():\n    \"\"\"\n    Defaults for image training.\n    \"\"\"\n    res = dict(\n        image_size=64,\n        num_channels=128,\n        num_res_blocks=2,\n        num_heads=4,\n        num_heads_upsample=-1,\n        num_head_channels=-1,\n        attention_resolutions=\"16,8\",\n        channel_mult=\"\",\n        dropout=0.0,\n        class_cond=False,\n        use_checkpoint=False,\n        use_scale_shift_norm=True,\n        resblock_updown=False,\n        use_fp16=False,\n        use_new_attention_order=False,\n    )\n    res.update(diffusion_defaults())\n    return res\n\n\ndef classifier_and_diffusion_defaults():\n    res = classifier_defaults()\n    res.update(diffusion_defaults())\n    return res\n\n\ndef create_model_and_diffusion(\n    image_size,\n    class_cond,\n    learn_sigma,\n    num_channels,\n    num_res_blocks,\n    channel_mult,\n    num_heads,\n    num_head_channels,\n    num_heads_upsample,\n    attention_resolutions,\n    dropout,\n    diffusion_steps,\n    noise_schedule,\n    timestep_respacing,\n    use_kl,\n    predict_xstart,\n    rescale_timesteps,\n    rescale_learned_sigmas,\n    use_checkpoint,\n    use_scale_shift_norm,\n    resblock_updown,\n    use_fp16,\n    use_new_attention_order,\n):\n    model = create_model(\n        image_size,\n        num_channels,\n        num_res_blocks,\n        channel_mult=channel_mult,\n        learn_sigma=learn_sigma,\n        class_cond=class_cond,\n        use_checkpoint=use_checkpoint,\n        attention_resolutions=attention_resolutions,\n        num_heads=num_heads,\n        num_head_channels=num_head_channels,\n        num_heads_upsample=num_heads_upsample,\n        use_scale_shift_norm=use_scale_shift_norm,\n        dropout=dropout,\n        resblock_updown=resblock_updown,\n        use_fp16=use_fp16,\n        use_new_attention_order=use_new_attention_order,\n    )\n    diffusion = create_gaussian_diffusion(\n        steps=diffusion_steps,\n        learn_sigma=learn_sigma,\n        noise_schedule=noise_schedule,\n        use_kl=use_kl,\n        predict_xstart=predict_xstart,\n        rescale_timesteps=rescale_timesteps,\n        rescale_learned_sigmas=rescale_learned_sigmas,\n        timestep_respacing=timestep_respacing,\n    )\n    return model, diffusion\n\n\ndef create_model(\n    image_size,\n    num_channels,\n    num_res_blocks,\n    channel_mult=\"\",\n    learn_sigma=False,\n    class_cond=False,\n    use_checkpoint=False,\n    attention_resolutions=\"16\",\n    num_heads=1,\n    num_head_channels=-1,\n    num_heads_upsample=-1,\n    use_scale_shift_norm=False,\n    dropout=0,\n    resblock_updown=False,\n    use_fp16=False,\n    use_new_attention_order=False,\n):\n    if channel_mult == \"\":\n        if image_size == 512:\n            channel_mult = (0.5, 1, 1, 2, 2, 4, 4)\n        elif image_size == 256:\n            channel_mult = (1, 1, 2, 2, 4, 4)\n        elif image_size == 128:\n            channel_mult = (1, 1, 2, 3, 4)\n        elif image_size == 64:\n            channel_mult = (1, 2, 3, 4)\n        else:\n            raise ValueError(f\"unsupported image size: {image_size}\")\n    else:\n        channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(\",\"))\n\n    attention_ds = []\n    for res in attention_resolutions.split(\",\"):\n        attention_ds.append(image_size // int(res))\n\n    return UNetModel(\n        image_size=image_size,\n        in_channels=3,\n        model_channels=num_channels,\n        out_channels=(3 if not learn_sigma else 6),\n        num_res_blocks=num_res_blocks,\n        attention_resolutions=tuple(attention_ds),\n        dropout=dropout,\n        channel_mult=channel_mult,\n        num_classes=(NUM_CLASSES if class_cond else None),\n        use_checkpoint=use_checkpoint,\n        use_fp16=use_fp16,\n        num_heads=num_heads,\n        num_head_channels=num_head_channels,\n        num_heads_upsample=num_heads_upsample,\n        use_scale_shift_norm=use_scale_shift_norm,\n        resblock_updown=resblock_updown,\n        use_new_attention_order=use_new_attention_order,\n    )\n\n\ndef create_classifier_and_diffusion(\n    image_size,\n    classifier_use_fp16,\n    classifier_width,\n    classifier_depth,\n    classifier_attention_resolutions,\n    classifier_use_scale_shift_norm,\n    classifier_resblock_updown,\n    classifier_pool,\n    learn_sigma,\n    diffusion_steps,\n    noise_schedule,\n    timestep_respacing,\n    use_kl,\n    predict_xstart,\n    rescale_timesteps,\n    rescale_learned_sigmas,\n):\n    classifier = create_classifier(\n        image_size,\n        classifier_use_fp16,\n        classifier_width,\n        classifier_depth,\n        classifier_attention_resolutions,\n        classifier_use_scale_shift_norm,\n        classifier_resblock_updown,\n        classifier_pool,\n    )\n    diffusion = create_gaussian_diffusion(\n        steps=diffusion_steps,\n        learn_sigma=learn_sigma,\n        noise_schedule=noise_schedule,\n        use_kl=use_kl,\n        predict_xstart=predict_xstart,\n        rescale_timesteps=rescale_timesteps,\n        rescale_learned_sigmas=rescale_learned_sigmas,\n        timestep_respacing=timestep_respacing,\n    )\n    return classifier, diffusion\n\n\ndef create_classifier(\n    image_size,\n    classifier_use_fp16,\n    classifier_width,\n    classifier_depth,\n    classifier_attention_resolutions,\n    classifier_use_scale_shift_norm,\n    classifier_resblock_updown,\n    classifier_pool,\n):\n    if image_size == 512:\n        channel_mult = (0.5, 1, 1, 2, 2, 4, 4)\n    elif image_size == 256:\n        channel_mult = (1, 1, 2, 2, 4, 4)\n    elif image_size == 128:\n        channel_mult = (1, 1, 2, 3, 4)\n    elif image_size == 64:\n        channel_mult = (1, 2, 3, 4)\n    else:\n        raise ValueError(f\"unsupported image size: {image_size}\")\n\n    attention_ds = []\n    for res in classifier_attention_resolutions.split(\",\"):\n        attention_ds.append(image_size // int(res))\n\n    return EncoderUNetModel(\n        image_size=image_size,\n        in_channels=3,\n        model_channels=classifier_width,\n        out_channels=1000,\n        num_res_blocks=classifier_depth,\n        attention_resolutions=tuple(attention_ds),\n        channel_mult=channel_mult,\n        use_fp16=classifier_use_fp16,\n        num_head_channels=64,\n        use_scale_shift_norm=classifier_use_scale_shift_norm,\n        resblock_updown=classifier_resblock_updown,\n        pool=classifier_pool,\n    )\n\n\ndef sr_model_and_diffusion_defaults():\n    res = model_and_diffusion_defaults()\n    res[\"large_size\"] = 256\n    res[\"small_size\"] = 64\n    arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0]\n    for k in res.copy().keys():\n        if k not in arg_names:\n            del res[k]\n    return res\n\n\ndef sr_create_model_and_diffusion(\n    large_size,\n    small_size,\n    class_cond,\n    learn_sigma,\n    num_channels,\n    num_res_blocks,\n    num_heads,\n    num_head_channels,\n    num_heads_upsample,\n    attention_resolutions,\n    dropout,\n    diffusion_steps,\n    noise_schedule,\n    timestep_respacing,\n    use_kl,\n    predict_xstart,\n    rescale_timesteps,\n    rescale_learned_sigmas,\n    use_checkpoint,\n    use_scale_shift_norm,\n    resblock_updown,\n    use_fp16,\n):\n    model = sr_create_model(\n        large_size,\n        small_size,\n        num_channels,\n        num_res_blocks,\n        learn_sigma=learn_sigma,\n        class_cond=class_cond,\n        use_checkpoint=use_checkpoint,\n        attention_resolutions=attention_resolutions,\n        num_heads=num_heads,\n        num_head_channels=num_head_channels,\n        num_heads_upsample=num_heads_upsample,\n        use_scale_shift_norm=use_scale_shift_norm,\n        dropout=dropout,\n        resblock_updown=resblock_updown,\n        use_fp16=use_fp16,\n    )\n    diffusion = create_gaussian_diffusion(\n        steps=diffusion_steps,\n        learn_sigma=learn_sigma,\n        noise_schedule=noise_schedule,\n        use_kl=use_kl,\n        predict_xstart=predict_xstart,\n        rescale_timesteps=rescale_timesteps,\n        rescale_learned_sigmas=rescale_learned_sigmas,\n        timestep_respacing=timestep_respacing,\n    )\n    return model, diffusion\n\n\ndef sr_create_model(\n    large_size,\n    small_size,\n    num_channels,\n    num_res_blocks,\n    learn_sigma,\n    class_cond,\n    use_checkpoint,\n    attention_resolutions,\n    num_heads,\n    num_head_channels,\n    num_heads_upsample,\n    use_scale_shift_norm,\n    dropout,\n    resblock_updown,\n    use_fp16,\n):\n    _ = small_size  # hack to prevent unused variable\n\n    if large_size == 512:\n        channel_mult = (1, 1, 2, 2, 4, 4)\n    elif large_size == 256:\n        channel_mult = (1, 1, 2, 2, 4, 4)\n    elif large_size == 64:\n        channel_mult = (1, 2, 3, 4)\n    else:\n        raise ValueError(f\"unsupported large size: {large_size}\")\n\n    attention_ds = []\n    for res in attention_resolutions.split(\",\"):\n        attention_ds.append(large_size // int(res))\n\n    return SuperResModel(\n        image_size=large_size,\n        in_channels=3,\n        model_channels=num_channels,\n        out_channels=(3 if not learn_sigma else 6),\n        num_res_blocks=num_res_blocks,\n        attention_resolutions=tuple(attention_ds),\n        dropout=dropout,\n        channel_mult=channel_mult,\n        num_classes=(NUM_CLASSES if class_cond else None),\n        use_checkpoint=use_checkpoint,\n        num_heads=num_heads,\n        num_head_channels=num_head_channels,\n        num_heads_upsample=num_heads_upsample,\n        use_scale_shift_norm=use_scale_shift_norm,\n        resblock_updown=resblock_updown,\n        use_fp16=use_fp16,\n    )\n\n\ndef create_gaussian_diffusion(\n    *,\n    steps=1000,\n    learn_sigma=False,\n    sigma_small=False,\n    noise_schedule=\"linear\",\n    use_kl=False,\n    predict_xstart=False,\n    rescale_timesteps=False,\n    rescale_learned_sigmas=False,\n    timestep_respacing=\"\",\n):\n    betas = gd.get_named_beta_schedule(noise_schedule, steps)\n    if use_kl:\n        loss_type = gd.LossType.RESCALED_KL\n    elif rescale_learned_sigmas:\n        loss_type = gd.LossType.RESCALED_MSE\n    else:\n        loss_type = gd.LossType.MSE\n    if not timestep_respacing:\n        timestep_respacing = [steps]\n    return SpacedDiffusion(\n        use_timesteps=space_timesteps(steps, timestep_respacing),\n        betas=betas,\n        model_mean_type=(\n            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X\n        ),\n        model_var_type=(\n            (\n                gd.ModelVarType.FIXED_LARGE\n                if not sigma_small\n                else gd.ModelVarType.FIXED_SMALL\n            )\n            if not learn_sigma\n            else gd.ModelVarType.LEARNED_RANGE\n        ),\n        loss_type=loss_type,\n        rescale_timesteps=rescale_timesteps,\n    )\n\n\ndef add_dict_to_argparser(parser, default_dict):\n    for k, v in default_dict.items():\n        v_type = type(v)\n        if v is None:\n            v_type = str\n        elif isinstance(v, bool):\n            v_type = str2bool\n        parser.add_argument(f\"--{k}\", default=v, type=v_type)\n\n\ndef args_to_dict(args, keys):\n    return {k: getattr(args, k) for k in keys}\n\n\ndef str2bool(v):\n    \"\"\"\n    https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse\n    \"\"\"\n    if isinstance(v, bool):\n        return v\n    if v.lower() in (\"yes\", \"true\", \"t\", \"y\", \"1\"):\n        return True\n    elif v.lower() in (\"no\", \"false\", \"f\", \"n\", \"0\"):\n        return False\n    else:\n        raise argparse.ArgumentTypeError(\"boolean value expected\")\n"
  },
  {
    "path": "guided_diffusion/train_util.py",
    "content": "import copy\nimport functools\nimport os\n\nimport blobfile as bf\nimport torch as th\nimport torch.distributed as dist\nfrom torch.nn.parallel.distributed import DistributedDataParallel as DDP\nfrom torch.optim import AdamW\n\nfrom . import dist_util, logger\nfrom .fp16_util import MixedPrecisionTrainer\nfrom .nn import update_ema\nfrom .resample import LossAwareSampler, UniformSampler\n\n# For ImageNet experiments, this was a good default value.\n# We found that the lg_loss_scale quickly climbed to\n# 20-21 within the first ~1K steps of training.\nINITIAL_LOG_LOSS_SCALE = 20.0\n\n\nclass TrainLoop:\n    def __init__(\n        self,\n        *,\n        model,\n        diffusion,\n        data,\n        batch_size,\n        microbatch,\n        lr,\n        ema_rate,\n        log_interval,\n        save_interval,\n        resume_checkpoint,\n        use_fp16=False,\n        fp16_scale_growth=1e-3,\n        schedule_sampler=None,\n        weight_decay=0.0,\n        lr_anneal_steps=0,\n    ):\n        self.model = model\n        self.diffusion = diffusion\n        self.data = data\n        self.batch_size = batch_size\n        self.microbatch = microbatch if microbatch > 0 else batch_size\n        self.lr = lr\n        self.ema_rate = (\n            [ema_rate]\n            if isinstance(ema_rate, float)\n            else [float(x) for x in ema_rate.split(\",\")]\n        )\n        self.log_interval = log_interval\n        self.save_interval = save_interval\n        self.resume_checkpoint = resume_checkpoint\n        self.use_fp16 = use_fp16\n        self.fp16_scale_growth = fp16_scale_growth\n        self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)\n        self.weight_decay = weight_decay\n        self.lr_anneal_steps = lr_anneal_steps\n\n        self.step = 0\n        self.resume_step = 0\n        self.global_batch = self.batch_size * dist.get_world_size()\n\n        self.sync_cuda = th.cuda.is_available()\n\n        self._load_and_sync_parameters()\n        self.mp_trainer = MixedPrecisionTrainer(\n            model=self.model,\n            use_fp16=self.use_fp16,\n            fp16_scale_growth=fp16_scale_growth,\n        )\n\n        self.opt = AdamW(\n            self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay\n        )\n        if self.resume_step:\n            self._load_optimizer_state()\n            # Model was resumed, either due to a restart or a checkpoint\n            # being specified at the command line.\n            self.ema_params = [\n                self._load_ema_parameters(rate) for rate in self.ema_rate\n            ]\n        else:\n            self.ema_params = [\n                copy.deepcopy(self.mp_trainer.master_params)\n                for _ in range(len(self.ema_rate))\n            ]\n\n        if th.cuda.is_available():\n            self.use_ddp = True\n            self.ddp_model = DDP(\n                self.model,\n                device_ids=[dist_util.dev()],\n                output_device=dist_util.dev(),\n                broadcast_buffers=False,\n                bucket_cap_mb=128,\n                find_unused_parameters=False,\n            )\n        else:\n            if dist.get_world_size() > 1:\n                logger.warn(\n                    \"Distributed training requires CUDA. \"\n                    \"Gradients will not be synchronized properly!\"\n                )\n            self.use_ddp = False\n            self.ddp_model = self.model\n\n    def _load_and_sync_parameters(self):\n        resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint\n\n        if resume_checkpoint:\n            self.resume_step = parse_resume_step_from_filename(resume_checkpoint)\n            if dist.get_rank() == 0:\n                logger.log(f\"loading model from checkpoint: {resume_checkpoint}...\")\n                self.model.load_state_dict(\n                    dist_util.load_state_dict(\n                        resume_checkpoint, map_location=dist_util.dev()\n                    )\n                )\n\n        dist_util.sync_params(self.model.parameters())\n\n    def _load_ema_parameters(self, rate):\n        ema_params = copy.deepcopy(self.mp_trainer.master_params)\n\n        main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint\n        ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate)\n        if ema_checkpoint:\n            if dist.get_rank() == 0:\n                logger.log(f\"loading EMA from checkpoint: {ema_checkpoint}...\")\n                state_dict = dist_util.load_state_dict(\n                    ema_checkpoint, map_location=dist_util.dev()\n                )\n                ema_params = self.mp_trainer.state_dict_to_master_params(state_dict)\n\n        dist_util.sync_params(ema_params)\n        return ema_params\n\n    def _load_optimizer_state(self):\n        main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint\n        opt_checkpoint = bf.join(\n            bf.dirname(main_checkpoint), f\"opt{self.resume_step:06}.pt\"\n        )\n        if bf.exists(opt_checkpoint):\n            logger.log(f\"loading optimizer state from checkpoint: {opt_checkpoint}\")\n            state_dict = dist_util.load_state_dict(\n                opt_checkpoint, map_location=dist_util.dev()\n            )\n            self.opt.load_state_dict(state_dict)\n\n    def run_loop(self):\n        while (\n            not self.lr_anneal_steps\n            or self.step + self.resume_step < self.lr_anneal_steps\n        ):\n            batch, cond = next(self.data)\n            self.run_step(batch, cond)\n            if self.step % self.log_interval == 0:\n                logger.dumpkvs()\n            if self.step % self.save_interval == 0:\n                self.save()\n                # Run for a finite amount of time in integration tests.\n                if os.environ.get(\"DIFFUSION_TRAINING_TEST\", \"\") and self.step > 0:\n                    return\n            self.step += 1\n        # Save the last checkpoint if it wasn't already saved.\n        if (self.step - 1) % self.save_interval != 0:\n            self.save()\n\n    def run_step(self, batch, cond):\n        self.forward_backward(batch, cond)\n        took_step = self.mp_trainer.optimize(self.opt)\n        if took_step:\n            self._update_ema()\n        self._anneal_lr()\n        self.log_step()\n\n    def forward_backward(self, batch, cond):\n        self.mp_trainer.zero_grad()\n        for i in range(0, batch.shape[0], self.microbatch):\n            micro = batch[i : i + self.microbatch].to(dist_util.dev())\n            micro_cond = {\n                k: v[i : i + self.microbatch].to(dist_util.dev())\n                for k, v in cond.items()\n            }\n            last_batch = (i + self.microbatch) >= batch.shape[0]\n            t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev())\n            # print(\"TrainLoop/forward_backward() t:\", t)\n            assert t[0] >= 0, f\"train_util.py t[0] : {t[0]}\"\n            compute_losses = functools.partial(\n                self.diffusion.training_losses,\n                self.ddp_model,\n                micro,\n                t,\n                model_kwargs=micro_cond,\n            )\n\n            if last_batch or not self.use_ddp:\n                losses = compute_losses()\n            else:\n                with self.ddp_model.no_sync():\n                    losses = compute_losses()\n\n            if isinstance(self.schedule_sampler, LossAwareSampler):\n                self.schedule_sampler.update_with_local_losses(\n                    t, losses[\"loss\"].detach()\n                )\n\n            loss = (losses[\"loss\"] * weights).mean()\n            log_loss_dict(\n                self.diffusion, t, {k: v * weights for k, v in losses.items()}\n            )\n            self.mp_trainer.backward(loss)\n\n    def _update_ema(self):\n        for rate, params in zip(self.ema_rate, self.ema_params):\n            update_ema(params, self.mp_trainer.master_params, rate=rate)\n\n    def _anneal_lr(self):\n        if not self.lr_anneal_steps:\n            return\n        frac_done = (self.step + self.resume_step) / self.lr_anneal_steps\n        lr = self.lr * (1 - frac_done)\n        for param_group in self.opt.param_groups:\n            param_group[\"lr\"] = lr\n\n    def log_step(self):\n        logger.logkv(\"step\", self.step + self.resume_step)\n        logger.logkv(\"samples\", (self.step + self.resume_step + 1) * self.global_batch)\n\n    def save(self):\n        def save_checkpoint(rate, params):\n            state_dict = self.mp_trainer.master_params_to_state_dict(params)\n            if dist.get_rank() == 0:\n                logger.log(f\"saving model {rate}...\")\n                if not rate:\n                    filename = f\"model{(self.step+self.resume_step):06d}.pt\"\n                else:\n                    filename = f\"ema_{rate}_{(self.step+self.resume_step):06d}.pt\"\n                with bf.BlobFile(bf.join(get_blob_logdir(), filename), \"wb\") as f:\n                    th.save(state_dict, f)\n\n        save_checkpoint(0, self.mp_trainer.master_params)\n        for rate, params in zip(self.ema_rate, self.ema_params):\n            save_checkpoint(rate, params)\n\n        if dist.get_rank() == 0:\n            with bf.BlobFile(\n                bf.join(get_blob_logdir(), f\"opt{(self.step+self.resume_step):06d}.pt\"),\n                \"wb\",\n            ) as f:\n                th.save(self.opt.state_dict(), f)\n\n        dist.barrier()\n\n\ndef parse_resume_step_from_filename(filename):\n    \"\"\"\n    Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the\n    checkpoint's number of steps.\n    \"\"\"\n    split = filename.split(\"model\")\n    if len(split) < 2:\n        return 0\n    split1 = split[-1].split(\".\")[0]\n    try:\n        return int(split1)\n    except ValueError:\n        return 0\n\n\ndef get_blob_logdir():\n    # You can change this to be a separate path to save checkpoints to\n    # a blobstore or some external drive.\n    return logger.get_dir()\n\n\ndef find_resume_checkpoint():\n    # On your infrastructure, you may want to override this to automatically\n    # discover the latest checkpoint on your blob storage, etc.\n    return None\n\n\ndef find_ema_checkpoint(main_checkpoint, step, rate):\n    if main_checkpoint is None:\n        return None\n    filename = f\"ema_{rate}_{(step):06d}.pt\"\n    path = bf.join(bf.dirname(main_checkpoint), filename)\n    if bf.exists(path):\n        return path\n    return None\n\n\ndef log_loss_dict(diffusion, ts, losses):\n    for key, values in losses.items():\n        logger.logkv_mean(key, values.mean().item())\n        # Log the quantiles (four quartiles, in particular).\n        for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):\n            quartile = int(4 * sub_t / diffusion.num_timesteps)\n            logger.logkv_mean(f\"{key}_q{quartile}\", sub_loss)\n"
  },
  {
    "path": "guided_diffusion/unet.py",
    "content": "from abc import abstractmethod\n\nimport math\n\nimport numpy as np\nimport torch as th\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom .fp16_util import convert_module_to_f16, convert_module_to_f32\nfrom .nn import (\n    checkpoint,\n    conv_nd,\n    linear,\n    avg_pool_nd,\n    zero_module,\n    normalization,\n    timestep_embedding,\n)\n\n\nclass AttentionPool2d(nn.Module):\n    \"\"\"\n    Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py\n    \"\"\"\n\n    def __init__(\n        self,\n        spacial_dim: int,\n        embed_dim: int,\n        num_heads_channels: int,\n        output_dim: int = None,\n    ):\n        super().__init__()\n        self.positional_embedding = nn.Parameter(\n            th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5\n        )\n        self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)\n        self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)\n        self.num_heads = embed_dim // num_heads_channels\n        self.attention = QKVAttention(self.num_heads)\n\n    def forward(self, x):\n        b, c, *_spatial = x.shape\n        x = x.reshape(b, c, -1)  # NC(HW)\n        x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(HW+1)\n        x = x + self.positional_embedding[None, :, :].to(x.dtype)  # NC(HW+1)\n        x = self.qkv_proj(x)\n        x = self.attention(x)\n        x = self.c_proj(x)\n        return x[:, :, 0]\n\n\nclass TimestepBlock(nn.Module):\n    \"\"\"\n    Any module where forward() takes timestep embeddings as a second argument.\n    \"\"\"\n\n    @abstractmethod\n    def forward(self, x, emb):\n        \"\"\"\n        Apply the module to `x` given `emb` timestep embeddings.\n        \"\"\"\n\n\nclass TimestepEmbedSequential(nn.Sequential, TimestepBlock):\n    \"\"\"\n    A sequential module that passes timestep embeddings to the children that\n    support it as an extra input.\n    \"\"\"\n\n    def forward(self, x, emb):\n        for layer in self:\n            if isinstance(layer, TimestepBlock):\n                x = layer(x, emb)\n            else:\n                x = layer(x)\n        return x\n\n\nclass Upsample(nn.Module):\n    \"\"\"\n    An upsampling layer with an optional convolution.\n\n    :param channels: channels in the inputs and outputs.\n    :param use_conv: a bool determining if a convolution is applied.\n    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then\n                 upsampling occurs in the inner-two dimensions.\n    \"\"\"\n\n    def __init__(self, channels, use_conv, dims=2, out_channels=None):\n        super().__init__()\n        self.channels = channels\n        self.out_channels = out_channels or channels\n        self.use_conv = use_conv\n        self.dims = dims\n        if use_conv:\n            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)\n\n    def forward(self, x):\n        assert x.shape[1] == self.channels\n        if self.dims == 3:\n            x = F.interpolate(\n                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode=\"nearest\"\n            )\n        else:\n            x = F.interpolate(x, scale_factor=2, mode=\"nearest\")\n        if self.use_conv:\n            x = self.conv(x)\n        return x\n\n\nclass Downsample(nn.Module):\n    \"\"\"\n    A downsampling layer with an optional convolution.\n\n    :param channels: channels in the inputs and outputs.\n    :param use_conv: a bool determining if a convolution is applied.\n    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then\n                 downsampling occurs in the inner-two dimensions.\n    \"\"\"\n\n    def __init__(self, channels, use_conv, dims=2, out_channels=None):\n        super().__init__()\n        self.channels = channels\n        self.out_channels = out_channels or channels\n        self.use_conv = use_conv\n        self.dims = dims\n        stride = 2 if dims != 3 else (1, 2, 2)\n        if use_conv:\n            self.op = conv_nd(\n                dims, self.channels, self.out_channels, 3, stride=stride, padding=1\n            )\n        else:\n            assert self.channels == self.out_channels\n            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)\n\n    def forward(self, x):\n        assert x.shape[1] == self.channels\n        return self.op(x)\n\n\nclass ResBlock(TimestepBlock):\n    \"\"\"\n    A residual block that can optionally change the number of channels.\n\n    :param channels: the number of input channels.\n    :param emb_channels: the number of timestep embedding channels.\n    :param dropout: the rate of dropout.\n    :param out_channels: if specified, the number of out channels.\n    :param use_conv: if True and out_channels is specified, use a spatial\n        convolution instead of a smaller 1x1 convolution to change the\n        channels in the skip connection.\n    :param dims: determines if the signal is 1D, 2D, or 3D.\n    :param use_checkpoint: if True, use gradient checkpointing on this module.\n    :param up: if True, use this block for upsampling.\n    :param down: if True, use this block for downsampling.\n    \"\"\"\n\n    def __init__(\n        self,\n        channels,\n        emb_channels,\n        dropout,\n        out_channels=None,\n        use_conv=False,\n        use_scale_shift_norm=False,\n        dims=2,\n        use_checkpoint=False,\n        up=False,\n        down=False,\n    ):\n        super().__init__()\n        self.channels = channels\n        self.emb_channels = emb_channels\n        self.dropout = dropout\n        self.out_channels = out_channels or channels\n        self.use_conv = use_conv\n        self.use_checkpoint = use_checkpoint\n        self.use_scale_shift_norm = use_scale_shift_norm\n\n        self.in_layers = nn.Sequential(\n            normalization(channels),\n            nn.SiLU(),\n            conv_nd(dims, channels, self.out_channels, 3, padding=1),\n        )\n\n        self.updown = up or down\n\n        if up:\n            self.h_upd = Upsample(channels, False, dims)\n            self.x_upd = Upsample(channels, False, dims)\n        elif down:\n            self.h_upd = Downsample(channels, False, dims)\n            self.x_upd = Downsample(channels, False, dims)\n        else:\n            self.h_upd = self.x_upd = nn.Identity()\n\n        self.emb_layers = nn.Sequential(\n            nn.SiLU(),\n            linear(\n                emb_channels,\n                2 * self.out_channels if use_scale_shift_norm else self.out_channels,\n            ),\n        )\n        self.out_layers = nn.Sequential(\n            normalization(self.out_channels),\n            nn.SiLU(),\n            nn.Dropout(p=dropout),\n            zero_module(\n                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)\n            ),\n        )\n\n        if self.out_channels == channels:\n            self.skip_connection = nn.Identity()\n        elif use_conv:\n            self.skip_connection = conv_nd(\n                dims, channels, self.out_channels, 3, padding=1\n            )\n        else:\n            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)\n\n    def forward(self, x, emb):\n        \"\"\"\n        Apply the block to a Tensor, conditioned on a timestep embedding.\n\n        :param x: an [N x C x ...] Tensor of features.\n        :param emb: an [N x emb_channels] Tensor of timestep embeddings.\n        :return: an [N x C x ...] Tensor of outputs.\n        \"\"\"\n        return checkpoint(\n            self._forward, (x, emb), self.parameters(), self.use_checkpoint\n        )\n\n    def _forward(self, x, emb):\n        if self.updown:\n            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]\n            h = in_rest(x)\n            h = self.h_upd(h)\n            x = self.x_upd(x)\n            h = in_conv(h)\n        else:\n            h = self.in_layers(x)\n        emb_out = self.emb_layers(emb).type(h.dtype)\n        while len(emb_out.shape) < len(h.shape):\n            emb_out = emb_out[..., None]\n        if self.use_scale_shift_norm:\n            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]\n            scale, shift = th.chunk(emb_out, 2, dim=1)\n            h = out_norm(h) * (1 + scale) + shift\n            h = out_rest(h)\n        else:\n            h = h + emb_out\n            h = self.out_layers(h)\n        return self.skip_connection(x) + h\n\n\nclass AttentionBlock(nn.Module):\n    \"\"\"\n    An attention block that allows spatial positions to attend to each other.\n\n    Originally ported from here, but adapted to the N-d case.\n    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.\n    \"\"\"\n\n    def __init__(\n        self,\n        channels,\n        num_heads=1,\n        num_head_channels=-1,\n        use_checkpoint=False,\n        use_new_attention_order=False,\n    ):\n        super().__init__()\n        self.channels = channels\n        if num_head_channels == -1:\n            self.num_heads = num_heads\n        else:\n            assert (\n                channels % num_head_channels == 0\n            ), f\"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}\"\n            self.num_heads = channels // num_head_channels\n        self.use_checkpoint = use_checkpoint\n        self.norm = normalization(channels)\n        self.qkv = conv_nd(1, channels, channels * 3, 1)\n        if use_new_attention_order:\n            # split qkv before split heads\n            self.attention = QKVAttention(self.num_heads)\n        else:\n            # split heads before split qkv\n            self.attention = QKVAttentionLegacy(self.num_heads)\n\n        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))\n\n    def forward(self, x):\n        return checkpoint(self._forward, (x,), self.parameters(), True)\n\n    def _forward(self, x):\n        b, c, *spatial = x.shape\n        x = x.reshape(b, c, -1)\n        qkv = self.qkv(self.norm(x))\n        h = self.attention(qkv)\n        h = self.proj_out(h)\n        return (x + h).reshape(b, c, *spatial)\n\n\ndef count_flops_attn(model, _x, y):\n    \"\"\"\n    A counter for the `thop` package to count the operations in an\n    attention operation.\n    Meant to be used like:\n        macs, params = thop.profile(\n            model,\n            inputs=(inputs, timestamps),\n            custom_ops={QKVAttention: QKVAttention.count_flops},\n        )\n    \"\"\"\n    b, c, *spatial = y[0].shape\n    num_spatial = int(np.prod(spatial))\n    # We perform two matmuls with the same number of ops.\n    # The first computes the weight matrix, the second computes\n    # the combination of the value vectors.\n    matmul_ops = 2 * b * (num_spatial ** 2) * c\n    model.total_ops += th.DoubleTensor([matmul_ops])\n\n\nclass QKVAttentionLegacy(nn.Module):\n    \"\"\"\n    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping\n    \"\"\"\n\n    def __init__(self, n_heads):\n        super().__init__()\n        self.n_heads = n_heads\n\n    def forward(self, qkv):\n        \"\"\"\n        Apply QKV attention.\n\n        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.\n        :return: an [N x (H * C) x T] tensor after attention.\n        \"\"\"\n        bs, width, length = qkv.shape\n        assert width % (3 * self.n_heads) == 0\n        ch = width // (3 * self.n_heads)\n        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)\n        scale = 1 / math.sqrt(math.sqrt(ch))\n        weight = th.einsum(\n            \"bct,bcs->bts\", q * scale, k * scale\n        )  # More stable with f16 than dividing afterwards\n        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)\n        a = th.einsum(\"bts,bcs->bct\", weight, v)\n        return a.reshape(bs, -1, length)\n\n    @staticmethod\n    def count_flops(model, _x, y):\n        return count_flops_attn(model, _x, y)\n\n\nclass QKVAttention(nn.Module):\n    \"\"\"\n    A module which performs QKV attention and splits in a different order.\n    \"\"\"\n\n    def __init__(self, n_heads):\n        super().__init__()\n        self.n_heads = n_heads\n\n    def forward(self, qkv):\n        \"\"\"\n        Apply QKV attention.\n\n        :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.\n        :return: an [N x (H * C) x T] tensor after attention.\n        \"\"\"\n        bs, width, length = qkv.shape\n        assert width % (3 * self.n_heads) == 0\n        ch = width // (3 * self.n_heads)\n        q, k, v = qkv.chunk(3, dim=1)\n        scale = 1 / math.sqrt(math.sqrt(ch))\n        weight = th.einsum(\n            \"bct,bcs->bts\",\n            (q * scale).view(bs * self.n_heads, ch, length),\n            (k * scale).view(bs * self.n_heads, ch, length),\n        )  # More stable with f16 than dividing afterwards\n        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)\n        a = th.einsum(\"bts,bcs->bct\", weight, v.reshape(bs * self.n_heads, ch, length))\n        return a.reshape(bs, -1, length)\n\n    @staticmethod\n    def count_flops(model, _x, y):\n        return count_flops_attn(model, _x, y)\n\n\nclass UNetModel(nn.Module):\n    \"\"\"\n    The full UNet model with attention and timestep embedding.\n\n    :param in_channels: channels in the input Tensor.\n    :param model_channels: base channel count for the model.\n    :param out_channels: channels in the output Tensor.\n    :param num_res_blocks: number of residual blocks per downsample.\n    :param attention_resolutions: a collection of downsample rates at which\n        attention will take place. May be a set, list, or tuple.\n        For example, if this contains 4, then at 4x downsampling, attention\n        will be used.\n    :param dropout: the dropout probability.\n    :param channel_mult: channel multiplier for each level of the UNet.\n    :param conv_resample: if True, use learned convolutions for upsampling and\n        downsampling.\n    :param dims: determines if the signal is 1D, 2D, or 3D.\n    :param num_classes: if specified (as an int), then this model will be\n        class-conditional with `num_classes` classes.\n    :param use_checkpoint: use gradient checkpointing to reduce memory usage.\n    :param num_heads: the number of attention heads in each attention layer.\n    :param num_heads_channels: if specified, ignore num_heads and instead use\n                               a fixed channel width per attention head.\n    :param num_heads_upsample: works with num_heads to set a different number\n                               of heads for upsampling. Deprecated.\n    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.\n    :param resblock_updown: use residual blocks for up/downsampling.\n    :param use_new_attention_order: use a different attention pattern for potentially\n                                    increased efficiency.\n    \"\"\"\n\n    def __init__(\n        self,\n        image_size,\n        in_channels,\n        model_channels,\n        out_channels,\n        blur,\n        num_res_blocks = 3,\n        attention_resolutions = [4,8],\n        dropout=0,\n        channel_mult=(1, 2, 3, 4),\n        conv_resample=True,\n        dims=2,\n        num_classes=None,\n        use_checkpoint=False,\n        use_fp16=False,\n        num_heads=4,\n        num_head_channels=-1,\n        num_heads_upsample=-1,\n        use_scale_shift_norm=True,\n        resblock_updown=False,\n        use_new_attention_order=False,\n        freq_feat = False\n    ):\n        super().__init__()\n\n        if num_heads_upsample == -1:\n            num_heads_upsample = num_heads\n\n        self.image_size = image_size\n        self.in_channels = in_channels\n        self.model_channels = model_channels\n        self.out_channels = out_channels\n        self.num_res_blocks = num_res_blocks\n        self.attention_resolutions = attention_resolutions\n        self.dropout = dropout\n        self.channel_mult = channel_mult\n        self.conv_resample = conv_resample\n        self.num_classes = num_classes\n        self.use_checkpoint = use_checkpoint\n        self.dtype = th.float16 if use_fp16 else th.float32\n        self.num_heads = num_heads\n        self.num_head_channels = num_head_channels\n        self.num_heads_upsample = num_heads_upsample\n        self.freq_feat = freq_feat\n        self.blur = blur\n        time_embed_dim = model_channels * 4\n        self.time_embed = nn.Sequential(\n            linear(model_channels, time_embed_dim),\n            nn.SiLU(),\n            linear(time_embed_dim, time_embed_dim),\n        )\n        if freq_feat:\n            self.in_channels += 3\n            in_channels += 3\n        if self.num_classes is not None:\n            self.label_emb = nn.Embedding(num_classes, time_embed_dim)\n        ch = input_ch = int(channel_mult[0] * model_channels)\n        self.input_blocks = nn.ModuleList(\n            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]\n        )\n        print(\"input_blocks\", self.input_blocks[0][0].weight.shape)\n        if freq_feat:\n            # Set the weights to the frequency features to zeros\n            self.input_blocks[0][0].weight.data[:, 3: :, :] = 0\n            # for j in range(6):\n            #     w = self.input_blocks[0][0].weight[:, j, :, :]\n            #     print(f\"j: {j} w.mean(): {w.mean()} w.std(): {w.std()}\")\n            # raise NotImplementedError\n\n        self._feature_size = ch\n        input_block_chans = [ch]\n        ds = 1\n        for level, mult in enumerate(channel_mult):\n            for _ in range(num_res_blocks):\n                layers = [\n                    ResBlock(\n                        ch,\n                        time_embed_dim,\n                        dropout,\n                        out_channels=int(mult * model_channels),\n                        dims=dims,\n                        use_checkpoint=use_checkpoint,\n                        use_scale_shift_norm=use_scale_shift_norm,\n                    )\n                ]\n                ch = int(mult * model_channels)\n                if ds in attention_resolutions:\n                    layers.append(\n                        AttentionBlock(\n                            ch,\n                            use_checkpoint=use_checkpoint,\n                            num_heads=num_heads,\n                            num_head_channels=num_head_channels,\n                            use_new_attention_order=use_new_attention_order,\n                        )\n                    )\n                self.input_blocks.append(TimestepEmbedSequential(*layers))\n                self._feature_size += ch\n                input_block_chans.append(ch)\n            if level != len(channel_mult) - 1:\n                out_ch = ch\n                self.input_blocks.append(\n                    TimestepEmbedSequential(\n                        ResBlock(\n                            ch,\n                            time_embed_dim,\n                            dropout,\n                            out_channels=out_ch,\n                            dims=dims,\n                            use_checkpoint=use_checkpoint,\n                            use_scale_shift_norm=use_scale_shift_norm,\n                            down=True,\n                        )\n                        if resblock_updown\n                        else Downsample(\n                            ch, conv_resample, dims=dims, out_channels=out_ch\n                        )\n                    )\n                )\n                ch = out_ch\n                input_block_chans.append(ch)\n                ds *= 2\n                self._feature_size += ch\n\n        self.middle_block = TimestepEmbedSequential(\n            ResBlock(\n                ch,\n                time_embed_dim,\n                dropout,\n                dims=dims,\n                use_checkpoint=use_checkpoint,\n                use_scale_shift_norm=use_scale_shift_norm,\n            ),\n            AttentionBlock(\n                ch,\n                use_checkpoint=use_checkpoint,\n                num_heads=num_heads,\n                num_head_channels=num_head_channels,\n                use_new_attention_order=use_new_attention_order,\n            ),\n            ResBlock(\n                ch,\n                time_embed_dim,\n                dropout,\n                dims=dims,\n                use_checkpoint=use_checkpoint,\n                use_scale_shift_norm=use_scale_shift_norm,\n            ),\n        )\n        self._feature_size += ch\n\n        self.output_blocks = nn.ModuleList([])\n        for level, mult in list(enumerate(channel_mult))[::-1]:\n            for i in range(num_res_blocks + 1):\n                ich = input_block_chans.pop()\n                layers = [\n                    ResBlock(\n                        ch + ich,\n                        time_embed_dim,\n                        dropout,\n                        out_channels=int(model_channels * mult),\n                        dims=dims,\n                        use_checkpoint=use_checkpoint,\n                        use_scale_shift_norm=use_scale_shift_norm,\n                    )\n                ]\n                ch = int(model_channels * mult)\n                if ds in attention_resolutions:\n                    layers.append(\n                        AttentionBlock(\n                            ch,\n                            use_checkpoint=use_checkpoint,\n                            num_heads=num_heads_upsample,\n                            num_head_channels=num_head_channels,\n                            use_new_attention_order=use_new_attention_order,\n                        )\n                    )\n                if level and i == num_res_blocks:\n                    out_ch = ch\n                    layers.append(\n                        ResBlock(\n                            ch,\n                            time_embed_dim,\n                            dropout,\n                            out_channels=out_ch,\n                            dims=dims,\n                            use_checkpoint=use_checkpoint,\n                            use_scale_shift_norm=use_scale_shift_norm,\n                            up=True,\n                        )\n                        if resblock_updown\n                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)\n                    )\n                    ds //= 2\n                self.output_blocks.append(TimestepEmbedSequential(*layers))\n                self._feature_size += ch\n\n        self.out = nn.Sequential(\n            normalization(ch),\n            nn.SiLU(),\n            zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),\n        )\n\n    def convert_to_fp16(self):\n        \"\"\"\n        Convert the torso of the model to float16.\n        \"\"\"\n        self.input_blocks.apply(convert_module_to_f16)\n        self.middle_block.apply(convert_module_to_f16)\n        self.output_blocks.apply(convert_module_to_f16)\n\n    def convert_to_fp32(self):\n        \"\"\"\n        Convert the torso of the model to float32.\n        \"\"\"\n        self.input_blocks.apply(convert_module_to_f32)\n        self.middle_block.apply(convert_module_to_f32)\n        self.output_blocks.apply(convert_module_to_f32)\n\n    def forward(self, x, timesteps, y=None):\n        \"\"\"\n        Apply the model to an input batch.\n\n        :param x: an [N x C x ...] Tensor of inputs.\n        :param timesteps: a 1-D batch of timesteps.\n        :param y: an [N] Tensor of labels, if class-conditional.\n        :return: an [N x C x ...] Tensor of outputs.\n        \"\"\"\n        assert (y is not None) == (\n            self.num_classes is not None\n        ), \"must specify y if and only if the model is class-conditional\"\n        # timesteps = [timesteps] * x.shape[0]\n        # timesteps = th.tensor(timesteps).to(x.device)\n        if self.freq_feat:\n            x_freq = self.blur.Ut(x).view(x.shape)\n            x = th.cat([x, x_freq], dim=1)\n        hs = []\n        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))\n\n        if self.num_classes is not None:\n            assert y.shape == (x.shape[0],)\n            emb = emb + self.label_emb(y)\n\n        h = x.type(self.dtype)\n        for module in self.input_blocks:\n            h = module(h, emb)\n            hs.append(h)\n        h = self.middle_block(h, emb)\n        for module in self.output_blocks:\n       \n            h = th.cat([h, hs.pop()], dim=1)\n            h = module(h, emb)\n        h = h.type(x.dtype)\n        return self.out(h)\n\n\nclass SuperResModel(UNetModel):\n    \"\"\"\n    A UNetModel that performs super-resolution.\n\n    Expects an extra kwarg `low_res` to condition on a low-resolution image.\n    \"\"\"\n\n    def __init__(self, image_size, in_channels, *args, **kwargs):\n        super().__init__(image_size, in_channels * 2, *args, **kwargs)\n\n    def forward(self, x, timesteps, low_res=None, **kwargs):\n        _, _, new_height, new_width = x.shape\n        upsampled = F.interpolate(low_res, (new_height, new_width), mode=\"bilinear\")\n        x = th.cat([x, upsampled], dim=1)\n        return super().forward(x, timesteps, **kwargs)\n\n\nclass EncoderUNetModel(nn.Module):\n    \"\"\"\n    The half UNet model with attention and timestep embedding.\n\n    For usage, see UNet.\n    \"\"\"\n\n    def __init__(\n        self,\n        image_size,\n        in_channels,\n        model_channels,\n        out_channels,\n        num_res_blocks,\n        attention_resolutions,\n        dropout=0,\n        channel_mult=(1, 2, 4, 8),\n        conv_resample=True,\n        dims=2,\n        use_checkpoint=False,\n        use_fp16=False,\n        num_heads=1,\n        num_head_channels=-1,\n        num_heads_upsample=-1,\n        use_scale_shift_norm=False,\n        resblock_updown=False,\n        use_new_attention_order=False,\n        pool=\"adaptive\",\n    ):\n        super().__init__()\n\n        if num_heads_upsample == -1:\n            num_heads_upsample = num_heads\n\n        self.in_channels = in_channels\n        self.model_channels = model_channels\n        self.out_channels = out_channels\n        self.num_res_blocks = num_res_blocks\n        self.attention_resolutions = attention_resolutions\n        self.dropout = dropout\n        self.channel_mult = channel_mult\n        self.conv_resample = conv_resample\n        self.use_checkpoint = use_checkpoint\n        self.dtype = th.float16 if use_fp16 else th.float32\n        self.num_heads = num_heads\n        self.num_head_channels = num_head_channels\n        self.num_heads_upsample = num_heads_upsample\n\n        time_embed_dim = model_channels * 4\n        self.time_embed = nn.Sequential(\n            linear(model_channels, time_embed_dim),\n            nn.SiLU(),\n            linear(time_embed_dim, time_embed_dim),\n        )\n\n        ch = int(channel_mult[0] * model_channels)\n        self.input_blocks = nn.ModuleList(\n            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]\n        )\n        self._feature_size = ch\n        input_block_chans = [ch]\n        ds = 1\n        for level, mult in enumerate(channel_mult):\n            for _ in range(num_res_blocks):\n                layers = [\n                    ResBlock(\n                        ch,\n                        time_embed_dim,\n                        dropout,\n                        out_channels=int(mult * model_channels),\n                        dims=dims,\n                        use_checkpoint=use_checkpoint,\n                        use_scale_shift_norm=use_scale_shift_norm,\n                    )\n                ]\n                ch = int(mult * model_channels)\n                if ds in attention_resolutions:\n                    layers.append(\n                        AttentionBlock(\n                            ch,\n                            use_checkpoint=use_checkpoint,\n                            num_heads=num_heads,\n                            num_head_channels=num_head_channels,\n                            use_new_attention_order=use_new_attention_order,\n                        )\n                    )\n                self.input_blocks.append(TimestepEmbedSequential(*layers))\n                self._feature_size += ch\n                input_block_chans.append(ch)\n            if level != len(channel_mult) - 1:\n                out_ch = ch\n                self.input_blocks.append(\n                    TimestepEmbedSequential(\n                        ResBlock(\n                            ch,\n                            time_embed_dim,\n                            dropout,\n                            out_channels=out_ch,\n                            dims=dims,\n                            use_checkpoint=use_checkpoint,\n                            use_scale_shift_norm=use_scale_shift_norm,\n                            down=True,\n                        )\n                        if resblock_updown\n                        else Downsample(\n                            ch, conv_resample, dims=dims, out_channels=out_ch\n                        )\n                    )\n                )\n                ch = out_ch\n                input_block_chans.append(ch)\n                ds *= 2\n                self._feature_size += ch\n\n        self.middle_block = TimestepEmbedSequential(\n            ResBlock(\n                ch,\n                time_embed_dim,\n                dropout,\n                dims=dims,\n                use_checkpoint=use_checkpoint,\n                use_scale_shift_norm=use_scale_shift_norm,\n            ),\n            AttentionBlock(\n                ch,\n                use_checkpoint=use_checkpoint,\n                num_heads=num_heads,\n                num_head_channels=num_head_channels,\n                use_new_attention_order=use_new_attention_order,\n            ),\n            ResBlock(\n                ch,\n                time_embed_dim,\n                dropout,\n                dims=dims,\n                use_checkpoint=use_checkpoint,\n                use_scale_shift_norm=use_scale_shift_norm,\n            ),\n        )\n        self._feature_size += ch\n        self.pool = pool\n        if pool == \"adaptive\":\n            self.out = nn.Sequential(\n                normalization(ch),\n                nn.SiLU(),\n                nn.AdaptiveAvgPool2d((1, 1)),\n                zero_module(conv_nd(dims, ch, out_channels, 1)),\n                nn.Flatten(),\n            )\n        elif pool == \"attention\":\n            assert num_head_channels != -1\n            self.out = nn.Sequential(\n                normalization(ch),\n                nn.SiLU(),\n                AttentionPool2d(\n                    (image_size // ds), ch, num_head_channels, out_channels\n                ),\n            )\n        elif pool == \"spatial\":\n            self.out = nn.Sequential(\n                nn.Linear(self._feature_size, 2048),\n                nn.ReLU(),\n                nn.Linear(2048, self.out_channels),\n            )\n        elif pool == \"spatial_v2\":\n            self.out = nn.Sequential(\n                nn.Linear(self._feature_size, 2048),\n                normalization(2048),\n                nn.SiLU(),\n                nn.Linear(2048, self.out_channels),\n            )\n        else:\n            raise NotImplementedError(f\"Unexpected {pool} pooling\")\n\n    def convert_to_fp16(self):\n        \"\"\"\n        Convert the torso of the model to float16.\n        \"\"\"\n        self.input_blocks.apply(convert_module_to_f16)\n        self.middle_block.apply(convert_module_to_f16)\n\n    def convert_to_fp32(self):\n        \"\"\"\n        Convert the torso of the model to float32.\n        \"\"\"\n        self.input_blocks.apply(convert_module_to_f32)\n        self.middle_block.apply(convert_module_to_f32)\n\n    def forward(self, x, timesteps):\n        \"\"\"\n        Apply the model to an input batch.\n\n        :param x: an [N x C x ...] Tensor of inputs.\n        :param timesteps: a 1-D batch of timesteps.\n        :return: an [N x K] Tensor of outputs.\n        \"\"\"\n        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))\n\n        results = []\n        h = x.type(self.dtype)\n        for module in self.input_blocks:\n            h = module(h, emb)\n            if self.pool.startswith(\"spatial\"):\n                results.append(h.type(x.dtype).mean(dim=(2, 3)))\n        h = self.middle_block(h, emb)\n        if self.pool.startswith(\"spatial\"):\n            results.append(h.type(x.dtype).mean(dim=(2, 3)))\n            h = th.cat(results, axis=-1)\n            return self.out(h)\n        else:\n            h = h.type(x.dtype)\n            return self.out(h)\n"
  },
  {
    "path": "main.py",
    "content": "import torch\nfrom PIL import Image\nfrom collections import defaultdict\nimport torch.nn.functional as TF\nimport torchvision.datasets as dsets\nfrom torchvision import transforms\nimport numpy as np\n\nimport torch.optim as optim\nimport torchvision\nimport matplotlib.pyplot as plt\nimport utils\nfrom guided_diffusion.unet import UNetModel\nimport math\nfrom tensorboardX import SummaryWriter\nimport os\nimport json\nfrom collections import namedtuple\nimport argparse\nfrom torchvision.utils import save_image\nfrom tqdm import tqdm\nfrom blur_diffusion import Deblurring, ForwardBlurIncreasing, gaussian_kernel_1d\nfrom utils import normalize_np, clear\nfrom EMA import EMA\nfrom torch.nn import DataParallel\nfrom fid import FID\nfrom scipy.integrate import solve_ivp\n\nparser = argparse.ArgumentParser(description='Configs')\nparser.add_argument('--gpu', type=str, help='gpu num')\nparser.add_argument('--dataset', type=str, help='cifar10 / mnist')\nparser.add_argument('--name', type=str, help='Saving directory name')\nparser.add_argument('--ckpt', default='', type=str, help='UNet checkpoint')\n\nparser.add_argument('--bsize', default=16, type=int, help='batchsize')\nparser.add_argument('--N', default=1000, type=int, help='Max diffusion timesteps')\nparser.add_argument('--sig', default=0.4, type=float, help='sigma value for blur kernel')\nparser.add_argument('--sig_min', default=0, type=float, help='sigma value for blur kernel')\nparser.add_argument('--sig_max', default=0.1, type=float, help='sigma value for blur kernel')\nparser.add_argument('--lr', default=0.00005, type=float, help='learning rate')\nparser.add_argument('--noise_schedule', default='linear', type=str, help='Type of noise schedule to use')\nparser.add_argument('--betamin', default=0.0001, type=float, help='beta (min). get_score(1) can diverge if this is too low.')\nparser.add_argument('--betamax', default=0.02, type=float, help='beta (max)')\nparser.add_argument('--fromprior', default=True, type=bool, help='start sampling from prior')\nparser.add_argument('--gtscore', action='store_true', help='Use ground truth score for reverse diffusion')\nparser.add_argument('--max_iter', default=1500000, type=int, help='max iterations')\nparser.add_argument('--eval_iter', default=10000, type=int, help='eval iterations')\nparser.add_argument('--fid_iter', default=50000, type=int, help='eval iterations')\nparser.add_argument('--fid_num_samples', default=10000, type=int, help='eval iterations')\nparser.add_argument('--fid_bsize', default=32, type=int, help='eval iterations')\nparser.add_argument('--loss_type', type=str, default = 'eps_simple', choices=['sm_simple', 'eps_simple', 'sm_exact', 'std_matching'])\nparser.add_argument('--f_type', type=str, default = 'linear', choices=['linear', 'log', 'quadratic', 'cubic', 'quartic', 'triangular'])\nparser.add_argument('--dropout', default=0, type=float, help='dropout')\n\n# EMA, save\nparser.add_argument('--use_ema', action='store_true',\n                    help='use EMA or not')\nparser.add_argument('--inference', action='store_true')\nparser.add_argument('--freq_feat', action='store_true', help = \"concat Utx_i\")\nparser.add_argument('--ode', action='store_true', help = \"ODE fast sampler\")\nparser.add_argument('--ema_decay', type=float, default=0.9999, help='decay rate for EMA')\nparser.add_argument('--save_every', type=int, default=50000, help='How often we wish to save ckpts')\nopt = parser.parse_args()\n\ndataset = opt.dataset\ndevice = torch.device(f'cuda:{opt.gpu}')\ndevice = torch.device('cuda')\nprint(\"N:\", opt.N)\nN = opt.N\nbsize = opt.bsize\nbeta_min = opt.betamin\nbeta_max = opt.betamax\nsig = opt.sig\n\nif opt.gtscore:\n    opt.fromprior = True\nif dataset == 'mnist':\n    train_dir = '/home/ubuntu/code/sangyoon/datasets/mnist_train'\n    dataset_train = dsets.MNIST(root='/home/ubuntu/code/sangyoon/datasets/mnist_train',\n                                train=True,\n                                transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()]),\n                                download=True)\n    dataset_test = dsets.MNIST(root='/home/ubuntu/code/sangyoon/datasets/mnist_test',\n                               train=False,\n                               transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()]),\n                               download=True)\n\nelif dataset == 'cifar10':\n    train_dir = \"/home/ubuntu/code/sangyoon/datasets/cifar10_train/png\"\n    dataset_train = dsets.CIFAR10(root='../datasets/cifar10_train',\n                                  train=True,\n                                  transform=transforms.Compose([transforms.RandomHorizontalFlip(0.5), transforms.ToTensor()]),\n                                  download=True\n                                  )\n    dataset_test = dsets.CIFAR10(root='../datasets/cifar10_test',\n                                 train=False,\n                                 transform=transforms.Compose([transforms.RandomHorizontalFlip(0.5), transforms.ToTensor()]),\n                                 download=True\n                                 )\nelif dataset == 'lsun-bedroom':\n    res = 64\n    train_dir = \"/home/ubuntu/code/sangyoon/forward-blur/datasets/bedroom_train/bedroom_train_lmdb/imgs\"\n    dataset_train = dsets.LSUN(root='../forward-blur/datasets/bedroom_train',\n                                 classes=['bedroom_train'],\n                                  transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])\n                                  )\n    dataset_test = dsets.LSUN(root='../forward-blur/datasets/bedroom_train',\n                                    classes=['bedroom_train'],\n                                 transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])\n                                 )\nelif dataset == 'lsun-church':\n    res = 64\n    train_dir = \"/home/ubuntu/code/sangyoon/datasets/church_train/church_outdoor_train_lmdb/imgs\"\n    dataset_train = dsets.LSUN(root='/home/ubuntu/code/sangyoon/datasets/church_train',\n                                 classes=['church_outdoor_train'],\n                                  transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])\n                                  )\n    dataset_test = dsets.LSUN(root='/home/ubuntu/code/sangyoon/datasets/church_test',\n                                    classes=['church_outdoor_val'],\n                                 transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])\n                                 )\n\nfid_eval = FID(real_dir = train_dir, device = device)\nresolution = dataset_train[0][0].shape[-1]\ninput_nc = dataset_train[0][0].shape[0]\nksize = resolution * 2 - 1\npad = 0\n\n# define forward blur\nkernel = gaussian_kernel_1d(ksize, sig)\nblur = Deblurring(kernel, input_nc, resolution, device=device)\nprint(\"blur.U_small.shape:\", blur.U_small.shape)\nD_diag = blur.singulars()\nfb = ForwardBlurIncreasing(N=N, beta_min=beta_min, beta_max=beta_max, sig=sig, sig_max = opt.sig_max, sig_min = opt.sig_min, D_diag=D_diag,\n                    blur=blur, channel=input_nc, device=device, noise_schedule=opt.noise_schedule, resolution=resolution, pad=pad, f_type=opt.f_type)\ndir = os.path.join('experiments', opt.name)\nwriter = SummaryWriter(dir)\n\n# # Bedroom config - ADM\n# model_channels = 256\n# num_res_blocks = 2\n# channel_mult = (1, 1, 2, 2, 4, 4)\n# attention_resolutions = [8,16,32]\n# out_channels = 6 # mean and varaince\n# dropout = 0.1\n# resblock_updown = True\n# # bsize 256\n# # iter -- cat (200K), horse (250K), bedroom (500K)\n# lr = 1e-5 # Fine tuning\n# model = UNetModel(image_size = resolution, in_channels = input_nc, model_channels = model_channels, out_channels = out_channels,\n#                   num_res_blocks=num_res_blocks, channel_mult=channel_mult, attention_resolutions=attention_resolutions,\n#                   dropout = dropout, resblock_updown = resblock_updown)\n\nmodel = UNetModel(resolution, input_nc, 128, input_nc, blur = blur, dropout=opt.dropout, freq_feat = opt.freq_feat)\nif not opt.ckpt == '' and os.path.exists(opt.ckpt):\n    model.load_state_dict(torch.load(opt.ckpt))\nif torch.cuda.device_count() > 1:\n  print(\"Let's use\", torch.cuda.device_count(), \"GPUs!\")\n  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs\n  model = DataParallel(model)\n\nmodel.to(device)\nprint(\"input_nc\", input_nc, \"resolution\", resolution)\n\ndata_loader = torch.utils.data.DataLoader(dataset=dataset_train,\n                                          batch_size=bsize,\n                                          shuffle=True,\n                                          drop_last=True)\ndata_loader_test = torch.utils.data.DataLoader(dataset=dataset_test,\n                                               batch_size=bsize,\n                                               shuffle=False,\n                                               drop_last=True)\noptimizer = optim.Adam(model.parameters(), lr=opt.lr)\nif opt.use_ema:\n    optimizer = EMA(optimizer, ema_decay=opt.ema_decay)\n\n# forward process visualization\nsample = dataset_train[1][0].unsqueeze(0)\n\nx_0 = sample[:4]\nx_0 = x_0.to(device)\ni = np.array([500] * x_0.shape[0])\ni = torch.from_numpy(i).to(device)\nfb.sanity(x_0, i)\n\nsample_list = []\nfor i in range(0, N+1, N//10):\n    if i == 0:\n        sample_list.append(x_0[0])\n        continue\n    i = np.array([i] * x_0.shape[0])\n    i = torch.from_numpy(i).to(device)\n    x_i = fb.get_x_i(x_0, i)\n    sample_list.append(x_i[0])\n    print(f\"x_{i.item()}.std() = {x_i.std()}\")\n    print(f\"x_{i.item()}.mean() = {x_i.mean()}\")\n\n\ngrid_sample = torch.cat(sample_list, dim=2)\nutils.tensor_imsave(grid_sample, \"./\" + dir, \"forward_process.jpg\")\nwith open(os.path.join(dir, \"config.json\"), \"w\") as json_file:\n    json.dump(vars(opt), json_file)\nimport time\nmeta_iter = 0\nfor step in range(opt.max_iter):\n    if not opt.inference:\n        elips = time.time()\n        try:\n            x_0, _ = train_iter.next()\n        except:\n            train_iter = iter(data_loader)\n            image, _ = train_iter.next()\n        \"\"\"\n        training\n        \"\"\"\n        assert x_0.shape[-1] == resolution, f\"{x_0.shape}\"\n        i = np.random.uniform(1 / N, 1, size = (x_0.shape[0])) * N\n        i = torch.from_numpy(i).to(device).type(torch.long)\n\n        x_0 = x_0.to(device)\n        x_i, eps = fb.get_x_i(x_0, i, return_eps = True)\n\n        if opt.loss_type == \"sm_simple\":\n            loss = fb.get_loss_i_simple(model, x_0, x_i, i)\n        elif opt.loss_type == \"eps_simple\":\n            loss = fb.get_loss_i_eps_simple(model, x_i, i, eps)\n        elif opt.loss_type == \"sm_exact\":\n            loss = fb.get_loss_i_exact(model, x_0, x_i, i)\n        elif opt.loss_type == \"std_matching\":\n            loss = fb.get_loss_i_std_matching(model, x_i, i, eps)\n        writer.add_scalar('loss_train', loss, step)\n        loss.backward()\n        optimizer.step()\n        optimizer.zero_grad()\n        if step % 100 == 0:\n            print(step, loss)\n        # print(f\"time: {time.time() - elips}\")\n    # Calcuate FID\n    if step > 240001:\n        fid_iter = opt.fid_iter\n    else:\n        fid_iter = 240000\n    if (step % fid_iter == 0 and step > 0):\n        id = 0\n        if not os.path.exists(os.path.join(\"./\",dir, f\"{step}\")):\n            os.mkdir(os.path.join(\"./\",dir, f\"{step}\"))\n        with torch.no_grad():\n            if opt.use_ema:\n                optimizer.swap_parameters_with_ema(store_params_in_ema=True)\n            model = model.eval()\n            for _ in range(opt.fid_num_samples // opt.fid_bsize):\n                i = np.array([opt.N - 1] * opt.fid_bsize)\n                i = torch.from_numpy(i).to(device)\n                pred = fb.get_x_N([opt.fid_bsize, input_nc, resolution, resolution], i)\n                for i in reversed(range(1, opt.N)):\n                    i = np.array([i] * opt.fid_bsize)\n                    i = torch.from_numpy(i).to(device)\n                    if opt.loss_type == \"sm_simple\":\n                        s = model(pred, i)\n                    elif opt.loss_type == \"eps_simple\":\n                        eps = model(pred, i)\n                        s = fb.get_score_from_eps(eps, i)\n                    elif opt.loss_type == \"sm_exact\":\n                        s = model(pred, i)\n                    elif opt.loss_type == \"std_matching\":\n                        std = model(pred, i)\n                        s = fb.get_score_from_std(std, i)\n                    else:\n                        raise NotImplementedError\n                    s = fb.U_I_minus_B_Ut(s, i)\n                    rms = lambda x: torch.sqrt(torch.mean(x ** 2))\n                    # print(f\"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}\")\n                    hf = pred - fb.W(pred, i)\n                    # Anderson theorem\n                    pred1 = pred + hf # unsharpening mask filtering\n                    pred2 = pred1 + s  # # denoising\n                    if i[0] > 2:\n                        pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise\n                    else:\n                        pred = pred2\n                    # print(f\"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}\" )\n                for sample in pred:\n                    save_image(sample, os.path.join(dir, f\"{step}\", f\"{id:05d}.png\"))\n                    id += 1\n        if opt.use_ema:\n            optimizer.swap_parameters_with_ema(store_params_in_ema=True)\n            model = model.train()\n        fid = fid_eval(os.path.join(dir, f\"{step}\"))\n        writer.add_scalar('fid', fid, step)\n        print(f\"step {step}, fid = {fid}\")\n    if (step % opt.eval_iter == 0 and step > 0) or opt.inference:\n        \"\"\"\n        sampling (eval)\n        \"\"\"\n        cnt = 0\n        loss = 0\n        \n\n        with torch.no_grad():\n            if opt.use_ema:\n                optimizer.swap_parameters_with_ema(store_params_in_ema=True)\n            model = model.eval()\n\n            if opt.ode:\n                raise NotImplementedError\n                def to_flattened_numpy(x):\n                    \"\"\"Flatten a torch tensor `x` and convert it to numpy.\"\"\"\n                    return x.detach().cpu().numpy().reshape((-1,))\n                def from_flattened_numpy(x, shape):\n                    \"\"\"Form a torch tensor with the given `shape` from a flattened numpy array `x`.\"\"\"\n                    return torch.from_numpy(x.reshape(shape))\n                def ode_func(i, y):\n                    i = int(i*N)\n                    print(f\"i = {i}\")\n                    y = from_flattened_numpy(y, [bsize, input_nc, resolution, resolution]).to(device).type(torch.float32)\n                    i = np.array([N - 1] * bsize)\n                    i = torch.from_numpy(i).to(device)\n                    if opt.loss_type == \"sm_simple\":\n                            s = model(y, i)\n                    elif opt.loss_type == \"eps_simple\":\n                        eps = model(y, i)\n                        s = fb.get_score_from_eps(eps, i)\n                    elif opt.loss_type == \"sm_exact\":\n                        s = model(y, i)\n                    elif opt.loss_type == \"std_matching\":\n                        std = model(y, i)\n                        s = fb.get_score_from_std(std, i)\n                    else:\n                        raise NotImplementedError\n                    s = fb.U_I_minus_B_Ut(s, i)\n                    hf = y - fb.W(y, i)\n                    dt = - 1.0 / N\n                    drift = (s/2 + hf) / dt\n                    drift = to_flattened_numpy(drift)\n                    return drift\n                x_N = fb.get_x_N([bsize, input_nc, resolution, resolution], N)\n                solution = solve_ivp(ode_func, (1, 1e-3), to_flattened_numpy(x_N),\n                                     rtol=1e-3, atol=1e-3, method=\"RK45\")\n                nfe = solution.nfev\n                solution = torch.tensor(solution.y[:, -1]).reshape(x_N.shape).to(device).type(torch.float32)\n                \n                save_image(solution, \"./solution.jpg\")\n                print(f\"nfe = {nfe}\")\n                raise NotImplementedError\n            for x_0, _ in data_loader_test:\n                x_0 = x_0.to(device)\n                # for v in range(0, 250, 20):\n                #     x_0[:, :, v, :] = 0\n                if opt.fromprior:\n                    i = np.array([N - 1] * x_0.shape[0])\n                    i = torch.from_numpy(i).to(device)\n                    pred = fb.get_x_N(x_0.shape, i)\n                    print(f\"pred.std() = {pred.std()}\")\n                else:\n                    i = np.array([N-1] * x_0.shape[0])\n                    i = torch.from_numpy(i).to(device)\n                    pred = fb.get_x_i(x_0, i)\n                preds = [pred]\n\n                for i in reversed(range(1, N)):\n                    i = np.array([i] * x_0.shape[0])\n                    i = torch.from_numpy(i).to(device)\n                    if opt.gtscore:\n                        s = fb.get_score_gt(pred, x_0, i)\n                    else:\n                        if opt.loss_type == \"sm_simple\":\n                            s = model(pred, i)\n                        elif opt.loss_type == \"eps_simple\":\n                            eps = model(pred, i)\n                            s = fb.get_score_from_eps(eps, i)\n                        elif opt.loss_type == \"sm_exact\":\n                            s = model(pred, i)\n                        elif opt.loss_type == \"std_matching\":\n                            std = model(pred, i)\n                            s = fb.get_score_from_std(std, i)\n                        else:\n                            raise NotImplementedError\n                    s = fb.U_I_minus_B_Ut(s, i)\n                    rms = lambda x: torch.sqrt(torch.mean(x ** 2))\n                    # print(f\"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}\")\n                    hf = pred - fb.W(pred, i)\n                    # Anderson theorem\n                    pred1 = pred + hf # unsharpening mask filtering\n                    pred2 = pred1 + s  # # denoising\n                    if i[0] > 2:\n                        pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise\n                    else:\n                        pred = pred2\n                    # print(f\"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}\")\n                    # assert rms(pred) < 100\n                    if (i[0]) % (N // 10) == 0:\n                        img = pred[0]\n                        preds.append(pred)\n\n                preds.append(pred)\n                assert x_0.shape == pred.shape\n                # visualize\n                grid = torch.cat(preds, dim=3) # grid_sample.shape: (bsize, channel, H, W * 12)\n                # (batch_size, channel, H, W * 12) -> (channel, H * bsize, W * 12)\n                grid = grid.permute(1, 0, 2, 3).contiguous().view(grid.shape[1], -1, grid.shape[3])\n                # (bsize, channel, H, W) -> (channel, H, W * bsize)\n                gt = x_0.permute(1, 2, 0, 3).contiguous().view(x_0.shape[1], -1, x_0.shape[3] * x_0.shape[0])\n                if cnt <= 2:\n                    utils.tensor_imsave(gt, \"./\" + dir, f\"{step}_{cnt}_GT.jpg\")\n                    utils.tensor_imsave(grid, \"./\" + dir, f\"{step}_{cnt}_pred.jpg\")\n              \n                cnt += 1\n                loss += TF.l1_loss(x_0, pred) / 2\n\n                if cnt == 2:\n                    break\n        print(f\"step: {step} loss: {loss}\")\n        writer.add_scalar('loss_val', loss, meta_iter)\n        f = open('./' + str(dir) + '/log.txt', 'a')\n\n        f.write(f\"Step: {step} loss: {loss}\" + '\\n')\n\n        f.close()\n        model.train()\n        if opt.use_ema:\n            optimizer.swap_parameters_with_ema(store_params_in_ema=True)\n    if step % opt.save_every == 1:\n        if opt.use_ema:\n            optimizer.swap_parameters_with_ema(store_params_in_ema=True)\n        if torch.cuda.device_count() > 1:\n            torch.save(model.module.state_dict(), os.path.join(dir, f\"model_{step}.ckpt\"))\n        else:\n            torch.save(model.state_dict(), os.path.join(dir, f\"model_{step}.ckpt\"))\n        if opt.use_ema:\n            optimizer.swap_parameters_with_ema(store_params_in_ema=True)\n"
  },
  {
    "path": "pytorch_fid/__init__.py",
    "content": "\"\"\"\nhttps://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py\n\"\"\"\n__version__ = '0.2.1'\n"
  },
  {
    "path": "pytorch_fid/__main__.py",
    "content": "import pytorch_fid.fid_score\n\npytorch_fid.fid_score.main()\n"
  },
  {
    "path": "pytorch_fid/fid_score.py",
    "content": "\"\"\"Calculates the Frechet Inception Distance (FID) to evalulate GANs\n\nThe FID metric calculates the distance between two distributions of images.\nTypically, we have summary statistics (mean & covariance matrix) of one\nof these distributions, while the 2nd distribution is given by a GAN.\n\nWhen run as a stand-alone program, it compares the distribution of\nimages that are stored as PNG/JPEG at a specified location with a\ndistribution given by summary statistics (in pickle format).\n\nThe FID is calculated by assuming that X_1 and X_2 are the activations of\nthe pool_3 layer of the inception net for generated samples and real world\nsamples respectively.\n\nSee --help to see further details.\n\nCode apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead\nof Tensorflow\n\nCopyright 2018 Institute of Bioinformatics, JKU Linz\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n   http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an \"AS IS\" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License.\n\"\"\"\nimport os\nimport pathlib\nfrom argparse import ArgumentDefaultsHelpFormatter, ArgumentParser\n\nimport numpy as np\nimport torch\nimport torchvision.transforms as TF\nfrom PIL import Image\nfrom scipy import linalg\nfrom torch.nn.functional import adaptive_avg_pool2d\n\ntry:\n    from tqdm import tqdm\nexcept ImportError:\n    # If tqdm is not available, provide a mock version of it\n    def tqdm(x):\n        return x\n\ntry:\n    from pytorch_fid.inception import InceptionV3\nexcept:\n    from inception import InceptionV3\n\nparser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)\nparser.add_argument('--batch-size', type=int, default=50,\n                    help='Batch size to use')\nparser.add_argument('--num-workers', type=int,\n                    help=('Number of processes to use for data loading. '\n                          'Defaults to `min(8, num_cpus)`'))\nparser.add_argument('--device', type=str, default=None,\n                    help='Device to use. Like cuda, cuda:0 or cpu')\nparser.add_argument('--dims', type=int, default=2048,\n                    choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),\n                    help=('Dimensionality of Inception features to use. '\n                          'By default, uses pool3 features'))\nparser.add_argument('path', type=str, nargs=2,\n                    help=('Paths to the generated images or '\n                          'to .npz statistic files'))\n\nIMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',\n                    'tif', 'tiff', 'webp'}\n\n\nclass ImagePathDataset(torch.utils.data.Dataset):\n    def __init__(self, files, transforms=None):\n        self.files = files\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, i):\n        path = self.files[i]\n        img = Image.open(path).convert('RGB')\n        if self.transforms is not None:\n            img = self.transforms(img)\n        return img\n\n\ndef get_activations(files, model, batch_size=50, dims=2048, device='cpu',\n                    num_workers=1):\n    \"\"\"Calculates the activations of the pool_3 layer for all images.\n\n    Params:\n    -- files       : List of image files paths\n    -- model       : Instance of inception model\n    -- batch_size  : Batch size of images for the model to process at once.\n                     Make sure that the number of samples is a multiple of\n                     the batch size, otherwise some samples are ignored. This\n                     behavior is retained to match the original FID score\n                     implementation.\n    -- dims        : Dimensionality of features returned by Inception\n    -- device      : Device to run calculations\n    -- num_workers : Number of parallel dataloader workers\n\n    Returns:\n    -- A numpy array of dimension (num images, dims) that contains the\n       activations of the given tensor when feeding inception with the\n       query tensor.\n    \"\"\"\n    model.eval()\n\n    if batch_size > len(files):\n        print(('Warning: batch size is bigger than the data size. '\n               'Setting batch size to data size'))\n        batch_size = len(files)\n\n    dataset = ImagePathDataset(files, transforms=TF.ToTensor())\n    dataloader = torch.utils.data.DataLoader(dataset,\n                                             batch_size=batch_size,\n                                             shuffle=False,\n                                             drop_last=False,\n                                             num_workers=num_workers)\n\n    pred_arr = np.empty((len(files), dims))\n\n    start_idx = 0\n\n    for batch in tqdm(dataloader):\n        batch = batch.to(device)\n\n        with torch.no_grad():\n            pred = model(batch)[0]\n\n        # If model output is not scalar, apply global spatial average pooling.\n        # This happens if you choose a dimensionality not equal 2048.\n        if pred.size(2) != 1 or pred.size(3) != 1:\n            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))\n\n        pred = pred.squeeze(3).squeeze(2).cpu().numpy()\n\n        pred_arr[start_idx:start_idx + pred.shape[0]] = pred\n\n        start_idx = start_idx + pred.shape[0]\n\n    return pred_arr\n\n\ndef calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):\n    \"\"\"Numpy implementation of the Frechet Distance.\n    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)\n    and X_2 ~ N(mu_2, C_2) is\n            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).\n\n    Stable version by Dougal J. Sutherland.\n\n    Params:\n    -- mu1   : Numpy array containing the activations of a layer of the\n               inception net (like returned by the function 'get_predictions')\n               for generated samples.\n    -- mu2   : The sample mean over activations, precalculated on an\n               representative data set.\n    -- sigma1: The covariance matrix over activations for generated samples.\n    -- sigma2: The covariance matrix over activations, precalculated on an\n               representative data set.\n\n    Returns:\n    --   : The Frechet Distance.\n    \"\"\"\n\n    mu1 = np.atleast_1d(mu1)\n    mu2 = np.atleast_1d(mu2)\n\n    sigma1 = np.atleast_2d(sigma1)\n    sigma2 = np.atleast_2d(sigma2)\n\n    assert mu1.shape == mu2.shape, \\\n        'Training and test mean vectors have different lengths'\n    assert sigma1.shape == sigma2.shape, \\\n        'Training and test covariances have different dimensions'\n\n    diff = mu1 - mu2\n\n    # Product might be almost singular\n    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)\n    if not np.isfinite(covmean).all():\n        msg = ('fid calculation produces singular product; '\n               'adding %s to diagonal of cov estimates') % eps\n        print(msg)\n        offset = np.eye(sigma1.shape[0]) * eps\n        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n\n    # Numerical error might give slight imaginary component\n    if np.iscomplexobj(covmean):\n        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):\n            m = np.max(np.abs(covmean.imag))\n            raise ValueError('Imaginary component {}'.format(m))\n        covmean = covmean.real\n\n    tr_covmean = np.trace(covmean)\n\n    return (diff.dot(diff) + np.trace(sigma1)\n            + np.trace(sigma2) - 2 * tr_covmean)\n\n\ndef calculate_activation_statistics(files, model, batch_size=50, dims=2048,\n                                    device='cpu', num_workers=1):\n    \"\"\"Calculation of the statistics used by the FID.\n    Params:\n    -- files       : List of image files paths\n    -- model       : Instance of inception model\n    -- batch_size  : The images numpy array is split into batches with\n                     batch size batch_size. A reasonable batch size\n                     depends on the hardware.\n    -- dims        : Dimensionality of features returned by Inception\n    -- device      : Device to run calculations\n    -- num_workers : Number of parallel dataloader workers\n\n    Returns:\n    -- mu    : The mean over samples of the activations of the pool_3 layer of\n               the inception model.\n    -- sigma : The covariance matrix of the activations of the pool_3 layer of\n               the inception model.\n    \"\"\"\n    act = get_activations(files, model, batch_size, dims, device, num_workers)\n    mu = np.mean(act, axis=0)\n    sigma = np.cov(act, rowvar=False)\n    return mu, sigma\n\n\ndef compute_statistics_of_path(path, model, batch_size, dims, device,\n                               num_workers=1):\n    if path.endswith('.npz'):\n        with np.load(path) as f:\n            m, s = f['mu'][:], f['sigma'][:]\n    else:\n        path = pathlib.Path(path)\n        files = sorted([file for ext in IMAGE_EXTENSIONS\n                       for file in path.glob('*.{}'.format(ext))])\n        print(f\"path: {path}\")\n        m, s = calculate_activation_statistics(files, model, batch_size,\n                                               dims, device, num_workers)\n        # Save to a npz file\n        np.savez_compressed(path/'inception_statistics.npz', mu=m, sigma=s)\n    return m, s\n\n\ndef calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1):\n    \"\"\"Calculates the FID of two paths\"\"\"\n    for p in paths:\n        if not os.path.exists(p):\n            raise RuntimeError('Invalid path: %s' % p)\n\n    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]\n\n    model = InceptionV3([block_idx]).to(device)\n\n    m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,\n                                        dims, device, num_workers)\n    m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,\n                                        dims, device, num_workers)\n    fid_value = calculate_frechet_distance(m1, s1, m2, s2)\n\n    return fid_value\n\n\ndef main():\n    args = parser.parse_args()\n\n    if args.device is None:\n        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')\n    else:\n        device = torch.device(args.device)\n\n    if args.num_workers is None:\n        num_avail_cpus = len(os.sched_getaffinity(0))\n        num_workers = min(num_avail_cpus, 8)\n    else:\n        num_workers = args.num_workers\n\n    fid_value = calculate_fid_given_paths(args.path,\n                                          args.batch_size,\n                                          device,\n                                          args.dims,\n                                          num_workers)\n    print('FID: ', fid_value)\n\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "pytorch_fid/inception.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\n\ntry:\n    from torchvision.models.utils import load_state_dict_from_url\nexcept ImportError:\n    from torch.utils.model_zoo import load_url as load_state_dict_from_url\n\n# Inception weights ported to Pytorch from\n# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz\nFID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'  # noqa: E501\n\n\nclass InceptionV3(nn.Module):\n    \"\"\"Pretrained InceptionV3 network returning feature maps\"\"\"\n\n    # Index of default block of inception to return,\n    # corresponds to output of final average pooling\n    DEFAULT_BLOCK_INDEX = 3\n\n    # Maps feature dimensionality to their output blocks indices\n    BLOCK_INDEX_BY_DIM = {\n        64: 0,   # First max pooling features\n        192: 1,  # Second max pooling featurs\n        768: 2,  # Pre-aux classifier features\n        2048: 3  # Final average pooling features\n    }\n\n    def __init__(self,\n                 output_blocks=(DEFAULT_BLOCK_INDEX,),\n                 resize_input=True,\n                 normalize_input=True,\n                 requires_grad=False,\n                 use_fid_inception=True):\n        \"\"\"Build pretrained InceptionV3\n\n        Parameters\n        ----------\n        output_blocks : list of int\n            Indices of blocks to return features of. Possible values are:\n                - 0: corresponds to output of first max pooling\n                - 1: corresponds to output of second max pooling\n                - 2: corresponds to output which is fed to aux classifier\n                - 3: corresponds to output of final average pooling\n        resize_input : bool\n            If true, bilinearly resizes input to width and height 299 before\n            feeding input to model. As the network without fully connected\n            layers is fully convolutional, it should be able to handle inputs\n            of arbitrary size, so resizing might not be strictly needed\n        normalize_input : bool\n            If true, scales the input from range (0, 1) to the range the\n            pretrained Inception network expects, namely (-1, 1)\n        requires_grad : bool\n            If true, parameters of the model require gradients. Possibly useful\n            for finetuning the network\n        use_fid_inception : bool\n            If true, uses the pretrained Inception model used in Tensorflow's\n            FID implementation. If false, uses the pretrained Inception model\n            available in torchvision. The FID Inception model has different\n            weights and a slightly different structure from torchvision's\n            Inception model. If you want to compute FID scores, you are\n            strongly advised to set this parameter to true to get comparable\n            results.\n        \"\"\"\n        super(InceptionV3, self).__init__()\n\n        self.resize_input = resize_input\n        self.normalize_input = normalize_input\n        self.output_blocks = sorted(output_blocks)\n        self.last_needed_block = max(output_blocks)\n\n        assert self.last_needed_block <= 3, \\\n            'Last possible output block index is 3'\n\n        self.blocks = nn.ModuleList()\n\n        if use_fid_inception:\n            inception = fid_inception_v3()\n        else:\n            inception = _inception_v3(pretrained=True)\n\n        # Block 0: input to maxpool1\n        block0 = [\n            inception.Conv2d_1a_3x3,\n            inception.Conv2d_2a_3x3,\n            inception.Conv2d_2b_3x3,\n            nn.MaxPool2d(kernel_size=3, stride=2)\n        ]\n        self.blocks.append(nn.Sequential(*block0))\n\n        # Block 1: maxpool1 to maxpool2\n        if self.last_needed_block >= 1:\n            block1 = [\n                inception.Conv2d_3b_1x1,\n                inception.Conv2d_4a_3x3,\n                nn.MaxPool2d(kernel_size=3, stride=2)\n            ]\n            self.blocks.append(nn.Sequential(*block1))\n\n        # Block 2: maxpool2 to aux classifier\n        if self.last_needed_block >= 2:\n            block2 = [\n                inception.Mixed_5b,\n                inception.Mixed_5c,\n                inception.Mixed_5d,\n                inception.Mixed_6a,\n                inception.Mixed_6b,\n                inception.Mixed_6c,\n                inception.Mixed_6d,\n                inception.Mixed_6e,\n            ]\n            self.blocks.append(nn.Sequential(*block2))\n\n        # Block 3: aux classifier to final avgpool\n        if self.last_needed_block >= 3:\n            block3 = [\n                inception.Mixed_7a,\n                inception.Mixed_7b,\n                inception.Mixed_7c,\n                nn.AdaptiveAvgPool2d(output_size=(1, 1))\n            ]\n            self.blocks.append(nn.Sequential(*block3))\n\n        for param in self.parameters():\n            param.requires_grad = requires_grad\n\n    def forward(self, inp):\n        \"\"\"Get Inception feature maps\n\n        Parameters\n        ----------\n        inp : torch.autograd.Variable\n            Input tensor of shape Bx3xHxW. Values are expected to be in\n            range (0, 1)\n\n        Returns\n        -------\n        List of torch.autograd.Variable, corresponding to the selected output\n        block, sorted ascending by index\n        \"\"\"\n        outp = []\n        x = inp\n\n        if self.resize_input:\n            x = F.interpolate(x,\n                              size=(299, 299),\n                              mode='bilinear',\n                              align_corners=False)\n\n        if self.normalize_input:\n            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)\n\n        for idx, block in enumerate(self.blocks):\n            x = block(x)\n            if idx in self.output_blocks:\n                outp.append(x)\n\n            if idx == self.last_needed_block:\n                break\n\n        return outp\n\n\ndef _inception_v3(*args, **kwargs):\n    \"\"\"Wraps `torchvision.models.inception_v3`\n\n    Skips default weight inititialization if supported by torchvision version.\n    See https://github.com/mseitzer/pytorch-fid/issues/28.\n    \"\"\"\n    try:\n        version = tuple(map(int, torchvision.__version__.split('.')[:2]))\n    except ValueError:\n        # Just a caution against weird version strings\n        version = (0,)\n\n    if version >= (0, 6):\n        kwargs['init_weights'] = False\n\n    return torchvision.models.inception_v3(*args, **kwargs)\n\n\ndef fid_inception_v3():\n    \"\"\"Build pretrained Inception model for FID computation\n\n    The Inception model for FID computation uses a different set of weights\n    and has a slightly different structure than torchvision's Inception.\n\n    This method first constructs torchvision's Inception and then patches the\n    necessary parts that are different in the FID Inception model.\n    \"\"\"\n    inception = _inception_v3(num_classes=1008,\n                              aux_logits=False,\n                              pretrained=False)\n    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)\n    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)\n    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)\n    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)\n    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)\n    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)\n    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)\n    inception.Mixed_7b = FIDInceptionE_1(1280)\n    inception.Mixed_7c = FIDInceptionE_2(2048)\n\n    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)\n    inception.load_state_dict(state_dict)\n    return inception\n\n\nclass FIDInceptionA(torchvision.models.inception.InceptionA):\n    \"\"\"InceptionA block patched for FID computation\"\"\"\n    def __init__(self, in_channels, pool_features):\n        super(FIDInceptionA, self).__init__(in_channels, pool_features)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch5x5 = self.branch5x5_1(x)\n        branch5x5 = self.branch5x5_2(branch5x5)\n\n        branch3x3dbl = self.branch3x3dbl_1(x)\n        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)\n        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)\n\n        # Patch: Tensorflow's average pool does not use the padded zero's in\n        # its average calculation\n        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,\n                                   count_include_pad=False)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]\n        return torch.cat(outputs, 1)\n\n\nclass FIDInceptionC(torchvision.models.inception.InceptionC):\n    \"\"\"InceptionC block patched for FID computation\"\"\"\n    def __init__(self, in_channels, channels_7x7):\n        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch7x7 = self.branch7x7_1(x)\n        branch7x7 = self.branch7x7_2(branch7x7)\n        branch7x7 = self.branch7x7_3(branch7x7)\n\n        branch7x7dbl = self.branch7x7dbl_1(x)\n        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)\n        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)\n        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)\n        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)\n\n        # Patch: Tensorflow's average pool does not use the padded zero's in\n        # its average calculation\n        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,\n                                   count_include_pad=False)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]\n        return torch.cat(outputs, 1)\n\n\nclass FIDInceptionE_1(torchvision.models.inception.InceptionE):\n    \"\"\"First InceptionE block patched for FID computation\"\"\"\n    def __init__(self, in_channels):\n        super(FIDInceptionE_1, self).__init__(in_channels)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch3x3 = self.branch3x3_1(x)\n        branch3x3 = [\n            self.branch3x3_2a(branch3x3),\n            self.branch3x3_2b(branch3x3),\n        ]\n        branch3x3 = torch.cat(branch3x3, 1)\n\n        branch3x3dbl = self.branch3x3dbl_1(x)\n        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)\n        branch3x3dbl = [\n            self.branch3x3dbl_3a(branch3x3dbl),\n            self.branch3x3dbl_3b(branch3x3dbl),\n        ]\n        branch3x3dbl = torch.cat(branch3x3dbl, 1)\n\n        # Patch: Tensorflow's average pool does not use the padded zero's in\n        # its average calculation\n        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,\n                                   count_include_pad=False)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]\n        return torch.cat(outputs, 1)\n\n\nclass FIDInceptionE_2(torchvision.models.inception.InceptionE):\n    \"\"\"Second InceptionE block patched for FID computation\"\"\"\n    def __init__(self, in_channels):\n        super(FIDInceptionE_2, self).__init__(in_channels)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch3x3 = self.branch3x3_1(x)\n        branch3x3 = [\n            self.branch3x3_2a(branch3x3),\n            self.branch3x3_2b(branch3x3),\n        ]\n        branch3x3 = torch.cat(branch3x3, 1)\n\n        branch3x3dbl = self.branch3x3dbl_1(x)\n        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)\n        branch3x3dbl = [\n            self.branch3x3dbl_3a(branch3x3dbl),\n            self.branch3x3dbl_3b(branch3x3dbl),\n        ]\n        branch3x3dbl = torch.cat(branch3x3dbl, 1)\n\n        # Patch: The FID Inception model uses max pooling instead of average\n        # pooling. This is likely an error in this specific Inception\n        # implementation, as other Inception models use average pooling here\n        # (which matches the description in the paper).\n        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]\n        return torch.cat(outputs, 1)\n"
  },
  {
    "path": "train.sh",
    "content": "export OMP_NUM_THREADS=1\nexport CUDA_VISIBLE_DEVICES=0\n\npython main.py \\\n--name test \\\n--dataset lsun-bedroom \\\n--use_ema \\\n--ema_decay 0.9999 \\\n--noise_schedule linear \\\n--f_type quartic \\\n--gpu 0 \\\n--bsize 16 \\\n--sig 0.4 \\\n--sig_min 0 \\\n--sig_max 0.15 \\\n--fid_bsize 16 \\\n--loss_type eps_simple \\\n--lr 0.00005 \\"
  },
  {
    "path": "utils.py",
    "content": "from torchvision import transforms\nfrom PIL import Image\nimport os\nfrom torch import nn\nimport torch\nimport glob\nimport numpy as np\nimport torch.nn.functional as F\ndef tensor2pil(t):\n    img = transforms.functional.to_pil_image(denorm(t))\n    return img\ndef denorm(t):\n    return t*0.5 + 0.5\ndef tensor_imsave(t, path, fname, denormalization=False, prt=False):\n    # Save tensor as .png file\n    check_folder(path)\n    if denormalization: t = denorm(t)\n    t = torch.clamp(input=t, min=0., max=1.)\n    img = transforms.functional.to_pil_image(t.detach().cpu())\n    img.save(os.path.join(path,fname))\n    if prt:\n        print(f\"Saved to {os.path.join(path,fname)}\")\ndef pil_loader(path):\n    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\n    with open(path, 'rb') as f:##Automatic closing using with. rb = read binary\n        img = Image.open(f)\n        return img.convert('RGB')\ndef check_folder(log_dir):\n    if not os.path.exists(log_dir):\n        os.makedirs(log_dir)\n    return log_dir\ndef freeze_model(model):\n    for p in model.parameters():\n        p.requires_grad = False\nclass TVLoss(nn.Module):\n    def __init__(self,TVLoss_weight=1):\n        super(TVLoss,self).__init__()\n        self.TVLoss_weight = TVLoss_weight\n\n    def forward(self,x):\n        batch_size = x.size()[0]\n        h_x = x.size()[2]\n        w_x = x.size()[3]\n        count_h = self._tensor_size(x[:,:,1:,:])\n        count_w = self._tensor_size(x[:,:,:,1:])\n        h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum()\n        w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum()\n        return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size\n\n    def _tensor_size(self,t):\n        return t.size()[1]*t.size()[2]*t.size()[3]\n\ndef pil_batch_loader(root_path, bsize):\n    # Returns the list of PIL.Images at given path. Length of list is bsize.\n    flist = sorted(glob.glob(root_path, '*.png'))\n    flist = flist[:bsize]\n    result = []\n    for p in flist:\n        img = pil_loader(p)\n        result.append(img)\n    return result\n\nclass PSNR(object):\n    def __init__(self, gpu, val_max=1, val_min=0, ycbcr=True):\n        super(PSNR,self).__init__()\n        self.val_max = val_max\n        self.val_min = val_min\n        self.gpu = gpu\n        self.ycbcr = ycbcr\n    \n    def __call__(self,x,y):\n        \"\"\"\n        if x.is_cuda:\n            x = x.detach().cpu()\n        if y.is_cuda:\n            y = y.detach().cpu()\n        \"\"\"\n        assert len(x.size()) == len(y.size())\n        with torch.no_grad():\n            x_lum = rgb_to_ycbcr(x)[:,0]\n            y_lum = rgb_to_ycbcr(y)[:,0]\n            # if len(x.size()) == 3:\n            #     mse = torch.mean((y-x)**2)\n            #     psnr = 20*torch.log10(torch.tensor(self.val_max-self.val_min, dtype=torch.float).cuda(self.gpu)) - 10*torch.log10(mse)\n            #     return psnr\n            # elif len(x.size()) == 4:\n        \n            mse = torch.mean((y_lum-x_lum)**2, dim=[1,2])\n            psnr = 20*torch.log10(torch.tensor(self.val_max-self.val_min, dtype=torch.float).cuda(self.gpu)) - 10*torch.log10(mse)\n            return torch.mean(psnr)\n\ndef rgb_to_ycbcr(image):\n    r\"\"\"Convert an RGB image to YCbCr.\n    Args:\n        image (torch.Tensor): RGB Image to be converted to YCbCr.\n    Returns:\n        torch.Tensor: YCbCr version of the image.\n    \"\"\"\n\n    if not torch.is_tensor(image):\n        raise TypeError(\"Input type is not a torch.Tensor. Got {}\".format(\n            type(image)))\n\n    if len(image.shape) < 3 or image.shape[-3] != 3:\n        raise ValueError(\"Input size must have a shape of (*, 3, H, W). Got {}\"\n                         .format(image.shape))\n\n    r: torch.Tensor = image[..., 0, :, :]\n    g: torch.Tensor = image[..., 1, :, :]\n    b: torch.Tensor = image[..., 2, :, :]\n\n    delta = .5\n    y: torch.Tensor = .299 * r + .587 * g + .114 * b\n    cb: torch.Tensor = (b - y) * .564 + delta\n    cr: torch.Tensor = (r - y) * .713 + delta\n    return torch.stack((y, cb, cr), -3)\n\ndef clear(x):\n    return x.detach().cpu().squeeze().numpy()\n\ndef normalize_np(x):\n    x -= x.min()\n    x /= x.max()\n    return x\n\ndef pad(x, pad):\n    return F.pad(x, (pad, pad, pad, pad), 'reflect')\n\ndef crop(x, pad):\n    if pad == 0:\n        return x\n    else:\n        return x[..., pad:-pad, pad:-pad]\n"
  }
]