wips

app.py

@@ -8,11 +8,9 @@ import numpy as np
 import argparse
 
 from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
-from point_e.diffusion.sampler import PointCloudSampler
+from .sampler import PointCloudSampler
 from point_e.models.download import load_checkpoint
 from point_e.models.configs import MODEL_CONFIGS, model_from_config
-from point_e.util.plotting import plot_point_cloud
-from point_e.util.ply_util import write_ply
 
 from diffusers import StableDiffusionPipeline
 
@@ -119,7 +117,7 @@ def generate_3D(input, model_name='base1B', guidance_scale=3.0, grid_size=128):
     set_state('Converting to mesh...')
 
     uniqid = uuid.uuid4()
-    file_path = f'/tmp/mesh-{uniqid}.ply'
+    file_path = f'/tmp/mesh-{uniqid}.npy'
     save_ply(pc, file_path)
 
     set_state('')
@@ -153,7 +151,7 @@ def ply_to_glb(ply_file, glb_file):
 def save_ply(pc, file_name):
     # Produce a mesh (with vertex colors)
     with open(file_name, 'wb') as f:
-        pc.write_ply(f)
+        pc.save(f)
 
 
 def create_gif(pc):

pc.py

@@ -0,0 +1,174 @@
+import random
+from dataclasses import dataclass
+from typing import BinaryIO, Dict, List, Optional, Union
+
+import numpy as np
+
+from .ply_util import write_ply
+
+COLORS = frozenset(["R", "G", "B", "A"])
+
+
+def preprocess(data, channel):
+    if channel in COLORS:
+        return np.round(data * 255.0)
+    return data
+
+
+@dataclass
+class PointCloud:
+    """
+    An array of points sampled on a surface. Each point may have zero or more
+    channel attributes.
+
+    :param coords: an [N x 3] array of point coordinates.
+    :param channels: a dict mapping names to [N] arrays of channel values.
+    """
+
+    coords: np.ndarray
+    channels: Dict[str, np.ndarray]
+
+    @classmethod
+    def load(cls, f: Union[str, BinaryIO]) -> "PointCloud":
+        """
+        Load the point cloud from a .npz file.
+        """
+        if isinstance(f, str):
+            with open(f, "rb") as reader:
+                return cls.load(reader)
+        else:
+            obj = np.load(f)
+            keys = list(obj.keys())
+            return PointCloud(
+                coords=obj["coords"],
+                channels={k: obj[k] for k in keys if k != "coords"},
+            )
+
+    def save(self, f: Union[str, BinaryIO]):
+        """
+        Save the point cloud to a .npz file.
+        """
+        if isinstance(f, str):
+            with open(f, "wb") as writer:
+                self.save(writer)
+        else:
+            np.save(f, coords=self.coords, **self.channels)
+
+    def write_ply(self, raw_f: BinaryIO):
+        write_ply(
+            raw_f,
+            coords=self.coords,
+            rgb=(
+                np.stack([self.channels[x] for x in "RGB"], axis=1)
+                if all(x in self.channels for x in "RGB")
+                else None
+            ),
+        )
+
+    def random_sample(self, num_points: int, **subsample_kwargs) -> "PointCloud":
+        """
+        Sample a random subset of this PointCloud.
+
+        :param num_points: maximum number of points to sample.
+        :param subsample_kwargs: arguments to self.subsample().
+        :return: a reduced PointCloud, or self if num_points is not less than
+                 the current number of points.
+        """
+        if len(self.coords) <= num_points:
+            return self
+        indices = np.random.choice(len(self.coords), size=(num_points,), replace=False)
+        return self.subsample(indices, **subsample_kwargs)
+
+    def farthest_point_sample(
+        self, num_points: int, init_idx: Optional[int] = None, **subsample_kwargs
+    ) -> "PointCloud":
+        """
+        Sample a subset of the point cloud that is evenly distributed in space.
+
+        First, a random point is selected. Then each successive point is chosen
+        such that it is furthest from the currently selected points.
+
+        The time complexity of this operation is O(NM), where N is the original
+        number of points and M is the reduced number. Therefore, performance
+        can be improved by randomly subsampling points with random_sample()
+        before running farthest_point_sample().
+
+        :param num_points: maximum number of points to sample.
+        :param init_idx: if specified, the first point to sample.
+        :param subsample_kwargs: arguments to self.subsample().
+        :return: a reduced PointCloud, or self if num_points is not less than
+                 the current number of points.
+        """
+        if len(self.coords) <= num_points:
+            return self
+        init_idx = random.randrange(len(self.coords)) if init_idx is None else init_idx
+        indices = np.zeros([num_points], dtype=np.int64)
+        indices[0] = init_idx
+        sq_norms = np.sum(self.coords**2, axis=-1)
+
+        def compute_dists(idx: int):
+            # Utilize equality: ||A-B||^2 = ||A||^2 + ||B||^2 - 2*(A @ B).
+            return sq_norms + sq_norms[idx] - 2 * (self.coords @ self.coords[idx])
+
+        cur_dists = compute_dists(init_idx)
+        for i in range(1, num_points):
+            idx = np.argmax(cur_dists)
+            indices[i] = idx
+            cur_dists = np.minimum(cur_dists, compute_dists(idx))
+        return self.subsample(indices, **subsample_kwargs)
+
+    def subsample(self, indices: np.ndarray, average_neighbors: bool = False) -> "PointCloud":
+        if not average_neighbors:
+            return PointCloud(
+                coords=self.coords[indices],
+                channels={k: v[indices] for k, v in self.channels.items()},
+            )
+
+        new_coords = self.coords[indices]
+        neighbor_indices = PointCloud(coords=new_coords, channels={}).nearest_points(self.coords)
+
+        # Make sure every point points to itself, which might not
+        # be the case if points are duplicated or there is rounding
+        # error.
+        neighbor_indices[indices] = np.arange(len(indices))
+
+        new_channels = {}
+        for k, v in self.channels.items():
+            v_sum = np.zeros_like(v[: len(indices)])
+            v_count = np.zeros_like(v[: len(indices)])
+            np.add.at(v_sum, neighbor_indices, v)
+            np.add.at(v_count, neighbor_indices, 1)
+            new_channels[k] = v_sum / v_count
+        return PointCloud(coords=new_coords, channels=new_channels)
+
+    def select_channels(self, channel_names: List[str]) -> np.ndarray:
+        data = np.stack([preprocess(self.channels[name], name) for name in channel_names], axis=-1)
+        return data
+
+    def nearest_points(self, points: np.ndarray, batch_size: int = 16384) -> np.ndarray:
+        """
+        For each point in another set of points, compute the point in this
+        pointcloud which is closest.
+
+        :param points: an [N x 3] array of points.
+        :param batch_size: the number of neighbor distances to compute at once.
+                           Smaller values save memory, while larger values may
+                           make the computation faster.
+        :return: an [N] array of indices into self.coords.
+        """
+        norms = np.sum(self.coords**2, axis=-1)
+        all_indices = []
+        for i in range(0, len(points), batch_size):
+            batch = points[i : i + batch_size]
+            dists = norms + np.sum(batch**2, axis=-1)[:, None] - 2 * (batch @ self.coords.T)
+            all_indices.append(np.argmin(dists, axis=-1))
+        return np.concatenate(all_indices, axis=0)
+
+    def combine(self, other: "PointCloud") -> "PointCloud":
+        assert self.channels.keys() == other.channels.keys()
+        return PointCloud(
+            coords=np.concatenate([self.coords, other.coords], axis=0),
+            channels={
+                k: np.concatenate([v, other.channels[k]], axis=0) for k, v in self.channels.items()
+            },
+        )

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 .pc import PointCloud
+
+from point_e.diffusion.gaussian_diffusion import GaussianDiffusion
+from point_e.diffusion.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,
+        )