Upload 98 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+point_e/examples/paper_banner.gif filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+*.egg-info/
+__pycache__/
+point_e_model_cache/
+.ipynb_checkpoints/
+.DS_Store
+

LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2022 OpenAI
+
+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.
+

README.md

@@ -1,3 +1,28 @@
----
-license: mit
----
+# Point·E
+
+![Animation of four 3D point clouds rotating](point_e/examples/paper_banner.gif)
+
+This is the official code and model release for [Point-E: A System for Generating 3D Point Clouds from Complex Prompts](https://arxiv.org/abs/2212.08751).
+
+# Usage
+
+Install with `pip install -e .`.
+
+To get started with examples, see the following notebooks:
+
+ * [image2pointcloud.ipynb](point_e/examples/image2pointcloud.ipynb) - sample a point cloud, conditioned on some example synthetic view images.
+ * [text2pointcloud.ipynb](point_e/examples/text2pointcloud.ipynb) - use our small, worse quality pure text-to-3D model to produce 3D point clouds directly from text descriptions. This model's capabilities are limited, but it does understand some simple categories and colors.
+ * [pointcloud2mesh.ipynb](point_e/examples/pointcloud2mesh.ipynb) - try our SDF regression model for producing meshes from point clouds.
+
+For our P-FID and P-IS evaluation scripts, see:
+
+ * [evaluate_pfid.py](point_e/evals/scripts/evaluate_pfid.py)
+ * [evaluate_pis.py](point_e/evals/scripts/evaluate_pis.py)
+
+For our Blender rendering code, see [blender_script.py](point_e/evals/scripts/blender_script.py)
+
+# Samples
+
+You can download the seed images and point clouds corresponding to the paper banner images [here](https://openaipublic.azureedge.net/main/point-e/banner_pcs.zip).
+
+You can download the seed images used for COCO CLIP R-Precision evaluations [here](https://openaipublic.azureedge.net/main/point-e/coco_images.zip).

model-card.md

@@ -0,0 +1,62 @@
+# Model Card: Point-E
+
+This is the official codebase for running the point cloud diffusion models and SDF regression models described in [Point-E: A System for Generating 3D Point Clouds from Complex Prompts](https://arxiv.org/abs/2212.08751). These models were trained and released by OpenAI.
+Following [Model Cards for Model Reporting (Mitchell et al.)](https://arxiv.org/abs/1810.03993), we're providing some information about how the models were trained and evaluated. 
+
+# Model Details
+
+The Point-E models are trained for use as point cloud diffusion models and SDF regression models.
+Our image-conditional models are often capable of producing coherent 3D point clouds, given a single rendering of a 3D object. However, the models sometimes fail to do so, either producing incorrect geometry where the rendering is occluded, or producing geometry that is inconsistent with visible parts of the rendering. The resulting point clouds are relatively low-resolution, and are often noisy and contain defects such as outliers or cracks.
+Our text-conditional model is sometimes capable of producing 3D point clouds which can be recognized as the provided text description, especially when the text description is simple. However, we find that this model fails to generalize to complex prompts or unusual objects.
+
+## Model Date
+
+December 2022
+
+## Model Versions
+
+ * `base40M-imagevec` - a 40 million parameter image to point cloud model that conditions on a single CLIP ViT-L/14 image vector. This model can be used to generate point clouds from rendered images, but does not perform as well as our other models for this task.
+ * `base40M-textvec` - a 40 million parameter text to point cloud model that conditions on a single CLIP ViT-L/14 text vector. This model can be used to directly generate point clouds from text descriptions, but only works for simple prompts.
+ * `base40M-uncond` - a 40 million parameter point cloud diffusion model that generates unconditional samples. This is included only as a baseline.
+ * `base40M` - a 40 million parameter image to point cloud diffusion model that conditions on the latent grid from a CLIP ViT-L/14 model. This model can be used to generate point clouds from rendered images, but is not as good as the larger models trained on the same task.
+ * `base300M` - a 300 million parameter image to point cloud diffusion model that conditions on the latent grid from a CLIP ViT-L/14 model. This model can be used to generate point clouds from rendered images, but it is slightly worse than base1B
+ * `base1B` - a 1 billion parameter image to point cloud diffusion model that conditions on the latent grid from a CLIP ViT-L/14 model.
+ * `upsample` - a 40 million parameter point cloud upsampling model that can optionally condition on an image as well. This takes a point cloud of 1024 points and upsamples it to 4096 points.
+ * `sdf` - a small model for predicting signed distance functions from 3D point clouds. This can be used to predict meshes from point clouds.
+ * `pointnet` - a small point cloud classification model used for our P-FID and P-IS evaluation metrics.
+
+## Paper & samples
+
+[Paper](https://arxiv.org/abs/2212.08751) / [Sample point clouds](point_e/examples/paper_banner.gif)
+
+# Training data
+
+These models were trained on a dataset of several million 3D models. We filtered the dataset to avoid flat objects, and used [CLIP](https://github.com/openai/CLIP/blob/main/model-card.md) to cluster the dataset and downweight clusters of 3D models which appeared to contain mostly unrecognizable objects. We additionally down-weighted clusters which appeared to consist of many similar-looking objects. We processed the resulting dataset into renders (RGB point clouds of 4K points each) and text captions from the associated metadata.
+Our SDF regression model was trained on a subset of the above dataset. In particular, we only retained 3D meshes which were manifold (i.e. watertight and free of singularities).
+
+# Evaluated Use
+
+We release these models to help advance research in generative modeling. Due to the limitations and biases of our models, we do not currently recommend it for commercial use. We understand that our models may be used in ways we haven't anticipated, and that it is difficult to define clear boundaries around what constitutes appropriate "research" use. In particular, we caution against using these models in applications where precision is critical, as subtle flaws in the outputs could lead to errors or inaccuracies.
+Functionally, these models are trained to be able to perform the following tasks for research purposes, and are evaluated on these tasks:
+
+ * Generate 3D point clouds conditioned on single rendered images
+ * Generate 3D point clouds conditioned on text
+ * Create 3D meshes from noisy 3D point clouds
+
+Our image-conditional models are intended to produce coherent point clouds, given a representative rendering of a 3D object. However, at their current level of capabilities, the models sometimes fail to generate coherent output, either producing incorrect geometry where the rendering is occluded, or producing geometry that is inconsistent with visible parts of the rendering. The resulting point clouds are relatively low-resolution, and are often noisy and contain defects such as outliers or cracks. 
+
+Our text-conditional model is sometimes capable of producing 3D point clouds which can be recognized as the provided text description, especially when the text description is simple. However, we find that this model fails to generalize to complex prompts or unusual objects.
+
+# Performance and Limitations
+
+Our image-conditional models are limited by the text-to-image model that is used to produce synthetic views. If the text-to-image model contains a bias or fails to understand a particular concept, these limitations will be passed down to the image-conditional point cloud model through conditioning images.
+While our main focus is on image-conditional models, we also experimented with a text-conditional model. We find that this model can sometimes produce 3D models of people that exhibit gender biases (for example, samples for "a man" tend to be wider and less narrow than samples for "a woman"). We additionally find that this model is sometimes capable of producing violent objects such as guns or tanks, although these generations are always low-quality and unrealistic.
+
+Since our dataset contains many simplistic, cartoonish 3D objects, our models are prone to mimicking this style.
+
+While these models were developed for research purposes, they have potential implications if used more broadly. For example, the ability to generate 3D point clouds from single images could help advance research in computer graphics, virtual reality, and robotics. The text-conditional model could allow for users to easily create 3D models from simple descriptions, which could be useful for rapid prototyping or 3D printing.
+ 
+The combination of these models with 3D printing could potentially be harmful, for example if used to prototype dangerous objects or when parts created by the model are trusted without external validation.
+
+Finally, point cloud models inherit many of the same risks and limitations as image-generation models, including the propensity to produce biased or otherwise harmful content or to carry dual-use risk. More research is needed on how these risks manifest themselves as capabilities improve.
+

point_e.egg-info/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.1
+Name: point-e
+Version: 0.0.0
+Author: OpenAI
+License-File: LICENSE

point_e.egg-info/SOURCES.txt

@@ -0,0 +1,35 @@
+LICENSE
+README.md
+setup.py
+point_e/__init__.py
+point_e.egg-info/PKG-INFO
+point_e.egg-info/SOURCES.txt
+point_e.egg-info/dependency_links.txt
+point_e.egg-info/requires.txt
+point_e.egg-info/top_level.txt
+point_e/diffusion/__init__.py
+point_e/diffusion/configs.py
+point_e/diffusion/gaussian_diffusion.py
+point_e/diffusion/k_diffusion.py
+point_e/diffusion/sampler.py
+point_e/evals/__init__.py
+point_e/evals/feature_extractor.py
+point_e/evals/fid_is.py
+point_e/evals/npz_stream.py
+point_e/evals/pointnet2_cls_ssg.py
+point_e/evals/pointnet2_utils.py
+point_e/models/__init__.py
+point_e/models/checkpoint.py
+point_e/models/configs.py
+point_e/models/download.py
+point_e/models/perceiver.py
+point_e/models/pretrained_clip.py
+point_e/models/sdf.py
+point_e/models/transformer.py
+point_e/models/util.py
+point_e/util/__init__.py
+point_e/util/mesh.py
+point_e/util/pc_to_mesh.py
+point_e/util/plotting.py
+point_e/util/ply_util.py
+point_e/util/point_cloud.py

point_e.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

point_e.egg-info/requires.txt

@@ -0,0 +1,12 @@
+filelock
+Pillow
+torch
+fire
+humanize
+requests
+tqdm
+matplotlib
+scikit-image
+scipy
+numpy
+clip@ git+https://github.com/openai/CLIP.git

point_e.egg-info/top_level.txt

@@ -0,0 +1 @@
+point_e

point_e/diffusion/configs.py

@@ -0,0 +1,64 @@
+"""
+Based on https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py
+"""
+
+from typing import Any, Dict
+
+import numpy as np
+
+from .gaussian_diffusion import (
+    GaussianDiffusion,
+    SpacedDiffusion,
+    get_named_beta_schedule,
+    space_timesteps,
+)
+
+BASE_DIFFUSION_CONFIG = {
+    "channel_biases": [0.0, 0.0, 0.0, -1.0, -1.0, -1.0],
+    "channel_scales": [2.0, 2.0, 2.0, 0.007843137255, 0.007843137255, 0.007843137255],
+    "mean_type": "epsilon",
+    "schedule": "cosine",
+    "timesteps": 1024,
+}
+
+DIFFUSION_CONFIGS = {
+    "base40M-imagevec": BASE_DIFFUSION_CONFIG,
+    "base40M-textvec": BASE_DIFFUSION_CONFIG,
+    "base40M-uncond": BASE_DIFFUSION_CONFIG,
+    "base40M": BASE_DIFFUSION_CONFIG,
+    "base300M": BASE_DIFFUSION_CONFIG,
+    "base1B": BASE_DIFFUSION_CONFIG,
+    "upsample": {
+        "channel_biases": [0.0, 0.0, 0.0, -1.0, -1.0, -1.0],
+        "channel_scales": [2.0, 2.0, 2.0, 0.007843137255, 0.007843137255, 0.007843137255],
+        "mean_type": "epsilon",
+        "schedule": "linear",
+        "timesteps": 1024,
+    },
+}
+
+
+def diffusion_from_config(config: Dict[str, Any]) -> GaussianDiffusion:
+    schedule = config["schedule"]
+    steps = config["timesteps"]
+    respace = config.get("respacing", None)
+    mean_type = config.get("mean_type", "epsilon")
+    betas = get_named_beta_schedule(schedule, steps)
+    channel_scales = config.get("channel_scales", None)
+    channel_biases = config.get("channel_biases", None)
+    if channel_scales is not None:
+        channel_scales = np.array(channel_scales)
+    if channel_biases is not None:
+        channel_biases = np.array(channel_biases)
+    kwargs = dict(
+        betas=betas,
+        model_mean_type=mean_type,
+        model_var_type="learned_range",
+        loss_type="mse",
+        channel_scales=channel_scales,
+        channel_biases=channel_biases,
+    )
+    if respace is None:
+        return GaussianDiffusion(**kwargs)
+    else:
+        return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace), **kwargs)

point_e/diffusion/gaussian_diffusion.py

@@ -0,0 +1,1091 @@
+"""
+Based on https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py
+"""
+
+import math
+from typing import Any, Dict, Iterable, Optional, Sequence, Union
+
+import numpy as np
+import torch as th
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    """
+    This is the deprecated API for creating beta schedules.
+
+    See get_named_beta_schedule() for the new library of schedules.
+    """
+    if beta_schedule == "linear":
+        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        return get_beta_schedule(
+            "linear",
+            beta_start=scale * 0.0001,
+            beta_end=scale * 0.02,
+            num_diffusion_timesteps=num_diffusion_timesteps,
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(f"cannot create exactly {num_timesteps} steps with an integer stride")
+        elif section_counts.startswith("exact"):
+            res = set(int(x) for x in section_counts[len("exact") :].split(","))
+            for x in res:
+                if x < 0 or x >= num_timesteps:
+                    raise ValueError(f"timestep out of bounds: {x}")
+            return res
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(f"cannot divide section of {size} steps into {section_count}")
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here:
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D array of betas for each diffusion timestep from T to 1.
+    :param model_mean_type: a string determining what the model outputs.
+    :param model_var_type: a string determining how variance is output.
+    :param loss_type: a string determining the loss function to use.
+    :param discretized_t0: if True, use discrete gaussian loss for t=0. Only
+                           makes sense for images.
+    :param channel_scales: a multiplier to apply to x_start in training_losses
+                           and sampling functions.
+    """
+
+    def __init__(
+        self,
+        *,
+        betas: Sequence[float],
+        model_mean_type: str,
+        model_var_type: str,
+        loss_type: str,
+        discretized_t0: bool = False,
+        channel_scales: Optional[np.ndarray] = None,
+        channel_biases: Optional[np.ndarray] = None,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.discretized_t0 = discretized_t0
+        self.channel_scales = channel_scales
+        self.channel_biases = channel_biases
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
+        )
+
+    def get_sigmas(self, t):
+        return _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, t.shape)
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=False, denoised_fn=None, model_kwargs=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, t, **model_kwargs)
+        if isinstance(model_output, tuple):
+            model_output, extra = model_output
+        else:
+            extra = None
+
+        if self.model_var_type in ["learned", "learned_range"]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.model_var_type == "learned":
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape)
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                "fixed_large": (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                "fixed_small": (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if self.model_mean_type == "x_prev":
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in ["x_start", "epsilon"]:
+            if self.model_mean_type == "x_start":
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
+        else:
+            raise NotImplementedError(self.model_mean_type)
+
+        assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+            "extra": extra,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, t, **model_kwargs)
+        new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t)
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        temp=1.0,
+    ):
+        """
+        Generate samples from the model.
+
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            temp=temp,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        temp=1.0,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device) * temp
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                )
+                yield self.unscale_out_dict(out)
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        temp=1.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+            temp=temp,
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        temp=1.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device) * temp
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield self.unscale_out_dict(out)
+                img = out["sample"]
+
+    def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False, model_kwargs=None):
+        """
+        Get a term for the variational lower-bound.
+
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(true_mean, true_log_variance_clipped, out["mean"], out["log_variance"])
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        if not self.discretized_t0:
+            decoder_nll = th.zeros_like(decoder_nll)
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {
+            "output": output,
+            "pred_xstart": out["pred_xstart"],
+            "extra": out["extra"],
+        }
+
+    def training_losses(
+        self, model, x_start, t, model_kwargs=None, noise=None
+    ) -> Dict[str, th.Tensor]:
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        x_start = self.scale_channels(x_start)
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == "kl" or self.loss_type == "rescaled_kl":
+            vb_terms = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )
+            terms["loss"] = vb_terms["output"]
+            if self.loss_type == "rescaled_kl":
+                terms["loss"] *= self.num_timesteps
+            extra = vb_terms["extra"]
+        elif self.loss_type == "mse" or self.loss_type == "rescaled_mse":
+            model_output = model(x_t, t, **model_kwargs)
+            if isinstance(model_output, tuple):
+                model_output, extra = model_output
+            else:
+                extra = {}
+
+            if self.model_var_type in [
+                "learned",
+                "learned_range",
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == "rescaled_mse":
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                "x_prev": self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[0],
+                "x_start": x_start,
+                "epsilon": noise,
+            }[self.model_mean_type]
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        if "losses" in extra:
+            terms.update({k: loss for k, (loss, _scale) in extra["losses"].items()})
+            for loss, scale in extra["losses"].values():
+                terms["loss"] = terms["loss"] + loss * scale
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+
+        This term can't be optimized, as it only depends on the encoder.
+
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+    def scale_channels(self, x: th.Tensor) -> th.Tensor:
+        if self.channel_scales is not None:
+            x = x * th.from_numpy(self.channel_scales).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        if self.channel_biases is not None:
+            x = x + th.from_numpy(self.channel_biases).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        return x
+
+    def unscale_channels(self, x: th.Tensor) -> th.Tensor:
+        if self.channel_biases is not None:
+            x = x - th.from_numpy(self.channel_biases).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        if self.channel_scales is not None:
+            x = x / th.from_numpy(self.channel_scales).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        return x
+
+    def unscale_out_dict(
+        self, out: Dict[str, Union[th.Tensor, Any]]
+    ) -> Dict[str, Union[th.Tensor, Any]]:
+        return {
+            k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v) for k, v in out.items()
+        }
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+    :param use_timesteps: (unordered) timesteps from the original diffusion
+                          process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps: Iterable[int], **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(self, model, *args, **kwargs):
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(self, model, *args, **kwargs):
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(model, self.timestep_map, self.original_num_steps)
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        return self.model(x, new_ts, **kwargs)
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + th.zeros(broadcast_shape, device=timesteps.device)
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.flatten(1).mean(1)

point_e/diffusion/k_diffusion.py

@@ -0,0 +1,332 @@
+"""
+Based on: https://github.com/crowsonkb/k-diffusion
+
+Copyright (c) 2022 Katherine Crowson
+
+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.
+"""
+
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion, mean_flat
+
+
+class KarrasDenoiser:
+    def __init__(self, sigma_data: float = 0.5):
+        self.sigma_data = sigma_data
+
+    def get_snr(self, sigmas):
+        return sigmas**-2
+
+    def get_sigmas(self, sigmas):
+        return sigmas
+
+    def get_scalings(self, sigma):
+        c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2)
+        c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5
+        c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
+        return c_skip, c_out, c_in
+
+    def training_losses(self, model, x_start, sigmas, model_kwargs=None, noise=None):
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+
+        terms = {}
+
+        dims = x_start.ndim
+        x_t = x_start + noise * append_dims(sigmas, dims)
+        c_skip, c_out, _ = [append_dims(x, dims) for x in self.get_scalings(sigmas)]
+        model_output, denoised = self.denoise(model, x_t, sigmas, **model_kwargs)
+        target = (x_start - c_skip * x_t) / c_out
+
+        terms["mse"] = mean_flat((model_output - target) ** 2)
+        terms["xs_mse"] = mean_flat((denoised - x_start) ** 2)
+
+        if "vb" in terms:
+            terms["loss"] = terms["mse"] + terms["vb"]
+        else:
+            terms["loss"] = terms["mse"]
+
+        return terms
+
+    def denoise(self, model, x_t, sigmas, **model_kwargs):
+        c_skip, c_out, c_in = [append_dims(x, x_t.ndim) for x in self.get_scalings(sigmas)]
+        rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
+        model_output = model(c_in * x_t, rescaled_t, **model_kwargs)
+        denoised = c_out * model_output + c_skip * x_t
+        return model_output, denoised
+
+
+class GaussianToKarrasDenoiser:
+    def __init__(self, model, diffusion):
+        from scipy import interpolate
+
+        self.model = model
+        self.diffusion = diffusion
+        self.alpha_cumprod_to_t = interpolate.interp1d(
+            diffusion.alphas_cumprod, np.arange(0, diffusion.num_timesteps)
+        )
+
+    def sigma_to_t(self, sigma):
+        alpha_cumprod = 1.0 / (sigma**2 + 1)
+        if alpha_cumprod > self.diffusion.alphas_cumprod[0]:
+            return 0
+        elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]:
+            return self.diffusion.num_timesteps - 1
+        else:
+            return float(self.alpha_cumprod_to_t(alpha_cumprod))
+
+    def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None):
+        t = th.tensor(
+            [self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()],
+            dtype=th.long,
+            device=sigmas.device,
+        )
+        c_in = append_dims(1.0 / (sigmas**2 + 1) ** 0.5, x_t.ndim)
+        out = self.diffusion.p_mean_variance(
+            self.model, x_t * c_in, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        return None, out["pred_xstart"]
+
+
+def karras_sample(*args, **kwargs):
+    last = None
+    for x in karras_sample_progressive(*args, **kwargs):
+        last = x["x"]
+    return last
+
+
+def karras_sample_progressive(
+    diffusion,
+    model,
+    shape,
+    steps,
+    clip_denoised=True,
+    progress=False,
+    model_kwargs=None,
+    device=None,
+    sigma_min=0.002,
+    sigma_max=80,  # higher for highres?
+    rho=7.0,
+    sampler="heun",
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+    guidance_scale=0.0,
+):
+    sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
+    x_T = th.randn(*shape, device=device) * sigma_max
+    sample_fn = {"heun": sample_heun, "dpm": sample_dpm, "ancestral": sample_euler_ancestral}[
+        sampler
+    ]
+
+    if sampler != "ancestral":
+        sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise)
+    else:
+        sampler_args = {}
+
+    if isinstance(diffusion, KarrasDenoiser):
+
+        def denoiser(x_t, sigma):
+            _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
+            if clip_denoised:
+                denoised = denoised.clamp(-1, 1)
+            return denoised
+
+    elif isinstance(diffusion, GaussianDiffusion):
+        model = GaussianToKarrasDenoiser(model, diffusion)
+
+        def denoiser(x_t, sigma):
+            _, denoised = model.denoise(
+                x_t, sigma, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+            )
+            return denoised
+
+    else:
+        raise NotImplementedError
+
+    if guidance_scale != 0 and guidance_scale != 1:
+
+        def guided_denoiser(x_t, sigma):
+            x_t = th.cat([x_t, x_t], dim=0)
+            sigma = th.cat([sigma, sigma], dim=0)
+            x_0 = denoiser(x_t, sigma)
+            cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0)
+            x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0)
+            return x_0
+
+    else:
+        guided_denoiser = denoiser
+
+    for obj in sample_fn(
+        guided_denoiser,
+        x_T,
+        sigmas,
+        progress=progress,
+        **sampler_args,
+    ):
+        if isinstance(diffusion, GaussianDiffusion):
+            yield diffusion.unscale_out_dict(obj)
+        else:
+            yield obj
+
+
+def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device="cpu"):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    ramp = th.linspace(0, 1, n)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+    return append_zero(sigmas).to(device)
+
+
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+
+def get_ancestral_step(sigma_from, sigma_to):
+    """Calculates the noise level (sigma_down) to step down to and the amount
+    of noise to add (sigma_up) when doing an ancestral sampling step."""
+    sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
+    sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
+    return sigma_down, sigma_up
+
+
+@th.no_grad()
+def sample_euler_ancestral(model, x, sigmas, progress=False):
+    """Ancestral sampling with Euler method steps."""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+
+        indices = tqdm(indices)
+
+    for i in indices:
+        denoised = model(x, sigmas[i] * s_in)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
+        yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "pred_xstart": denoised}
+        d = to_d(x, sigmas[i], denoised)
+        # Euler method
+        dt = sigma_down - sigmas[i]
+        x = x + d * dt
+        x = x + th.randn_like(x) * sigma_up
+    yield {"x": x, "pred_xstart": x}
+
+
+@th.no_grad()
+def sample_heun(
+    denoiser,
+    x,
+    sigmas,
+    progress=False,
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+):
+    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+
+        indices = tqdm(indices)
+
+    for i in indices:
+        gamma = (
+            min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+        )
+        eps = th.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
+        denoised = denoiser(x, sigma_hat * s_in)
+        d = to_d(x, sigma_hat, denoised)
+        yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "pred_xstart": denoised}
+        dt = sigmas[i + 1] - sigma_hat
+        if sigmas[i + 1] == 0:
+            # Euler method
+            x = x + d * dt
+        else:
+            # Heun's method
+            x_2 = x + d * dt
+            denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
+            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+    yield {"x": x, "pred_xstart": denoised}
+
+
+@th.no_grad()
+def sample_dpm(
+    denoiser,
+    x,
+    sigmas,
+    progress=False,
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+):
+    """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+
+        indices = tqdm(indices)
+
+    for i in indices:
+        gamma = (
+            min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+        )
+        eps = th.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
+        denoised = denoiser(x, sigma_hat * s_in)
+        d = to_d(x, sigma_hat, denoised)
+        yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "denoised": denoised}
+        # Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule
+        sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2) ** 3
+        dt_1 = sigma_mid - sigma_hat
+        dt_2 = sigmas[i + 1] - sigma_hat
+        x_2 = x + d * dt_1
+        denoised_2 = denoiser(x_2, sigma_mid * s_in)
+        d_2 = to_d(x_2, sigma_mid, denoised_2)
+        x = x + d_2 * dt_2
+    yield {"x": x, "pred_xstart": denoised}
+
+
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
+    return x[(...,) + (None,) * dims_to_append]
+
+
+def append_zero(x):
+    return th.cat([x, x.new_zeros([1])])

point_e/diffusion/sampler.py

@@ -0,0 +1,263 @@
+"""
+Helpers for sampling from a single- or multi-stage point cloud diffusion model.
+"""
+
+from typing import Any, Callable, Dict, Iterator, List, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+
+from point_e.util.point_cloud import PointCloud
+
+from .gaussian_diffusion import GaussianDiffusion
+from .k_diffusion import karras_sample_progressive
+
+
+class PointCloudSampler:
+    """
+    A wrapper around a model or stack of models that produces conditional or
+    unconditional sample tensors.
+
+    By default, this will load models and configs from files.
+    If you want to modify the sampler arguments of an existing sampler, call
+    with_options() or with_args().
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        models: Sequence[nn.Module],
+        diffusions: Sequence[GaussianDiffusion],
+        num_points: Sequence[int],
+        aux_channels: Sequence[str],
+        model_kwargs_key_filter: Sequence[str] = ("*",),
+        guidance_scale: Sequence[float] = (3.0, 3.0),
+        clip_denoised: bool = True,
+        use_karras: Sequence[bool] = (True, True),
+        karras_steps: Sequence[int] = (64, 64),
+        sigma_min: Sequence[float] = (1e-3, 1e-3),
+        sigma_max: Sequence[float] = (120, 160),
+        s_churn: Sequence[float] = (3, 0),
+    ):
+        n = len(models)
+        assert n > 0
+
+        if n > 1:
+            if len(guidance_scale) == 1:
+                # Don't guide the upsamplers by default.
+                guidance_scale = list(guidance_scale) + [1.0] * (n - 1)
+            if len(use_karras) == 1:
+                use_karras = use_karras * n
+            if len(karras_steps) == 1:
+                karras_steps = karras_steps * n
+            if len(sigma_min) == 1:
+                sigma_min = sigma_min * n
+            if len(sigma_max) == 1:
+                sigma_max = sigma_max * n
+            if len(s_churn) == 1:
+                s_churn = s_churn * n
+            if len(model_kwargs_key_filter) == 1:
+                model_kwargs_key_filter = model_kwargs_key_filter * n
+        if len(model_kwargs_key_filter) == 0:
+            model_kwargs_key_filter = ["*"] * n
+        assert len(guidance_scale) == n
+        assert len(use_karras) == n
+        assert len(karras_steps) == n
+        assert len(sigma_min) == n
+        assert len(sigma_max) == n
+        assert len(s_churn) == n
+        assert len(model_kwargs_key_filter) == n
+
+        self.device = device
+        self.num_points = num_points
+        self.aux_channels = aux_channels
+        self.model_kwargs_key_filter = model_kwargs_key_filter
+        self.guidance_scale = guidance_scale
+        self.clip_denoised = clip_denoised
+        self.use_karras = use_karras
+        self.karras_steps = karras_steps
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.s_churn = s_churn
+
+        self.models = models
+        self.diffusions = diffusions
+
+    @property
+    def num_stages(self) -> int:
+        return len(self.models)
+
+    def sample_batch(self, batch_size: int, model_kwargs: Dict[str, Any]) -> torch.Tensor:
+        samples = None
+        for x in self.sample_batch_progressive(batch_size, model_kwargs):
+            samples = x
+        return samples
+
+    def sample_batch_progressive(
+        self, batch_size: int, model_kwargs: Dict[str, Any]
+    ) -> Iterator[torch.Tensor]:
+        samples = None
+        for (
+            model,
+            diffusion,
+            stage_num_points,
+            stage_guidance_scale,
+            stage_use_karras,
+            stage_karras_steps,
+            stage_sigma_min,
+            stage_sigma_max,
+            stage_s_churn,
+            stage_key_filter,
+        ) in zip(
+            self.models,
+            self.diffusions,
+            self.num_points,
+            self.guidance_scale,
+            self.use_karras,
+            self.karras_steps,
+            self.sigma_min,
+            self.sigma_max,
+            self.s_churn,
+            self.model_kwargs_key_filter,
+        ):
+            stage_model_kwargs = model_kwargs.copy()
+            if stage_key_filter != "*":
+                use_keys = set(stage_key_filter.split(","))
+                stage_model_kwargs = {k: v for k, v in stage_model_kwargs.items() if k in use_keys}
+            if samples is not None:
+                stage_model_kwargs["low_res"] = samples
+            if hasattr(model, "cached_model_kwargs"):
+                stage_model_kwargs = model.cached_model_kwargs(batch_size, stage_model_kwargs)
+            sample_shape = (batch_size, 3 + len(self.aux_channels), stage_num_points)
+
+            if stage_guidance_scale != 1 and stage_guidance_scale != 0:
+                for k, v in stage_model_kwargs.copy().items():
+                    stage_model_kwargs[k] = torch.cat([v, torch.zeros_like(v)], dim=0)
+
+            if stage_use_karras:
+                samples_it = karras_sample_progressive(
+                    diffusion=diffusion,
+                    model=model,
+                    shape=sample_shape,
+                    steps=stage_karras_steps,
+                    clip_denoised=self.clip_denoised,
+                    model_kwargs=stage_model_kwargs,
+                    device=self.device,
+                    sigma_min=stage_sigma_min,
+                    sigma_max=stage_sigma_max,
+                    s_churn=stage_s_churn,
+                    guidance_scale=stage_guidance_scale,
+                )
+            else:
+                internal_batch_size = batch_size
+                if stage_guidance_scale:
+                    model = self._uncond_guide_model(model, stage_guidance_scale)
+                    internal_batch_size *= 2
+                samples_it = diffusion.p_sample_loop_progressive(
+                    model,
+                    shape=(internal_batch_size, *sample_shape[1:]),
+                    model_kwargs=stage_model_kwargs,
+                    device=self.device,
+                    clip_denoised=self.clip_denoised,
+                )
+            for x in samples_it:
+                samples = x["pred_xstart"][:batch_size]
+                if "low_res" in stage_model_kwargs:
+                    samples = torch.cat(
+                        [stage_model_kwargs["low_res"][: len(samples)], samples], dim=-1
+                    )
+                yield samples
+
+    @classmethod
+    def combine(cls, *samplers: "PointCloudSampler") -> "PointCloudSampler":
+        assert all(x.device == samplers[0].device for x in samplers[1:])
+        assert all(x.aux_channels == samplers[0].aux_channels for x in samplers[1:])
+        assert all(x.clip_denoised == samplers[0].clip_denoised for x in samplers[1:])
+        return cls(
+            device=samplers[0].device,
+            models=[x for y in samplers for x in y.models],
+            diffusions=[x for y in samplers for x in y.diffusions],
+            num_points=[x for y in samplers for x in y.num_points],
+            aux_channels=samplers[0].aux_channels,
+            model_kwargs_key_filter=[x for y in samplers for x in y.model_kwargs_key_filter],
+            guidance_scale=[x for y in samplers for x in y.guidance_scale],
+            clip_denoised=samplers[0].clip_denoised,
+            use_karras=[x for y in samplers for x in y.use_karras],
+            karras_steps=[x for y in samplers for x in y.karras_steps],
+            sigma_min=[x for y in samplers for x in y.sigma_min],
+            sigma_max=[x for y in samplers for x in y.sigma_max],
+            s_churn=[x for y in samplers for x in y.s_churn],
+        )
+
+    def _uncond_guide_model(
+        self, model: Callable[..., torch.Tensor], scale: float
+    ) -> Callable[..., torch.Tensor]:
+        def model_fn(x_t, ts, **kwargs):
+            half = x_t[: len(x_t) // 2]
+            combined = torch.cat([half, half], dim=0)
+            model_out = model(combined, ts, **kwargs)
+            eps, rest = model_out[:, :3], model_out[:, 3:]
+            cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)
+            half_eps = uncond_eps + scale * (cond_eps - uncond_eps)
+            eps = torch.cat([half_eps, half_eps], dim=0)
+            return torch.cat([eps, rest], dim=1)
+
+        return model_fn
+
+    def split_model_output(
+        self,
+        output: torch.Tensor,
+        rescale_colors: bool = False,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        assert (
+            len(self.aux_channels) + 3 == output.shape[1]
+        ), "there must be three spatial channels before aux"
+        pos, joined_aux = output[:, :3], output[:, 3:]
+
+        aux = {}
+        for i, name in enumerate(self.aux_channels):
+            v = joined_aux[:, i]
+            if name in {"R", "G", "B", "A"}:
+                v = v.clamp(0, 255).round()
+                if rescale_colors:
+                    v = v / 255.0
+            aux[name] = v
+        return pos, aux
+
+    def output_to_point_clouds(self, output: torch.Tensor) -> List[PointCloud]:
+        res = []
+        for sample in output:
+            xyz, aux = self.split_model_output(sample[None], rescale_colors=True)
+            res.append(
+                PointCloud(
+                    coords=xyz[0].t().cpu().numpy(),
+                    channels={k: v[0].cpu().numpy() for k, v in aux.items()},
+                )
+            )
+        return res
+
+    def with_options(
+        self,
+        guidance_scale: float,
+        clip_denoised: bool,
+        use_karras: Sequence[bool] = (True, True),
+        karras_steps: Sequence[int] = (64, 64),
+        sigma_min: Sequence[float] = (1e-3, 1e-3),
+        sigma_max: Sequence[float] = (120, 160),
+        s_churn: Sequence[float] = (3, 0),
+    ) -> "PointCloudSampler":
+        return PointCloudSampler(
+            device=self.device,
+            models=self.models,
+            diffusions=self.diffusions,
+            num_points=self.num_points,
+            aux_channels=self.aux_channels,
+            model_kwargs_key_filter=self.model_kwargs_key_filter,
+            guidance_scale=guidance_scale,
+            clip_denoised=clip_denoised,
+            use_karras=use_karras,
+            karras_steps=karras_steps,
+            sigma_min=sigma_min,
+            sigma_max=sigma_max,
+            s_churn=s_churn,
+        )

point_e/evals/feature_extractor.py

@@ -0,0 +1,119 @@
+from abc import ABC, abstractmethod
+from multiprocessing.pool import ThreadPool
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from point_e.models.download import load_checkpoint
+
+from .npz_stream import NpzStreamer
+from .pointnet2_cls_ssg import get_model
+
+
+def get_torch_devices() -> List[Union[str, torch.device]]:
+    if torch.cuda.is_available():
+        return [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
+    else:
+        return ["cpu"]
+
+
+class FeatureExtractor(ABC):
+    @property
+    @abstractmethod
+    def supports_predictions(self) -> bool:
+        pass
+
+    @property
+    @abstractmethod
+    def feature_dim(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def num_classes(self) -> int:
+        pass
+
+    @abstractmethod
+    def features_and_preds(self, streamer: NpzStreamer) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        For a stream of point cloud batches, compute feature vectors and class
+        predictions.
+
+        :param point_clouds: a streamer for a sample batch. Typically, arr_0
+                             will contain the XYZ coordinates.
+        :return: a tuple (features, predictions)
+                 - features: a [B x feature_dim] array of feature vectors.
+                 - predictions: a [B x num_classes] array of probabilities.
+        """
+
+
+class PointNetClassifier(FeatureExtractor):
+    def __init__(
+        self,
+        devices: List[Union[str, torch.device]],
+        device_batch_size: int = 64,
+        cache_dir: Optional[str] = None,
+    ):
+        state_dict = load_checkpoint("pointnet", device=torch.device("cpu"), cache_dir=cache_dir)[
+            "model_state_dict"
+        ]
+
+        self.device_batch_size = device_batch_size
+        self.devices = devices
+        self.models = []
+        for device in devices:
+            model = get_model(num_class=40, normal_channel=False, width_mult=2)
+            model.load_state_dict(state_dict)
+            model.to(device)
+            model.eval()
+            self.models.append(model)
+
+    @property
+    def supports_predictions(self) -> bool:
+        return True
+
+    @property
+    def feature_dim(self) -> int:
+        return 256
+
+    @property
+    def num_classes(self) -> int:
+        return 40
+
+    def features_and_preds(self, streamer: NpzStreamer) -> Tuple[np.ndarray, np.ndarray]:
+        batch_size = self.device_batch_size * len(self.devices)
+        point_clouds = (x["arr_0"] for x in streamer.stream(batch_size, ["arr_0"]))
+
+        output_features = []
+        output_predictions = []
+
+        with ThreadPool(len(self.devices)) as pool:
+            for batch in point_clouds:
+                batch = normalize_point_clouds(batch)
+                batches = []
+                for i, device in zip(range(0, len(batch), self.device_batch_size), self.devices):
+                    batches.append(
+                        torch.from_numpy(batch[i : i + self.device_batch_size])
+                        .permute(0, 2, 1)
+                        .to(dtype=torch.float32, device=device)
+                    )
+
+                def compute_features(i_batch):
+                    i, batch = i_batch
+                    with torch.no_grad():
+                        return self.models[i](batch, features=True)
+
+                for logits, _, features in pool.imap(compute_features, enumerate(batches)):
+                    output_features.append(features.cpu().numpy())
+                    output_predictions.append(logits.exp().cpu().numpy())
+
+        return np.concatenate(output_features, axis=0), np.concatenate(output_predictions, axis=0)
+
+
+def normalize_point_clouds(pc: np.ndarray) -> np.ndarray:
+    centroids = np.mean(pc, axis=1, keepdims=True)
+    pc = pc - centroids
+    m = np.max(np.sqrt(np.sum(pc**2, axis=-1, keepdims=True)), axis=1, keepdims=True)
+    pc = pc / m
+    return pc

point_e/evals/fid_is.py

@@ -0,0 +1,81 @@
+"""
+Adapted from https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/evaluations/evaluator.py
+"""
+
+
+import warnings
+
+import numpy as np
+from scipy import linalg
+
+
+class InvalidFIDException(Exception):
+    pass
+
+
+class FIDStatistics:
+    def __init__(self, mu: np.ndarray, sigma: np.ndarray):
+        self.mu = mu
+        self.sigma = sigma
+
+    def frechet_distance(self, other, eps=1e-6):
+        """
+        Compute the Frechet distance between two sets of statistics.
+        """
+        # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L132
+        mu1, sigma1 = self.mu, self.sigma
+        mu2, sigma2 = other.mu, other.sigma
+
+        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
+        ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}"
+        assert (
+            sigma1.shape == sigma2.shape
+        ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}"
+
+        diff = mu1 - mu2
+
+        # product might be almost singular
+        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+        if not np.isfinite(covmean).all():
+            msg = (
+                "fid calculation produces singular product; adding %s to diagonal of cov estimates"
+                % eps
+            )
+            warnings.warn(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 compute_statistics(feats: np.ndarray) -> FIDStatistics:
+    mu = np.mean(feats, axis=0)
+    sigma = np.cov(feats, rowvar=False)
+    return FIDStatistics(mu, sigma)
+
+
+def compute_inception_score(preds: np.ndarray, split_size: int = 5000) -> float:
+    # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46
+    scores = []
+    for i in range(0, len(preds), split_size):
+        part = preds[i : i + split_size]
+        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 float(np.mean(scores))

point_e/evals/npz_stream.py

@@ -0,0 +1,270 @@
+import glob
+import io
+import os
+import re
+import zipfile
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Dict, Iterator, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+
+@dataclass
+class NumpyArrayInfo:
+    """
+    Information about an array in an npz file.
+    """
+
+    name: str
+    dtype: np.dtype
+    shape: Tuple[int]
+
+    @classmethod
+    def infos_from_first_file(cls, glob_path: str) -> Dict[str, "NumpyArrayInfo"]:
+        paths, _ = _npz_paths_and_length(glob_path)
+        return cls.infos_from_file(paths[0])
+
+    @classmethod
+    def infos_from_file(cls, npz_path: str) -> Dict[str, "NumpyArrayInfo"]:
+        """
+        Extract the info of every array in an npz file.
+        """
+        if not os.path.exists(npz_path):
+            raise FileNotFoundError(f"batch of samples was not found: {npz_path}")
+        results = {}
+        with open(npz_path, "rb") as f:
+            with zipfile.ZipFile(f, "r") as zip_f:
+                for name in zip_f.namelist():
+                    if not name.endswith(".npy"):
+                        continue
+                    key_name = name[: -len(".npy")]
+                    with zip_f.open(name, "r") as arr_f:
+                        version = np.lib.format.read_magic(arr_f)
+                        if version == (1, 0):
+                            header = np.lib.format.read_array_header_1_0(arr_f)
+                        elif version == (2, 0):
+                            header = np.lib.format.read_array_header_2_0(arr_f)
+                        else:
+                            raise ValueError(f"unknown numpy array version: {version}")
+                        shape, _, dtype = header
+                        results[key_name] = cls(name=key_name, dtype=dtype, shape=shape)
+        return results
+
+    @property
+    def elem_shape(self) -> Tuple[int]:
+        return self.shape[1:]
+
+    def validate(self):
+        if self.name in {"R", "G", "B"}:
+            if len(self.shape) != 2:
+                raise ValueError(
+                    f"expecting exactly 2-D shape for '{self.name}' but got: {self.shape}"
+                )
+        elif self.name == "arr_0":
+            if len(self.shape) < 2:
+                raise ValueError(f"expecting at least 2-D shape but got: {self.shape}")
+            elif len(self.shape) == 3:
+                # For audio, we require continuous samples.
+                if not np.issubdtype(self.dtype, np.floating):
+                    raise ValueError(
+                        f"invalid dtype for audio batch: {self.dtype} (expected float)"
+                    )
+            elif self.dtype != np.uint8:
+                raise ValueError(f"invalid dtype for image batch: {self.dtype} (expected uint8)")
+
+
+class NpzStreamer:
+    def __init__(self, glob_path: str):
+        self.paths, self.trunc_length = _npz_paths_and_length(glob_path)
+        self.infos = NumpyArrayInfo.infos_from_file(self.paths[0])
+
+    def keys(self) -> List[str]:
+        return list(self.infos.keys())
+
+    def stream(self, batch_size: int, keys: Sequence[str]) -> Iterator[Dict[str, np.ndarray]]:
+        cur_batch = None
+        num_remaining = self.trunc_length
+        for path in self.paths:
+            if num_remaining is not None and num_remaining <= 0:
+                break
+            with open_npz_arrays(path, keys) as readers:
+                combined_reader = CombinedReader(keys, readers)
+                while num_remaining is None or num_remaining > 0:
+                    read_bs = batch_size
+                    if cur_batch is not None:
+                        read_bs -= _dict_batch_size(cur_batch)
+                    if num_remaining is not None:
+                        read_bs = min(read_bs, num_remaining)
+
+                    batch = combined_reader.read_batch(read_bs)
+                    if batch is None:
+                        break
+                    if num_remaining is not None:
+                        num_remaining -= _dict_batch_size(batch)
+                    if cur_batch is None:
+                        cur_batch = batch
+                    else:
+                        cur_batch = {
+                            # pylint: disable=unsubscriptable-object
+                            k: np.concatenate([cur_batch[k], v], axis=0)
+                            for k, v in batch.items()
+                        }
+                    if _dict_batch_size(cur_batch) == batch_size:
+                        yield cur_batch
+                        cur_batch = None
+        if cur_batch is not None:
+            yield cur_batch
+
+
+def _npz_paths_and_length(glob_path: str) -> Tuple[List[str], Optional[int]]:
+    # Match slice syntax like path[:100].
+    count_match = re.match("^(.*)\\[:([0-9]*)\\]$", glob_path)
+    if count_match:
+        raw_path = count_match[1]
+        max_count = int(count_match[2])
+    else:
+        raw_path = glob_path
+        max_count = None
+    paths = sorted(glob.glob(raw_path))
+    if not len(paths):
+        raise ValueError(f"no paths found matching: {glob_path}")
+    return paths, max_count
+
+
+class NpzArrayReader(ABC):
+    @abstractmethod
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        pass
+
+
+class StreamingNpzArrayReader(NpzArrayReader):
+    def __init__(self, arr_f, shape, dtype):
+        self.arr_f = arr_f
+        self.shape = shape
+        self.dtype = dtype
+        self.idx = 0
+
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        if self.idx >= self.shape[0]:
+            return None
+
+        bs = min(batch_size, self.shape[0] - self.idx)
+        self.idx += bs
+
+        if self.dtype.itemsize == 0:
+            return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
+
+        read_count = bs * np.prod(self.shape[1:])
+        read_size = int(read_count * self.dtype.itemsize)
+        data = _read_bytes(self.arr_f, read_size, "array data")
+        return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
+
+
+class MemoryNpzArrayReader(NpzArrayReader):
+    def __init__(self, arr):
+        self.arr = arr
+        self.idx = 0
+
+    @classmethod
+    def load(cls, path: str, arr_name: str):
+        with open(path, "rb") as f:
+            arr = np.load(f)[arr_name]
+        return cls(arr)
+
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        if self.idx >= self.arr.shape[0]:
+            return None
+
+        res = self.arr[self.idx : self.idx + batch_size]
+        self.idx += batch_size
+        return res
+
+
+@contextmanager
+def open_npz_arrays(path: str, arr_names: Sequence[str]) -> List[NpzArrayReader]:
+    if not len(arr_names):
+        yield []
+        return
+    arr_name = arr_names[0]
+    with open_array(path, arr_name) as arr_f:
+        version = np.lib.format.read_magic(arr_f)
+        header = None
+        if version == (1, 0):
+            header = np.lib.format.read_array_header_1_0(arr_f)
+        elif version == (2, 0):
+            header = np.lib.format.read_array_header_2_0(arr_f)
+
+        if header is None:
+            reader = MemoryNpzArrayReader.load(path, arr_name)
+        else:
+            shape, fortran, dtype = header
+            if fortran or dtype.hasobject:
+                reader = MemoryNpzArrayReader.load(path, arr_name)
+            else:
+                reader = StreamingNpzArrayReader(arr_f, shape, dtype)
+
+        with open_npz_arrays(path, arr_names[1:]) as next_readers:
+            yield [reader] + next_readers
+
+
+class CombinedReader:
+    def __init__(self, keys: List[str], readers: List[NpzArrayReader]):
+        self.keys = keys
+        self.readers = readers
+
+    def read_batch(self, batch_size: int) -> Optional[Dict[str, np.ndarray]]:
+        batches = [r.read_batch(batch_size) for r in self.readers]
+        any_none = any(x is None for x in batches)
+        all_none = all(x is None for x in batches)
+        if any_none != all_none:
+            raise RuntimeError("different keys had different numbers of elements")
+        if any_none:
+            return None
+        if any(len(x) != len(batches[0]) for x in batches):
+            raise RuntimeError("different keys had different numbers of elements")
+        return dict(zip(self.keys, batches))
+
+
+def _read_bytes(fp, size, error_template="ran out of data"):
+    """
+    Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886
+
+    Read from file-like object until size bytes are read.
+    Raises ValueError if not EOF is encountered before size bytes are read.
+    Non-blocking objects only supported if they derive from io objects.
+    Required as e.g. ZipExtFile in python 2.6 can return less data than
+    requested.
+    """
+    data = bytes()
+    while True:
+        # io files (default in python3) return None or raise on
+        # would-block, python2 file will truncate, probably nothing can be
+        # done about that.  note that regular files can't be non-blocking
+        try:
+            r = fp.read(size - len(data))
+            data += r
+            if len(r) == 0 or len(data) == size:
+                break
+        except io.BlockingIOError:
+            pass
+    if len(data) != size:
+        msg = "EOF: reading %s, expected %d bytes got %d"
+        raise ValueError(msg % (error_template, size, len(data)))
+    else:
+        return data
+
+
+@contextmanager
+def open_array(path: str, arr_name: str):
+    with open(path, "rb") as f:
+        with zipfile.ZipFile(f, "r") as zip_f:
+            if f"{arr_name}.npy" not in zip_f.namelist():
+                raise ValueError(f"missing {arr_name} in npz file")
+            with zip_f.open(f"{arr_name}.npy", "r") as arr_f:
+                yield arr_f
+
+
+def _dict_batch_size(objs: Dict[str, np.ndarray]) -> int:
+    return len(next(iter(objs.values())))

point_e/evals/pointnet2_cls_ssg.py

@@ -0,0 +1,101 @@
+"""
+Based on: https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/eb64fe0b4c24055559cea26299cb485dcb43d8dd/models/pointnet2_cls_ssg.py
+
+MIT License
+
+Copyright (c) 2019 benny
+
+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.
+"""
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .pointnet2_utils import PointNetSetAbstraction
+
+
+class get_model(nn.Module):
+    def __init__(self, num_class, normal_channel=True, width_mult=1):
+        super(get_model, self).__init__()
+        self.width_mult = width_mult
+        in_channel = 6 if normal_channel else 3
+        self.normal_channel = normal_channel
+        self.sa1 = PointNetSetAbstraction(
+            npoint=512,
+            radius=0.2,
+            nsample=32,
+            in_channel=in_channel,
+            mlp=[64 * width_mult, 64 * width_mult, 128 * width_mult],
+            group_all=False,
+        )
+        self.sa2 = PointNetSetAbstraction(
+            npoint=128,
+            radius=0.4,
+            nsample=64,
+            in_channel=128 * width_mult + 3,
+            mlp=[128 * width_mult, 128 * width_mult, 256 * width_mult],
+            group_all=False,
+        )
+        self.sa3 = PointNetSetAbstraction(
+            npoint=None,
+            radius=None,
+            nsample=None,
+            in_channel=256 * width_mult + 3,
+            mlp=[256 * width_mult, 512 * width_mult, 1024 * width_mult],
+            group_all=True,
+        )
+        self.fc1 = nn.Linear(1024 * width_mult, 512 * width_mult)
+        self.bn1 = nn.BatchNorm1d(512 * width_mult)
+        self.drop1 = nn.Dropout(0.4)
+        self.fc2 = nn.Linear(512 * width_mult, 256 * width_mult)
+        self.bn2 = nn.BatchNorm1d(256 * width_mult)
+        self.drop2 = nn.Dropout(0.4)
+        self.fc3 = nn.Linear(256 * width_mult, num_class)
+
+    def forward(self, xyz, features=False):
+        B, _, _ = xyz.shape
+        if self.normal_channel:
+            norm = xyz[:, 3:, :]
+            xyz = xyz[:, :3, :]
+        else:
+            norm = None
+        l1_xyz, l1_points = self.sa1(xyz, norm)
+        l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
+        l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
+        x = l3_points.view(B, 1024 * self.width_mult)
+        x = self.drop1(F.relu(self.bn1(self.fc1(x))))
+        result_features = self.bn2(self.fc2(x))
+        x = self.drop2(F.relu(result_features))
+        x = self.fc3(x)
+        x = F.log_softmax(x, -1)
+
+        if features:
+            return x, l3_points, result_features
+        else:
+            return x, l3_points
+
+
+class get_loss(nn.Module):
+    def __init__(self):
+        super(get_loss, self).__init__()
+
+    def forward(self, pred, target, trans_feat):
+        total_loss = F.nll_loss(pred, target)
+
+        return total_loss

point_e/evals/pointnet2_utils.py

@@ -0,0 +1,356 @@
+"""
+Based on: https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/eb64fe0b4c24055559cea26299cb485dcb43d8dd/models/pointnet_utils.py
+
+MIT License
+
+Copyright (c) 2019 benny
+
+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.
+"""
+
+from time import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def timeit(tag, t):
+    print("{}: {}s".format(tag, time() - t))
+    return time()
+
+
+def pc_normalize(pc):
+    l = pc.shape[0]
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+def square_distance(src, dst):
+    """
+    Calculate Euclid distance between each two points.
+
+    src^T * dst = xn * xm + yn * ym + zn * zm;
+    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+
+    Input:
+        src: source points, [B, N, C]
+        dst: target points, [B, M, C]
+    Output:
+        dist: per-point square distance, [B, N, M]
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src**2, -1).view(B, N, 1)
+    dist += torch.sum(dst**2, -1).view(B, 1, M)
+    return dist
+
+
+def index_points(points, idx):
+    """
+
+    Input:
+        points: input points data, [B, N, C]
+        idx: sample index data, [B, S]
+    Return:
+        new_points:, indexed points data, [B, S, C]
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = (
+        torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
+    )
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+
+def farthest_point_sample(xyz, npoint, deterministic=False):
+    """
+    Input:
+        xyz: pointcloud data, [B, N, 3]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [B, npoint]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+    distance = torch.ones(B, N).to(device) * 1e10
+    if deterministic:
+        farthest = torch.arange(0, B, dtype=torch.long).to(device)
+    else:
+        farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
+    batch_indices = torch.arange(B, dtype=torch.long).to(device)
+    for i in range(npoint):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+    return centroids
+
+
+def query_ball_point(radius, nsample, xyz, new_xyz):
+    """
+    Input:
+        radius: local region radius
+        nsample: max sample number in local region
+        xyz: all points, [B, N, 3]
+        new_xyz: query points, [B, S, 3]
+    Return:
+        group_idx: grouped points index, [B, S, nsample]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    _, S, _ = new_xyz.shape
+    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
+    sqrdists = square_distance(new_xyz, xyz)
+    group_idx[sqrdists > radius**2] = N
+    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
+    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
+    mask = group_idx == N
+    group_idx[mask] = group_first[mask]
+    return group_idx
+
+
+def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False, deterministic=False):
+    """
+    Input:
+        npoint:
+        radius:
+        nsample:
+        xyz: input points position data, [B, N, 3]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, npoint, nsample, 3]
+        new_points: sampled points data, [B, npoint, nsample, 3+D]
+    """
+    B, N, C = xyz.shape
+    S = npoint
+    fps_idx = farthest_point_sample(xyz, npoint, deterministic=deterministic)  # [B, npoint, C]
+    new_xyz = index_points(xyz, fps_idx)
+    idx = query_ball_point(radius, nsample, xyz, new_xyz)
+    grouped_xyz = index_points(xyz, idx)  # [B, npoint, nsample, C]
+    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
+
+    if points is not None:
+        grouped_points = index_points(points, idx)
+        new_points = torch.cat(
+            [grouped_xyz_norm, grouped_points], dim=-1
+        )  # [B, npoint, nsample, C+D]
+    else:
+        new_points = grouped_xyz_norm
+    if returnfps:
+        return new_xyz, new_points, grouped_xyz, fps_idx
+    else:
+        return new_xyz, new_points
+
+
+def sample_and_group_all(xyz, points):
+    """
+    Input:
+        xyz: input points position data, [B, N, 3]
+        points: input points data, [B, N, D]
+    Return:
+        new_xyz: sampled points position data, [B, 1, 3]
+        new_points: sampled points data, [B, 1, N, 3+D]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    new_xyz = torch.zeros(B, 1, C).to(device)
+    grouped_xyz = xyz.view(B, 1, N, C)
+    if points is not None:
+        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
+    else:
+        new_points = grouped_xyz
+    return new_xyz, new_points
+
+
+class PointNetSetAbstraction(nn.Module):
+    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
+        super(PointNetSetAbstraction, self).__init__()
+        self.npoint = npoint
+        self.radius = radius
+        self.nsample = nsample
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        last_channel = in_channel
+        for out_channel in mlp:
+            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
+            last_channel = out_channel
+        self.group_all = group_all
+
+    def forward(self, xyz, points):
+        """
+        Input:
+            xyz: input points position data, [B, C, N]
+            points: input points data, [B, D, N]
+        Return:
+            new_xyz: sampled points position data, [B, C, S]
+            new_points_concat: sample points feature data, [B, D', S]
+        """
+        xyz = xyz.permute(0, 2, 1)
+        if points is not None:
+            points = points.permute(0, 2, 1)
+
+        if self.group_all:
+            new_xyz, new_points = sample_and_group_all(xyz, points)
+        else:
+            new_xyz, new_points = sample_and_group(
+                self.npoint, self.radius, self.nsample, xyz, points, deterministic=not self.training
+            )
+        # new_xyz: sampled points position data, [B, npoint, C]
+        # new_points: sampled points data, [B, npoint, nsample, C+D]
+        new_points = new_points.permute(0, 3, 2, 1)  # [B, C+D, nsample,npoint]
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            new_points = F.relu(bn(conv(new_points)))
+
+        new_points = torch.max(new_points, 2)[0]
+        new_xyz = new_xyz.permute(0, 2, 1)
+        return new_xyz, new_points
+
+
+class PointNetSetAbstractionMsg(nn.Module):
+    def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
+        super(PointNetSetAbstractionMsg, self).__init__()
+        self.npoint = npoint
+        self.radius_list = radius_list
+        self.nsample_list = nsample_list
+        self.conv_blocks = nn.ModuleList()
+        self.bn_blocks = nn.ModuleList()
+        for i in range(len(mlp_list)):
+            convs = nn.ModuleList()
+            bns = nn.ModuleList()
+            last_channel = in_channel + 3
+            for out_channel in mlp_list[i]:
+                convs.append(nn.Conv2d(last_channel, out_channel, 1))
+                bns.append(nn.BatchNorm2d(out_channel))
+                last_channel = out_channel
+            self.conv_blocks.append(convs)
+            self.bn_blocks.append(bns)
+
+    def forward(self, xyz, points):
+        """
+        Input:
+            xyz: input points position data, [B, C, N]
+            points: input points data, [B, D, N]
+        Return:
+            new_xyz: sampled points position data, [B, C, S]
+            new_points_concat: sample points feature data, [B, D', S]
+        """
+        xyz = xyz.permute(0, 2, 1)
+        if points is not None:
+            points = points.permute(0, 2, 1)
+
+        B, N, C = xyz.shape
+        S = self.npoint
+        new_xyz = index_points(xyz, farthest_point_sample(xyz, S, deterministic=not self.training))
+        new_points_list = []
+        for i, radius in enumerate(self.radius_list):
+            K = self.nsample_list[i]
+            group_idx = query_ball_point(radius, K, xyz, new_xyz)
+            grouped_xyz = index_points(xyz, group_idx)
+            grouped_xyz -= new_xyz.view(B, S, 1, C)
+            if points is not None:
+                grouped_points = index_points(points, group_idx)
+                grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
+            else:
+                grouped_points = grouped_xyz
+
+            grouped_points = grouped_points.permute(0, 3, 2, 1)  # [B, D, K, S]
+            for j in range(len(self.conv_blocks[i])):
+                conv = self.conv_blocks[i][j]
+                bn = self.bn_blocks[i][j]
+                grouped_points = F.relu(bn(conv(grouped_points)))
+            new_points = torch.max(grouped_points, 2)[0]  # [B, D', S]
+            new_points_list.append(new_points)
+
+        new_xyz = new_xyz.permute(0, 2, 1)
+        new_points_concat = torch.cat(new_points_list, dim=1)
+        return new_xyz, new_points_concat
+
+
+class PointNetFeaturePropagation(nn.Module):
+    def __init__(self, in_channel, mlp):
+        super(PointNetFeaturePropagation, self).__init__()
+        self.mlp_convs = nn.ModuleList()
+        self.mlp_bns = nn.ModuleList()
+        last_channel = in_channel
+        for out_channel in mlp:
+            self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
+            self.mlp_bns.append(nn.BatchNorm1d(out_channel))
+            last_channel = out_channel
+
+    def forward(self, xyz1, xyz2, points1, points2):
+        """
+        Input:
+            xyz1: input points position data, [B, C, N]
+            xyz2: sampled input points position data, [B, C, S]
+            points1: input points data, [B, D, N]
+            points2: input points data, [B, D, S]
+        Return:
+            new_points: upsampled points data, [B, D', N]
+        """
+        xyz1 = xyz1.permute(0, 2, 1)
+        xyz2 = xyz2.permute(0, 2, 1)
+
+        points2 = points2.permute(0, 2, 1)
+        B, N, C = xyz1.shape
+        _, S, _ = xyz2.shape
+
+        if S == 1:
+            interpolated_points = points2.repeat(1, N, 1)
+        else:
+            dists = square_distance(xyz1, xyz2)
+            dists, idx = dists.sort(dim=-1)
+            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]
+
+            dist_recip = 1.0 / (dists + 1e-8)
+            norm = torch.sum(dist_recip, dim=2, keepdim=True)
+            weight = dist_recip / norm
+            interpolated_points = torch.sum(
+                index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2
+            )
+
+        if points1 is not None:
+            points1 = points1.permute(0, 2, 1)
+            new_points = torch.cat([points1, interpolated_points], dim=-1)
+        else:
+            new_points = interpolated_points
+
+        new_points = new_points.permute(0, 2, 1)
+        for i, conv in enumerate(self.mlp_convs):
+            bn = self.mlp_bns[i]
+            new_points = F.relu(bn(conv(new_points)))
+        return new_points

point_e/evals/scripts/blender_script.py

@@ -0,0 +1,533 @@
+"""
+Script to run within Blender to render a 3D model as RGBAD images.
+
+Example usage
+
+    blender -b -P blender_script.py -- \
+        --input_path ../../examples/example_data/corgi.ply \
+        --output_path render_out
+
+Pass `--camera_pose z-circular-elevated` for the rendering used to compute
+CLIP R-Precision results.
+
+The output directory will include metadata json files for each rendered view,
+as well as a global metadata file for the render. Each image will be saved as
+a collection of 16-bit PNG files for each channel (rgbad), as well as a full
+grayscale render of the view.
+"""
+
+import argparse
+import json
+import math
+import os
+import random
+import sys
+
+import bpy
+from mathutils import Vector
+from mathutils.noise import random_unit_vector
+
+MAX_DEPTH = 5.0
+FORMAT_VERSION = 6
+UNIFORM_LIGHT_DIRECTION = [0.09387503, -0.63953443, -0.7630093]
+
+
+def clear_scene():
+    bpy.ops.object.select_all(action="SELECT")
+    bpy.ops.object.delete()
+
+
+def clear_lights():
+    bpy.ops.object.select_all(action="DESELECT")
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, bpy.types.Light):
+            obj.select_set(True)
+    bpy.ops.object.delete()
+
+
+def import_model(path):
+    clear_scene()
+    _, ext = os.path.splitext(path)
+    ext = ext.lower()
+    if ext == ".obj":
+        bpy.ops.import_scene.obj(filepath=path)
+    elif ext in [".glb", ".gltf"]:
+        bpy.ops.import_scene.gltf(filepath=path)
+    elif ext == ".stl":
+        bpy.ops.import_mesh.stl(filepath=path)
+    elif ext == ".fbx":
+        bpy.ops.import_scene.fbx(filepath=path)
+    elif ext == ".dae":
+        bpy.ops.wm.collada_import(filepath=path)
+    elif ext == ".ply":
+        bpy.ops.import_mesh.ply(filepath=path)
+    else:
+        raise RuntimeError(f"unexpected extension: {ext}")
+
+
+def scene_root_objects():
+    for obj in bpy.context.scene.objects.values():
+        if not obj.parent:
+            yield obj
+
+
+def scene_bbox(single_obj=None, ignore_matrix=False):
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    for obj in scene_meshes() if single_obj is None else [single_obj]:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            if not ignore_matrix:
+                coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+    return Vector(bbox_min), Vector(bbox_max)
+
+
+def scene_meshes():
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, (bpy.types.Mesh)):
+            yield obj
+
+
+def normalize_scene():
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+
+    for obj in scene_root_objects():
+        obj.scale = obj.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    for obj in scene_root_objects():
+        obj.matrix_world.translation += offset
+
+    bpy.ops.object.select_all(action="DESELECT")
+
+
+def create_camera():
+    # https://b3d.interplanety.org/en/how-to-create-camera-through-the-blender-python-api/
+    camera_data = bpy.data.cameras.new(name="Camera")
+    camera_object = bpy.data.objects.new("Camera", camera_data)
+    bpy.context.scene.collection.objects.link(camera_object)
+    bpy.context.scene.camera = camera_object
+
+
+def set_camera(direction, camera_dist=2.0):
+    camera_pos = -camera_dist * direction
+    bpy.context.scene.camera.location = camera_pos
+
+    # https://blender.stackexchange.com/questions/5210/pointing-the-camera-in-a-particular-direction-programmatically
+    rot_quat = direction.to_track_quat("-Z", "Y")
+    bpy.context.scene.camera.rotation_euler = rot_quat.to_euler()
+
+    bpy.context.view_layer.update()
+
+
+def randomize_camera(camera_dist=2.0):
+    direction = random_unit_vector()
+    set_camera(direction, camera_dist=camera_dist)
+
+
+def pan_camera(time, axis="Z", camera_dist=2.0, elevation=-0.1):
+    angle = time * math.pi * 2
+    direction = [-math.cos(angle), -math.sin(angle), -elevation]
+    assert axis in ["X", "Y", "Z"]
+    if axis == "X":
+        direction = [direction[2], *direction[:2]]
+    elif axis == "Y":
+        direction = [direction[0], -elevation, direction[1]]
+    direction = Vector(direction).normalized()
+    set_camera(direction, camera_dist=camera_dist)
+
+
+def place_camera(time, camera_pose_mode="random", camera_dist_min=2.0, camera_dist_max=2.0):
+    camera_dist = random.uniform(camera_dist_min, camera_dist_max)
+    if camera_pose_mode == "random":
+        randomize_camera(camera_dist=camera_dist)
+    elif camera_pose_mode == "z-circular":
+        pan_camera(time, axis="Z", camera_dist=camera_dist)
+    elif camera_pose_mode == "z-circular-elevated":
+        pan_camera(time, axis="Z", camera_dist=camera_dist, elevation=0.2617993878)
+    else:
+        raise ValueError(f"Unknown camera pose mode: {camera_pose_mode}")
+
+
+def create_light(location, energy=1.0, angle=0.5 * math.pi / 180):
+    # https://blender.stackexchange.com/questions/215624/how-to-create-a-light-with-the-python-api-in-blender-2-92
+    light_data = bpy.data.lights.new(name="Light", type="SUN")
+    light_data.energy = energy
+    light_data.angle = angle
+    light_object = bpy.data.objects.new(name="Light", object_data=light_data)
+
+    direction = -location
+    rot_quat = direction.to_track_quat("-Z", "Y")
+    light_object.rotation_euler = rot_quat.to_euler()
+    bpy.context.view_layer.update()
+
+    bpy.context.collection.objects.link(light_object)
+    light_object.location = location
+
+
+def create_random_lights(count=4, distance=2.0, energy=1.5):
+    clear_lights()
+    for _ in range(count):
+        create_light(random_unit_vector() * distance, energy=energy)
+
+
+def create_camera_light():
+    clear_lights()
+    create_light(bpy.context.scene.camera.location, energy=5.0)
+
+
+def create_uniform_light(backend):
+    clear_lights()
+    # Random direction to decorrelate axis-aligned sides.
+    pos = Vector(UNIFORM_LIGHT_DIRECTION)
+    angle = 0.0092 if backend == "CYCLES" else math.pi
+    create_light(pos, energy=5.0, angle=angle)
+    create_light(-pos, energy=5.0, angle=angle)
+
+
+def create_vertex_color_shaders():
+    # By default, Blender will ignore vertex colors in both the
+    # Eevee and Cycles backends, since these colors aren't
+    # associated with a material.
+    #
+    # What we do here is create a simple material shader and link
+    # the vertex color to the material color.
+    for obj in bpy.context.scene.objects.values():
+        if not isinstance(obj.data, (bpy.types.Mesh)):
+            continue
+
+        if len(obj.data.materials):
+            # We don't want to override any existing materials.
+            continue
+
+        color_keys = (obj.data.vertex_colors or {}).keys()
+        if not len(color_keys):
+            # Many objects will have no materials *or* vertex colors.
+            continue
+
+        mat = bpy.data.materials.new(name="VertexColored")
+        mat.use_nodes = True
+
+        # There should be a Principled BSDF by default.
+        bsdf_node = None
+        for node in mat.node_tree.nodes:
+            if node.type == "BSDF_PRINCIPLED":
+                bsdf_node = node
+        assert bsdf_node is not None, "material has no Principled BSDF node to modify"
+
+        socket_map = {}
+        for input in bsdf_node.inputs:
+            socket_map[input.name] = input
+
+        # Make sure nothing lights the object except for the diffuse color.
+        socket_map["Specular"].default_value = 0.0
+        socket_map["Roughness"].default_value = 1.0
+
+        v_color = mat.node_tree.nodes.new("ShaderNodeVertexColor")
+        v_color.layer_name = color_keys[0]
+
+        mat.node_tree.links.new(v_color.outputs[0], socket_map["Base Color"])
+
+        obj.data.materials.append(mat)
+
+
+def create_default_materials():
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, (bpy.types.Mesh)):
+            if not len(obj.data.materials):
+                mat = bpy.data.materials.new(name="DefaultMaterial")
+                mat.use_nodes = True
+                obj.data.materials.append(mat)
+
+
+def find_materials():
+    all_materials = set()
+    for obj in bpy.context.scene.objects.values():
+        if not isinstance(obj.data, (bpy.types.Mesh)):
+            continue
+        for mat in obj.data.materials:
+            all_materials.add(mat)
+    return all_materials
+
+
+def get_socket_value(tree, socket):
+    default = socket.default_value
+    if not isinstance(default, float):
+        default = list(default)
+    for link in tree.links:
+        if link.to_socket == socket:
+            return (link.from_socket, default)
+    return (None, default)
+
+
+def clear_socket_input(tree, socket):
+    for link in list(tree.links):
+        if link.to_socket == socket:
+            tree.links.remove(link)
+
+
+def set_socket_value(tree, socket, socket_and_default):
+    clear_socket_input(tree, socket)
+    old_source_socket, default = socket_and_default
+    if isinstance(default, float) and not isinstance(socket.default_value, float):
+        # Codepath for setting Emission to a previous alpha value.
+        socket.default_value = [default] * 3 + [1.0]
+    else:
+        socket.default_value = default
+    if old_source_socket is not None:
+        tree.links.new(old_source_socket, socket)
+
+
+def setup_nodes(output_path, capturing_material_alpha: bool = False):
+    tree = bpy.context.scene.node_tree
+    links = tree.links
+
+    for node in tree.nodes:
+        tree.nodes.remove(node)
+
+    # Helpers to perform math on links and constants.
+    def node_op(op: str, *args, clamp=False):
+        node = tree.nodes.new(type="CompositorNodeMath")
+        node.operation = op
+        if clamp:
+            node.use_clamp = True
+        for i, arg in enumerate(args):
+            if isinstance(arg, (int, float)):
+                node.inputs[i].default_value = arg
+            else:
+                links.new(arg, node.inputs[i])
+        return node.outputs[0]
+
+    def node_clamp(x, maximum=1.0):
+        return node_op("MINIMUM", x, maximum)
+
+    def node_mul(x, y, **kwargs):
+        return node_op("MULTIPLY", x, y, **kwargs)
+
+    input_node = tree.nodes.new(type="CompositorNodeRLayers")
+    input_node.scene = bpy.context.scene
+
+    input_sockets = {}
+    for output in input_node.outputs:
+        input_sockets[output.name] = output
+
+    if capturing_material_alpha:
+        color_socket = input_sockets["Image"]
+    else:
+        raw_color_socket = input_sockets["Image"]
+
+        # We apply sRGB here so that our fixed-point depth map and material
+        # alpha values are not sRGB, and so that we perform ambient+diffuse
+        # lighting in linear RGB space.
+        color_node = tree.nodes.new(type="CompositorNodeConvertColorSpace")
+        color_node.from_color_space = "Linear"
+        color_node.to_color_space = "sRGB"
+        tree.links.new(raw_color_socket, color_node.inputs[0])
+        color_socket = color_node.outputs[0]
+    split_node = tree.nodes.new(type="CompositorNodeSepRGBA")
+    tree.links.new(color_socket, split_node.inputs[0])
+    # Create separate file output nodes for every channel we care about.
+    # The process calling this script must decide how to recombine these
+    # channels, possibly into a single image.
+    for i, channel in enumerate("rgba") if not capturing_material_alpha else [(0, "MatAlpha")]:
+        output_node = tree.nodes.new(type="CompositorNodeOutputFile")
+        output_node.base_path = f"{output_path}_{channel}"
+        links.new(split_node.outputs[i], output_node.inputs[0])
+
+    if capturing_material_alpha:
+        # No need to re-write depth here.
+        return
+
+    depth_out = node_clamp(node_mul(input_sockets["Depth"], 1 / MAX_DEPTH))
+    output_node = tree.nodes.new(type="CompositorNodeOutputFile")
+    output_node.base_path = f"{output_path}_depth"
+    links.new(depth_out, output_node.inputs[0])
+
+
+def render_scene(output_path, fast_mode: bool):
+    use_workbench = bpy.context.scene.render.engine == "BLENDER_WORKBENCH"
+    if use_workbench:
+        # We must use a different engine to compute depth maps.
+        bpy.context.scene.render.engine = "BLENDER_EEVEE"
+        bpy.context.scene.eevee.taa_render_samples = 1  # faster, since we discard image.
+    if fast_mode:
+        if bpy.context.scene.render.engine == "BLENDER_EEVEE":
+            bpy.context.scene.eevee.taa_render_samples = 1
+        elif bpy.context.scene.render.engine == "CYCLES":
+            bpy.context.scene.cycles.samples = 256
+    else:
+        if bpy.context.scene.render.engine == "CYCLES":
+            # We should still impose a per-frame time limit
+            # so that we don't timeout completely.
+            bpy.context.scene.cycles.time_limit = 40
+    bpy.context.view_layer.update()
+    bpy.context.scene.use_nodes = True
+    bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
+    bpy.context.scene.view_settings.view_transform = "Raw"  # sRGB done in graph nodes
+    bpy.context.scene.render.film_transparent = True
+    bpy.context.scene.render.resolution_x = 512
+    bpy.context.scene.render.resolution_y = 512
+    bpy.context.scene.render.image_settings.file_format = "PNG"
+    bpy.context.scene.render.image_settings.color_mode = "BW"
+    bpy.context.scene.render.image_settings.color_depth = "16"
+    bpy.context.scene.render.filepath = output_path
+    setup_nodes(output_path)
+    bpy.ops.render.render(write_still=True)
+
+    # The output images must be moved from their own sub-directories, or
+    # discarded if we are using workbench for the color.
+    for channel_name in ["r", "g", "b", "a", "depth"]:
+        sub_dir = f"{output_path}_{channel_name}"
+        image_path = os.path.join(sub_dir, os.listdir(sub_dir)[0])
+        name, ext = os.path.splitext(output_path)
+        if channel_name == "depth" or not use_workbench:
+            os.rename(image_path, f"{name}_{channel_name}{ext}")
+        else:
+            os.remove(image_path)
+        os.removedirs(sub_dir)
+
+    if use_workbench:
+        # Re-render RGBA using workbench with texture mode, since this seems
+        # to show the most reasonable colors when lighting is broken.
+        bpy.context.scene.use_nodes = False
+        bpy.context.scene.render.engine = "BLENDER_WORKBENCH"
+        bpy.context.scene.render.image_settings.color_mode = "RGBA"
+        bpy.context.scene.render.image_settings.color_depth = "8"
+        bpy.context.scene.display.shading.color_type = "TEXTURE"
+        bpy.context.scene.display.shading.light = "FLAT"
+        if fast_mode:
+            # Single pass anti-aliasing.
+            bpy.context.scene.display.render_aa = "FXAA"
+        os.remove(output_path)
+        bpy.ops.render.render(write_still=True)
+        bpy.context.scene.render.image_settings.color_mode = "BW"
+        bpy.context.scene.render.image_settings.color_depth = "16"
+
+
+def scene_fov():
+    x_fov = bpy.context.scene.camera.data.angle_x
+    y_fov = bpy.context.scene.camera.data.angle_y
+    width = bpy.context.scene.render.resolution_x
+    height = bpy.context.scene.render.resolution_y
+    if bpy.context.scene.camera.data.angle == x_fov:
+        y_fov = 2 * math.atan(math.tan(x_fov / 2) * height / width)
+    else:
+        x_fov = 2 * math.atan(math.tan(y_fov / 2) * width / height)
+    return x_fov, y_fov
+
+
+def write_camera_metadata(path):
+    x_fov, y_fov = scene_fov()
+    bbox_min, bbox_max = scene_bbox()
+    matrix = bpy.context.scene.camera.matrix_world
+    with open(path, "w") as f:
+        json.dump(
+            dict(
+                format_version=FORMAT_VERSION,
+                max_depth=MAX_DEPTH,
+                bbox=[list(bbox_min), list(bbox_max)],
+                origin=list(matrix.col[3])[:3],
+                x_fov=x_fov,
+                y_fov=y_fov,
+                x=list(matrix.col[0])[:3],
+                y=list(-matrix.col[1])[:3],
+                z=list(-matrix.col[2])[:3],
+            ),
+            f,
+        )
+
+
+def save_rendering_dataset(
+    input_path: str,
+    output_path: str,
+    num_images: int,
+    backend: str,
+    light_mode: str,
+    camera_pose: str,
+    camera_dist_min: float,
+    camera_dist_max: float,
+    fast_mode: bool,
+):
+    assert light_mode in ["random", "uniform", "camera"]
+    assert camera_pose in ["random", "z-circular", "z-circular-elevated"]
+
+    import_model(input_path)
+    bpy.context.scene.render.engine = backend
+    normalize_scene()
+    if light_mode == "random":
+        create_random_lights()
+    elif light_mode == "uniform":
+        create_uniform_light(backend)
+    create_camera()
+    create_vertex_color_shaders()
+    for i in range(num_images):
+        t = i / max(num_images - 1, 1)  # same as np.linspace(0, 1, num_images)
+        place_camera(
+            t,
+            camera_pose_mode=camera_pose,
+            camera_dist_min=camera_dist_min,
+            camera_dist_max=camera_dist_max,
+        )
+        if light_mode == "camera":
+            create_camera_light()
+        render_scene(
+            os.path.join(output_path, f"{i:05}.png"),
+            fast_mode=fast_mode,
+        )
+        write_camera_metadata(os.path.join(output_path, f"{i:05}.json"))
+    with open(os.path.join(output_path, "info.json"), "w") as f:
+        info = dict(
+            backend=backend,
+            light_mode=light_mode,
+            fast_mode=fast_mode,
+            format_version=FORMAT_VERSION,
+            channels=["R", "G", "B", "A", "D"],
+            scale=0.5,  # The scene is bounded by [-scale, scale].
+        )
+        json.dump(info, f)
+
+
+def main():
+    try:
+        dash_index = sys.argv.index("--")
+    except ValueError as exc:
+        raise ValueError("arguments must be preceded by '--'") from exc
+
+    raw_args = sys.argv[dash_index + 1 :]
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_path", required=True, type=str)
+    parser.add_argument("--output_path", required=True, type=str)
+    parser.add_argument("--num_images", type=int, default=20)
+    parser.add_argument("--backend", type=str, default="BLENDER_EEVEE")
+    parser.add_argument("--light_mode", type=str, default="uniform")
+    parser.add_argument("--camera_pose", type=str, default="random")
+    parser.add_argument("--camera_dist_min", type=float, default=2.0)
+    parser.add_argument("--camera_dist_max", type=float, default=2.0)
+    parser.add_argument("--fast_mode", action="store_true")
+    args = parser.parse_args(raw_args)
+
+    save_rendering_dataset(
+        input_path=args.input_path,
+        output_path=args.output_path,
+        num_images=args.num_images,
+        backend=args.backend,
+        light_mode=args.light_mode,
+        camera_pose=args.camera_pose,
+        camera_dist_min=args.camera_dist_min,
+        camera_dist_max=args.camera_dist_max,
+        fast_mode=args.fast_mode,
+    )
+
+
+main()

point_e/evals/scripts/evaluate_pfid.py

@@ -0,0 +1,40 @@
+"""
+Evaluate P-FID between two batches of point clouds.
+
+The point cloud batches should be saved to two npz files, where there
+is an arr_0 key of shape [N x K x 3], where K is the dimensionality of
+each point cloud and N is the number of clouds.
+"""
+
+import argparse
+
+from point_e.evals.feature_extractor import PointNetClassifier, get_torch_devices
+from point_e.evals.fid_is import compute_statistics
+from point_e.evals.npz_stream import NpzStreamer
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cache_dir", type=str, default=None)
+    parser.add_argument("batch_1", type=str)
+    parser.add_argument("batch_2", type=str)
+    args = parser.parse_args()
+
+    print("creating classifier...")
+    clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.cache_dir)
+
+    print("computing first batch activations")
+
+    features_1, _ = clf.features_and_preds(NpzStreamer(args.batch_1))
+    stats_1 = compute_statistics(features_1)
+    del features_1
+
+    features_2, _ = clf.features_and_preds(NpzStreamer(args.batch_2))
+    stats_2 = compute_statistics(features_2)
+    del features_2
+
+    print(f"P-FID: {stats_1.frechet_distance(stats_2)}")
+
+
+if __name__ == "__main__":
+    main()

point_e/evals/scripts/evaluate_pis.py

@@ -0,0 +1,31 @@
+"""
+Evaluate P-IS of a batch of point clouds.
+
+The point cloud batch should be saved to an npz file, where there is an
+arr_0 key of shape [N x K x 3], where K is the dimensionality of each
+point cloud and N is the number of clouds.
+"""
+
+import argparse
+
+from point_e.evals.feature_extractor import PointNetClassifier, get_torch_devices
+from point_e.evals.fid_is import compute_inception_score
+from point_e.evals.npz_stream import NpzStreamer
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cache_dir", type=str, default=None)
+    parser.add_argument("batch", type=str)
+    args = parser.parse_args()
+
+    print("creating classifier...")
+    clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.cache_dir)
+
+    print("computing batch predictions")
+    _, preds = clf.features_and_preds(NpzStreamer(args.batch))
+    print(f"P-IS: {compute_inception_score(preds)}")
+
+
+if __name__ == "__main__":
+    main()

point_e/examples/.ipynb_checkpoints/Test-checkpoint.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+from tensorflow import keras
+print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
+tf.test.is_built_with_cuda()
+print(tf.version.VERSION)
+import sys
+sys.version

point_e/examples/.ipynb_checkpoints/pointcloud2mesh-checkpoint.ipynb

@@ -0,0 +1,106 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import torch\n",
+    "import matplotlib.pyplot as plt\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "from point_e.models.download import load_checkpoint\n",
+    "from point_e.models.configs import MODEL_CONFIGS, model_from_config\n",
+    "from point_e.util.pc_to_mesh import marching_cubes_mesh\n",
+    "from point_e.util.plotting import plot_point_cloud\n",
+    "from point_e.util.point_cloud import PointCloud"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "\n",
+    "print('creating SDF model...')\n",
+    "name = 'sdf'\n",
+    "model = model_from_config(MODEL_CONFIGS[name], device)\n",
+    "model.eval()\n",
+    "\n",
+    "print('loading SDF model...')\n",
+    "model.load_state_dict(load_checkpoint(name, device))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load a point cloud we want to convert into a mesh.\n",
+    "pc = PointCloud.load('example_data/pc_corgi.npz')\n",
+    "\n",
+    "# Plot the point cloud as a sanity check.\n",
+    "fig = plot_point_cloud(pc, grid_size=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Produce a mesh (with vertex colors)\n",
+    "mesh = marching_cubes_mesh(\n",
+    "    pc=pc,\n",
+    "    model=model,\n",
+    "    batch_size=4096,\n",
+    "    grid_size=32, # increase to 128 for resolution used in evals\n",
+    "    progress=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Write the mesh to a PLY file to import into some other program.\n",
+    "with open('mesh.ply', 'wb') as f:\n",
+    "    mesh.write_ply(f)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.9.9 64-bit ('3.9.9')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.9"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "b270b0f43bc427bcab7703c037711644cc480aac7c1cc8d2940cfaf0b447ee2e"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

point_e/examples/.ipynb_checkpoints/text2pointcloud-checkpoint.ipynb

@@ -0,0 +1,150 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config\n",
+    "from point_e.diffusion.sampler import PointCloudSampler\n",
+    "from point_e.models.download import load_checkpoint\n",
+    "from point_e.models.configs import MODEL_CONFIGS, model_from_config\n",
+    "from point_e.util.plotting import plot_point_cloud"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "creating base model...\n",
+      "creating upsample model...\n",
+      "downloading base checkpoint...\n",
+      "downloading upsampler checkpoint...\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "device = torch.device('cuda')\n",
+    "\n",
+    "print('creating base model...')\n",
+    "base_name = 'base40M-textvec'\n",
+    "base_model = model_from_config(MODEL_CONFIGS[base_name], device)\n",
+    "base_model.eval()\n",
+    "base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name])\n",
+    "\n",
+    "print('creating upsample model...')\n",
+    "upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)\n",
+    "upsampler_model.eval()\n",
+    "upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])\n",
+    "\n",
+    "print('downloading base checkpoint...')\n",
+    "base_model.load_state_dict(load_checkpoint(base_name, device))\n",
+    "\n",
+    "print('downloading upsampler checkpoint...')\n",
+    "upsampler_model.load_state_dict(load_checkpoint('upsample', device))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sampler = PointCloudSampler(\n",
+    "    device=device,\n",
+    "    models=[base_model, upsampler_model],\n",
+    "    diffusions=[base_diffusion, upsampler_diffusion],\n",
+    "    num_points=[1024, 4096 - 1024],\n",
+    "    aux_channels=['R', 'G', 'B'],\n",
+    "    guidance_scale=[3.0, 0.0],\n",
+    "    model_kwargs_key_filter=('texts', ''), # Do not condition the upsampler at all\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2777bd89bbef428aaae750480cbdf123",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Set a prompt to condition on.\n",
+    "prompt = 'a yellow dinosaur'\n",
+    "\n",
+    "# Produce a sample from the model.\n",
+    "samples = None\n",
+    "for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=[prompt]))):\n",
+    "    samples = x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pc = sampler.output_to_point_clouds(samples)[0]\n",
+    "fig = plot_point_cloud(pc, grid_size=3, fixed_bounds=((-0.75, -0.75, -0.75),(0.75, 0.75, 0.75)))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (GPU)",
+   "language": "python",
+   "name": "gpu_env"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "b270b0f43bc427bcab7703c037711644cc480aac7c1cc8d2940cfaf0b447ee2e"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

point_e/examples/GPUtest.py

@@ -0,0 +1,56 @@
+import numpy as np
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+from datetime import datetime
+
+# Choose which device you want to test on: either 'cpu' or 'gpu'
+devices = ['cpu', 'gpu']
+
+# Choose size of the matrix to be used.
+# Make it bigger to see bigger benefits of parallel computation
+shapes = [(50, 50), (100, 100), (500, 500), (1000, 1000)]
+
+
+def compute_operations(device, shape):
+    """Run a simple set of operations on a matrix of given shape on given device
+
+    Parameters
+    ----------
+    device : the type of device to use, either 'cpu' or 'gpu' 
+    shape : a tuple for the shape of a 2d tensor, e.g. (10, 10)
+
+    Returns
+    -------
+    out : results of the operations as the time taken
+    """
+
+    # Define operations to be computed on selected device
+    with tf.device(device):
+        random_matrix = tf.random_uniform(shape=shape, minval=0, maxval=1)
+        dot_operation = tf.matmul(random_matrix, tf.transpose(random_matrix))
+        sum_operation = tf.reduce_sum(dot_operation)
+
+    # Time the actual runtime of the operations
+    start_time = datetime.now()
+    with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as session:
+            result = session.run(sum_operation)
+    elapsed_time = datetime.now() - start_time
+
+    return result, elapsed_time
+
+
+
+if __name__ == '__main__':
+
+    # Run the computations and print summary of each run
+    for device in devices:
+        print("--" * 20)
+
+        for shape in shapes:
+            _, time_taken = compute_operations(device, shape)
+
+            # Print the result and also the time taken on the selected device
+            print("Input shape:", shape, "using Device:", device, "took: {:.2f}".format(time_taken.seconds + time_taken.microseconds/1e6))
+            #print("Computation on shape:", shape, "using Device:", device, "took:")
+
+    print("--" * 20)

point_e/examples/Saving Model Code.txt

@@ -0,0 +1,6 @@
+#Attempts to save the point cloud (pc) to a ply file
+
+pc = sampler.output_to_point_clouds(samples)[0]
+with open('example_data/blue_bird.ply','wb') as f:
+    pc.write_ply(f)
+#pc.save('example_data/blue_bird.npz')

point_e/examples/Test.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+from tensorflow import keras
+print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
+tf.test.is_built_with_cuda()
+print(tf.version.VERSION)
+import sys
+sys.version

point_e/examples/example_data/blue_bird.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bb37522cb2781a44b6272f7c826485f5e218f6b8a16c3445ae1573e4118ca8e
+size 99272

point_e/examples/example_data/pc_corgi.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8adfe52905e03d6e090e52d500f25afb59722153325e05decf4c80002ef76180
+size 99272

point_e/examples/example_data/pc_cube_stack.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41126e73e3a6eaf7cb68d35f6594a6ee221c074ee51b921ce2e8ccce52009c22
+size 99272