| "pipeline_source": "class TripoSplatPipeline:\n def __init__(self, ckpt_path: str, decoder_path: str, dinov3_path: str,\n flux2_vae_encoder_path: str, rmbg_path: str, device: str = \"cuda\"):\n self._device = torch.device(device)\n self.dinov3 = load_dinov3 (dinov3_path, device=self._device, dtype=torch.bfloat16)\n self.vae_encoder = load_vae_encoder (flux2_vae_encoder_path, device=self._device, dtype=torch.bfloat16)\n self.rmbg = load_rmbg (rmbg_path, device=self._device, dtype=torch.float16)\n self.flow_model = load_flow_model (ckpt_path, device=self._device, dtype=torch.float16)\n self.decoder = load_decoder (decoder_path, device=self._device, dtype=torch.float16)\n\n def preprocess_image(self, image, erode_radius: int = 1) -> Image.Image:\n return preprocess_image(image, self.rmbg, erode_radius=erode_radius)\n\n def encode_image(self, image: Image.Image, generator: torch.Generator = None) -> dict:\n return encode_image(image, self.dinov3, self.vae_encoder, generator=generator)\n\n def sample_latent(self, cond: dict, steps: int = 50, guidance_scale: float = 7.0,\n shift: float = 3.0, generator: torch.Generator = None,\n show_progress: bool = False, callback=None) -> dict:\n return sample_latent(self.flow_model, cond, steps=steps, guidance_scale=guidance_scale,\n shift=shift, generator=generator,\n show_progress=show_progress, callback=callback)\n\n def decode_latent(self, latent: torch.Tensor, num_gaussians: int = 262144):\n return self.decoder.decode(latent, num_gaussians=num_gaussians)\n\n _NUM_GAUSSIANS_MIN = 32768\n _NUM_GAUSSIANS_MAX = 262144\n\n def _validate_num_gaussians(self, n: int) -> int:\n assert self._NUM_GAUSSIANS_MIN <= n <= self._NUM_GAUSSIANS_MAX, (\n f\"num_gaussians must be in [{self._NUM_GAUSSIANS_MIN}, {self._NUM_GAUSSIANS_MAX}], got {n}\"\n )\n gpp = self.decoder.gaussians_per_point\n if n % gpp == 0:\n return n\n rounded = round(n / gpp) * gpp\n print(f\"[TripoSplatPipeline] num_gaussians={n} is not a multiple of {gpp}; rounding to {rounded}\")\n return rounded\n\n @torch.no_grad()\n def run(self, image, seed: int = 42, steps: int = 20, guidance_scale: float = 3.0,\n shift: float = 3.0, num_gaussians=262144, erode_radius: int = 1,\n show_progress: bool = False, callback=None):\n \"\"\"\n Args:\n image: Input image. Accepts a file path / PIL.Image / torch.Tensor\n (`[1,H,W,C]` or `[H,W,C]`, float in `[0, 1]`, optional alpha\n channel as the 4th channel).\n seed: RNG seed for the VAE encoder's stochastic latent sampling and\n the initial flow-matching noise. Same seed \u2192 same output.\n steps: Number of Euler integrator steps in the flow-matching sampler.\n More steps \u2192 better fidelity, linear runtime cost.\n Recommend: 10~20.\n guidance_scale: Classifier-free-guidance strength (diffusers\n convention). `\u2264 1.0` disables CFG. Higher \u2192 more detail,\n stronger adherence to the input image; too high can cause color\n oversaturation.\n Recommend: 3.0.\n shift: Flow-matching timestep schedule shift. `1.0` gives a uniform\n schedule; `>1.0` allocates more steps to the early/high-noise end.\n Recommend: 3.0.\n num_gaussians: Target Gaussian-splat count. An `int` returns a\n single `Gaussian`. A `list` / `tuple` of ints returns a\n `list[Gaussian]`. Each count is rounded to the nearest multiple\n of 32. More gaussians \u2192 more detail but higher rendering and\n storage cost.\n Recommend: 32768~262144.\n erode_radius: Pixel radius used to erode the alpha matte after\n background removal, to avoid segmentation-border bleed before\n compositing on black. `0` disables; `1` is a 3\u00d73 minimum filter.\n Recommend: 1.\n show_progress: Print a `tqdm` progress bar over sampler steps.\n callback: Optional `fn(step, total)` invoked after each sampler step.\n Useful for external progress UIs (e.g. ComfyUI's\n `ProgressBar.update`).\n\n Returns:\n `(gaussian, prepared_image)` for an `int` `num_gaussians`, or\n `(list_of_gaussians, prepared_image)` for a `list` / `tuple`. The\n second element is the RGB composite the encoders actually saw \u2014\n useful for display / debugging.\n \"\"\"\n if isinstance(num_gaussians, (list, tuple)):\n counts = [self._validate_num_gaussians(n) for n in num_gaussians]\n else:\n counts = [self._validate_num_gaussians(num_gaussians)]\n\n gen = torch.Generator(device=self._device).manual_seed(seed)\n prepared = self.preprocess_image(image, erode_radius=erode_radius)\n cond = self.encode_image(prepared, generator=gen)\n out = self.sample_latent(cond, steps=steps, guidance_scale=guidance_scale, shift=shift,\n generator=gen, show_progress=show_progress, callback=callback)\n gaussians = [self.decode_latent(out['latent'], num_gaussians=n) for n in counts]\n if isinstance(num_gaussians, (list, tuple)):\n return gaussians, prepared\n return gaussians[0], prepared\n" |
