wipz

app.py

@@ -117,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}.npy'
+    file_path = f'/tmp/mesh-{uniqid}.npz'
     save_ply(pc, file_path)
 
     set_state('')

pc.py

@@ -1,174 +0,0 @@
-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

@@ -1,263 +0,0 @@
-"""
-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,
-        )