""" 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, )