initial code release

.gitignore

@@ -0,0 +1,46 @@
+# Node artifact files
+node_modules/
+dist/
+ngc_scripts/
+
+# Compiled Java class files
+*.class
+
+# Compiled Python bytecode
+*.py[cod]
+
+# Log files
+*.log
+
+# Package files
+*.jar
+
+# Maven
+target/
+dist/
+
+# JetBrains IDE
+.idea/
+
+# Unit test reports
+TEST*.xml
+
+# Generated by MacOS
+.DS_Store
+
+# Generated by Windows
+Thumbs.db
+
+# Applications
+*.app
+*.exe
+*.war
+
+# Large media files
+*.mp4
+*.tiff
+*.avi
+*.flv
+*.mov
+*.wmv
+

EMA.py

@@ -0,0 +1,90 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
+# ---------------------------------------------------------------
+
+'''
+Codes adapted from https://github.com/NVlabs/LSGM/blob/main/util/ema.py
+'''
+import warnings
+
+import torch
+from torch.optim import Optimizer
+
+
+class EMA(Optimizer):
+    def __init__(self, opt, ema_decay):
+        self.ema_decay = ema_decay
+        self.apply_ema = self.ema_decay > 0.
+        self.optimizer = opt
+        self.state = opt.state
+        self.param_groups = opt.param_groups
+
+    def step(self, *args, **kwargs):
+        retval = self.optimizer.step(*args, **kwargs)
+
+        # stop here if we are not applying EMA
+        if not self.apply_ema:
+            return retval
+
+        ema, params = {}, {}
+        for group in self.optimizer.param_groups:
+            for i, p in enumerate(group['params']):
+                if p.grad is None:
+                    continue
+                state = self.optimizer.state[p]
+
+                # State initialization
+                if 'ema' not in state:
+                    state['ema'] = p.data.clone()
+
+                if p.shape not in params:
+                    params[p.shape] = {'idx': 0, 'data': []}
+                    ema[p.shape] = []
+
+                params[p.shape]['data'].append(p.data)
+                ema[p.shape].append(state['ema'])
+
+            for i in params:
+                params[i]['data'] = torch.stack(params[i]['data'], dim=0)
+                ema[i] = torch.stack(ema[i], dim=0)
+                ema[i].mul_(self.ema_decay).add_(params[i]['data'], alpha=1. - self.ema_decay)
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                idx = params[p.shape]['idx']
+                self.optimizer.state[p]['ema'] = ema[p.shape][idx, :]
+                params[p.shape]['idx'] += 1
+
+        return retval
+
+    def load_state_dict(self, state_dict):
+        super(EMA, self).load_state_dict(state_dict)
+        # load_state_dict loads the data to self.state and self.param_groups. We need to pass this data to
+        # the underlying optimizer too.
+        self.optimizer.state = self.state
+        self.optimizer.param_groups = self.param_groups
+
+    def swap_parameters_with_ema(self, store_params_in_ema):
+        """ This function swaps parameters with their ema values. It records original parameters in the ema
+        parameters, if store_params_in_ema is true."""
+
+        # stop here if we are not applying EMA
+        if not self.apply_ema:
+            warnings.warn('swap_parameters_with_ema was called when there is no EMA weights.')
+            return
+
+        for group in self.optimizer.param_groups:
+            for i, p in enumerate(group['params']):
+                if not p.requires_grad:
+                    continue
+                ema = self.optimizer.state[p]['ema']
+                if store_params_in_ema:
+                    tmp = p.data.detach()
+                    p.data = ema.detach()
+                    self.optimizer.state[p]['ema'] = tmp
+                else:
+                    p.data = ema.detach()

LICENSE

@@ -0,0 +1,63 @@
+NVIDIA License
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under this license, which may include software,
+documentation, or other files, and (b) any additions to or derivative works thereof that are made available under this license.
+
+The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the meaning as provided under U.S.
+copyright law; provided, however, that for the purposes of this license, derivative works shall not include works that
+remain separable from, or merely link (or bind by name) to the interfaces of, the Work.
+
+Works are “made available” under this license by including in or with the Work either (a) a copyright notice
+referencing the applicability of this license to the Work, or (b) a copy of this license.
+
+2. License Grant
+
+2.1 Copyright Grant. Subject to the terms and conditions of this license, each Licensor grants to you a perpetual,
+worldwide, non-exclusive, royalty-free, copyright license to use, reproduce, prepare derivative works of, publicly
+display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form.
+
+3. Limitations
+
+3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under this license, (b) you
+include a complete copy of this license with your distribution, and (c) you retain without modification any
+copyright, patent, trademark, or attribution notices that are present in the Work.
+
+3.2 Derivative Works. You may specify that additional or different terms apply to the use, reproduction, and
+distribution of your derivative works of the Work (“Your Terms”) only if (a) Your Terms provide that the use
+limitation in Section 3.3 applies to your derivative works, and (b) you identify the specific derivative works
+that are subject to Your Terms. Notwithstanding Your Terms, this license (including the redistribution
+requirements in Section 3.1) will continue to apply to the Work itself.
+
+3.3 Use Limitation. The Work and any derivative works thereof only may be used or intended for use non-commercially.
+Notwithstanding the foregoing, NVIDIA Corporation and its affiliates may use the Work and any derivative
+works commercially. As used herein, “non-commercially” means for research or evaluation purposes only.
+
+3.4 Patent Claims. If you bring or threaten to bring a patent claim against any Licensor (including any claim,
+cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then
+your rights under this license from such Licensor (including the grant in Section 2.1) will terminate immediately.
+
+3.5 Trademarks. This license does not grant any rights to use any Licensor’s or its affiliates’ names, logos,
+or trademarks, except as necessary to reproduce the notices described in this license.
+
+3.6 Termination. If you violate any term of this license, then your rights under this license (including the
+grant in Section 2.1) will terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
+WARRANTIES OR CONDITIONS OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU
+BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING
+NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR
+INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR
+DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN
+ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+

README.md

@@ -1,16 +0,0 @@
-## <p align="center">Tackling the Generative Learning Trilemma with Denoising Diffusion GANs</p>
-
-<div align="center">
-  <a href="https://xavierxiao.github.io/" target="_blank">Zhisheng&nbsp;Xiao</a> &emsp; <b>&middot;</b> &emsp;
-  <a href="https://karstenkreis.github.io/" target="_blank">Karsten&nbsp;Kreis</a> &emsp; <b>&middot;</b> &emsp;
-  <a href="http://latentspace.cc/" target="_blank">Arash&nbsp;Vahdat</a>
-  <br> <br>
-  <a href="https://nvlabs.github.io/denoising-diffusion-gan" target="_blank">Project&nbsp;Page</a>
-</div>
-<br><br>
-<p align="center">:construction: :pick: :hammer_and_wrench: :construction_worker:</p>
-<p align="center">Code coming soon (the current expected release is in March 2022). Stay tuned!</p>
-<br><br>
-<p align="center">
-    <img width="800" alt="teaser" src="assets/teaser.png"/>
-</p>

datasets_prep/LICENSE_PyTorch

@@ -0,0 +1,70 @@
+From PyTorch:
+
+Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+
+From Caffe2:
+
+Copyright (c) 2016-present, Facebook Inc. All rights reserved.
+
+All contributions by Facebook:
+Copyright (c) 2016 Facebook Inc.
+
+All contributions by Google:
+Copyright (c) 2015 Google Inc.
+All rights reserved.
+
+All contributions by Yangqing Jia:
+Copyright (c) 2015 Yangqing Jia
+All rights reserved.
+
+All contributions from Caffe:
+Copyright(c) 2013, 2014, 2015, the respective contributors
+All rights reserved.
+
+All other contributions:
+Copyright(c) 2015, 2016 the respective contributors
+All rights reserved.
+
+Caffe2 uses a copyright model similar to Caffe: each contributor holds
+copyright over their contributions to Caffe2. The project versioning records
+all such contribution and copyright details. If a contributor wants to further
+mark their specific copyright on a particular contribution, they should
+indicate their copyright solely in the commit message of the change when it is
+committed.
+
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+
+3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
+   and IDIAP Research Institute nor the names of its contributors may be
+   used to endorse or promote products derived from this software without
+   specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.

datasets_prep/LICENSE_torchvision

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) Soumith Chintala 2016,
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

datasets_prep/lmdb_datasets.py

@@ -0,0 +1,58 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
+# ---------------------------------------------------------------
+
+import torch.utils.data as data
+import numpy as np
+import lmdb
+import os
+import io
+from PIL import Image
+
+
+def num_samples(dataset, train):
+    if dataset == 'celeba':
+        return 27000 if train else 3000
+    
+    else:
+        raise NotImplementedError('dataset %s is unknown' % dataset)
+
+
+class LMDBDataset(data.Dataset):
+    def __init__(self, root, name='', train=True, transform=None, is_encoded=False):
+        self.train = train
+        self.name = name
+        self.transform = transform
+        if self.train:
+            lmdb_path = os.path.join(root, 'train.lmdb')
+        else:
+            lmdb_path = os.path.join(root, 'validation.lmdb')
+        self.data_lmdb = lmdb.open(lmdb_path, readonly=True, max_readers=1,
+                                   lock=False, readahead=False, meminit=False)
+        self.is_encoded = is_encoded
+
+    def __getitem__(self, index):
+        target = [0]
+        with self.data_lmdb.begin(write=False, buffers=True) as txn:
+            data = txn.get(str(index).encode())
+            if self.is_encoded:
+                img = Image.open(io.BytesIO(data))
+                img = img.convert('RGB')
+            else:
+                img = np.asarray(data, dtype=np.uint8)
+                # assume data is RGB
+                size = int(np.sqrt(len(img) / 3))
+                img = np.reshape(img, (size, size, 3))
+                img = Image.fromarray(img, mode='RGB')
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        return img, target
+
+    def __len__(self):
+        return num_samples(self.name, self.train)
+

datasets_prep/lsun.py

@@ -0,0 +1,170 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from a file in the torchvision library
+# which was released under the BSD 3-Clause License.
+#
+# Source:
+# https://github.com/pytorch/vision/blob/ea6b879e90459006e71a164dc76b7e2cc3bff9d9/torchvision/datasets/lsun.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_torchvision). The modifications
+# to this file are subject to the same BSD 3-Clause License.
+# ---------------------------------------------------------------
+
+from torchvision.datasets.vision import VisionDataset
+from PIL import Image
+import os
+import os.path
+import io
+import string
+from collections.abc import Iterable
+import pickle
+from torchvision.datasets.utils import verify_str_arg, iterable_to_str
+
+
+class LSUNClass(VisionDataset):
+    def __init__(self, root, transform=None, target_transform=None):
+        import lmdb
+        super(LSUNClass, self).__init__(root, transform=transform,
+                                        target_transform=target_transform)
+
+        self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False,
+                             readahead=False, meminit=False)
+        with self.env.begin(write=False) as txn:
+            self.length = txn.stat()['entries']
+        # cache_file = '_cache_' + ''.join(c for c in root if c in string.ascii_letters)
+        # av begin
+        # We only modified the location of cache_file.
+        cache_file = os.path.join(self.root, '_cache_')
+        # av end
+        if os.path.isfile(cache_file):
+            self.keys = pickle.load(open(cache_file, "rb"))
+        else:
+            with self.env.begin(write=False) as txn:
+                self.keys = [key for key, _ in txn.cursor()]
+            pickle.dump(self.keys, open(cache_file, "wb"))
+
+    def __getitem__(self, index):
+        img, target = None, -1
+        env = self.env
+        with env.begin(write=False) as txn:
+            imgbuf = txn.get(self.keys[index])
+
+        buf = io.BytesIO()
+        buf.write(imgbuf)
+        buf.seek(0)
+        img = Image.open(buf).convert('RGB')
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+    def __len__(self):
+        return self.length
+
+
+class LSUN(VisionDataset):
+    """
+    `LSUN <https://www.yf.io/p/lsun>`_ dataset.
+
+    Args:
+        root (string): Root directory for the database files.
+        classes (string or list): One of {'train', 'val', 'test'} or a list of
+            categories to load. e,g. ['bedroom_train', 'church_outdoor_train'].
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+    """
+
+    def __init__(self, root, classes='train', transform=None, target_transform=None):
+        super(LSUN, self).__init__(root, transform=transform,
+                                   target_transform=target_transform)
+        self.classes = self._verify_classes(classes)
+
+        # for each class, create an LSUNClassDataset
+        self.dbs = []
+        for c in self.classes:
+            self.dbs.append(LSUNClass(
+                root=root + '/' + c + '_lmdb',
+                transform=transform))
+
+        self.indices = []
+        count = 0
+        for db in self.dbs:
+            count += len(db)
+            self.indices.append(count)
+
+        self.length = count
+
+    def _verify_classes(self, classes):
+        categories = ['bedroom', 'bridge', 'church_outdoor', 'classroom',
+                      'conference_room', 'dining_room', 'kitchen',
+                      'living_room', 'restaurant', 'tower', 'cat']
+        dset_opts = ['train', 'val', 'test']
+
+        try:
+            verify_str_arg(classes, "classes", dset_opts)
+            if classes == 'test':
+                classes = [classes]
+            else:
+                classes = [c + '_' + classes for c in categories]
+        except ValueError:
+            if not isinstance(classes, Iterable):
+                msg = ("Expected type str or Iterable for argument classes, "
+                       "but got type {}.")
+                raise ValueError(msg.format(type(classes)))
+
+            classes = list(classes)
+            msg_fmtstr = ("Expected type str for elements in argument classes, "
+                          "but got type {}.")
+            for c in classes:
+                verify_str_arg(c, custom_msg=msg_fmtstr.format(type(c)))
+                c_short = c.split('_')
+                category, dset_opt = '_'.join(c_short[:-1]), c_short[-1]
+
+                msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}."
+                msg = msg_fmtstr.format(category, "LSUN class",
+                                        iterable_to_str(categories))
+                verify_str_arg(category, valid_values=categories, custom_msg=msg)
+
+                msg = msg_fmtstr.format(dset_opt, "postfix", iterable_to_str(dset_opts))
+                verify_str_arg(dset_opt, valid_values=dset_opts, custom_msg=msg)
+
+        return classes
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+
+        Returns:
+            tuple: Tuple (image, target) where target is the index of the target category.
+        """
+        target = 0
+        sub = 0
+        for ind in self.indices:
+            if index < ind:
+                break
+            target += 1
+            sub = ind
+
+        db = self.dbs[target]
+        index = index - sub
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        img, _ = db[index]
+        return img, target
+
+    def __len__(self):
+        return self.length
+
+    def extra_repr(self):
+        return "Classes: {classes}".format(**self.__dict__)

datasets_prep/stackmnist_data.py

@@ -0,0 +1,65 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
+# ---------------------------------------------------------------
+
+
+import numpy as np
+from PIL import Image
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+
+
+class StackedMNIST(dset.MNIST):
+    def __init__(self, root, train=True, transform=None, target_transform=None,
+                 download=False):
+        super(StackedMNIST, self).__init__(root=root, train=train, transform=transform,
+                                           target_transform=target_transform, download=download)
+
+        index1 = np.hstack([np.random.permutation(len(self.data)), np.random.permutation(len(self.data))])
+        index2 = np.hstack([np.random.permutation(len(self.data)), np.random.permutation(len(self.data))])
+        index3 = np.hstack([np.random.permutation(len(self.data)), np.random.permutation(len(self.data))])
+        self.num_images = 2 * len(self.data)
+
+        self.index = []
+        for i in range(self.num_images):
+            self.index.append((index1[i], index2[i], index3[i]))
+
+    def __len__(self):
+        return self.num_images
+
+    def __getitem__(self, index):
+        img = np.zeros((28, 28, 3), dtype=np.uint8)
+        target = 0
+        for i in range(3):
+            img_, target_ = self.data[self.index[index][i]], int(self.targets[self.index[index][i]])
+            img[:, :, i] = img_
+            target += target_ * 10 ** (2 - i)
+
+        img = Image.fromarray(img, mode="RGB")
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+def _data_transforms_stacked_mnist():
+    """Get data transforms for cifar10."""
+    train_transform = transforms.Compose([
+        transforms.Pad(padding=2),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+    ])
+
+    valid_transform = transforms.Compose([
+        transforms.Pad(padding=2),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+    ])
+
+    return train_transform, valid_transform

pytorch_fid/LICENSE_MIT

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Zhifeng Kong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pytorch_fid/LICENSE_inception

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

pytorch_fid/LICENSE_pytorch_fid

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

pytorch_fid/fid_score.py

@@ -0,0 +1,305 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from FAST_DPM.
+#
+# Source:
+# https://github.com/FengNiMa/FastDPM_pytorch/blob/6540c1cdac3799aff8a5f7b9de430269bbd0b7c3/pytorch_fid/fid_score.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_MIT).
+# The modifications to this file are subject to the same license.
+# ---------------------------------------------------------------
+
+"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
+
+The FID metric calculates the distance between two distributions of images.
+Typically, we have summary statistics (mean & covariance matrix) of one
+of these distributions, while the 2nd distribution is given by a GAN.
+
+When run as a stand-alone program, it compares the distribution of
+images that are stored as PNG/JPEG at a specified location with a
+distribution given by summary statistics (in pickle format).
+
+The FID is calculated by assuming that X_1 and X_2 are the activations of
+the pool_3 layer of the inception net for generated samples and real world
+samples respectively.
+
+See --help to see further details.
+
+Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
+of Tensorflow
+
+Copyright 2018 Institute of Bioinformatics, JKU Linz
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+import pathlib
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+from multiprocessing import cpu_count
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as TF
+from PIL import Image
+from scipy import linalg
+from torch.nn.functional import adaptive_avg_pool2d
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+try:
+    from inception import InceptionV3
+except ImportError:
+    from .inception import InceptionV3
+
+parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+parser.add_argument('--batch-size', type=int, default=50,
+                    help='Batch size to use')
+parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+parser.add_argument('--dims', type=int, default=2048,
+                    choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
+                    help=('Dimensionality of Inception features to use. '
+                          'By default, uses pool3 features'))
+parser.add_argument('path', type=str, nargs=2,
+                    help=('Paths to the generated images or '
+                          'to .npz statistic files'))
+
+IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
+                    'tif', 'tiff', 'webp'}
+
+
+
+
+class ImagePathDataset(torch.utils.data.Dataset):
+    def __init__(self, files, transforms=None):
+        self.files = files
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, i):
+        path = self.files[i]
+        img = Image.open(path).convert('RGB')
+        if self.transforms is not None:
+            img = self.transforms(img)
+        return img
+
+
+def get_activations(files, model, batch_size=50, dims=2048, device='cpu', resize=0):
+    """Calculates the activations of the pool_3 layer for all images.
+
+    Params:
+    -- files       : List of image files paths
+    -- model       : Instance of inception model
+    -- batch_size  : Batch size of images for the model to process at once.
+                     Make sure that the number of samples is a multiple of
+                     the batch size, otherwise some samples are ignored. This
+                     behavior is retained to match the original FID score
+                     implementation.
+    -- dims        : Dimensionality of features returned by Inception
+    -- device      : Device to run calculations
+
+    Returns:
+    -- A numpy array of dimension (num images, dims) that contains the
+       activations of the given tensor when feeding inception with the
+       query tensor.
+    """
+    model.eval()
+
+    if batch_size > len(files):
+        print(('Warning: batch size is bigger than the data size. '
+               'Setting batch size to data size'))
+        batch_size = len(files)
+
+    if resize > 0:
+        print('Resized to ({}, {})'.format(resize, resize))
+        dataset = ImagePathDataset(files, transforms=TF.Compose([TF.Resize(size=(resize, resize)), 
+                                                                 TF.ToTensor()]))
+    else:
+        dataset = ImagePathDataset(files, transforms=TF.ToTensor())
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=batch_size,
+                                             shuffle=False,
+                                             drop_last=False,
+                                             num_workers=cpu_count())
+
+    pred_arr = np.empty((len(files), dims))
+
+    start_idx = 0
+
+    for batch in tqdm(dataloader):
+        batch = batch.to(device)
+
+        with torch.no_grad():
+            pred = model(batch)[0]
+
+        # If model output is not scalar, apply global spatial average pooling.
+        # This happens if you choose a dimensionality not equal 2048.
+        if pred.size(2) != 1 or pred.size(3) != 1:
+            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+
+        pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+
+        pred_arr[start_idx:start_idx + pred.shape[0]] = pred
+
+        start_idx = start_idx + pred.shape[0]
+
+    return pred_arr
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean = linalg.sqrtm(sigma1.dot(sigma2))
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1)
+            + np.trace(sigma2) - 2 * tr_covmean)
+
+
+def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
+                                    device='cpu', resize=0):
+    """Calculation of the statistics used by the FID.
+    Params:
+    -- files       : List of image files paths
+    -- model       : Instance of inception model
+    -- batch_size  : The images numpy array is split into batches with
+                     batch size batch_size. A reasonable batch size
+                     depends on the hardware.
+    -- dims        : Dimensionality of features returned by Inception
+    -- device      : Device to run calculations
+    -- resize      : resize image to this shape
+
+    Returns:
+    -- mu    : The mean over samples of the activations of the pool_3 layer of
+               the inception model.
+    -- sigma : The covariance matrix of the activations of the pool_3 layer of
+               the inception model.
+    """
+    act = get_activations(files, model, batch_size, dims, device, resize)
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+def compute_statistics_of_path(path, model, batch_size, dims, device, resize=0):
+    if path.endswith('.npz') or path.endswith('.npy'):
+        f = np.load(path, allow_pickle=True)
+        try:
+            m, s = f['mu'][:], f['sigma'][:]
+        except:
+            m, s = f.item()['mu'][:], f.item()['sigma'][:]
+    else:
+        path_str = path[:]
+        path = pathlib.Path(path)
+        files = sorted([file for ext in IMAGE_EXTENSIONS
+                       for file in path.glob('*.{}'.format(ext))])
+        m, s = calculate_activation_statistics(files, model, batch_size,
+                                               dims, device, resize)
+    return m, s
+
+
+def calculate_fid_given_paths(paths, batch_size, device, dims, resize=0):
+    """Calculates the FID of two paths"""
+    for p in paths:
+        if not os.path.exists(p):
+            raise RuntimeError('Invalid path: %s' % p)
+
+    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+
+    model = InceptionV3([block_idx]).to(device)
+
+    m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
+                                        dims, device, resize)
+    m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,
+                                        dims, device, resize)
+    
+    del model
+    fid_value = calculate_frechet_distance(m1, s1, m2, s2)
+    return fid_value
+
+
+
+def main():
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    fid_value = calculate_fid_given_paths(args.path,
+                                          args.batch_size,
+                                          device,
+                                          args.dims)
+    print('FID: ', fid_value)
+
+
+if __name__ == '__main__':
+    main()

pytorch_fid/inception.py

@@ -0,0 +1,337 @@
+# ---------------------------------------------------------------
+# Taken from the following link as is from:
+# https://github.com/mseitzer/pytorch-fid/blob/3d604a25516746c3a4a5548c8610e99010b2c819/src/pytorch_fid/inception.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_pytorch_fid).
+# ---------------------------------------------------------------
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# Inception weights ported to Pytorch from
+# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'  # noqa: E501
+
+
+class InceptionV3(nn.Module):
+    """Pretrained InceptionV3 network returning feature maps"""
+
+    # Index of default block of inception to return,
+    # corresponds to output of final average pooling
+    DEFAULT_BLOCK_INDEX = 3
+
+    # Maps feature dimensionality to their output blocks indices
+    BLOCK_INDEX_BY_DIM = {
+        64: 0,   # First max pooling features
+        192: 1,  # Second max pooling featurs
+        768: 2,  # Pre-aux classifier features
+        2048: 3  # Final average pooling features
+    }
+
+    def __init__(self,
+                 output_blocks=(DEFAULT_BLOCK_INDEX,),
+                 resize_input=True,
+                 normalize_input=True,
+                 requires_grad=False,
+                 use_fid_inception=True):
+        """Build pretrained InceptionV3
+
+        Parameters
+        ----------
+        output_blocks : list of int
+            Indices of blocks to return features of. Possible values are:
+                - 0: corresponds to output of first max pooling
+                - 1: corresponds to output of second max pooling
+                - 2: corresponds to output which is fed to aux classifier
+                - 3: corresponds to output of final average pooling
+        resize_input : bool
+            If true, bilinearly resizes input to width and height 299 before
+            feeding input to model. As the network without fully connected
+            layers is fully convolutional, it should be able to handle inputs
+            of arbitrary size, so resizing might not be strictly needed
+        normalize_input : bool
+            If true, scales the input from range (0, 1) to the range the
+            pretrained Inception network expects, namely (-1, 1)
+        requires_grad : bool
+            If true, parameters of the model require gradients. Possibly useful
+            for finetuning the network
+        use_fid_inception : bool
+            If true, uses the pretrained Inception model used in Tensorflow's
+            FID implementation. If false, uses the pretrained Inception model
+            available in torchvision. The FID Inception model has different
+            weights and a slightly different structure from torchvision's
+            Inception model. If you want to compute FID scores, you are
+            strongly advised to set this parameter to true to get comparable
+            results.
+        """
+        super(InceptionV3, self).__init__()
+
+        self.resize_input = resize_input
+        self.normalize_input = normalize_input
+        self.output_blocks = sorted(output_blocks)
+        self.last_needed_block = max(output_blocks)
+
+        assert self.last_needed_block <= 3, \
+            'Last possible output block index is 3'
+
+        self.blocks = nn.ModuleList()
+
+        if use_fid_inception:
+            inception = fid_inception_v3()
+        else:
+            inception = _inception_v3(pretrained=True)
+
+        # Block 0: input to maxpool1
+        block0 = [
+            inception.Conv2d_1a_3x3,
+            inception.Conv2d_2a_3x3,
+            inception.Conv2d_2b_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2)
+        ]
+        self.blocks.append(nn.Sequential(*block0))
+
+        # Block 1: maxpool1 to maxpool2
+        if self.last_needed_block >= 1:
+            block1 = [
+                inception.Conv2d_3b_1x1,
+                inception.Conv2d_4a_3x3,
+                nn.MaxPool2d(kernel_size=3, stride=2)
+            ]
+            self.blocks.append(nn.Sequential(*block1))
+
+        # Block 2: maxpool2 to aux classifier
+        if self.last_needed_block >= 2:
+            block2 = [
+                inception.Mixed_5b,
+                inception.Mixed_5c,
+                inception.Mixed_5d,
+                inception.Mixed_6a,
+                inception.Mixed_6b,
+                inception.Mixed_6c,
+                inception.Mixed_6d,
+                inception.Mixed_6e,
+            ]
+            self.blocks.append(nn.Sequential(*block2))
+
+        # Block 3: aux classifier to final avgpool
+        if self.last_needed_block >= 3:
+            block3 = [
+                inception.Mixed_7a,
+                inception.Mixed_7b,
+                inception.Mixed_7c,
+                nn.AdaptiveAvgPool2d(output_size=(1, 1))
+            ]
+            self.blocks.append(nn.Sequential(*block3))
+
+        for param in self.parameters():
+            param.requires_grad = requires_grad
+
+    def forward(self, inp):
+        """Get Inception feature maps
+
+        Parameters
+        ----------
+        inp : torch.autograd.Variable
+            Input tensor of shape Bx3xHxW. Values are expected to be in
+            range (0, 1)
+
+        Returns
+        -------
+        List of torch.autograd.Variable, corresponding to the selected output
+        block, sorted ascending by index
+        """
+        outp = []
+        x = inp
+
+        if self.resize_input:
+            x = F.interpolate(x,
+                              size=(299, 299),
+                              mode='bilinear',
+                              align_corners=False)
+
+        if self.normalize_input:
+            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x)
+            if idx in self.output_blocks:
+                outp.append(x)
+
+            if idx == self.last_needed_block:
+                break
+
+        return outp
+
+
+def _inception_v3(*args, **kwargs):
+    """Wraps `torchvision.models.inception_v3`
+
+    Skips default weight inititialization if supported by torchvision version.
+    See https://github.com/mseitzer/pytorch-fid/issues/28.
+    """
+    try:
+        version = tuple(map(int, torchvision.__version__.split('.')[:2]))
+    except ValueError:
+        # Just a caution against weird version strings
+        version = (0,)
+
+    if version >= (0, 6):
+        kwargs['init_weights'] = False
+
+    return torchvision.models.inception_v3(*args, **kwargs)
+
+
+def fid_inception_v3():
+    """Build pretrained Inception model for FID computation
+
+    The Inception model for FID computation uses a different set of weights
+    and has a slightly different structure than torchvision's Inception.
+
+    This method first constructs torchvision's Inception and then patches the
+    necessary parts that are different in the FID Inception model.
+    """
+    inception = _inception_v3(num_classes=1008,
+                              aux_logits=False,
+                              pretrained=False)
+    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
+    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
+    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
+    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
+    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
+    inception.Mixed_7b = FIDInceptionE_1(1280)
+    inception.Mixed_7c = FIDInceptionE_2(2048)
+
+    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
+    inception.load_state_dict(state_dict)
+    return inception
+
+
+class FIDInceptionA(torchvision.models.inception.InceptionA):
+    """InceptionA block patched for FID computation"""
+    def __init__(self, in_channels, pool_features):
+        super(FIDInceptionA, self).__init__(in_channels, pool_features)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch5x5 = self.branch5x5_1(x)
+        branch5x5 = self.branch5x5_2(branch5x5)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionC(torchvision.models.inception.InceptionC):
+    """InceptionC block patched for FID computation"""
+    def __init__(self, in_channels, channels_7x7):
+        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch7x7 = self.branch7x7_1(x)
+        branch7x7 = self.branch7x7_2(branch7x7)
+        branch7x7 = self.branch7x7_3(branch7x7)
+
+        branch7x7dbl = self.branch7x7dbl_1(x)
+        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_1(torchvision.models.inception.InceptionE):
+    """First InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_1, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_2(torchvision.models.inception.InceptionE):
+    """Second InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_2, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: The FID Inception model uses max pooling instead of average
+        # pooling. This is likely an error in this specific Inception
+        # implementation, as other Inception models use average pooling here
+        # (which matches the description in the paper).
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)

pytorch_fid/inception_score.py

@@ -0,0 +1,103 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from inception score
+#
+# Source:
+# https://github.com/tsc2017/Inception-Score/blob/04390da2ebb3c9a3860337a33297c4f270bd906d/inception_score.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_inception).
+# The modifications to this file are subject to the same license.
+# ---------------------------------------------------------------
+
+
+'''
+Usage:
+    Call get_inception_score(images, splits=10)
+Args:
+    images: A numpy array with values ranging from 0 to 255 and shape in the form [N, 3, HEIGHT, WIDTH] where N, HEIGHT and WIDTH can be arbitrary. A dtype of np.uint8 is recommended to save CPU memory.
+    splits: The number of splits of the images, default is 10.
+Returns:
+    Mean and standard deviation of the Inception Score across the splits.
+'''
+import argparse
+
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+import tensorflow_gan as tfgan
+import os
+import functools
+import numpy as np
+import time
+from tensorflow.python.ops import array_ops
+# pip install tensorflow-gan
+import tensorflow_gan as tfgan
+session=tf.compat.v1.InteractiveSession()
+# A smaller BATCH_SIZE reduces GPU memory usage, but at the cost of a slight slowdown
+BATCH_SIZE = 64
+INCEPTION_TFHUB = 'https://tfhub.dev/tensorflow/tfgan/eval/inception/1'
+INCEPTION_OUTPUT = 'logits'
+
+# Run images through Inception.
+inception_images = tf.compat.v1.placeholder(tf.float32, [None, 3, None, None], name = 'inception_images')
+def inception_logits(images = inception_images, num_splits = 1):
+    images = tf.transpose(images, [0, 2, 3, 1])
+    size = 299
+    images = tf.compat.v1.image.resize_bilinear(images, [size, size])
+    generated_images_list = array_ops.split(images, num_or_size_splits = num_splits)
+    logits = tf.map_fn(
+        fn = tfgan.eval.classifier_fn_from_tfhub(INCEPTION_TFHUB, INCEPTION_OUTPUT, True),
+        elems = array_ops.stack(generated_images_list),
+        parallel_iterations = 8,
+        back_prop = False,
+        swap_memory = True,
+        name = 'RunClassifier')
+    logits = array_ops.concat(array_ops.unstack(logits), 0)
+    return logits
+
+logits=inception_logits()
+
+def get_inception_probs(inps):
+    session=tf.get_default_session()
+    n_batches = int(np.ceil(float(inps.shape[0]) / BATCH_SIZE))
+    preds = np.zeros([inps.shape[0], 1000], dtype = np.float32)
+    for i in range(n_batches):
+        inp = inps[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] / 255. * 2 - 1
+        preds[i * BATCH_SIZE : i * BATCH_SIZE + min(BATCH_SIZE, inp.shape[0])] = session.run(logits,{inception_images: inp})[:, :1000]
+    preds = np.exp(preds) / np.sum(np.exp(preds), 1, keepdims=True)
+    return preds
+
+def preds2score(preds, splits=10):
+    scores = []
+    for i in range(splits):
+        part = preds[(i * preds.shape[0] // splits):((i + 1) * preds.shape[0] // splits), :]
+        kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
+        kl = np.mean(np.sum(kl, 1))
+        scores.append(np.exp(kl))
+    return np.mean(scores), np.std(scores)
+
+def get_inception_score(images, splits=10):
+    assert(type(images) == np.ndarray)
+    assert(len(images.shape) == 4)
+    assert(images.shape[1] == 3)
+    assert(np.min(images[0]) >= 0 and np.max(images[0]) > 10), 'Image values should be in the range [0, 255]'
+    print('Calculating Inception Score with %i images in %i splits' % (images.shape[0], splits))
+    start_time=time.time()
+    preds = get_inception_probs(images)
+    mean, std = preds2score(preds, splits)
+    print('Inception Score calculation time: %f s' % (time.time() - start_time))
+    return mean, std  # Reference values: 11.38 for 50000 CIFAR-10 training set images, or mean=11.31, std=0.10 if in 10 splits.
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--sample_dir', default='./saved_samples/', help='path to saved images')
+    opt = parser.parse_args()
+
+    data = np.load(opt.sample_dir)
+    data = np.clip(data, 0, 255)
+    m, s = get_inception_score(data, splits=1)
+    
+    print('mean: ', m)
+    print('std: ', s)

readme.md

@@ -0,0 +1,113 @@
+# Official PyTorch implementation of "Tackling the Generative Learning Trilemma with Denoising Diffusion GANs" [(ICLR 2022 Spotlight Paper)](https://arxiv.org/abs/2112.07804) #
+
+<div align="center">
+  <a href="https://xavierxiao.github.io/" target="_blank">Zhisheng&nbsp;Xiao</a> &emsp; <b>&middot;</b> &emsp;
+  <a href="https://karstenkreis.github.io/" target="_blank">Karsten&nbsp;Kreis</a> &emsp; <b>&middot;</b> &emsp;
+  <a href="http://latentspace.cc/arash_vahdat/" target="_blank">Arash&nbsp;Vahdat</a>
+  <br> <br>
+  <a href="https://nvlabs.github.io/denoising-diffusion-gan/" target="_blank">Project&nbsp;Page</a>
+</div>
+<br>
+<br>
+
+<div align="center">
+    <img width="800" alt="teaser" src="assets/teaser.png"/>
+</div>
+
+Generative denoising diffusion models typically assume that the denoising distribution can be modeled by a Gaussian distribution. This assumption holds only for small denoising steps, which in practice translates to thousands of denoising steps in the synthesis process. In our denoising diffusion GANs, we represent the denoising model using multimodal and complex conditional GANs, enabling us to efficiently generate data in as few as two steps.
+
+## Set up datasets ##
+We trained on several datasets, including CIFAR10, LSUN Church Outdoor 256 and CelebA HQ 256. 
+For large datasets, we store the data in LMDB datasets for I/O efficiency. Check [here](https://github.com/NVlabs/NVAE#set-up-file-paths-and-data) for information regarding dataset preparation.
+
+
+## Training Denoising Diffusion GANs ##
+We use the following commands on each dataset for training denoising diffusion GANs.
+
+#### CIFAR-10 ####
+
+We train Denoising Diffusion GANs on CIFAR-10 using 4 32-GB V100 GPU. 
+```
+python3 train_ddgan.py --dataset cifar10 --exp ddgan_cifar10_exp1 --num_channels 3 --num_channels_dae 128 --num_timesteps 4 \
+--num_res_blocks 2 --batch_size 64 --num_epoch 1800 --ngf 64 --nz 100 --z_emb_dim 256 --n_mlp 4 --embedding_type positional \
+--use_ema --ema_decay 0.9999 --r1_gamma 0.02 --lr_d 1.25e-4 --lr_g 1.6e-4 --lazy_reg 15 --num_process_per_node 4 \
+--ch_mult 1 2 2 2 --save_content
+```
+
+#### LSUN Church Outdoor 256 ####
+
+We train Denoising Diffusion GANs on LSUN Church Outdoor 256 using 8 32-GB V100 GPU. 
+```
+python3 train_ddgan.py --dataset lsun --image_size 256 --exp ddgan_lsun_exp1 --num_channels 3 --num_channels_dae 64 --ch_mult 1 1 2 2 4 4 --num_timesteps 4 \
+--num_res_blocks 2 --batch_size 8 --num_epoch 500 --ngf 64 --embedding_type positional --use_ema --ema_decay 0.999 --r1_gamma 1. \
+--z_emb_dim 256 --lr_d 1e-4 --lr_g 1.6e-4 --lazy_reg 10 --num_process_per_node 8 --save_content
+```
+
+#### CelebA HQ 256 ####
+
+We train Denoising Diffusion GANs on CelebA HQ 256 using 8 32-GB V100 GPUs. 
+```
+python3 train_ddgan.py --dataset celeba_256 --image_size 256 --exp ddgan_celebahq_exp1 --num_channels 3 --num_channels_dae 64 --ch_mult 1 1 2 2 4 4 --num_timesteps 2 \
+--num_res_blocks 2 --batch_size 4 --num_epoch 800 --ngf 64 --embedding_type positional --use_ema --r1_gamma 2. \
+--z_emb_dim 256 --lr_d 1e-4 --lr_g 2e-4 --lazy_reg 10  --num_process_per_node 8 --save_content
+```
+
+## Pretrained Checkpoints ##
+We have released pretrained checkpoints on CIFAR-10 and CelebA HQ 256 at this 
+[Google drive directory](https://drive.google.com/drive/folders/1UkzsI0SwBRstMYysRdR76C1XdSv5rQNz?usp=sharing).
+Simply download the `saved_info` directory to the code directory. Use `--epoch_id 1200` for CIFAR-10 and `--epoch_id 550`
+for CelebA HQ 256 in the commands below.
+
+## Evaluation ##
+After training, samples can be generated by calling ```test_ddgan.py```. We evaluate the models with single V100 GPU.
+Below, we use `--epoch_id` to specify the checkpoint saved at a particular epoch.
+Specifically, for models trained by above commands, the scripts for generating samples on CIFAR-10 is
+```
+python3 test_ddgan.py --dataset cifar10 --exp ddgan_cifar10_exp1 --num_channels 3 --num_channels_dae 128 --num_timesteps 4 \
+--num_res_blocks 2 --nz 100 --z_emb_dim 256 --n_mlp 4 --ch_mult 1 2 2 2 --epoch_id $EPOCH
+```
+The scripts for generating samples on CelebA HQ is 
+```
+python3 test_ddgan.py --dataset celeba_256 --image_size 256 --exp ddgan_celebahq_exp1 --num_channels 3 --num_channels_dae 64 \
+--ch_mult 1 1 2 2 4 4 --num_timesteps 2 --num_res_blocks 2  --epoch_id $EPOCH
+```
+The scripts for generating samples on LSUN Church Outdoor is 
+```
+python3 test_ddgan.py --dataset lsun --image_size 256 --exp ddgan_lsun_exp1 --num_channels 3 --num_channels_dae 64 \
+--ch_mult 1 1 2 2 4 4  --num_timesteps 4 --num_res_blocks 2  --epoch_id $EPOCH
+```
+
+We use the [PyTorch](https://github.com/mseitzer/pytorch-fid) implementation to compute the FID scores, and in particular, codes for computing the FID are adapted from [FastDPM](https://github.com/FengNiMa/FastDPM_pytorch).
+
+To compute FID, run the same scripts above for sampling, with additional arguments ```--compute_fid``` and ```--real_img_dir /path/to/real/images```.
+
+For Inception Score, save samples in a single numpy array with pixel values in range [0, 255] and simply run 
+```
+python ./pytorch_fid/inception_score.py --sample_dir /path/to/sampled_images
+```
+where the code for computing Inception Score is adapted from [here](https://github.com/tsc2017/Inception-Score).
+
+For Improved Precision and Recall, follow the instruction [here](https://github.com/kynkaat/improved-precision-and-recall-metric).
+
+
+## License ##
+Please check the LICENSE file. Denoising diffusion GAN may be used non-commercially, meaning for research or 
+evaluation purposes only. For business inquiries, please contact 
+[researchinquiries@nvidia.com](mailto:researchinquiries@nvidia.com).
+
+## Bibtex ##
+Cite our paper using the following bibtex item:
+
+```
+@inproceedings{
+xiao2022tackling,
+title={Tackling the Generative Learning Trilemma with Denoising Diffusion GANs},
+author={Zhisheng Xiao and Karsten Kreis and Arash Vahdat},
+booktitle={International Conference on Learning Representations},
+year={2022}
+}
+```
+
+## Contributors ##
+Denoising Diffusion GAN was built primarily by [Zhisheng Xiao](https://xavierxiao.github.io/) during a summer 
+internship at NVIDIA research.

requirements.txt

@@ -0,0 +1,10 @@
+torch==1.8.0
+torchvision==0.9.0
+pillow
+matplotlib
+tensorboard
+tensorboardX
+lmdb
+matplotlib
+scipy
+ninja

score_sde/LICENSE_Apache

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

score_sde/models/LICENSE_MIT

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Chin-Wei Huang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

score_sde/models/__init__.py

@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+

score_sde/models/dense_layer.py

@@ -0,0 +1,83 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from a file released under the MIT License.
+#
+# Source:
+# https://github.com/CW-Huang/sdeflow-light/blob/524650bc5ad69522b3e0905672deef0650374512/lib/models/unet.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_MIT). The modifications
+# to this file are subject to the same MIT License.
+# ---------------------------------------------------------------
+
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import _calculate_fan_in_and_fan_out
+import numpy as np
+
+
+def _calculate_correct_fan(tensor, mode):
+    """
+    copied and modified from https://github.com/pytorch/pytorch/blob/master/torch/nn/init.py#L337
+    """
+    mode = mode.lower()
+    valid_modes = ['fan_in', 'fan_out', 'fan_avg']
+    if mode not in valid_modes:
+        raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes))
+
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    return fan_in if mode == 'fan_in' else fan_out
+
+
+def kaiming_uniform_(tensor, gain=1., mode='fan_in'):
+    r"""Fills the input `Tensor` with values according to the method
+    described in `Delving deep into rectifiers: Surpassing human-level
+    performance on ImageNet classification` - He, K. et al. (2015), using a
+    uniform distribution. The resulting tensor will have values sampled from
+    :math:`\mathcal{U}(-\text{bound}, \text{bound})` where
+    .. math::
+        \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}}
+    Also known as He initialization.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        gain: multiplier to the dispersion
+        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
+            preserves the magnitude of the variance of the weights in the
+            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
+            backwards pass.
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.kaiming_uniform_(w, mode='fan_in')
+    """
+    fan = _calculate_correct_fan(tensor, mode)
+    var = gain / max(1., fan)
+    bound = math.sqrt(3.0 * var)  # Calculate uniform bounds from standard deviation
+    with torch.no_grad():
+        return tensor.uniform_(-bound, bound)
+
+
+def variance_scaling_init_(tensor, scale):
+    return kaiming_uniform_(tensor, gain=1e-10 if scale == 0 else scale, mode='fan_avg')
+
+
+def dense(in_channels, out_channels, init_scale=1.):
+    lin = nn.Linear(in_channels, out_channels)
+    variance_scaling_init_(lin.weight, scale=init_scale)
+    nn.init.zeros_(lin.bias)
+    return lin
+
+def conv2d(in_planes, out_planes, kernel_size=(3, 3), stride=1, dilation=1, padding=1, bias=True, padding_mode='zeros',
+           init_scale=1.):
+    conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation,
+                     bias=bias, padding_mode=padding_mode)
+    variance_scaling_init_(conv.weight, scale=init_scale)
+    if bias:
+        nn.init.zeros_(conv.bias)
+    return conv
+
+
+    

score_sde/models/discriminator.py

@@ -0,0 +1,239 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
+# ---------------------------------------------------------------
+import torch
+import torch.nn as nn
+import numpy as np
+
+from . import up_or_down_sampling
+from . import dense_layer
+from . import layers
+
+dense = dense_layer.dense
+conv2d = dense_layer.conv2d
+get_sinusoidal_positional_embedding = layers.get_timestep_embedding
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, embedding_dim, hidden_dim, output_dim, act=nn.LeakyReLU(0.2)):
+        super().__init__()
+
+        self.embedding_dim = embedding_dim
+        self.output_dim = output_dim
+        self.hidden_dim = hidden_dim
+
+        self.main = nn.Sequential(
+            dense(embedding_dim, hidden_dim),
+            act,
+            dense(hidden_dim, output_dim),
+        )
+
+    def forward(self, temp):
+        temb = get_sinusoidal_positional_embedding(temp, self.embedding_dim)
+        temb = self.main(temb)
+        return temb
+#%%
+class DownConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size=3,
+        padding=1,
+        t_emb_dim = 128,
+        downsample=False,
+        act = nn.LeakyReLU(0.2),
+        fir_kernel=(1, 3, 3, 1)
+    ):
+        super().__init__()
+     
+        
+        self.fir_kernel = fir_kernel
+        self.downsample = downsample
+        
+        self.conv1 = nn.Sequential(
+                    conv2d(in_channel, out_channel, kernel_size, padding=padding),
+                    )
+
+        
+        self.conv2 = nn.Sequential(
+                    conv2d(out_channel, out_channel, kernel_size, padding=padding,init_scale=0.)
+                    )
+        self.dense_t1= dense(t_emb_dim, out_channel)
+
+
+        self.act = act
+        
+            
+        self.skip = nn.Sequential(
+                    conv2d(in_channel, out_channel, 1, padding=0, bias=False),
+                    )
+        
+            
+
+    def forward(self, input, t_emb):
+        
+        out = self.act(input)
+        out = self.conv1(out)
+        out += self.dense_t1(t_emb)[..., None, None]
+       
+        out = self.act(out)
+       
+        if self.downsample:
+            out = up_or_down_sampling.downsample_2d(out, self.fir_kernel, factor=2)
+            input = up_or_down_sampling.downsample_2d(input, self.fir_kernel, factor=2)
+        out = self.conv2(out)
+        
+        
+        skip = self.skip(input)
+        out = (out + skip) / np.sqrt(2)
+
+
+        return out
+    
+class Discriminator_small(nn.Module):
+  """A time-dependent discriminator for small images (CIFAR10, StackMNIST)."""
+
+  def __init__(self, nc = 3, ngf = 64, t_emb_dim = 128, act=nn.LeakyReLU(0.2)):
+    super().__init__()
+    # Gaussian random feature embedding layer for time
+    self.act = act
+    
+    
+    self.t_embed = TimestepEmbedding(
+        embedding_dim=t_emb_dim,
+        hidden_dim=t_emb_dim,
+        output_dim=t_emb_dim,
+        act=act,
+        )
+    
+    
+     
+    # Encoding layers where the resolution decreases
+    self.start_conv = conv2d(nc,ngf*2,1, padding=0)
+    self.conv1 = DownConvBlock(ngf*2, ngf*2, t_emb_dim = t_emb_dim,act=act)
+    
+    self.conv2 = DownConvBlock(ngf*2, ngf*4,  t_emb_dim = t_emb_dim, downsample=True,act=act)
+    
+    
+    self.conv3 = DownConvBlock(ngf*4, ngf*8,  t_emb_dim = t_emb_dim, downsample=True,act=act)
+
+    
+    self.conv4 = DownConvBlock(ngf*8, ngf*8, t_emb_dim = t_emb_dim, downsample=True,act=act)
+    
+    
+    self.final_conv = conv2d(ngf*8 + 1, ngf*8, 3,padding=1, init_scale=0.)
+    self.end_linear = dense(ngf*8, 1)
+    
+    self.stddev_group = 4
+    self.stddev_feat = 1
+    
+        
+  def forward(self, x, t, x_t):
+    t_embed = self.act(self.t_embed(t))  
+    
+  
+    input_x = torch.cat((x, x_t), dim = 1)
+    
+    h0 = self.start_conv(input_x)
+    h1 = self.conv1(h0,t_embed)    
+    
+    h2 = self.conv2(h1,t_embed)   
+   
+    h3 = self.conv3(h2,t_embed)
+   
+    
+    out = self.conv4(h3,t_embed)
+    
+    batch, channel, height, width = out.shape
+    group = min(batch, self.stddev_group)
+    stddev = out.view(
+            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
+        )
+    stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
+    stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
+    stddev = stddev.repeat(group, 1, height, width)
+    out = torch.cat([out, stddev], 1)
+    
+    out = self.final_conv(out)
+    out = self.act(out)
+   
+    
+    out = out.view(out.shape[0], out.shape[1], -1).sum(2)
+    out = self.end_linear(out)
+    
+    return out
+
+
+class Discriminator_large(nn.Module):
+  """A time-dependent discriminator for large images (CelebA, LSUN)."""
+
+  def __init__(self, nc = 1, ngf = 32, t_emb_dim = 128, act=nn.LeakyReLU(0.2)):
+    super().__init__()
+    # Gaussian random feature embedding layer for time
+    self.act = act
+    
+    self.t_embed = TimestepEmbedding(
+            embedding_dim=t_emb_dim,
+            hidden_dim=t_emb_dim,
+            output_dim=t_emb_dim,
+            act=act,
+        )
+      
+    self.start_conv = conv2d(nc,ngf*2,1, padding=0)
+    self.conv1 = DownConvBlock(ngf*2, ngf*4, t_emb_dim = t_emb_dim, downsample = True, act=act)
+    
+    self.conv2 = DownConvBlock(ngf*4, ngf*8,  t_emb_dim = t_emb_dim, downsample=True,act=act)
+
+    self.conv3 = DownConvBlock(ngf*8, ngf*8,  t_emb_dim = t_emb_dim, downsample=True,act=act)
+    
+    
+    self.conv4 = DownConvBlock(ngf*8, ngf*8, t_emb_dim = t_emb_dim, downsample=True,act=act)
+    self.conv5 = DownConvBlock(ngf*8, ngf*8, t_emb_dim = t_emb_dim, downsample=True,act=act)
+    self.conv6 = DownConvBlock(ngf*8, ngf*8, t_emb_dim = t_emb_dim, downsample=True,act=act)
+
+  
+    self.final_conv = conv2d(ngf*8 + 1, ngf*8, 3,padding=1)
+    self.end_linear = dense(ngf*8, 1)
+    
+    self.stddev_group = 4
+    self.stddev_feat = 1
+    
+        
+  def forward(self, x, t, x_t):
+    t_embed = self.act(self.t_embed(t))  
+    
+    input_x = torch.cat((x, x_t), dim = 1)
+    
+    h = self.start_conv(input_x)
+    h = self.conv1(h,t_embed)    
+   
+    h = self.conv2(h,t_embed)
+   
+    h = self.conv3(h,t_embed)
+    h = self.conv4(h,t_embed)
+    h = self.conv5(h,t_embed)
+   
+    
+    out = self.conv6(h,t_embed)
+    
+    batch, channel, height, width = out.shape
+    group = min(batch, self.stddev_group)
+    stddev = out.view(
+            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
+        )
+    stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
+    stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
+    stddev = stddev.repeat(group, 1, height, width)
+    out = torch.cat([out, stddev], 1)
+    
+    out = self.final_conv(out)
+    out = self.act(out)
+    
+    out = out.view(out.shape[0], out.shape[1], -1).sum(2)
+    out = self.end_linear(out)
+    
+    return out
+

score_sde/models/layers.py

@@ -0,0 +1,619 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from a file in the Score SDE library
+# which was released under the Apache License.
+#
+# Source:
+# https://github.com/yang-song/score_sde_pytorch/blob/main/models/layers.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_Apache). The modifications
+# to this file are subject to the same Apache License.
+# ---------------------------------------------------------------
+
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: skip-file
+"""Common layers for defining score networks. 
+    Adapted from https://github.com/yang-song/score_sde_pytorch/blob/main/models/layers.py
+"""
+import math
+import string
+from functools import partial
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+def get_act(config):
+  """Get activation functions from the config file."""
+
+  if config.model.nonlinearity.lower() == 'elu':
+    return nn.ELU()
+  elif config.model.nonlinearity.lower() == 'relu':
+    return nn.ReLU()
+  elif config.model.nonlinearity.lower() == 'lrelu':
+    return nn.LeakyReLU(negative_slope=0.2)
+  elif config.model.nonlinearity.lower() == 'swish':
+    return nn.SiLU()
+  else:
+    raise NotImplementedError('activation function does not exist!')
+
+
+def ncsn_conv1x1(in_planes, out_planes, stride=1, bias=True, dilation=1, init_scale=1., padding=0):
+  """1x1 convolution. Same as NCSNv1/v2."""
+  conv = nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=bias, dilation=dilation,
+                   padding=padding)
+  init_scale = 1e-10 if init_scale == 0 else init_scale
+  conv.weight.data *= init_scale
+  conv.bias.data *= init_scale
+  return conv
+
+
+def variance_scaling(scale, mode, distribution,
+                     in_axis=1, out_axis=0,
+                     dtype=torch.float32,
+                     device='cpu'):
+  """Ported from JAX. """
+
+  def _compute_fans(shape, in_axis=1, out_axis=0):
+    receptive_field_size = np.prod(shape) / shape[in_axis] / shape[out_axis]
+    fan_in = shape[in_axis] * receptive_field_size
+    fan_out = shape[out_axis] * receptive_field_size
+    return fan_in, fan_out
+
+  def init(shape, dtype=dtype, device=device):
+    fan_in, fan_out = _compute_fans(shape, in_axis, out_axis)
+    if mode == "fan_in":
+      denominator = fan_in
+    elif mode == "fan_out":
+      denominator = fan_out
+    elif mode == "fan_avg":
+      denominator = (fan_in + fan_out) / 2
+    else:
+      raise ValueError(
+        "invalid mode for variance scaling initializer: {}".format(mode))
+    variance = scale / denominator
+    if distribution == "normal":
+      return torch.randn(*shape, dtype=dtype, device=device) * np.sqrt(variance)
+    elif distribution == "uniform":
+      return (torch.rand(*shape, dtype=dtype, device=device) * 2. - 1.) * np.sqrt(3 * variance)
+    else:
+      raise ValueError("invalid distribution for variance scaling initializer")
+
+  return init
+
+
+def default_init(scale=1.):
+  """The same initialization used in DDPM."""
+  scale = 1e-10 if scale == 0 else scale
+  return variance_scaling(scale, 'fan_avg', 'uniform')
+
+
+class Dense(nn.Module):
+  """Linear layer with `default_init`."""
+  def __init__(self):
+    super().__init__()
+
+
+def ddpm_conv1x1(in_planes, out_planes, stride=1, bias=True, init_scale=1., padding=0):
+  """1x1 convolution with DDPM initialization."""
+  conv = nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=padding, bias=bias)
+  conv.weight.data = default_init(init_scale)(conv.weight.data.shape)
+  nn.init.zeros_(conv.bias)
+  return conv
+
+
+def ncsn_conv3x3(in_planes, out_planes, stride=1, bias=True, dilation=1, init_scale=1., padding=1):
+  """3x3 convolution with PyTorch initialization. Same as NCSNv1/NCSNv2."""
+  init_scale = 1e-10 if init_scale == 0 else init_scale
+  conv = nn.Conv2d(in_planes, out_planes, stride=stride, bias=bias,
+                   dilation=dilation, padding=padding, kernel_size=3)
+  conv.weight.data *= init_scale
+  conv.bias.data *= init_scale
+  return conv
+
+
+def ddpm_conv3x3(in_planes, out_planes, stride=1, bias=True, dilation=1, init_scale=1., padding=1):
+  """3x3 convolution with DDPM initialization."""
+  conv = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=padding,
+                   dilation=dilation, bias=bias)
+  conv.weight.data = default_init(init_scale)(conv.weight.data.shape)
+  nn.init.zeros_(conv.bias)
+  return conv
+
+  ###########################################################################
+  # Functions below are ported over from the NCSNv1/NCSNv2 codebase:
+  # https://github.com/ermongroup/ncsn
+  # https://github.com/ermongroup/ncsnv2
+  ###########################################################################
+
+
+class CRPBlock(nn.Module):
+  def __init__(self, features, n_stages, act=nn.ReLU(), maxpool=True):
+    super().__init__()
+    self.convs = nn.ModuleList()
+    for i in range(n_stages):
+      self.convs.append(ncsn_conv3x3(features, features, stride=1, bias=False))
+    self.n_stages = n_stages
+    if maxpool:
+      self.pool = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)
+    else:
+      self.pool = nn.AvgPool2d(kernel_size=5, stride=1, padding=2)
+
+    self.act = act
+
+  def forward(self, x):
+    x = self.act(x)
+    path = x
+    for i in range(self.n_stages):
+      path = self.pool(path)
+      path = self.convs[i](path)
+      x = path + x
+    return x
+
+
+class CondCRPBlock(nn.Module):
+  def __init__(self, features, n_stages, num_classes, normalizer, act=nn.ReLU()):
+    super().__init__()
+    self.convs = nn.ModuleList()
+    self.norms = nn.ModuleList()
+    self.normalizer = normalizer
+    for i in range(n_stages):
+      self.norms.append(normalizer(features, num_classes, bias=True))
+      self.convs.append(ncsn_conv3x3(features, features, stride=1, bias=False))
+
+    self.n_stages = n_stages
+    self.pool = nn.AvgPool2d(kernel_size=5, stride=1, padding=2)
+    self.act = act
+
+  def forward(self, x, y):
+    x = self.act(x)
+    path = x
+    for i in range(self.n_stages):
+      path = self.norms[i](path, y)
+      path = self.pool(path)
+      path = self.convs[i](path)
+
+      x = path + x
+    return x
+
+
+class RCUBlock(nn.Module):
+  def __init__(self, features, n_blocks, n_stages, act=nn.ReLU()):
+    super().__init__()
+
+    for i in range(n_blocks):
+      for j in range(n_stages):
+        setattr(self, '{}_{}_conv'.format(i + 1, j + 1), ncsn_conv3x3(features, features, stride=1, bias=False))
+
+    self.stride = 1
+    self.n_blocks = n_blocks
+    self.n_stages = n_stages
+    self.act = act
+
+  def forward(self, x):
+    for i in range(self.n_blocks):
+      residual = x
+      for j in range(self.n_stages):
+        x = self.act(x)
+        x = getattr(self, '{}_{}_conv'.format(i + 1, j + 1))(x)
+
+      x += residual
+    return x
+
+
+class CondRCUBlock(nn.Module):
+  def __init__(self, features, n_blocks, n_stages, num_classes, normalizer, act=nn.ReLU()):
+    super().__init__()
+
+    for i in range(n_blocks):
+      for j in range(n_stages):
+        setattr(self, '{}_{}_norm'.format(i + 1, j + 1), normalizer(features, num_classes, bias=True))
+        setattr(self, '{}_{}_conv'.format(i + 1, j + 1), ncsn_conv3x3(features, features, stride=1, bias=False))
+
+    self.stride = 1
+    self.n_blocks = n_blocks
+    self.n_stages = n_stages
+    self.act = act
+    self.normalizer = normalizer
+
+  def forward(self, x, y):
+    for i in range(self.n_blocks):
+      residual = x
+      for j in range(self.n_stages):
+        x = getattr(self, '{}_{}_norm'.format(i + 1, j + 1))(x, y)
+        x = self.act(x)
+        x = getattr(self, '{}_{}_conv'.format(i + 1, j + 1))(x)
+
+      x += residual
+    return x
+
+
+class MSFBlock(nn.Module):
+  def __init__(self, in_planes, features):
+    super().__init__()
+    assert isinstance(in_planes, list) or isinstance(in_planes, tuple)
+    self.convs = nn.ModuleList()
+    self.features = features
+
+    for i in range(len(in_planes)):
+      self.convs.append(ncsn_conv3x3(in_planes[i], features, stride=1, bias=True))
+
+  def forward(self, xs, shape):
+    sums = torch.zeros(xs[0].shape[0], self.features, *shape, device=xs[0].device)
+    for i in range(len(self.convs)):
+      h = self.convs[i](xs[i])
+      h = F.interpolate(h, size=shape, mode='bilinear', align_corners=True)
+      sums += h
+    return sums
+
+
+class CondMSFBlock(nn.Module):
+  def __init__(self, in_planes, features, num_classes, normalizer):
+    super().__init__()
+    assert isinstance(in_planes, list) or isinstance(in_planes, tuple)
+
+    self.convs = nn.ModuleList()
+    self.norms = nn.ModuleList()
+    self.features = features
+    self.normalizer = normalizer
+
+    for i in range(len(in_planes)):
+      self.convs.append(ncsn_conv3x3(in_planes[i], features, stride=1, bias=True))
+      self.norms.append(normalizer(in_planes[i], num_classes, bias=True))
+
+  def forward(self, xs, y, shape):
+    sums = torch.zeros(xs[0].shape[0], self.features, *shape, device=xs[0].device)
+    for i in range(len(self.convs)):
+      h = self.norms[i](xs[i], y)
+      h = self.convs[i](h)
+      h = F.interpolate(h, size=shape, mode='bilinear', align_corners=True)
+      sums += h
+    return sums
+
+
+class RefineBlock(nn.Module):
+  def __init__(self, in_planes, features, act=nn.ReLU(), start=False, end=False, maxpool=True):
+    super().__init__()
+
+    assert isinstance(in_planes, tuple) or isinstance(in_planes, list)
+    self.n_blocks = n_blocks = len(in_planes)
+
+    self.adapt_convs = nn.ModuleList()
+    for i in range(n_blocks):
+      self.adapt_convs.append(RCUBlock(in_planes[i], 2, 2, act))
+
+    self.output_convs = RCUBlock(features, 3 if end else 1, 2, act)
+
+    if not start:
+      self.msf = MSFBlock(in_planes, features)
+
+    self.crp = CRPBlock(features, 2, act, maxpool=maxpool)
+
+  def forward(self, xs, output_shape):
+    assert isinstance(xs, tuple) or isinstance(xs, list)
+    hs = []
+    for i in range(len(xs)):
+      h = self.adapt_convs[i](xs[i])
+      hs.append(h)
+
+    if self.n_blocks > 1:
+      h = self.msf(hs, output_shape)
+    else:
+      h = hs[0]
+
+    h = self.crp(h)
+    h = self.output_convs(h)
+
+    return h
+
+
+class CondRefineBlock(nn.Module):
+  def __init__(self, in_planes, features, num_classes, normalizer, act=nn.ReLU(), start=False, end=False):
+    super().__init__()
+
+    assert isinstance(in_planes, tuple) or isinstance(in_planes, list)
+    self.n_blocks = n_blocks = len(in_planes)
+
+    self.adapt_convs = nn.ModuleList()
+    for i in range(n_blocks):
+      self.adapt_convs.append(
+        CondRCUBlock(in_planes[i], 2, 2, num_classes, normalizer, act)
+      )
+
+    self.output_convs = CondRCUBlock(features, 3 if end else 1, 2, num_classes, normalizer, act)
+
+    if not start:
+      self.msf = CondMSFBlock(in_planes, features, num_classes, normalizer)
+
+    self.crp = CondCRPBlock(features, 2, num_classes, normalizer, act)
+
+  def forward(self, xs, y, output_shape):
+    assert isinstance(xs, tuple) or isinstance(xs, list)
+    hs = []
+    for i in range(len(xs)):
+      h = self.adapt_convs[i](xs[i], y)
+      hs.append(h)
+
+    if self.n_blocks > 1:
+      h = self.msf(hs, y, output_shape)
+    else:
+      h = hs[0]
+
+    h = self.crp(h, y)
+    h = self.output_convs(h, y)
+
+    return h
+
+
+class ConvMeanPool(nn.Module):
+  def __init__(self, input_dim, output_dim, kernel_size=3, biases=True, adjust_padding=False):
+    super().__init__()
+    if not adjust_padding:
+      conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride=1, padding=kernel_size // 2, bias=biases)
+      self.conv = conv
+    else:
+      conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride=1, padding=kernel_size // 2, bias=biases)
+
+      self.conv = nn.Sequential(
+        nn.ZeroPad2d((1, 0, 1, 0)),
+        conv
+      )
+
+  def forward(self, inputs):
+    output = self.conv(inputs)
+    output = sum([output[:, :, ::2, ::2], output[:, :, 1::2, ::2],
+                  output[:, :, ::2, 1::2], output[:, :, 1::2, 1::2]]) / 4.
+    return output
+
+
+class MeanPoolConv(nn.Module):
+  def __init__(self, input_dim, output_dim, kernel_size=3, biases=True):
+    super().__init__()
+    self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride=1, padding=kernel_size // 2, bias=biases)
+
+  def forward(self, inputs):
+    output = inputs
+    output = sum([output[:, :, ::2, ::2], output[:, :, 1::2, ::2],
+                  output[:, :, ::2, 1::2], output[:, :, 1::2, 1::2]]) / 4.
+    return self.conv(output)
+
+
+class UpsampleConv(nn.Module):
+  def __init__(self, input_dim, output_dim, kernel_size=3, biases=True):
+    super().__init__()
+    self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride=1, padding=kernel_size // 2, bias=biases)
+    self.pixelshuffle = nn.PixelShuffle(upscale_factor=2)
+
+  def forward(self, inputs):
+    output = inputs
+    output = torch.cat([output, output, output, output], dim=1)
+    output = self.pixelshuffle(output)
+    return self.conv(output)
+
+
+
+
+class ResidualBlock(nn.Module):
+  def __init__(self, input_dim, output_dim, resample=None, act=nn.ELU(),
+               normalization=nn.InstanceNorm2d, adjust_padding=False, dilation=1):
+    super().__init__()
+    self.non_linearity = act
+    self.input_dim = input_dim
+    self.output_dim = output_dim
+    self.resample = resample
+    self.normalization = normalization
+    if resample == 'down':
+      if dilation > 1:
+        self.conv1 = ncsn_conv3x3(input_dim, input_dim, dilation=dilation)
+        self.normalize2 = normalization(input_dim)
+        self.conv2 = ncsn_conv3x3(input_dim, output_dim, dilation=dilation)
+        conv_shortcut = partial(ncsn_conv3x3, dilation=dilation)
+      else:
+        self.conv1 = ncsn_conv3x3(input_dim, input_dim)
+        self.normalize2 = normalization(input_dim)
+        self.conv2 = ConvMeanPool(input_dim, output_dim, 3, adjust_padding=adjust_padding)
+        conv_shortcut = partial(ConvMeanPool, kernel_size=1, adjust_padding=adjust_padding)
+
+    elif resample is None:
+      if dilation > 1:
+        conv_shortcut = partial(ncsn_conv3x3, dilation=dilation)
+        self.conv1 = ncsn_conv3x3(input_dim, output_dim, dilation=dilation)
+        self.normalize2 = normalization(output_dim)
+        self.conv2 = ncsn_conv3x3(output_dim, output_dim, dilation=dilation)
+      else:
+        # conv_shortcut = nn.Conv2d ### Something wierd here.
+        conv_shortcut = partial(ncsn_conv1x1)
+        self.conv1 = ncsn_conv3x3(input_dim, output_dim)
+        self.normalize2 = normalization(output_dim)
+        self.conv2 = ncsn_conv3x3(output_dim, output_dim)
+    else:
+      raise Exception('invalid resample value')
+
+    if output_dim != input_dim or resample is not None:
+      self.shortcut = conv_shortcut(input_dim, output_dim)
+
+    self.normalize1 = normalization(input_dim)
+
+  def forward(self, x):
+    output = self.normalize1(x)
+    output = self.non_linearity(output)
+    output = self.conv1(output)
+    output = self.normalize2(output)
+    output = self.non_linearity(output)
+    output = self.conv2(output)
+
+    if self.output_dim == self.input_dim and self.resample is None:
+      shortcut = x
+    else:
+      shortcut = self.shortcut(x)
+
+    return shortcut + output
+
+
+###########################################################################
+# Functions below are ported over from the DDPM codebase:
+#  https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/nn.py
+###########################################################################
+
+def get_timestep_embedding(timesteps, embedding_dim, max_positions=10000):
+  assert len(timesteps.shape) == 1  # and timesteps.dtype == tf.int32
+  half_dim = embedding_dim // 2
+  # magic number 10000 is from transformers
+  emb = math.log(max_positions) / (half_dim - 1)
+  emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb)
+  emb = timesteps.float()[:, None] * emb[None, :]
+  emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+  if embedding_dim % 2 == 1:  # zero pad
+    emb = F.pad(emb, (0, 1), mode='constant')
+  assert emb.shape == (timesteps.shape[0], embedding_dim)
+  return emb
+
+
+def _einsum(a, b, c, x, y):
+  einsum_str = '{},{}->{}'.format(''.join(a), ''.join(b), ''.join(c))
+  return torch.einsum(einsum_str, x, y)
+
+
+def contract_inner(x, y):
+  """tensordot(x, y, 1)."""
+  x_chars = list(string.ascii_lowercase[:len(x.shape)])
+  y_chars = list(string.ascii_lowercase[len(x.shape):len(y.shape) + len(x.shape)])
+  y_chars[0] = x_chars[-1]  # first axis of y and last of x get summed
+  out_chars = x_chars[:-1] + y_chars[1:]
+  return _einsum(x_chars, y_chars, out_chars, x, y)
+
+
+class NIN(nn.Module):
+  def __init__(self, in_dim, num_units, init_scale=0.1):
+    super().__init__()
+    self.W = nn.Parameter(default_init(scale=init_scale)((in_dim, num_units)), requires_grad=True)
+    self.b = nn.Parameter(torch.zeros(num_units), requires_grad=True)
+
+  def forward(self, x):
+    x = x.permute(0, 2, 3, 1)
+    y = contract_inner(x, self.W) + self.b
+    return y.permute(0, 3, 1, 2)
+
+
+class AttnBlock(nn.Module):
+  """Channel-wise self-attention block."""
+  def __init__(self, channels):
+    super().__init__()
+    self.GroupNorm_0 = nn.GroupNorm(num_groups=32, num_channels=channels, eps=1e-6)
+    self.NIN_0 = NIN(channels, channels)
+    self.NIN_1 = NIN(channels, channels)
+    self.NIN_2 = NIN(channels, channels)
+    self.NIN_3 = NIN(channels, channels, init_scale=0.)
+
+  def forward(self, x):
+    B, C, H, W = x.shape
+    h = self.GroupNorm_0(x)
+    q = self.NIN_0(h)
+    k = self.NIN_1(h)
+    v = self.NIN_2(h)
+
+    w = torch.einsum('bchw,bcij->bhwij', q, k) * (int(C) ** (-0.5))
+    w = torch.reshape(w, (B, H, W, H * W))
+    w = F.softmax(w, dim=-1)
+    w = torch.reshape(w, (B, H, W, H, W))
+    h = torch.einsum('bhwij,bcij->bchw', w, v)
+    h = self.NIN_3(h)
+    return x + h
+
+
+class Upsample(nn.Module):
+  def __init__(self, channels, with_conv=False):
+    super().__init__()
+    if with_conv:
+      self.Conv_0 = ddpm_conv3x3(channels, channels)
+    self.with_conv = with_conv
+
+  def forward(self, x):
+    B, C, H, W = x.shape
+    h = F.interpolate(x, (H * 2, W * 2), mode='nearest')
+    if self.with_conv:
+      h = self.Conv_0(h)
+    return h
+
+
+class Downsample(nn.Module):
+  def __init__(self, channels, with_conv=False):
+    super().__init__()
+    if with_conv:
+      self.Conv_0 = ddpm_conv3x3(channels, channels, stride=2, padding=0)
+    self.with_conv = with_conv
+
+  def forward(self, x):
+    B, C, H, W = x.shape
+    # Emulate 'SAME' padding
+    if self.with_conv:
+      x = F.pad(x, (0, 1, 0, 1))
+      x = self.Conv_0(x)
+    else:
+      x = F.avg_pool2d(x, kernel_size=2, stride=2, padding=0)
+
+    assert x.shape == (B, C, H // 2, W // 2)
+    return x
+
+
+class ResnetBlockDDPM(nn.Module):
+  """The ResNet Blocks used in DDPM."""
+  def __init__(self, act, in_ch, out_ch=None, temb_dim=None, conv_shortcut=False, dropout=0.1):
+    super().__init__()
+    if out_ch is None:
+      out_ch = in_ch
+    self.GroupNorm_0 = nn.GroupNorm(num_groups=32, num_channels=in_ch, eps=1e-6)
+    self.act = act
+    self.Conv_0 = ddpm_conv3x3(in_ch, out_ch)
+    if temb_dim is not None:
+      self.Dense_0 = nn.Linear(temb_dim, out_ch)
+      self.Dense_0.weight.data = default_init()(self.Dense_0.weight.data.shape)
+      nn.init.zeros_(self.Dense_0.bias)
+
+    self.GroupNorm_1 = nn.GroupNorm(num_groups=32, num_channels=out_ch, eps=1e-6)
+    self.Dropout_0 = nn.Dropout(dropout)
+    self.Conv_1 = ddpm_conv3x3(out_ch, out_ch, init_scale=0.)
+    if in_ch != out_ch:
+      if conv_shortcut:
+        self.Conv_2 = ddpm_conv3x3(in_ch, out_ch)
+      else:
+        self.NIN_0 = NIN(in_ch, out_ch)
+    self.out_ch = out_ch
+    self.in_ch = in_ch
+    self.conv_shortcut = conv_shortcut
+
+  def forward(self, x, temb=None):
+    B, C, H, W = x.shape
+    assert C == self.in_ch
+    out_ch = self.out_ch if self.out_ch else self.in_ch
+    h = self.act(self.GroupNorm_0(x))
+    h = self.Conv_0(h)
+    # Add bias to each feature map conditioned on the time embedding
+    if temb is not None:
+      h += self.Dense_0(self.act(temb))[:, :, None, None]
+    h = self.act(self.GroupNorm_1(h))
+    h = self.Dropout_0(h)
+    h = self.Conv_1(h)
+    if C != out_ch:
+      if self.conv_shortcut:
+        x = self.Conv_2(x)
+      else:
+        x = self.NIN_0(x)
+    return x + h

score_sde/models/layerspp.py

@@ -0,0 +1,380 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from a file in the Score SDE library
+# which was released under the Apache License.
+#
+# Source:
+# https://github.com/yang-song/score_sde_pytorch/blob/main/models/layerspp.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_Apache). The modifications
+# to this file are subject to the same Apache License.
+# ---------------------------------------------------------------
+
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: skip-file
+
+from . import layers
+from . import up_or_down_sampling, dense_layer
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+conv1x1 = layers.ddpm_conv1x1
+conv3x3 = layers.ddpm_conv3x3
+NIN = layers.NIN
+default_init = layers.default_init
+dense = dense_layer.dense
+
+class AdaptiveGroupNorm(nn.Module):
+    def __init__(self, num_groups,in_channel, style_dim):
+        super().__init__()
+
+        self.norm = nn.GroupNorm(num_groups, in_channel, affine=False, eps=1e-6)
+        self.style = dense(style_dim, in_channel * 2)
+
+        self.style.bias.data[:in_channel] = 1
+        self.style.bias.data[in_channel:] = 0
+
+    def forward(self, input, style):
+        style = self.style(style).unsqueeze(2).unsqueeze(3)
+        gamma, beta = style.chunk(2, 1)
+
+        out = self.norm(input)
+        out = gamma * out + beta
+
+        return out
+
+class GaussianFourierProjection(nn.Module):
+  """Gaussian Fourier embeddings for noise levels."""
+
+  def __init__(self, embedding_size=256, scale=1.0):
+    super().__init__()
+    self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)
+
+  def forward(self, x):
+    x_proj = x[:, None] * self.W[None, :] * 2 * np.pi
+    return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)
+
+
+class Combine(nn.Module):
+  """Combine information from skip connections."""
+
+  def __init__(self, dim1, dim2, method='cat'):
+    super().__init__()
+    self.Conv_0 = conv1x1(dim1, dim2)
+    self.method = method
+
+  def forward(self, x, y):
+    h = self.Conv_0(x)
+    if self.method == 'cat':
+      return torch.cat([h, y], dim=1)
+    elif self.method == 'sum':
+      return h + y
+    else:
+      raise ValueError(f'Method {self.method} not recognized.')
+
+
+class AttnBlockpp(nn.Module):
+  """Channel-wise self-attention block. Modified from DDPM."""
+
+  def __init__(self, channels, skip_rescale=False, init_scale=0.):
+    super().__init__()
+    self.GroupNorm_0 = nn.GroupNorm(num_groups=min(channels // 4, 32), num_channels=channels,
+                                  eps=1e-6)
+    self.NIN_0 = NIN(channels, channels)
+    self.NIN_1 = NIN(channels, channels)
+    self.NIN_2 = NIN(channels, channels)
+    self.NIN_3 = NIN(channels, channels, init_scale=init_scale)
+    self.skip_rescale = skip_rescale
+
+  def forward(self, x):
+    B, C, H, W = x.shape
+    h = self.GroupNorm_0(x)
+    q = self.NIN_0(h)
+    k = self.NIN_1(h)
+    v = self.NIN_2(h)
+
+    w = torch.einsum('bchw,bcij->bhwij', q, k) * (int(C) ** (-0.5))
+    w = torch.reshape(w, (B, H, W, H * W))
+    w = F.softmax(w, dim=-1)
+    w = torch.reshape(w, (B, H, W, H, W))
+    h = torch.einsum('bhwij,bcij->bchw', w, v)
+    h = self.NIN_3(h)
+    if not self.skip_rescale:
+      return x + h
+    else:
+      return (x + h) / np.sqrt(2.)
+
+
+class Upsample(nn.Module):
+  def __init__(self, in_ch=None, out_ch=None, with_conv=False, fir=False,
+               fir_kernel=(1, 3, 3, 1)):
+    super().__init__()
+    out_ch = out_ch if out_ch else in_ch
+    if not fir:
+      if with_conv:
+        self.Conv_0 = conv3x3(in_ch, out_ch)
+    else:
+      if with_conv:
+        self.Conv2d_0 = up_or_down_sampling.Conv2d(in_ch, out_ch,
+                                                   kernel=3, up=True,
+                                                   resample_kernel=fir_kernel,
+                                                   use_bias=True,
+                                                   kernel_init=default_init())
+    self.fir = fir
+    self.with_conv = with_conv
+    self.fir_kernel = fir_kernel
+    self.out_ch = out_ch
+
+  def forward(self, x):
+    B, C, H, W = x.shape
+    if not self.fir:
+      h = F.interpolate(x, (H * 2, W * 2), 'nearest')
+      if self.with_conv:
+        h = self.Conv_0(h)
+    else:
+      if not self.with_conv:
+        h = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
+      else:
+        h = self.Conv2d_0(x)
+
+    return h
+
+
+class Downsample(nn.Module):
+  def __init__(self, in_ch=None, out_ch=None, with_conv=False, fir=False,
+               fir_kernel=(1, 3, 3, 1)):
+    super().__init__()
+    out_ch = out_ch if out_ch else in_ch
+    if not fir:
+      if with_conv:
+        self.Conv_0 = conv3x3(in_ch, out_ch, stride=2, padding=0)
+    else:
+      if with_conv:
+        self.Conv2d_0 = up_or_down_sampling.Conv2d(in_ch, out_ch,
+                                                   kernel=3, down=True,
+                                                   resample_kernel=fir_kernel,
+                                                   use_bias=True,
+                                                   kernel_init=default_init())
+    self.fir = fir
+    self.fir_kernel = fir_kernel
+    self.with_conv = with_conv
+    self.out_ch = out_ch
+
+  def forward(self, x):
+    B, C, H, W = x.shape
+    if not self.fir:
+      if self.with_conv:
+        x = F.pad(x, (0, 1, 0, 1))
+        x = self.Conv_0(x)
+      else:
+        x = F.avg_pool2d(x, 2, stride=2)
+    else:
+      if not self.with_conv:
+        x = up_or_down_sampling.downsample_2d(x, self.fir_kernel, factor=2)
+      else:
+        x = self.Conv2d_0(x)
+
+    return x
+
+
+class ResnetBlockDDPMpp_Adagn(nn.Module):
+  """ResBlock adapted from DDPM."""
+
+  def __init__(self, act, in_ch, out_ch=None, temb_dim=None, zemb_dim=None, conv_shortcut=False,
+               dropout=0.1, skip_rescale=False, init_scale=0.):
+    super().__init__()
+    out_ch = out_ch if out_ch else in_ch
+    self.GroupNorm_0 = AdaptiveGroupNorm(min(in_ch // 4, 32), in_ch, zemb_dim)
+    self.Conv_0 = conv3x3(in_ch, out_ch)
+    if temb_dim is not None:
+      self.Dense_0 = nn.Linear(temb_dim, out_ch)
+      self.Dense_0.weight.data = default_init()(self.Dense_0.weight.data.shape)
+      nn.init.zeros_(self.Dense_0.bias)
+      
+      
+    self.GroupNorm_1 = AdaptiveGroupNorm(min(out_ch // 4, 32), out_ch, zemb_dim)
+    self.Dropout_0 = nn.Dropout(dropout)
+    self.Conv_1 = conv3x3(out_ch, out_ch, init_scale=init_scale)
+    if in_ch != out_ch:
+      if conv_shortcut:
+        self.Conv_2 = conv3x3(in_ch, out_ch)
+      else:
+        self.NIN_0 = NIN(in_ch, out_ch)
+
+    self.skip_rescale = skip_rescale
+    self.act = act
+    self.out_ch = out_ch
+    self.conv_shortcut = conv_shortcut
+
+  def forward(self, x, temb=None, zemb=None):
+    h = self.act(self.GroupNorm_0(x, zemb))
+    h = self.Conv_0(h)
+    if temb is not None:
+      h += self.Dense_0(self.act(temb))[:, :, None, None]
+    h = self.act(self.GroupNorm_1(h, zemb))
+    h = self.Dropout_0(h)
+    h = self.Conv_1(h)
+    if x.shape[1] != self.out_ch:
+      if self.conv_shortcut:
+        x = self.Conv_2(x)
+      else:
+        x = self.NIN_0(x)
+    if not self.skip_rescale:
+      return x + h
+    else:
+      return (x + h) / np.sqrt(2.)
+
+
+class ResnetBlockBigGANpp_Adagn(nn.Module):
+  def __init__(self, act, in_ch, out_ch=None, temb_dim=None, zemb_dim=None, up=False, down=False,
+               dropout=0.1, fir=False, fir_kernel=(1, 3, 3, 1),
+               skip_rescale=True, init_scale=0.):
+    super().__init__()
+
+    out_ch = out_ch if out_ch else in_ch
+    self.GroupNorm_0 = AdaptiveGroupNorm(min(in_ch // 4, 32), in_ch, zemb_dim)
+    
+    self.up = up
+    self.down = down
+    self.fir = fir
+    self.fir_kernel = fir_kernel
+
+    self.Conv_0 = conv3x3(in_ch, out_ch)
+    if temb_dim is not None:
+      self.Dense_0 = nn.Linear(temb_dim, out_ch)
+      self.Dense_0.weight.data = default_init()(self.Dense_0.weight.shape)
+      nn.init.zeros_(self.Dense_0.bias)
+   
+    self.GroupNorm_1 = AdaptiveGroupNorm(min(out_ch // 4, 32), out_ch, zemb_dim)
+    self.Dropout_0 = nn.Dropout(dropout)
+    self.Conv_1 = conv3x3(out_ch, out_ch, init_scale=init_scale)
+    if in_ch != out_ch or up or down:
+      self.Conv_2 = conv1x1(in_ch, out_ch)
+
+    self.skip_rescale = skip_rescale
+    self.act = act
+    self.in_ch = in_ch
+    self.out_ch = out_ch
+
+  def forward(self, x, temb=None, zemb=None):
+    h = self.act(self.GroupNorm_0(x, zemb))
+
+    if self.up:
+      if self.fir:
+        h = up_or_down_sampling.upsample_2d(h, self.fir_kernel, factor=2)
+        x = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
+      else:
+        h = up_or_down_sampling.naive_upsample_2d(h, factor=2)
+        x = up_or_down_sampling.naive_upsample_2d(x, factor=2)
+    elif self.down:
+      if self.fir:
+        h = up_or_down_sampling.downsample_2d(h, self.fir_kernel, factor=2)
+        x = up_or_down_sampling.downsample_2d(x, self.fir_kernel, factor=2)
+      else:
+        h = up_or_down_sampling.naive_downsample_2d(h, factor=2)
+        x = up_or_down_sampling.naive_downsample_2d(x, factor=2)
+
+    h = self.Conv_0(h)
+    # Add bias to each feature map conditioned on the time embedding
+    if temb is not None:
+      h += self.Dense_0(self.act(temb))[:, :, None, None]
+    h = self.act(self.GroupNorm_1(h, zemb))
+    h = self.Dropout_0(h)
+    h = self.Conv_1(h)
+   
+    if self.in_ch != self.out_ch or self.up or self.down:
+      x = self.Conv_2(x)
+
+    if not self.skip_rescale:
+      return x + h
+    else:
+      return (x + h) / np.sqrt(2.)
+  
+
+class ResnetBlockBigGANpp_Adagn_one(nn.Module):
+  def __init__(self, act, in_ch, out_ch=None, temb_dim=None, zemb_dim=None, up=False, down=False,
+               dropout=0.1, fir=False, fir_kernel=(1, 3, 3, 1),
+               skip_rescale=True, init_scale=0.):
+    super().__init__()
+
+    out_ch = out_ch if out_ch else in_ch
+    self.GroupNorm_0 = AdaptiveGroupNorm(min(in_ch // 4, 32), in_ch, zemb_dim)
+   
+    self.up = up
+    self.down = down
+    self.fir = fir
+    self.fir_kernel = fir_kernel
+
+    self.Conv_0 = conv3x3(in_ch, out_ch)
+    if temb_dim is not None:
+      self.Dense_0 = nn.Linear(temb_dim, out_ch)
+      self.Dense_0.weight.data = default_init()(self.Dense_0.weight.shape)
+      nn.init.zeros_(self.Dense_0.bias)
+    
+
+    self.GroupNorm_1 = nn.GroupNorm(num_groups=min(out_ch // 4, 32), num_channels=out_ch, eps=1e-6)
+    
+    self.Dropout_0 = nn.Dropout(dropout)
+    self.Conv_1 = conv3x3(out_ch, out_ch, init_scale=init_scale)
+    if in_ch != out_ch or up or down:
+      self.Conv_2 = conv1x1(in_ch, out_ch)
+
+    self.skip_rescale = skip_rescale
+    self.act = act
+    self.in_ch = in_ch
+    self.out_ch = out_ch
+
+  def forward(self, x, temb=None, zemb=None):
+    h = self.act(self.GroupNorm_0(x, zemb))
+
+    if self.up:
+      if self.fir:
+        h = up_or_down_sampling.upsample_2d(h, self.fir_kernel, factor=2)
+        x = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
+      else:
+        h = up_or_down_sampling.naive_upsample_2d(h, factor=2)
+        x = up_or_down_sampling.naive_upsample_2d(x, factor=2)
+    elif self.down:
+      if self.fir:
+        h = up_or_down_sampling.downsample_2d(h, self.fir_kernel, factor=2)
+        x = up_or_down_sampling.downsample_2d(x, self.fir_kernel, factor=2)
+      else:
+        h = up_or_down_sampling.naive_downsample_2d(h, factor=2)
+        x = up_or_down_sampling.naive_downsample_2d(x, factor=2)
+
+    h = self.Conv_0(h)
+    # Add bias to each feature map conditioned on the time embedding
+    if temb is not None:
+      h += self.Dense_0(self.act(temb))[:, :, None, None]
+    h = self.act(self.GroupNorm_1(h))
+    h = self.Dropout_0(h)
+    h = self.Conv_1(h)
+    
+
+    if self.in_ch != self.out_ch or self.up or self.down:
+      x = self.Conv_2(x)
+
+    if not self.skip_rescale:
+      return x + h
+    else:
+      return (x + h) / np.sqrt(2.)
+  

score_sde/models/ncsnpp_generator_adagn.py

@@ -0,0 +1,431 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from a file in the Score SDE library
+# which was released under the Apache License.
+#
+# Source:
+# https://github.com/yang-song/score_sde_pytorch/blob/main/models/layerspp.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_Apache). The modifications
+# to this file are subject to the same Apache License.
+# ---------------------------------------------------------------
+
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: skip-file
+''' Codes adapted from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ncsnpp.py
+'''
+
+from . import utils, layers, layerspp, dense_layer
+import torch.nn as nn
+import functools
+import torch
+import numpy as np
+
+
+ResnetBlockDDPM = layerspp.ResnetBlockDDPMpp_Adagn
+ResnetBlockBigGAN = layerspp.ResnetBlockBigGANpp_Adagn
+ResnetBlockBigGAN_one = layerspp.ResnetBlockBigGANpp_Adagn_one
+Combine = layerspp.Combine
+conv3x3 = layerspp.conv3x3
+conv1x1 = layerspp.conv1x1
+get_act = layers.get_act
+default_initializer = layers.default_init
+dense = dense_layer.dense
+
+class PixelNorm(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
+
+
+@utils.register_model(name='ncsnpp')
+class NCSNpp(nn.Module):
+  """NCSN++ model"""
+
+  def __init__(self, config):
+    super().__init__()
+    self.config = config
+    self.not_use_tanh = config.not_use_tanh
+    self.act = act = nn.SiLU()
+    self.z_emb_dim = z_emb_dim = config.z_emb_dim
+    
+    self.nf = nf = config.num_channels_dae
+    ch_mult = config.ch_mult
+    self.num_res_blocks = num_res_blocks = config.num_res_blocks
+    self.attn_resolutions = attn_resolutions = config.attn_resolutions
+    dropout = config.dropout
+    resamp_with_conv = config.resamp_with_conv
+    self.num_resolutions = num_resolutions = len(ch_mult)
+    self.all_resolutions = all_resolutions = [config.image_size // (2 ** i) for i in range(num_resolutions)]
+
+    self.conditional = conditional = config.conditional  # noise-conditional
+    fir = config.fir
+    fir_kernel = config.fir_kernel
+    self.skip_rescale = skip_rescale = config.skip_rescale
+    self.resblock_type = resblock_type = config.resblock_type.lower()
+    self.progressive = progressive = config.progressive.lower()
+    self.progressive_input = progressive_input = config.progressive_input.lower()
+    self.embedding_type = embedding_type = config.embedding_type.lower()
+    init_scale = 0.
+    assert progressive in ['none', 'output_skip', 'residual']
+    assert progressive_input in ['none', 'input_skip', 'residual']
+    assert embedding_type in ['fourier', 'positional']
+    combine_method = config.progressive_combine.lower()
+    combiner = functools.partial(Combine, method=combine_method)
+
+    modules = []
+    # timestep/noise_level embedding; only for continuous training
+    if embedding_type == 'fourier':
+      # Gaussian Fourier features embeddings.
+      #assert config.training.continuous, "Fourier features are only used for continuous training."
+
+      modules.append(layerspp.GaussianFourierProjection(
+        embedding_size=nf, scale=config.fourier_scale
+      ))
+      embed_dim = 2 * nf
+
+    elif embedding_type == 'positional':
+      embed_dim = nf
+
+    else:
+      raise ValueError(f'embedding type {embedding_type} unknown.')
+
+    if conditional:
+      modules.append(nn.Linear(embed_dim, nf * 4))
+      modules[-1].weight.data = default_initializer()(modules[-1].weight.shape)
+      nn.init.zeros_(modules[-1].bias)
+      modules.append(nn.Linear(nf * 4, nf * 4))
+      modules[-1].weight.data = default_initializer()(modules[-1].weight.shape)
+      nn.init.zeros_(modules[-1].bias)
+
+    AttnBlock = functools.partial(layerspp.AttnBlockpp,
+                                  init_scale=init_scale,
+                                  skip_rescale=skip_rescale)
+
+    Upsample = functools.partial(layerspp.Upsample,
+                                 with_conv=resamp_with_conv, fir=fir, fir_kernel=fir_kernel)
+
+    if progressive == 'output_skip':
+      self.pyramid_upsample = layerspp.Upsample(fir=fir, fir_kernel=fir_kernel, with_conv=False)
+    elif progressive == 'residual':
+      pyramid_upsample = functools.partial(layerspp.Upsample,
+                                           fir=fir, fir_kernel=fir_kernel, with_conv=True)
+
+    Downsample = functools.partial(layerspp.Downsample,
+                                   with_conv=resamp_with_conv, fir=fir, fir_kernel=fir_kernel)
+
+    if progressive_input == 'input_skip':
+      self.pyramid_downsample = layerspp.Downsample(fir=fir, fir_kernel=fir_kernel, with_conv=False)
+    elif progressive_input == 'residual':
+      pyramid_downsample = functools.partial(layerspp.Downsample,
+                                             fir=fir, fir_kernel=fir_kernel, with_conv=True)
+
+    if resblock_type == 'ddpm':
+      ResnetBlock = functools.partial(ResnetBlockDDPM,
+                                      act=act,
+                                      dropout=dropout,
+                                      init_scale=init_scale,
+                                      skip_rescale=skip_rescale,
+                                      temb_dim=nf * 4,
+                                      zemb_dim = z_emb_dim)
+
+    elif resblock_type == 'biggan':
+      ResnetBlock = functools.partial(ResnetBlockBigGAN,
+                                      act=act,
+                                      dropout=dropout,
+                                      fir=fir,
+                                      fir_kernel=fir_kernel,
+                                      init_scale=init_scale,
+                                      skip_rescale=skip_rescale,
+                                      temb_dim=nf * 4,
+                                      zemb_dim = z_emb_dim)
+    elif resblock_type == 'biggan_oneadagn':
+      ResnetBlock = functools.partial(ResnetBlockBigGAN_one,
+                                      act=act,
+                                      dropout=dropout,
+                                      fir=fir,
+                                      fir_kernel=fir_kernel,
+                                      init_scale=init_scale,
+                                      skip_rescale=skip_rescale,
+                                      temb_dim=nf * 4,
+                                      zemb_dim = z_emb_dim)
+
+    else:
+      raise ValueError(f'resblock type {resblock_type} unrecognized.')
+
+    # Downsampling block
+
+    channels = config.num_channels
+    if progressive_input != 'none':
+      input_pyramid_ch = channels
+
+    modules.append(conv3x3(channels, nf))
+    hs_c = [nf]
+
+    in_ch = nf
+    for i_level in range(num_resolutions):
+      # Residual blocks for this resolution
+      for i_block in range(num_res_blocks):
+        out_ch = nf * ch_mult[i_level]
+        modules.append(ResnetBlock(in_ch=in_ch, out_ch=out_ch))
+        in_ch = out_ch
+
+        if all_resolutions[i_level] in attn_resolutions:
+          modules.append(AttnBlock(channels=in_ch))
+        hs_c.append(in_ch)
+
+      if i_level != num_resolutions - 1:
+        if resblock_type == 'ddpm':
+          modules.append(Downsample(in_ch=in_ch))
+        else:
+          modules.append(ResnetBlock(down=True, in_ch=in_ch))
+
+        if progressive_input == 'input_skip':
+          modules.append(combiner(dim1=input_pyramid_ch, dim2=in_ch))
+          if combine_method == 'cat':
+            in_ch *= 2
+
+        elif progressive_input == 'residual':
+          modules.append(pyramid_downsample(in_ch=input_pyramid_ch, out_ch=in_ch))
+          input_pyramid_ch = in_ch
+
+        hs_c.append(in_ch)
+
+    in_ch = hs_c[-1]
+    modules.append(ResnetBlock(in_ch=in_ch))
+    modules.append(AttnBlock(channels=in_ch))
+    modules.append(ResnetBlock(in_ch=in_ch))
+
+    pyramid_ch = 0
+    # Upsampling block
+    for i_level in reversed(range(num_resolutions)):
+      for i_block in range(num_res_blocks + 1):
+        out_ch = nf * ch_mult[i_level]
+        modules.append(ResnetBlock(in_ch=in_ch + hs_c.pop(),
+                                   out_ch=out_ch))
+        in_ch = out_ch
+
+      if all_resolutions[i_level] in attn_resolutions:
+        modules.append(AttnBlock(channels=in_ch))
+
+      if progressive != 'none':
+        if i_level == num_resolutions - 1:
+          if progressive == 'output_skip':
+            modules.append(nn.GroupNorm(num_groups=min(in_ch // 4, 32),
+                                        num_channels=in_ch, eps=1e-6))
+            modules.append(conv3x3(in_ch, channels, init_scale=init_scale))
+            pyramid_ch = channels
+          elif progressive == 'residual':
+            modules.append(nn.GroupNorm(num_groups=min(in_ch // 4, 32),
+                                        num_channels=in_ch, eps=1e-6))
+            modules.append(conv3x3(in_ch, in_ch, bias=True))
+            pyramid_ch = in_ch
+          else:
+            raise ValueError(f'{progressive} is not a valid name.')
+        else:
+          if progressive == 'output_skip':
+            modules.append(nn.GroupNorm(num_groups=min(in_ch // 4, 32),
+                                        num_channels=in_ch, eps=1e-6))
+            modules.append(conv3x3(in_ch, channels, bias=True, init_scale=init_scale))
+            pyramid_ch = channels
+          elif progressive == 'residual':
+            modules.append(pyramid_upsample(in_ch=pyramid_ch, out_ch=in_ch))
+            pyramid_ch = in_ch
+          else:
+            raise ValueError(f'{progressive} is not a valid name')
+
+      if i_level != 0:
+        if resblock_type == 'ddpm':
+          modules.append(Upsample(in_ch=in_ch))
+        else:
+          modules.append(ResnetBlock(in_ch=in_ch, up=True))
+
+    assert not hs_c
+
+    if progressive != 'output_skip':
+      modules.append(nn.GroupNorm(num_groups=min(in_ch // 4, 32),
+                                  num_channels=in_ch, eps=1e-6))
+      modules.append(conv3x3(in_ch, channels, init_scale=init_scale))
+
+    self.all_modules = nn.ModuleList(modules)
+    
+    
+    mapping_layers = [PixelNorm(),
+                      dense(config.nz, z_emb_dim),
+                      self.act,]
+    for _ in range(config.n_mlp):
+        mapping_layers.append(dense(z_emb_dim, z_emb_dim))
+        mapping_layers.append(self.act)
+    self.z_transform = nn.Sequential(*mapping_layers)
+    
+
+  def forward(self, x, time_cond, z):
+    # timestep/noise_level embedding; only for continuous training
+    zemb = self.z_transform(z)
+    modules = self.all_modules
+    m_idx = 0
+    if self.embedding_type == 'fourier':
+      # Gaussian Fourier features embeddings.
+      used_sigmas = time_cond
+      temb = modules[m_idx](torch.log(used_sigmas))
+      m_idx += 1
+
+    elif self.embedding_type == 'positional':
+      # Sinusoidal positional embeddings.
+      timesteps = time_cond
+     
+      temb = layers.get_timestep_embedding(timesteps, self.nf)
+
+    else:
+      raise ValueError(f'embedding type {self.embedding_type} unknown.')
+
+    if self.conditional:
+      temb = modules[m_idx](temb)
+      m_idx += 1
+      temb = modules[m_idx](self.act(temb))
+      m_idx += 1
+    else:
+      temb = None
+
+    if not self.config.centered:
+      # If input data is in [0, 1]
+      x = 2 * x - 1.
+
+    # Downsampling block
+    input_pyramid = None
+    if self.progressive_input != 'none':
+      input_pyramid = x
+
+    hs = [modules[m_idx](x)]
+    m_idx += 1
+    for i_level in range(self.num_resolutions):
+      # Residual blocks for this resolution
+      for i_block in range(self.num_res_blocks):
+        h = modules[m_idx](hs[-1], temb, zemb)
+        m_idx += 1
+        if h.shape[-1] in self.attn_resolutions:
+          h = modules[m_idx](h)
+          m_idx += 1
+
+        hs.append(h)
+
+      if i_level != self.num_resolutions - 1:
+        if self.resblock_type == 'ddpm':
+          h = modules[m_idx](hs[-1])
+          m_idx += 1
+        else:
+          h = modules[m_idx](hs[-1], temb, zemb)
+          m_idx += 1
+
+        if self.progressive_input == 'input_skip':
+          input_pyramid = self.pyramid_downsample(input_pyramid)
+          h = modules[m_idx](input_pyramid, h)
+          m_idx += 1
+
+        elif self.progressive_input == 'residual':
+          input_pyramid = modules[m_idx](input_pyramid)
+          m_idx += 1
+          if self.skip_rescale:
+            input_pyramid = (input_pyramid + h) / np.sqrt(2.)
+          else:
+            input_pyramid = input_pyramid + h
+          h = input_pyramid
+
+        hs.append(h)
+
+    h = hs[-1]
+    h = modules[m_idx](h, temb, zemb)
+    m_idx += 1
+    h = modules[m_idx](h)
+    m_idx += 1
+    h = modules[m_idx](h, temb, zemb)
+    m_idx += 1
+
+    pyramid = None
+
+    # Upsampling block
+    for i_level in reversed(range(self.num_resolutions)):
+      for i_block in range(self.num_res_blocks + 1):
+        h = modules[m_idx](torch.cat([h, hs.pop()], dim=1), temb, zemb)
+        m_idx += 1
+
+      if h.shape[-1] in self.attn_resolutions:
+        h = modules[m_idx](h)
+        m_idx += 1
+
+      if self.progressive != 'none':
+        if i_level == self.num_resolutions - 1:
+          if self.progressive == 'output_skip':
+            pyramid = self.act(modules[m_idx](h))
+            m_idx += 1
+            pyramid = modules[m_idx](pyramid)
+            m_idx += 1
+          elif self.progressive == 'residual':
+            pyramid = self.act(modules[m_idx](h))
+            m_idx += 1
+            pyramid = modules[m_idx](pyramid)
+            m_idx += 1
+          else:
+            raise ValueError(f'{self.progressive} is not a valid name.')
+        else:
+          if self.progressive == 'output_skip':
+            pyramid = self.pyramid_upsample(pyramid)
+            pyramid_h = self.act(modules[m_idx](h))
+            m_idx += 1
+            pyramid_h = modules[m_idx](pyramid_h)
+            m_idx += 1
+            pyramid = pyramid + pyramid_h
+          elif self.progressive == 'residual':
+            pyramid = modules[m_idx](pyramid)
+            m_idx += 1
+            if self.skip_rescale:
+              pyramid = (pyramid + h) / np.sqrt(2.)
+            else:
+              pyramid = pyramid + h
+            h = pyramid
+          else:
+            raise ValueError(f'{self.progressive} is not a valid name')
+
+      if i_level != 0:
+        if self.resblock_type == 'ddpm':
+          h = modules[m_idx](h)
+          m_idx += 1
+        else:
+          h = modules[m_idx](h, temb, zemb)
+          m_idx += 1
+
+    assert not hs
+
+    if self.progressive == 'output_skip':
+      h = pyramid
+    else:
+      h = self.act(modules[m_idx](h))
+      m_idx += 1
+      h = modules[m_idx](h)
+      m_idx += 1
+
+    assert m_idx == len(modules)
+    
+    if not self.not_use_tanh:
+
+        return torch.tanh(h)
+    else:
+        return h

score_sde/models/up_or_down_sampling.py

@@ -0,0 +1,262 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# ---------------------------------------------------------------
+
+
+"""Layers used for up-sampling or down-sampling images.
+
+Many functions are ported from https://github.com/NVlabs/stylegan2.
+"""
+
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+from score_sde.op import upfirdn2d
+
+
+# Function ported from StyleGAN2
+def get_weight(module,
+               shape,
+               weight_var='weight',
+               kernel_init=None):
+  """Get/create weight tensor for a convolution or fully-connected layer."""
+
+  return module.param(weight_var, kernel_init, shape)
+
+
+class Conv2d(nn.Module):
+  """Conv2d layer with optimal upsampling and downsampling (StyleGAN2)."""
+
+  def __init__(self, in_ch, out_ch, kernel, up=False, down=False,
+               resample_kernel=(1, 3, 3, 1),
+               use_bias=True,
+               kernel_init=None):
+    super().__init__()
+    assert not (up and down)
+    assert kernel >= 1 and kernel % 2 == 1
+    self.weight = nn.Parameter(torch.zeros(out_ch, in_ch, kernel, kernel))
+    if kernel_init is not None:
+      self.weight.data = kernel_init(self.weight.data.shape)
+    if use_bias:
+      self.bias = nn.Parameter(torch.zeros(out_ch))
+
+    self.up = up
+    self.down = down
+    self.resample_kernel = resample_kernel
+    self.kernel = kernel
+    self.use_bias = use_bias
+
+  def forward(self, x):
+    if self.up:
+      x = upsample_conv_2d(x, self.weight, k=self.resample_kernel)
+    elif self.down:
+      x = conv_downsample_2d(x, self.weight, k=self.resample_kernel)
+    else:
+      x = F.conv2d(x, self.weight, stride=1, padding=self.kernel // 2)
+
+    if self.use_bias:
+      x = x + self.bias.reshape(1, -1, 1, 1)
+
+    return x
+
+
+def naive_upsample_2d(x, factor=2):
+  _N, C, H, W = x.shape
+  x = torch.reshape(x, (-1, C, H, 1, W, 1))
+  x = x.repeat(1, 1, 1, factor, 1, factor)
+  return torch.reshape(x, (-1, C, H * factor, W * factor))
+
+
+def naive_downsample_2d(x, factor=2):
+  _N, C, H, W = x.shape
+  x = torch.reshape(x, (-1, C, H // factor, factor, W // factor, factor))
+  return torch.mean(x, dim=(3, 5))
+
+
+def upsample_conv_2d(x, w, k=None, factor=2, gain=1):
+  """Fused `upsample_2d()` followed by `tf.nn.conv2d()`.
+
+     Padding is performed only once at the beginning, not between the
+     operations.
+     The fused op is considerably more efficient than performing the same
+     calculation
+     using standard TensorFlow ops. It supports gradients of arbitrary order.
+     Args:
+       x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W,
+         C]`.
+       w:            Weight tensor of the shape `[filterH, filterW, inChannels,
+         outChannels]`. Grouped convolution can be performed by `inChannels =
+         x.shape[0] // numGroups`.
+       k:            FIR filter of the shape `[firH, firW]` or `[firN]`
+         (separable). The default is `[1] * factor`, which corresponds to
+         nearest-neighbor upsampling.
+       factor:       Integer upsampling factor (default: 2).
+       gain:         Scaling factor for signal magnitude (default: 1.0).
+
+     Returns:
+       Tensor of the shape `[N, C, H * factor, W * factor]` or
+       `[N, H * factor, W * factor, C]`, and same datatype as `x`.
+  """
+
+  assert isinstance(factor, int) and factor >= 1
+
+  # Check weight shape.
+  assert len(w.shape) == 4
+  convH = w.shape[2]
+  convW = w.shape[3]
+  inC = w.shape[1]
+  outC = w.shape[0]
+
+  assert convW == convH
+
+  # Setup filter kernel.
+  if k is None:
+    k = [1] * factor
+  k = _setup_kernel(k) * (gain * (factor ** 2))
+  p = (k.shape[0] - factor) - (convW - 1)
+
+  stride = (factor, factor)
+
+  # Determine data dimensions.
+  stride = [1, 1, factor, factor]
+  output_shape = ((_shape(x, 2) - 1) * factor + convH, (_shape(x, 3) - 1) * factor + convW)
+  output_padding = (output_shape[0] - (_shape(x, 2) - 1) * stride[0] - convH,
+                    output_shape[1] - (_shape(x, 3) - 1) * stride[1] - convW)
+  assert output_padding[0] >= 0 and output_padding[1] >= 0
+  num_groups = _shape(x, 1) // inC
+
+  # Transpose weights.
+  w = torch.reshape(w, (num_groups, -1, inC, convH, convW))
+  w = w[..., ::-1, ::-1].permute(0, 2, 1, 3, 4)
+  w = torch.reshape(w, (num_groups * inC, -1, convH, convW))
+
+  x = F.conv_transpose2d(x, w, stride=stride, output_padding=output_padding, padding=0)
+  ## Original TF code.
+  # x = tf.nn.conv2d_transpose(
+  #     x,
+  #     w,
+  #     output_shape=output_shape,
+  #     strides=stride,
+  #     padding='VALID',
+  #     data_format=data_format)
+  ## JAX equivalent
+
+  return upfirdn2d(x, torch.tensor(k, device=x.device),
+                   pad=((p + 1) // 2 + factor - 1, p // 2 + 1))
+
+
+def conv_downsample_2d(x, w, k=None, factor=2, gain=1):
+  """Fused `tf.nn.conv2d()` followed by `downsample_2d()`.
+
+    Padding is performed only once at the beginning, not between the operations.
+    The fused op is considerably more efficient than performing the same
+    calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W,
+          C]`.
+        w:            Weight tensor of the shape `[filterH, filterW, inChannels,
+          outChannels]`. Grouped convolution can be performed by `inChannels =
+          x.shape[0] // numGroups`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to
+          average pooling.
+        factor:       Integer downsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        Tensor of the shape `[N, C, H // factor, W // factor]` or
+        `[N, H // factor, W // factor, C]`, and same datatype as `x`.
+  """
+
+  assert isinstance(factor, int) and factor >= 1
+  _outC, _inC, convH, convW = w.shape
+  assert convW == convH
+  if k is None:
+    k = [1] * factor
+  k = _setup_kernel(k) * gain
+  p = (k.shape[0] - factor) + (convW - 1)
+  s = [factor, factor]
+  x = upfirdn2d(x, torch.tensor(k, device=x.device),
+                pad=((p + 1) // 2, p // 2))
+  return F.conv2d(x, w, stride=s, padding=0)
+
+
+def _setup_kernel(k):
+  k = np.asarray(k, dtype=np.float32)
+  if k.ndim == 1:
+    k = np.outer(k, k)
+  k /= np.sum(k)
+  assert k.ndim == 2
+  assert k.shape[0] == k.shape[1]
+  return k
+
+
+def _shape(x, dim):
+  return x.shape[dim]
+
+
+def upsample_2d(x, k=None, factor=2, gain=1):
+  r"""Upsample a batch of 2D images with the given filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and upsamples each image with the given filter. The filter is normalized so
+    that
+    if the input pixels are constant, they will be scaled by the specified
+    `gain`.
+    Pixels outside the image are assumed to be zero, and the filter is padded
+    with
+    zeros so that its shape is a multiple of the upsampling factor.
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W,
+          C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to
+          nearest-neighbor upsampling.
+        factor:       Integer upsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        Tensor of the shape `[N, C, H * factor, W * factor]`
+  """
+  assert isinstance(factor, int) and factor >= 1
+  if k is None:
+    k = [1] * factor
+  k = _setup_kernel(k) * (gain * (factor ** 2))
+  p = k.shape[0] - factor
+  return upfirdn2d(x, torch.tensor(k, device=x.device),
+                   up=factor, pad=((p + 1) // 2 + factor - 1, p // 2))
+
+
+def downsample_2d(x, k=None, factor=2, gain=1):
+  r"""Downsample a batch of 2D images with the given filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and downsamples each image with the given filter. The filter is normalized
+    so that
+    if the input pixels are constant, they will be scaled by the specified
+    `gain`.
+    Pixels outside the image are assumed to be zero, and the filter is padded
+    with
+    zeros so that its shape is a multiple of the downsampling factor.
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W,
+          C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to
+          average pooling.
+        factor:       Integer downsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        Tensor of the shape `[N, C, H // factor, W // factor]`
+  """
+
+  assert isinstance(factor, int) and factor >= 1
+  if k is None:
+    k = [1] * factor
+  k = _setup_kernel(k) * gain
+  p = k.shape[0] - factor
+  return upfirdn2d(x, torch.tensor(k, device=x.device),
+                   down=factor, pad=((p + 1) // 2, p // 2))

score_sde/models/utils.py

@@ -0,0 +1,148 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from a file in the Score SDE library
+# which was released under the Apache License.
+#
+# Source:
+# https://github.com/yang-song/score_sde_pytorch/blob/main/models/utils.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_Apache). The modifications
+# to this file are subject to the same Apache License.
+# ---------------------------------------------------------------
+
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import numpy as np
+
+
+_MODELS = {}
+
+
+def register_model(cls=None, *, name=None):
+  """A decorator for registering model classes."""
+
+  def _register(cls):
+    if name is None:
+      local_name = cls.__name__
+    else:
+      local_name = name
+    if local_name in _MODELS:
+      raise ValueError(f'Already registered model with name: {local_name}')
+    _MODELS[local_name] = cls
+    return cls
+
+  if cls is None:
+    return _register
+  else:
+    return _register(cls)
+
+
+def get_model(name):
+  return _MODELS[name]
+
+
+def get_sigmas(config):
+  """Get sigmas --- the set of noise levels for SMLD from config files.
+  Args:
+    config: A ConfigDict object parsed from the config file
+  Returns:
+    sigmas: a jax numpy arrary of noise levels
+  """
+  sigmas = np.exp(
+    np.linspace(np.log(config.model.sigma_max), np.log(config.model.sigma_min), config.model.num_scales))
+
+  return sigmas
+
+
+def get_ddpm_params(config):
+  """Get betas and alphas --- parameters used in the original DDPM paper."""
+  num_diffusion_timesteps = 1000
+  # parameters need to be adapted if number of time steps differs from 1000
+  beta_start = config.model.beta_min / config.model.num_scales
+  beta_end = config.model.beta_max / config.model.num_scales
+  betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+
+  alphas = 1. - betas
+  alphas_cumprod = np.cumprod(alphas, axis=0)
+  sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
+  sqrt_1m_alphas_cumprod = np.sqrt(1. - alphas_cumprod)
+
+  return {
+    'betas': betas,
+    'alphas': alphas,
+    'alphas_cumprod': alphas_cumprod,
+    'sqrt_alphas_cumprod': sqrt_alphas_cumprod,
+    'sqrt_1m_alphas_cumprod': sqrt_1m_alphas_cumprod,
+    'beta_min': beta_start * (num_diffusion_timesteps - 1),
+    'beta_max': beta_end * (num_diffusion_timesteps - 1),
+    'num_diffusion_timesteps': num_diffusion_timesteps
+  }
+
+
+def create_model(config):
+  """Create the score model."""
+  model_name = config.model.name
+  score_model = get_model(model_name)(config)
+  score_model = score_model.to(config.device)
+  score_model = torch.nn.DataParallel(score_model)
+  return score_model
+
+
+def get_model_fn(model, train=False):
+  """Create a function to give the output of the score-based model.
+
+  Args:
+    model: The score model.
+    train: `True` for training and `False` for evaluation.
+
+  Returns:
+    A model function.
+  """
+
+  def model_fn(x, labels):
+    """Compute the output of the score-based model.
+
+    Args:
+      x: A mini-batch of input data.
+      labels: A mini-batch of conditioning variables for time steps. Should be interpreted differently
+        for different models.
+
+    Returns:
+      A tuple of (model output, new mutable states)
+    """
+    if not train:
+      model.eval()
+      return model(x, labels)
+    else:
+      model.train()
+      return model(x, labels)
+
+  return model_fn
+
+
+
+
+def to_flattened_numpy(x):
+  """Flatten a torch tensor `x` and convert it to numpy."""
+  return x.detach().cpu().numpy().reshape((-1,))
+
+
+def from_flattened_numpy(x, shape):
+  """Form a torch tensor with the given `shape` from a flattened numpy array `x`."""
+  return torch.from_numpy(x.reshape(shape))

score_sde/op/LICENSE_MIT

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Kim Seonghyeon
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

score_sde/op/__init__.py

@@ -0,0 +1,2 @@
+from .fused_act import FusedLeakyReLU, fused_leaky_relu
+from .upfirdn2d import upfirdn2d

score_sde/op/fused_act.py

@@ -0,0 +1,105 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# ---------------------------------------------------------------
+
+""" Originated from https://github.com/rosinality/stylegan2-pytorch
+The license for the original version of this file can be found in this directory (LICENSE_MIT).
+"""
+
+import os
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+
+module_path = os.path.dirname(__file__)
+fused = load(
+    "fused",
+    sources=[
+        os.path.join(module_path, "fused_bias_act.cpp"),
+        os.path.join(module_path, "fused_bias_act_kernel.cu"),
+    ],
+)
+
+
+class FusedLeakyReLUFunctionBackward(Function):
+    @staticmethod
+    def forward(ctx, grad_output, out, negative_slope, scale):
+        ctx.save_for_backward(out)
+        ctx.negative_slope = negative_slope
+        ctx.scale = scale
+
+        empty = grad_output.new_empty(0)
+
+        grad_input = fused.fused_bias_act(
+            grad_output, empty, out, 3, 1, negative_slope, scale
+        )
+
+        dim = [0]
+
+        if grad_input.ndim > 2:
+            dim += list(range(2, grad_input.ndim))
+
+        grad_bias = grad_input.sum(dim).detach()
+
+        return grad_input, grad_bias
+
+    @staticmethod
+    def backward(ctx, gradgrad_input, gradgrad_bias):
+        out, = ctx.saved_tensors
+        gradgrad_out = fused.fused_bias_act(
+            gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
+        )
+
+        return gradgrad_out, None, None, None
+
+
+class FusedLeakyReLUFunction(Function):
+    @staticmethod
+    def forward(ctx, input, bias, negative_slope, scale):
+        empty = input.new_empty(0)
+        out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
+        ctx.save_for_backward(out)
+        ctx.negative_slope = negative_slope
+        ctx.scale = scale
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        out, = ctx.saved_tensors
+
+        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
+            grad_output, out, ctx.negative_slope, ctx.scale
+        )
+
+        return grad_input, grad_bias, None, None
+
+
+class FusedLeakyReLU(nn.Module):
+    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
+        super().__init__()
+
+        self.bias = nn.Parameter(torch.zeros(channel))
+        self.negative_slope = negative_slope
+        self.scale = scale
+
+    def forward(self, input):
+        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
+
+
+def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
+    if input.device.type == "cpu":
+        rest_dim = [1] * (input.ndim - bias.ndim - 1)
+        return (
+            F.leaky_relu(
+                input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
+            )
+            * scale
+        )
+
+    else:
+        return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)

score_sde/op/fused_bias_act.cpp

@@ -0,0 +1,28 @@
+// ---------------------------------------------------------------
+// Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+// ---------------------------------------------------------------
+
+// Originated from https://github.com/rosinality/stylegan2-pytorch
+// The license for the original version of this file can be found in this directory (LICENSE_MIT).
+
+#include <torch/extension.h>
+
+
+torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+    int act, int grad, float alpha, float scale);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+    int act, int grad, float alpha, float scale) {
+    CHECK_CUDA(input);
+    CHECK_CUDA(bias);
+
+    return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
+}

score_sde/op/fused_bias_act_kernel.cu

@@ -0,0 +1,101 @@
+// ---------------------------------------------------------------
+// Copyright (c) 2019-2022, NVIDIA Corporation. All rights reserved.
+// ---------------------------------------------------------------
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include <torch/types.h>
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAApplyUtils.cuh>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+
+template <typename scalar_t>
+static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,
+    int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) {
+    int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;
+
+    scalar_t zero = 0.0;
+
+    for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {
+        scalar_t x = p_x[xi];
+
+        if (use_bias) {
+            x += p_b[(xi / step_b) % size_b];
+        }
+
+        scalar_t ref = use_ref ? p_ref[xi] : zero;
+
+        scalar_t y;
+
+        switch (act * 10 + grad) {
+            default:
+            case 10: y = x; break;
+            case 11: y = x; break;
+            case 12: y = 0.0; break;
+
+            case 30: y = (x > 0.0) ? x : x * alpha; break;
+            case 31: y = (ref > 0.0) ? x : x * alpha; break;
+            case 32: y = 0.0; break;
+        }
+
+        out[xi] = y * scale;
+    }
+}
+
+
+torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+    int act, int grad, float alpha, float scale) {
+    int curDevice = -1;
+    cudaGetDevice(&curDevice);
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
+
+    auto x = input.contiguous();
+    auto b = bias.contiguous();
+    auto ref = refer.contiguous();
+
+    int use_bias = b.numel() ? 1 : 0;
+    int use_ref = ref.numel() ? 1 : 0;
+
+    int size_x = x.numel();
+    int size_b = b.numel();
+    int step_b = 1;
+
+    for (int i = 1 + 1; i < x.dim(); i++) {
+        step_b *= x.size(i);
+    }
+
+    int loop_x = 4;
+    int block_size = 4 * 32;
+    int grid_size = (size_x - 1) / (loop_x * block_size) + 1;
+
+    auto y = torch::empty_like(x);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] {
+        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(
+            y.data_ptr<scalar_t>(),
+            x.data_ptr<scalar_t>(),
+            b.data_ptr<scalar_t>(),
+            ref.data_ptr<scalar_t>(),
+            act,
+            grad,
+            alpha,
+            scale,
+            loop_x,
+            size_x,
+            step_b,
+            size_b,
+            use_bias,
+            use_ref
+        );
+    });
+
+    return y;
+}

score_sde/op/upfirdn2d.cpp

@@ -0,0 +1,31 @@
+// ---------------------------------------------------------------
+// Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+// ---------------------------------------------------------------
+
+// Originated from https://github.com/rosinality/stylegan2-pytorch
+// The license for the original version of this file can be found in this directory (LICENSE_MIT).
+
+
+#include <torch/extension.h>
+
+
+torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
+                            int up_x, int up_y, int down_x, int down_y,
+                            int pad_x0, int pad_x1, int pad_y0, int pad_y1);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel,
+                        int up_x, int up_y, int down_x, int down_y,
+                        int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
+    CHECK_CUDA(input);
+    CHECK_CUDA(kernel);
+
+    return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
+}

score_sde/op/upfirdn2d.py

@@ -0,0 +1,225 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# ---------------------------------------------------------------
+
+""" Originated from https://github.com/rosinality/stylegan2-pytorch
+The license for the original version of this file can be found in this directory (LICENSE_MIT).
+"""
+
+import os
+
+import torch
+from torch.nn import functional as F
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+from collections import abc
+
+module_path = os.path.dirname(__file__)
+upfirdn2d_op = load(
+    "upfirdn2d",
+    sources=[
+        os.path.join(module_path, "upfirdn2d.cpp"),
+        os.path.join(module_path, "upfirdn2d_kernel.cu"),
+    ],
+)
+
+
+class UpFirDn2dBackward(Function):
+    @staticmethod
+    def forward(
+        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
+    ):
+
+        up_x, up_y = up
+        down_x, down_y = down
+        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
+
+        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
+
+        grad_input = upfirdn2d_op.upfirdn2d(
+            grad_output,
+            grad_kernel,
+            down_x,
+            down_y,
+            up_x,
+            up_y,
+            g_pad_x0,
+            g_pad_x1,
+            g_pad_y0,
+            g_pad_y1,
+        )
+        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
+
+        ctx.save_for_backward(kernel)
+
+        pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+        ctx.up_x = up_x
+        ctx.up_y = up_y
+        ctx.down_x = down_x
+        ctx.down_y = down_y
+        ctx.pad_x0 = pad_x0
+        ctx.pad_x1 = pad_x1
+        ctx.pad_y0 = pad_y0
+        ctx.pad_y1 = pad_y1
+        ctx.in_size = in_size
+        ctx.out_size = out_size
+
+        return grad_input
+
+    @staticmethod
+    def backward(ctx, gradgrad_input):
+        kernel, = ctx.saved_tensors
+
+        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
+
+        gradgrad_out = upfirdn2d_op.upfirdn2d(
+            gradgrad_input,
+            kernel,
+            ctx.up_x,
+            ctx.up_y,
+            ctx.down_x,
+            ctx.down_y,
+            ctx.pad_x0,
+            ctx.pad_x1,
+            ctx.pad_y0,
+            ctx.pad_y1,
+        )
+        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
+        gradgrad_out = gradgrad_out.view(
+            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
+        )
+
+        return gradgrad_out, None, None, None, None, None, None, None, None
+
+
+class UpFirDn2d(Function):
+    @staticmethod
+    def forward(ctx, input, kernel, up, down, pad):
+        up_x, up_y = up
+        down_x, down_y = down
+        pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+        kernel_h, kernel_w = kernel.shape
+        batch, channel, in_h, in_w = input.shape
+        ctx.in_size = input.shape
+
+        input = input.reshape(-1, in_h, in_w, 1)
+
+        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
+
+        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+        ctx.out_size = (out_h, out_w)
+
+        ctx.up = (up_x, up_y)
+        ctx.down = (down_x, down_y)
+        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
+
+        g_pad_x0 = kernel_w - pad_x0 - 1
+        g_pad_y0 = kernel_h - pad_y0 - 1
+        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
+        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
+
+        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
+
+        out = upfirdn2d_op.upfirdn2d(
+            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+        )
+        # out = out.view(major, out_h, out_w, minor)
+        out = out.view(-1, channel, out_h, out_w)
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        kernel, grad_kernel = ctx.saved_tensors
+
+        grad_input = UpFirDn2dBackward.apply(
+            grad_output,
+            kernel,
+            grad_kernel,
+            ctx.up,
+            ctx.down,
+            ctx.pad,
+            ctx.g_pad,
+            ctx.in_size,
+            ctx.out_size,
+        )
+
+        return grad_input, None, None, None, None
+
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+    if input.device.type == "cpu":
+        out = upfirdn2d_native(
+            input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]
+        )
+
+    else:
+        out = UpFirDn2d.apply(
+            input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
+        )
+
+    return out
+
+def upfirdn2d_ada(input, kernel, up=1, down=1, pad=(0, 0)):
+    if not isinstance(up, abc.Iterable):
+        up = (up, up)
+
+    if not isinstance(down, abc.Iterable):
+        down = (down, down)
+
+    if len(pad) == 2:
+        pad = (pad[0], pad[1], pad[0], pad[1])
+
+    if input.device.type == "cpu":
+        out = upfirdn2d_native(input, kernel, *up, *down, *pad)
+
+    else:
+        out = UpFirDn2d.apply(input, kernel, up, down, pad)
+
+    return out
+
+def upfirdn2d_native(
+    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+):
+    _, channel, in_h, in_w = input.shape
+    input = input.reshape(-1, in_h, in_w, 1)
+
+    _, in_h, in_w, minor = input.shape
+    kernel_h, kernel_w = kernel.shape
+
+    out = input.view(-1, in_h, 1, in_w, 1, minor)
+    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+    out = F.pad(
+        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
+    )
+    out = out[
+        :,
+        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+        :,
+    ]
+
+    out = out.permute(0, 3, 1, 2)
+    out = out.reshape(
+        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
+    )
+    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+    out = F.conv2d(out, w)
+    out = out.reshape(
+        -1,
+        minor,
+        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+    )
+    out = out.permute(0, 2, 3, 1)
+    out = out[:, ::down_y, ::down_x, :]
+
+    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+
+    return out.view(-1, channel, out_h, out_w)

score_sde/op/upfirdn2d_kernel.cu

@@ -0,0 +1,371 @@
+// ---------------------------------------------------------------
+// Copyright (c) 2019-2022, NVIDIA Corporation. All rights reserved.
+// ---------------------------------------------------------------
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include <torch/types.h>
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAApplyUtils.cuh>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
+  int c = a / b;
+
+  if (c * b > a) {
+    c--;
+  }
+
+  return c;
+}
+
+struct UpFirDn2DKernelParams {
+  int up_x;
+  int up_y;
+  int down_x;
+  int down_y;
+  int pad_x0;
+  int pad_x1;
+  int pad_y0;
+  int pad_y1;
+
+  int major_dim;
+  int in_h;
+  int in_w;
+  int minor_dim;
+  int kernel_h;
+  int kernel_w;
+  int out_h;
+  int out_w;
+  int loop_major;
+  int loop_x;
+};
+
+template <typename scalar_t>
+__global__ void upfirdn2d_kernel_large(scalar_t *out, const scalar_t *input,
+                                       const scalar_t *kernel,
+                                       const UpFirDn2DKernelParams p) {
+  int minor_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int out_y = minor_idx / p.minor_dim;
+  minor_idx -= out_y * p.minor_dim;
+  int out_x_base = blockIdx.y * p.loop_x * blockDim.y + threadIdx.y;
+  int major_idx_base = blockIdx.z * p.loop_major;
+
+  if (out_x_base >= p.out_w || out_y >= p.out_h ||
+      major_idx_base >= p.major_dim) {
+    return;
+  }
+
+  int mid_y = out_y * p.down_y + p.up_y - 1 - p.pad_y0;
+  int in_y = min(max(floor_div(mid_y, p.up_y), 0), p.in_h);
+  int h = min(max(floor_div(mid_y + p.kernel_h, p.up_y), 0), p.in_h) - in_y;
+  int kernel_y = mid_y + p.kernel_h - (in_y + 1) * p.up_y;
+
+  for (int loop_major = 0, major_idx = major_idx_base;
+       loop_major < p.loop_major && major_idx < p.major_dim;
+       loop_major++, major_idx++) {
+    for (int loop_x = 0, out_x = out_x_base;
+         loop_x < p.loop_x && out_x < p.out_w; loop_x++, out_x += blockDim.y) {
+      int mid_x = out_x * p.down_x + p.up_x - 1 - p.pad_x0;
+      int in_x = min(max(floor_div(mid_x, p.up_x), 0), p.in_w);
+      int w = min(max(floor_div(mid_x + p.kernel_w, p.up_x), 0), p.in_w) - in_x;
+      int kernel_x = mid_x + p.kernel_w - (in_x + 1) * p.up_x;
+
+      const scalar_t *x_p =
+          &input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim +
+                 minor_idx];
+      const scalar_t *k_p = &kernel[kernel_y * p.kernel_w + kernel_x];
+      int x_px = p.minor_dim;
+      int k_px = -p.up_x;
+      int x_py = p.in_w * p.minor_dim;
+      int k_py = -p.up_y * p.kernel_w;
+
+      scalar_t v = 0.0f;
+
+      for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+          v += static_cast<scalar_t>(*x_p) * static_cast<scalar_t>(*k_p);
+          x_p += x_px;
+          k_p += k_px;
+        }
+
+        x_p += x_py - w * x_px;
+        k_p += k_py - w * k_px;
+      }
+
+      out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
+          minor_idx] = v;
+    }
+  }
+}
+
+template <typename scalar_t, int up_x, int up_y, int down_x, int down_y,
+          int kernel_h, int kernel_w, int tile_out_h, int tile_out_w>
+__global__ void upfirdn2d_kernel(scalar_t *out, const scalar_t *input,
+                                 const scalar_t *kernel,
+                                 const UpFirDn2DKernelParams p) {
+  const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
+  const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;
+
+  __shared__ volatile float sk[kernel_h][kernel_w];
+  __shared__ volatile float sx[tile_in_h][tile_in_w];
+
+  int minor_idx = blockIdx.x;
+  int tile_out_y = minor_idx / p.minor_dim;
+  minor_idx -= tile_out_y * p.minor_dim;
+  tile_out_y *= tile_out_h;
+  int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
+  int major_idx_base = blockIdx.z * p.loop_major;
+
+  if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h |
+      major_idx_base >= p.major_dim) {
+    return;
+  }
+
+  for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w;
+       tap_idx += blockDim.x) {
+    int ky = tap_idx / kernel_w;
+    int kx = tap_idx - ky * kernel_w;
+    scalar_t v = 0.0;
+
+    if (kx < p.kernel_w & ky < p.kernel_h) {
+      v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
+    }
+
+    sk[ky][kx] = v;
+  }
+
+  for (int loop_major = 0, major_idx = major_idx_base;
+       loop_major < p.loop_major & major_idx < p.major_dim;
+       loop_major++, major_idx++) {
+    for (int loop_x = 0, tile_out_x = tile_out_x_base;
+         loop_x < p.loop_x & tile_out_x < p.out_w;
+         loop_x++, tile_out_x += tile_out_w) {
+      int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
+      int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
+      int tile_in_x = floor_div(tile_mid_x, up_x);
+      int tile_in_y = floor_div(tile_mid_y, up_y);
+
+      __syncthreads();
+
+      for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w;
+           in_idx += blockDim.x) {
+        int rel_in_y = in_idx / tile_in_w;
+        int rel_in_x = in_idx - rel_in_y * tile_in_w;
+        int in_x = rel_in_x + tile_in_x;
+        int in_y = rel_in_y + tile_in_y;
+
+        scalar_t v = 0.0;
+
+        if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
+          v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) *
+                        p.minor_dim +
+                    minor_idx];
+        }
+
+        sx[rel_in_y][rel_in_x] = v;
+      }
+
+      __syncthreads();
+      for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w;
+           out_idx += blockDim.x) {
+        int rel_out_y = out_idx / tile_out_w;
+        int rel_out_x = out_idx - rel_out_y * tile_out_w;
+        int out_x = rel_out_x + tile_out_x;
+        int out_y = rel_out_y + tile_out_y;
+
+        int mid_x = tile_mid_x + rel_out_x * down_x;
+        int mid_y = tile_mid_y + rel_out_y * down_y;
+        int in_x = floor_div(mid_x, up_x);
+        int in_y = floor_div(mid_y, up_y);
+        int rel_in_x = in_x - tile_in_x;
+        int rel_in_y = in_y - tile_in_y;
+        int kernel_x = (in_x + 1) * up_x - mid_x - 1;
+        int kernel_y = (in_y + 1) * up_y - mid_y - 1;
+
+        scalar_t v = 0.0;
+
+#pragma unroll
+        for (int y = 0; y < kernel_h / up_y; y++)
+#pragma unroll
+          for (int x = 0; x < kernel_w / up_x; x++)
+            v += sx[rel_in_y + y][rel_in_x + x] *
+                 sk[kernel_y + y * up_y][kernel_x + x * up_x];
+
+        if (out_x < p.out_w & out_y < p.out_h) {
+          out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
+              minor_idx] = v;
+        }
+      }
+    }
+  }
+}
+
+torch::Tensor upfirdn2d_op(const torch::Tensor &input,
+                           const torch::Tensor &kernel, int up_x, int up_y,
+                           int down_x, int down_y, int pad_x0, int pad_x1,
+                           int pad_y0, int pad_y1) {
+  int curDevice = -1;
+  cudaGetDevice(&curDevice);
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
+
+  UpFirDn2DKernelParams p;
+
+  auto x = input.contiguous();
+  auto k = kernel.contiguous();
+
+  p.major_dim = x.size(0);
+  p.in_h = x.size(1);
+  p.in_w = x.size(2);
+  p.minor_dim = x.size(3);
+  p.kernel_h = k.size(0);
+  p.kernel_w = k.size(1);
+  p.up_x = up_x;
+  p.up_y = up_y;
+  p.down_x = down_x;
+  p.down_y = down_y;
+  p.pad_x0 = pad_x0;
+  p.pad_x1 = pad_x1;
+  p.pad_y0 = pad_y0;
+  p.pad_y1 = pad_y1;
+
+  p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) /
+            p.down_y;
+  p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) /
+            p.down_x;
+
+  auto out =
+      at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());
+
+  int mode = -1;
+
+  int tile_out_h = -1;
+  int tile_out_w = -1;
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 4 && p.kernel_w <= 4) {
+    mode = 1;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 3 && p.kernel_w <= 3) {
+    mode = 2;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 4 && p.kernel_w <= 4) {
+    mode = 3;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 2 && p.kernel_w <= 2) {
+    mode = 4;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
+      p.kernel_h <= 4 && p.kernel_w <= 4) {
+    mode = 5;
+    tile_out_h = 8;
+    tile_out_w = 32;
+  }
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
+      p.kernel_h <= 2 && p.kernel_w <= 2) {
+    mode = 6;
+    tile_out_h = 8;
+    tile_out_w = 32;
+  }
+
+  dim3 block_size;
+  dim3 grid_size;
+
+  if (tile_out_h > 0 && tile_out_w > 0) {
+    p.loop_major = (p.major_dim - 1) / 16384 + 1;
+    p.loop_x = 1;
+    block_size = dim3(32 * 8, 1, 1);
+    grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
+                     (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
+                     (p.major_dim - 1) / p.loop_major + 1);
+  } else {
+    p.loop_major = (p.major_dim - 1) / 16384 + 1;
+    p.loop_x = 4;
+    block_size = dim3(4, 32, 1);
+    grid_size = dim3((p.out_h * p.minor_dim - 1) / block_size.x + 1,
+                     (p.out_w - 1) / (p.loop_x * block_size.y) + 1,
+                     (p.major_dim - 1) / p.loop_major + 1);
+  }
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
+    switch (mode) {
+    case 1:
+      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 4, 4, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 2:
+      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 3, 3, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 3:
+      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 4, 4, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 4:
+      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 2, 2, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 5:
+      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 6:
+      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    default:
+      upfirdn2d_kernel_large<scalar_t><<<grid_size, block_size, 0, stream>>>(
+          out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),
+          k.data_ptr<scalar_t>(), p);
+    }
+  });
+
+  return out;
+}

test_ddgan.py

@@ -0,0 +1,275 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
+# ---------------------------------------------------------------
+import argparse
+import torch
+import numpy as np
+
+import os
+
+import torchvision
+from score_sde.models.ncsnpp_generator_adagn import NCSNpp
+from pytorch_fid.fid_score import calculate_fid_given_paths
+
+#%% Diffusion coefficients 
+def var_func_vp(t, beta_min, beta_max):
+    log_mean_coeff = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min
+    var = 1. - torch.exp(2. * log_mean_coeff)
+    return var
+
+def var_func_geometric(t, beta_min, beta_max):
+    return beta_min * ((beta_max / beta_min) ** t)
+
+def extract(input, t, shape):
+    out = torch.gather(input, 0, t)
+    reshape = [shape[0]] + [1] * (len(shape) - 1)
+    out = out.reshape(*reshape)
+
+    return out
+
+def get_time_schedule(args, device):
+    n_timestep = args.num_timesteps
+    eps_small = 1e-3
+    t = np.arange(0, n_timestep + 1, dtype=np.float64)
+    t = t / n_timestep
+    t = torch.from_numpy(t) * (1. - eps_small)  + eps_small
+    return t.to(device)
+
+def get_sigma_schedule(args, device):
+    n_timestep = args.num_timesteps
+    beta_min = args.beta_min
+    beta_max = args.beta_max
+    eps_small = 1e-3
+   
+    t = np.arange(0, n_timestep + 1, dtype=np.float64)
+    t = t / n_timestep
+    t = torch.from_numpy(t) * (1. - eps_small) + eps_small
+    
+    if args.use_geometric:
+        var = var_func_geometric(t, beta_min, beta_max)
+    else:
+        var = var_func_vp(t, beta_min, beta_max)
+    alpha_bars = 1.0 - var
+    betas = 1 - alpha_bars[1:] / alpha_bars[:-1]
+    
+    first = torch.tensor(1e-8)
+    betas = torch.cat((first[None], betas)).to(device)
+    betas = betas.type(torch.float32)
+    sigmas = betas**0.5
+    a_s = torch.sqrt(1-betas)
+    return sigmas, a_s, betas
+
+#%% posterior sampling
+class Posterior_Coefficients():
+    def __init__(self, args, device):
+        
+        _, _, self.betas = get_sigma_schedule(args, device=device)
+        
+        #we don't need the zeros
+        self.betas = self.betas.type(torch.float32)[1:]
+        
+        self.alphas = 1 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, 0)
+        self.alphas_cumprod_prev = torch.cat(
+                                    (torch.tensor([1.], dtype=torch.float32,device=device), self.alphas_cumprod[:-1]), 0
+                                        )               
+        self.posterior_variance = self.betas * (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod)
+        
+        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1 / self.alphas_cumprod - 1)
+        
+        self.posterior_mean_coef1 = (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1 - self.alphas_cumprod))
+        self.posterior_mean_coef2 = ((1 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1 - self.alphas_cumprod))
+        
+        self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min=1e-20))
+        
+def sample_posterior(coefficients, x_0,x_t, t):
+    
+    def q_posterior(x_0, x_t, t):
+        mean = (
+            extract(coefficients.posterior_mean_coef1, t, x_t.shape) * x_0
+            + extract(coefficients.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        var = extract(coefficients.posterior_variance, t, x_t.shape)
+        log_var_clipped = extract(coefficients.posterior_log_variance_clipped, t, x_t.shape)
+        return mean, var, log_var_clipped
+    
+  
+    def p_sample(x_0, x_t, t):
+        mean, _, log_var = q_posterior(x_0, x_t, t)
+        
+        noise = torch.randn_like(x_t)
+        
+        nonzero_mask = (1 - (t == 0).type(torch.float32))
+
+        return mean + nonzero_mask[:,None,None,None] * torch.exp(0.5 * log_var) * noise
+            
+    sample_x_pos = p_sample(x_0, x_t, t)
+    
+    return sample_x_pos
+
+def sample_from_model(coefficients, generator, n_time, x_init, T, opt):
+    x = x_init
+    with torch.no_grad():
+        for i in reversed(range(n_time)):
+            t = torch.full((x.size(0),), i, dtype=torch.int64).to(x.device)
+            
+            t_time = t
+            latent_z = torch.randn(x.size(0), opt.nz, device=x.device)#.to(x.device)
+            x_0 = generator(x, t_time, latent_z)
+            x_new = sample_posterior(coefficients, x_0, x, t)
+            x = x_new.detach()
+        
+    return x
+
+#%%
+def sample_and_test(args):
+    torch.manual_seed(42)
+    device = 'cuda:0'
+    
+    if args.dataset == 'cifar10':
+        real_img_dir = 'pytorch_fid/cifar10_train_stat.npy'
+    elif args.dataset == 'celeba_256':
+        real_img_dir = 'pytorch_fid/celeba_256_stat.npy'
+    elif args.dataset == 'lsun':
+        real_img_dir = 'pytorch_fid/lsun_church_stat.npy'
+    else:
+        real_img_dir = args.real_img_dir
+    
+    to_range_0_1 = lambda x: (x + 1.) / 2.
+
+    
+    netG = NCSNpp(args).to(device)
+    ckpt = torch.load('./saved_info/dd_gan/{}/{}/netG_{}.pth'.format(args.dataset, args.exp, args.epoch_id), map_location=device)
+    
+    #loading weights from ddp in single gpu
+    for key in list(ckpt.keys()):
+        ckpt[key[7:]] = ckpt.pop(key)
+    netG.load_state_dict(ckpt)
+    netG.eval()
+    
+    
+    T = get_time_schedule(args, device)
+    
+    pos_coeff = Posterior_Coefficients(args, device)
+        
+    iters_needed = 50000 //args.batch_size
+    
+    save_dir = "./generated_samples/{}".format(args.dataset)
+    
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    
+    if args.compute_fid:
+        for i in range(iters_needed):
+            with torch.no_grad():
+                x_t_1 = torch.randn(args.batch_size, args.num_channels,args.image_size, args.image_size).to(device)
+                fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args)
+                
+                fake_sample = to_range_0_1(fake_sample)
+                for j, x in enumerate(fake_sample):
+                    index = i * args.batch_size + j 
+                    torchvision.utils.save_image(x, './generated_samples/{}/{}.jpg'.format(args.dataset, index))
+                print('generating batch ', i)
+        
+        paths = [save_dir, real_img_dir]
+    
+        kwargs = {'batch_size': 100, 'device': device, 'dims': 2048}
+        fid = calculate_fid_given_paths(paths=paths, **kwargs)
+        print('FID = {}'.format(fid))
+    else:
+        x_t_1 = torch.randn(args.batch_size, args.num_channels,args.image_size, args.image_size).to(device)
+        fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args)
+        fake_sample = to_range_0_1(fake_sample)
+        torchvision.utils.save_image(fake_sample, './samples_{}.jpg'.format(args.dataset))
+
+    
+    
+            
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('ddgan parameters')
+    parser.add_argument('--seed', type=int, default=1024,
+                        help='seed used for initialization')
+    parser.add_argument('--compute_fid', action='store_true', default=False,
+                            help='whether or not compute FID')
+    parser.add_argument('--epoch_id', type=int,default=1000)
+    parser.add_argument('--num_channels', type=int, default=3,
+                            help='channel of image')
+    parser.add_argument('--centered', action='store_false', default=True,
+                            help='-1,1 scale')
+    parser.add_argument('--use_geometric', action='store_true',default=False)
+    parser.add_argument('--beta_min', type=float, default= 0.1,
+                            help='beta_min for diffusion')
+    parser.add_argument('--beta_max', type=float, default=20.,
+                            help='beta_max for diffusion')
+    
+    
+    parser.add_argument('--num_channels_dae', type=int, default=128,
+                            help='number of initial channels in denosing model')
+    parser.add_argument('--n_mlp', type=int, default=3,
+                            help='number of mlp layers for z')
+    parser.add_argument('--ch_mult', nargs='+', type=int,
+                            help='channel multiplier')
+
+    parser.add_argument('--num_res_blocks', type=int, default=2,
+                            help='number of resnet blocks per scale')
+    parser.add_argument('--attn_resolutions', default=(16,),
+                            help='resolution of applying attention')
+    parser.add_argument('--dropout', type=float, default=0.,
+                            help='drop-out rate')
+    parser.add_argument('--resamp_with_conv', action='store_false', default=True,
+                            help='always up/down sampling with conv')
+    parser.add_argument('--conditional', action='store_false', default=True,
+                            help='noise conditional')
+    parser.add_argument('--fir', action='store_false', default=True,
+                            help='FIR')
+    parser.add_argument('--fir_kernel', default=[1, 3, 3, 1],
+                            help='FIR kernel')
+    parser.add_argument('--skip_rescale', action='store_false', default=True,
+                            help='skip rescale')
+    parser.add_argument('--resblock_type', default='biggan',
+                            help='tyle of resnet block, choice in biggan and ddpm')
+    parser.add_argument('--progressive', type=str, default='none', choices=['none', 'output_skip', 'residual'],
+                            help='progressive type for output')
+    parser.add_argument('--progressive_input', type=str, default='residual', choices=['none', 'input_skip', 'residual'],
+                        help='progressive type for input')
+    parser.add_argument('--progressive_combine', type=str, default='sum', choices=['sum', 'cat'],
+                        help='progressive combine method.')
+
+    parser.add_argument('--embedding_type', type=str, default='positional', choices=['positional', 'fourier'],
+                        help='type of time embedding')
+    parser.add_argument('--fourier_scale', type=float, default=16.,
+                            help='scale of fourier transform')
+    parser.add_argument('--not_use_tanh', action='store_true',default=False)
+    
+    #geenrator and training
+    parser.add_argument('--exp', default='experiment_cifar_default', help='name of experiment')
+    parser.add_argument('--real_img_dir', default='./pytorch_fid/cifar10_train_stat.npy', help='directory to real images for FID computation')
+
+    parser.add_argument('--dataset', default='cifar10', help='name of dataset')
+    parser.add_argument('--image_size', type=int, default=32,
+                            help='size of image')
+
+    parser.add_argument('--nz', type=int, default=100)
+    parser.add_argument('--num_timesteps', type=int, default=4)
+    
+    
+    parser.add_argument('--z_emb_dim', type=int, default=256)
+    parser.add_argument('--t_emb_dim', type=int, default=256)
+    parser.add_argument('--batch_size', type=int, default=200, help='sample generating batch size')
+        
+
+
+
+   
+    args = parser.parse_args()
+    
+    sample_and_test(args)
+    
+   
+                

train_ddgan.py

@@ -0,0 +1,602 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the NVIDIA Source Code License
+# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
+# ---------------------------------------------------------------
+
+
+import argparse
+import torch
+import numpy as np
+
+import os
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR10
+from datasets_prep.lsun import LSUN
+from datasets_prep.stackmnist_data import StackedMNIST, _data_transforms_stacked_mnist
+from datasets_prep.lmdb_datasets import LMDBDataset
+
+
+from torch.multiprocessing import Process
+import torch.distributed as dist
+import shutil
+
+def copy_source(file, output_dir):
+    shutil.copyfile(file, os.path.join(output_dir, os.path.basename(file)))
+            
+def broadcast_params(params):
+    for param in params:
+        dist.broadcast(param.data, src=0)
+
+
+#%% Diffusion coefficients 
+def var_func_vp(t, beta_min, beta_max):
+    log_mean_coeff = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min
+    var = 1. - torch.exp(2. * log_mean_coeff)
+    return var
+
+def var_func_geometric(t, beta_min, beta_max):
+    return beta_min * ((beta_max / beta_min) ** t)
+
+def extract(input, t, shape):
+    out = torch.gather(input, 0, t)
+    reshape = [shape[0]] + [1] * (len(shape) - 1)
+    out = out.reshape(*reshape)
+
+    return out
+
+def get_time_schedule(args, device):
+    n_timestep = args.num_timesteps
+    eps_small = 1e-3
+    t = np.arange(0, n_timestep + 1, dtype=np.float64)
+    t = t / n_timestep
+    t = torch.from_numpy(t) * (1. - eps_small)  + eps_small
+    return t.to(device)
+
+def get_sigma_schedule(args, device):
+    n_timestep = args.num_timesteps
+    beta_min = args.beta_min
+    beta_max = args.beta_max
+    eps_small = 1e-3
+   
+    t = np.arange(0, n_timestep + 1, dtype=np.float64)
+    t = t / n_timestep
+    t = torch.from_numpy(t) * (1. - eps_small) + eps_small
+    
+    if args.use_geometric:
+        var = var_func_geometric(t, beta_min, beta_max)
+    else:
+        var = var_func_vp(t, beta_min, beta_max)
+    alpha_bars = 1.0 - var
+    betas = 1 - alpha_bars[1:] / alpha_bars[:-1]
+    
+    first = torch.tensor(1e-8)
+    betas = torch.cat((first[None], betas)).to(device)
+    betas = betas.type(torch.float32)
+    sigmas = betas**0.5
+    a_s = torch.sqrt(1-betas)
+    return sigmas, a_s, betas
+
+class Diffusion_Coefficients():
+    def __init__(self, args, device):
+                
+        self.sigmas, self.a_s, _ = get_sigma_schedule(args, device=device)
+        self.a_s_cum = np.cumprod(self.a_s.cpu())
+        self.sigmas_cum = np.sqrt(1 - self.a_s_cum ** 2)
+        self.a_s_prev = self.a_s.clone()
+        self.a_s_prev[-1] = 1
+        
+        self.a_s_cum = self.a_s_cum.to(device)
+        self.sigmas_cum = self.sigmas_cum.to(device)
+        self.a_s_prev = self.a_s_prev.to(device)
+    
+def q_sample(coeff, x_start, t, *, noise=None):
+    """
+    Diffuse the data (t == 0 means diffused for t step)
+    """
+    if noise is None:
+      noise = torch.randn_like(x_start)
+      
+    x_t = extract(coeff.a_s_cum, t, x_start.shape) * x_start + \
+          extract(coeff.sigmas_cum, t, x_start.shape) * noise
+    
+    return x_t
+
+def q_sample_pairs(coeff, x_start, t):
+    """
+    Generate a pair of disturbed images for training
+    :param x_start: x_0
+    :param t: time step t
+    :return: x_t, x_{t+1}
+    """
+    noise = torch.randn_like(x_start)
+    x_t = q_sample(coeff, x_start, t)
+    x_t_plus_one = extract(coeff.a_s, t+1, x_start.shape) * x_t + \
+                   extract(coeff.sigmas, t+1, x_start.shape) * noise
+    
+    return x_t, x_t_plus_one
+#%% posterior sampling
+class Posterior_Coefficients():
+    def __init__(self, args, device):
+        
+        _, _, self.betas = get_sigma_schedule(args, device=device)
+        
+        #we don't need the zeros
+        self.betas = self.betas.type(torch.float32)[1:]
+        
+        self.alphas = 1 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, 0)
+        self.alphas_cumprod_prev = torch.cat(
+                                    (torch.tensor([1.], dtype=torch.float32,device=device), self.alphas_cumprod[:-1]), 0
+                                        )               
+        self.posterior_variance = self.betas * (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod)
+        
+        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1 / self.alphas_cumprod - 1)
+        
+        self.posterior_mean_coef1 = (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1 - self.alphas_cumprod))
+        self.posterior_mean_coef2 = ((1 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1 - self.alphas_cumprod))
+        
+        self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min=1e-20))
+        
+def sample_posterior(coefficients, x_0,x_t, t):
+    
+    def q_posterior(x_0, x_t, t):
+        mean = (
+            extract(coefficients.posterior_mean_coef1, t, x_t.shape) * x_0
+            + extract(coefficients.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        var = extract(coefficients.posterior_variance, t, x_t.shape)
+        log_var_clipped = extract(coefficients.posterior_log_variance_clipped, t, x_t.shape)
+        return mean, var, log_var_clipped
+    
+  
+    def p_sample(x_0, x_t, t):
+        mean, _, log_var = q_posterior(x_0, x_t, t)
+        
+        noise = torch.randn_like(x_t)
+        
+        nonzero_mask = (1 - (t == 0).type(torch.float32))
+
+        return mean + nonzero_mask[:,None,None,None] * torch.exp(0.5 * log_var) * noise
+            
+    sample_x_pos = p_sample(x_0, x_t, t)
+    
+    return sample_x_pos
+
+def sample_from_model(coefficients, generator, n_time, x_init, T, opt):
+    x = x_init
+    with torch.no_grad():
+        for i in reversed(range(n_time)):
+            t = torch.full((x.size(0),), i, dtype=torch.int64).to(x.device)
+          
+            t_time = t
+            latent_z = torch.randn(x.size(0), opt.nz, device=x.device)
+            x_0 = generator(x, t_time, latent_z)
+            x_new = sample_posterior(coefficients, x_0, x, t)
+            x = x_new.detach()
+        
+    return x
+
+#%%
+def train(rank, gpu, args):
+    from score_sde.models.discriminator import Discriminator_small, Discriminator_large
+    from score_sde.models.ncsnpp_generator_adagn import NCSNpp
+    from EMA import EMA
+    
+    torch.manual_seed(args.seed + rank)
+    torch.cuda.manual_seed(args.seed + rank)
+    torch.cuda.manual_seed_all(args.seed + rank)
+    device = torch.device('cuda:{}'.format(gpu))
+    
+    batch_size = args.batch_size
+    
+    nz = args.nz #latent dimension
+    
+    if args.dataset == 'cifar10':
+        dataset = CIFAR10('./data', train=True, transform=transforms.Compose([
+                        transforms.Resize(32),
+                        transforms.RandomHorizontalFlip(),
+                        transforms.ToTensor(),
+                        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))]), download=True)
+       
+    
+    elif args.dataset == 'stackmnist':
+        train_transform, valid_transform = _data_transforms_stacked_mnist()
+        dataset = StackedMNIST(root='./data', train=True, download=False, transform=train_transform)
+        
+    elif args.dataset == 'lsun':
+        
+        train_transform = transforms.Compose([
+                        transforms.Resize(args.image_size),
+                        transforms.CenterCrop(args.image_size),
+                        transforms.RandomHorizontalFlip(),
+                        transforms.ToTensor(),
+                        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+                    ])
+
+        train_data = LSUN(root='/datasets/LSUN/', classes=['church_outdoor_train'], transform=train_transform)
+        subset = list(range(0, 120000))
+        dataset = torch.utils.data.Subset(train_data, subset)
+      
+    
+    elif args.dataset == 'celeba_256':
+        train_transform = transforms.Compose([
+                transforms.Resize(args.image_size),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+            ])
+        dataset = LMDBDataset(root='/datasets/celeba-lmdb/', name='celeba', train=True, transform=train_transform)
+      
+    
+    
+    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset,
+                                                                    num_replicas=args.world_size,
+                                                                    rank=rank)
+    data_loader = torch.utils.data.DataLoader(dataset,
+                                               batch_size=batch_size,
+                                               shuffle=False,
+                                               num_workers=4,
+                                               pin_memory=True,
+                                               sampler=train_sampler,
+                                               drop_last = True)
+    
+    netG = NCSNpp(args).to(device)
+    
+
+    if args.dataset == 'cifar10' or args.dataset == 'stackmnist':    
+        netD = Discriminator_small(nc = 2*args.num_channels, ngf = args.ngf,
+                               t_emb_dim = args.t_emb_dim,
+                               act=nn.LeakyReLU(0.2)).to(device)
+    else:
+        netD = Discriminator_large(nc = 2*args.num_channels, ngf = args.ngf, 
+                                   t_emb_dim = args.t_emb_dim,
+                                   act=nn.LeakyReLU(0.2)).to(device)
+    
+    broadcast_params(netG.parameters())
+    broadcast_params(netD.parameters())
+    
+    optimizerD = optim.Adam(netD.parameters(), lr=args.lr_d, betas = (args.beta1, args.beta2))
+    
+    optimizerG = optim.Adam(netG.parameters(), lr=args.lr_g, betas = (args.beta1, args.beta2))
+    
+    if args.use_ema:
+        optimizerG = EMA(optimizerG, ema_decay=args.ema_decay)
+    
+    schedulerG = torch.optim.lr_scheduler.CosineAnnealingLR(optimizerG, args.num_epoch, eta_min=1e-5)
+    schedulerD = torch.optim.lr_scheduler.CosineAnnealingLR(optimizerD, args.num_epoch, eta_min=1e-5)
+    
+    
+    
+    #ddp
+    netG = nn.parallel.DistributedDataParallel(netG, device_ids=[gpu])
+    netD = nn.parallel.DistributedDataParallel(netD, device_ids=[gpu])
+
+    
+    exp = args.exp
+    parent_dir = "./saved_info/dd_gan/{}".format(args.dataset)
+
+    exp_path = os.path.join(parent_dir,exp)
+    if rank == 0:
+        if not os.path.exists(exp_path):
+            os.makedirs(exp_path)
+            copy_source(__file__, exp_path)
+            shutil.copytree('score_sde/models', os.path.join(exp_path, 'score_sde/models'))
+    
+    
+    coeff = Diffusion_Coefficients(args, device)
+    pos_coeff = Posterior_Coefficients(args, device)
+    T = get_time_schedule(args, device)
+    
+    if args.resume:
+        checkpoint_file = os.path.join(exp_path, 'content.pth')
+        checkpoint = torch.load(checkpoint_file, map_location=device)
+        init_epoch = checkpoint['epoch']
+        epoch = init_epoch
+        netG.load_state_dict(checkpoint['netG_dict'])
+        # load G
+        
+        optimizerG.load_state_dict(checkpoint['optimizerG'])
+        schedulerG.load_state_dict(checkpoint['schedulerG'])
+        # load D
+        netD.load_state_dict(checkpoint['netD_dict'])
+        optimizerD.load_state_dict(checkpoint['optimizerD'])
+        schedulerD.load_state_dict(checkpoint['schedulerD'])
+        global_step = checkpoint['global_step']
+        print("=> loaded checkpoint (epoch {})"
+                  .format(checkpoint['epoch']))
+    else:
+        global_step, epoch, init_epoch = 0, 0, 0
+    
+    
+    for epoch in range(init_epoch, args.num_epoch+1):
+        train_sampler.set_epoch(epoch)
+       
+        for iteration, (x, y) in enumerate(data_loader):
+            for p in netD.parameters():  
+                p.requires_grad = True  
+        
+            
+            netD.zero_grad()
+            
+            #sample from p(x_0)
+            real_data = x.to(device, non_blocking=True)
+            
+            #sample t
+            t = torch.randint(0, args.num_timesteps, (real_data.size(0),), device=device)
+            
+            x_t, x_tp1 = q_sample_pairs(coeff, real_data, t)
+            x_t.requires_grad = True
+            
+    
+            # train with real
+            D_real = netD(x_t, t, x_tp1.detach()).view(-1)
+            
+            errD_real = F.softplus(-D_real)
+            errD_real = errD_real.mean()
+            
+            errD_real.backward(retain_graph=True)
+            
+            
+            if args.lazy_reg is None:
+                grad_real = torch.autograd.grad(
+                            outputs=D_real.sum(), inputs=x_t, create_graph=True
+                            )[0]
+                grad_penalty = (
+                                grad_real.view(grad_real.size(0), -1).norm(2, dim=1) ** 2
+                                ).mean()
+                
+                
+                grad_penalty = args.r1_gamma / 2 * grad_penalty
+                grad_penalty.backward()
+            else:
+                if global_step % args.lazy_reg == 0:
+                    grad_real = torch.autograd.grad(
+                            outputs=D_real.sum(), inputs=x_t, create_graph=True
+                            )[0]
+                    grad_penalty = (
+                                grad_real.view(grad_real.size(0), -1).norm(2, dim=1) ** 2
+                                ).mean()
+                
+                
+                    grad_penalty = args.r1_gamma / 2 * grad_penalty
+                    grad_penalty.backward()
+
+            # train with fake
+            latent_z = torch.randn(batch_size, nz, device=device)
+            
+         
+            x_0_predict = netG(x_tp1.detach(), t, latent_z)
+            x_pos_sample = sample_posterior(pos_coeff, x_0_predict, x_tp1, t)
+            
+            output = netD(x_pos_sample, t, x_tp1.detach()).view(-1)
+                
+            
+            errD_fake = F.softplus(output)
+            errD_fake = errD_fake.mean()
+            errD_fake.backward()
+    
+            
+            errD = errD_real + errD_fake
+            # Update D
+            optimizerD.step()
+            
+        
+            #update G
+            for p in netD.parameters():
+                p.requires_grad = False
+            netG.zero_grad()
+            
+            
+            t = torch.randint(0, args.num_timesteps, (real_data.size(0),), device=device)
+            
+            
+            x_t, x_tp1 = q_sample_pairs(coeff, real_data, t)
+                
+            
+            latent_z = torch.randn(batch_size, nz,device=device)
+            
+            
+                
+           
+            x_0_predict = netG(x_tp1.detach(), t, latent_z)
+            x_pos_sample = sample_posterior(pos_coeff, x_0_predict, x_tp1, t)
+            
+            output = netD(x_pos_sample, t, x_tp1.detach()).view(-1)
+               
+            
+            errG = F.softplus(-output)
+            errG = errG.mean()
+            
+            errG.backward()
+            optimizerG.step()
+                
+           
+            
+            global_step += 1
+            if iteration % 100 == 0:
+                if rank == 0:
+                    print('epoch {} iteration{}, G Loss: {}, D Loss: {}'.format(epoch,iteration, errG.item(), errD.item()))
+        
+        if not args.no_lr_decay:
+            
+            schedulerG.step()
+            schedulerD.step()
+        
+        if rank == 0:
+            if epoch % 10 == 0:
+                torchvision.utils.save_image(x_pos_sample, os.path.join(exp_path, 'xpos_epoch_{}.png'.format(epoch)), normalize=True)
+            
+            x_t_1 = torch.randn_like(real_data)
+            fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1, T, args)
+            torchvision.utils.save_image(fake_sample, os.path.join(exp_path, 'sample_discrete_epoch_{}.png'.format(epoch)), normalize=True)
+            
+            if args.save_content:
+                if epoch % args.save_content_every == 0:
+                    print('Saving content.')
+                    content = {'epoch': epoch + 1, 'global_step': global_step, 'args': args,
+                               'netG_dict': netG.state_dict(), 'optimizerG': optimizerG.state_dict(),
+                               'schedulerG': schedulerG.state_dict(), 'netD_dict': netD.state_dict(),
+                               'optimizerD': optimizerD.state_dict(), 'schedulerD': schedulerD.state_dict()}
+                    
+                    torch.save(content, os.path.join(exp_path, 'content.pth'))
+                
+            if epoch % args.save_ckpt_every == 0:
+                if args.use_ema:
+                    optimizerG.swap_parameters_with_ema(store_params_in_ema=True)
+                    
+                torch.save(netG.state_dict(), os.path.join(exp_path, 'netG_{}.pth'.format(epoch)))
+                if args.use_ema:
+                    optimizerG.swap_parameters_with_ema(store_params_in_ema=True)
+            
+
+
+def init_processes(rank, size, fn, args):
+    """ Initialize the distributed environment. """
+    os.environ['MASTER_ADDR'] = args.master_address
+    os.environ['MASTER_PORT'] = '6020'
+    torch.cuda.set_device(args.local_rank)
+    gpu = args.local_rank
+    dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=size)
+    fn(rank, gpu, args)
+    dist.barrier()
+    cleanup()  
+
+def cleanup():
+    dist.destroy_process_group()    
+#%%
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('ddgan parameters')
+    parser.add_argument('--seed', type=int, default=1024,
+                        help='seed used for initialization')
+    
+    parser.add_argument('--resume', action='store_true',default=False)
+    
+    parser.add_argument('--image_size', type=int, default=32,
+                            help='size of image')
+    parser.add_argument('--num_channels', type=int, default=3,
+                            help='channel of image')
+    parser.add_argument('--centered', action='store_false', default=True,
+                            help='-1,1 scale')
+    parser.add_argument('--use_geometric', action='store_true',default=False)
+    parser.add_argument('--beta_min', type=float, default= 0.1,
+                            help='beta_min for diffusion')
+    parser.add_argument('--beta_max', type=float, default=20.,
+                            help='beta_max for diffusion')
+    
+    
+    parser.add_argument('--num_channels_dae', type=int, default=128,
+                            help='number of initial channels in denosing model')
+    parser.add_argument('--n_mlp', type=int, default=3,
+                            help='number of mlp layers for z')
+    parser.add_argument('--ch_mult', nargs='+', type=int,
+                            help='channel multiplier')
+    parser.add_argument('--num_res_blocks', type=int, default=2,
+                            help='number of resnet blocks per scale')
+    parser.add_argument('--attn_resolutions', default=(16,),
+                            help='resolution of applying attention')
+    parser.add_argument('--dropout', type=float, default=0.,
+                            help='drop-out rate')
+    parser.add_argument('--resamp_with_conv', action='store_false', default=True,
+                            help='always up/down sampling with conv')
+    parser.add_argument('--conditional', action='store_false', default=True,
+                            help='noise conditional')
+    parser.add_argument('--fir', action='store_false', default=True,
+                            help='FIR')
+    parser.add_argument('--fir_kernel', default=[1, 3, 3, 1],
+                            help='FIR kernel')
+    parser.add_argument('--skip_rescale', action='store_false', default=True,
+                            help='skip rescale')
+    parser.add_argument('--resblock_type', default='biggan',
+                            help='tyle of resnet block, choice in biggan and ddpm')
+    parser.add_argument('--progressive', type=str, default='none', choices=['none', 'output_skip', 'residual'],
+                            help='progressive type for output')
+    parser.add_argument('--progressive_input', type=str, default='residual', choices=['none', 'input_skip', 'residual'],
+                        help='progressive type for input')
+    parser.add_argument('--progressive_combine', type=str, default='sum', choices=['sum', 'cat'],
+                        help='progressive combine method.')
+    
+    parser.add_argument('--embedding_type', type=str, default='positional', choices=['positional', 'fourier'],
+                        help='type of time embedding')
+    parser.add_argument('--fourier_scale', type=float, default=16.,
+                            help='scale of fourier transform')
+    parser.add_argument('--not_use_tanh', action='store_true',default=False)
+    
+    #geenrator and training
+    parser.add_argument('--exp', default='experiment_cifar_default', help='name of experiment')
+    parser.add_argument('--dataset', default='cifar10', help='name of dataset')
+    parser.add_argument('--nz', type=int, default=100)
+    parser.add_argument('--num_timesteps', type=int, default=4)
+
+    parser.add_argument('--z_emb_dim', type=int, default=256)
+    parser.add_argument('--t_emb_dim', type=int, default=256)
+    parser.add_argument('--batch_size', type=int, default=128, help='input batch size')
+    parser.add_argument('--num_epoch', type=int, default=1200)
+    parser.add_argument('--ngf', type=int, default=64)
+
+    parser.add_argument('--lr_g', type=float, default=1.5e-4, help='learning rate g')
+    parser.add_argument('--lr_d', type=float, default=1e-4, help='learning rate d')
+    parser.add_argument('--beta1', type=float, default=0.5,
+                            help='beta1 for adam')
+    parser.add_argument('--beta2', type=float, default=0.9,
+                            help='beta2 for adam')
+    parser.add_argument('--no_lr_decay',action='store_true', default=False)
+    
+    parser.add_argument('--use_ema', action='store_true', default=False,
+                            help='use EMA or not')
+    parser.add_argument('--ema_decay', type=float, default=0.9999, help='decay rate for EMA')
+    
+    parser.add_argument('--r1_gamma', type=float, default=0.05, help='coef for r1 reg')
+    parser.add_argument('--lazy_reg', type=int, default=None,
+                        help='lazy regulariation.')
+
+    parser.add_argument('--save_content', action='store_true',default=False)
+    parser.add_argument('--save_content_every', type=int, default=50, help='save content for resuming every x epochs')
+    parser.add_argument('--save_ckpt_every', type=int, default=25, help='save ckpt every x epochs')
+   
+    ###ddp
+    parser.add_argument('--num_proc_node', type=int, default=1,
+                        help='The number of nodes in multi node env.')
+    parser.add_argument('--num_process_per_node', type=int, default=1,
+                        help='number of gpus')
+    parser.add_argument('--node_rank', type=int, default=0,
+                        help='The index of node.')
+    parser.add_argument('--local_rank', type=int, default=0,
+                        help='rank of process in the node')
+    parser.add_argument('--master_address', type=str, default='127.0.0.1',
+                        help='address for master')
+
+   
+    args = parser.parse_args()
+    args.world_size = args.num_proc_node * args.num_process_per_node
+    size = args.num_process_per_node
+
+    if size > 1:
+        processes = []
+        for rank in range(size):
+            args.local_rank = rank
+            global_rank = rank + args.node_rank * args.num_process_per_node
+            global_size = args.num_proc_node * args.num_process_per_node
+            args.global_rank = global_rank
+            print('Node rank %d, local proc %d, global proc %d' % (args.node_rank, rank, global_rank))
+            p = Process(target=init_processes, args=(global_rank, global_size, train, args))
+            p.start()
+            processes.append(p)
+            
+        for p in processes:
+            p.join()
+    else:
+        print('starting in debug mode')
+        
+        init_processes(0, size, train, args)
+   
+