trellis/pipelines/dvd_voxel.py

from dataclasses import dataclass
from typing import Callable, Literal

import torch

from .trellis_image_to_3d import TrellisImageTo3DPipeline
from .trellis_text_to_3d import TrellisTextTo3DPipeline
from .samplers import USDMAncestralGuidanceFunctionSampler

CfgSchedule = float | Callable[[float], float]

DVD_PRETRAINED_FILENAMES = {
    "image": {
        "base": ("dvd_img.json", "dvd_img.safetensors"),
        "bsp": ("dvd_img_BSP_ft.json", "dvd_img_BSP_ft.safetensors"),
    },
    "text": {
        "base": ("dvd_text.json", "dvd_text.safetensors"),
        "bsp": ("dvd_text_BSP_ft.json", "dvd_text_BSP_ft.safetensors"),
    },
}
DVD_VARIANT_ALIASES = {
    "default": "base",
    "generation": "base",
    "gen": "base",
    "edit": "bsp",
    "editing": "bsp",
    "bsp_ft": "bsp",
    "bsp-ft": "bsp",
    "bsp_finetune": "bsp",
    "bsp-finetune": "bsp",
}


def _dvd_pretrained_filenames(
    modality: Literal["image", "text"],
    variant: str = "base",
    config_name: str | None = None,
    checkpoint_name: str | None = None,
) -> tuple[str, str]:
    if (config_name is None) != (checkpoint_name is None):
        raise ValueError("Pass both config_name and checkpoint_name, or neither.")
    if config_name is not None and checkpoint_name is not None:
        return config_name, checkpoint_name

    variant_key = DVD_VARIANT_ALIASES.get(variant.lower(), variant.lower())
    try:
        return DVD_PRETRAINED_FILENAMES[modality][variant_key]
    except KeyError as exc:
        valid = ", ".join(sorted(DVD_PRETRAINED_FILENAMES[modality]))
        raise ValueError(f"Unknown DVD {modality} variant '{variant}'. Valid variants: {valid}.") from exc


def _resolve_dvd_pretrained_files(
    pretrained_model_name_or_path: str,
    config_name: str,
    checkpoint_name: str,
    subfolder: str | None = None,
    revision: str | None = None,
    repo_type: str | None = None,
    cache_dir: str | None = None,
    local_files_only: bool = False,
    token: bool | str | None = None,
) -> tuple[str, str]:
    from pathlib import Path

    local_root = Path(pretrained_model_name_or_path)
    if local_root.is_dir():
        local_root = local_root / subfolder if subfolder else local_root
        config_path = local_root / config_name
        checkpoint_path = local_root / checkpoint_name
        if not config_path.exists():
            raise FileNotFoundError(f"DVD config not found: {config_path}")
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"DVD checkpoint not found: {checkpoint_path}")
        return str(config_path), str(checkpoint_path)

    from huggingface_hub import hf_hub_download

    kwargs = {
        "repo_id": pretrained_model_name_or_path,
        "subfolder": subfolder,
        "revision": revision,
        "repo_type": repo_type,
        "cache_dir": cache_dir,
        "local_files_only": local_files_only,
        "token": token,
    }
    kwargs = {key: value for key, value in kwargs.items() if value is not None}
    config_path = hf_hub_download(filename=config_name, **kwargs)
    checkpoint_path = hf_hub_download(filename=checkpoint_name, **kwargs)
    return config_path, checkpoint_path


@dataclass
class DVDVoxelOutput:
    """
    DVD voxel samples in DVD convention.

    Attributes:
        samples: Occupancy tensor with shape [B, D, H, W].
        coords: Occupied coordinates with shape [N, 4] in [batch, x, y, z]
            DVD convention.
    """
    samples: torch.Tensor
    coords: torch.Tensor

    @property
    def coords_without_batch(self) -> torch.Tensor:
        return self.coords[:, 1:]

    def to_trellis_coords(self, resolution: int = 64, device: torch.device | str | None = None) -> torch.Tensor:
        coords = dvd_coords_to_trellis_coords(self.coords, resolution=resolution)
        if device is not None:
            coords = coords.to(device)
        return coords


def load_dvd_denoiser(config_path: str, checkpoint_path: str, device: torch.device | str | None = None) -> torch.nn.Module:
    import json
    from safetensors.torch import load_file
    from trellis import models

    with open(config_path, "r") as f:
        config = json.load(f)

    denoiser = getattr(models, config["name"])(**config["args"])
    denoiser.load_state_dict(load_file(checkpoint_path))
    denoiser.eval()
    if device is not None:
        denoiser.to(device)
    return denoiser


def image_cfg_schedule(t: float) -> float:
    return 0.4 if t < 0.5 else 0.7


def text_cfg_schedule(t: float) -> float:
    return 0.4 if t < 0.8 else 1.2


def bsp_edit_cfg_schedule(t: float) -> float:
    return 0.45


def normalize_cfg_schedule(cfg_strength: CfgSchedule):
    if callable(cfg_strength):
        return cfg_strength
    return lambda _t: float(cfg_strength)


def coords_to_samples(coords: torch.Tensor, resolution: int = 64, batch_size: int | None = None) -> torch.Tensor:
    coords = coords.detach().cpu().long()
    if coords.ndim != 2 or coords.shape[1] not in (3, 4):
        raise ValueError("coords must have shape [N, 3] or [N, 4]")

    if coords.shape[1] == 3:
        coords = torch.cat([torch.zeros((coords.shape[0], 1), dtype=coords.dtype), coords], dim=1)

    if batch_size is None:
        batch_size = int(coords[:, 0].max().item()) + 1 if coords.numel() > 0 else 1

    samples = torch.zeros((batch_size, resolution, resolution, resolution), dtype=torch.long)
    samples[coords[:, 0], coords[:, 1], coords[:, 2], coords[:, 3]] = 1
    return samples


def samples_to_coords(samples: torch.Tensor) -> torch.Tensor:
    if samples.ndim == 5 and samples.shape[1] == 1:
        samples = samples.squeeze(1)
    if samples.ndim != 4:
        raise ValueError("samples must have shape [B, D, H, W] or [B, 1, D, H, W]")
    return torch.argwhere(samples.detach().cpu() > 0).int()


def as_voxel_output(voxels: DVDVoxelOutput | torch.Tensor, resolution: int = 64) -> DVDVoxelOutput:
    if isinstance(voxels, DVDVoxelOutput):
        return voxels
    if not isinstance(voxels, torch.Tensor):
        raise TypeError("voxels must be DVDVoxelOutput or torch.Tensor")

    if voxels.ndim in (4, 5):
        samples = voxels.detach().cpu()
        if samples.ndim == 5 and samples.shape[1] == 1:
            samples = samples.squeeze(1)
        return DVDVoxelOutput(samples=samples.long(), coords=samples_to_coords(samples))

    if voxels.ndim == 2 and voxels.shape[1] in (3, 4):
        coords = voxels.detach().cpu().int()
        if coords.shape[1] == 3:
            coords = torch.cat([torch.zeros((coords.shape[0], 1), dtype=coords.dtype), coords], dim=1)
        return DVDVoxelOutput(samples=coords_to_samples(coords, resolution=resolution), coords=coords)

    raise ValueError("voxel tensor must be samples [B, D, H, W] or coords [N, 3/4]")


def dvd_coords_to_trellis_coords(coords: torch.Tensor, resolution: int = 64) -> torch.Tensor:
    coords = coords.detach().clone().int()
    has_batch = coords.shape[1] == 4
    if not has_batch:
        coords = torch.cat([torch.zeros((coords.shape[0], 1), dtype=coords.dtype), coords], dim=1)
    coords = coords[:, [0, 1, 3, 2]]
    coords[:, 1] = resolution - 1 - coords[:, 1].clone()
    return coords if has_batch else coords[:, 1:]


def trellis_coords_to_dvd_coords(coords: torch.Tensor, resolution: int = 64) -> torch.Tensor:
    coords = coords.detach().clone().int()
    has_batch = coords.shape[1] == 4
    if not has_batch:
        coords = torch.cat([torch.zeros((coords.shape[0], 1), dtype=coords.dtype), coords], dim=1)
    coords = coords[:, [0, 1, 3, 2]]
    coords[:, 1] = resolution - 1 - coords[:, 1].clone()
    return coords if has_batch else coords[:, 1:]


def make_box_mask(
    resolution: int = 64,
    batch_size: int = 1,
    x: tuple[int, int] | None = None,
    y: tuple[int, int] | None = None,
    z: tuple[int, int] | None = None,
    mode: Literal["keep", "edit"] = "keep",
) -> torch.Tensor:
    """
    Build a boolean mask with shape [B, D, H, W].

    The sampler expects True for known voxels to keep. Use mode="edit" when
    describing the region to regenerate; the returned mask is inverted for you.
    """
    edit_mask = torch.zeros((batch_size, resolution, resolution, resolution), dtype=torch.bool)
    xs = x or (0, resolution)
    ys = y or (0, resolution)
    zs = z or (0, resolution)
    edit_mask[:, xs[0]:xs[1], ys[0]:ys[1], zs[0]:zs[1]] = True
    return ~edit_mask if mode == "edit" else edit_mask


def export_cubified_voxels(
    voxels: DVDVoxelOutput | torch.Tensor,
    output_path: str,
    voxel_size: float = 0.5,
    resolution: int = 64,
    color: tuple[int, int, int, int] = (150, 150, 150, 255),
) -> str:
    import trimesh
    from pytorch3d.ops import cubify

    output = as_voxel_output(voxels, resolution=resolution)
    samples = output.samples.float()
    cubified_meshes = cubify(samples.flip(dims=[-3]), voxel_size, align="center")
    vertices = cubified_meshes.verts_packed().cpu().numpy()
    faces = cubified_meshes.faces_packed().cpu().numpy()
    if vertices.size:
        vmin = vertices.min(axis=0)
        vmax = vertices.max(axis=0)
        center = (vmin + vmax) / 2.0
        scale = (vmax - vmin).max() / 2.0
        if scale > 0:
            vertices = (vertices - center) / scale
    mesh = trimesh.Trimesh(
        vertices=vertices,
        faces=faces,
    )
    mesh.visual.face_colors = color
    mesh.export(output_path)
    return output_path


class DVDVoxelPipelineMixin:
    resolution: int
    dvd_sampler: USDMAncestralGuidanceFunctionSampler

    @property
    def denoiser(self) -> torch.nn.Module:
        return self.models["denoiser"]

    def _init_dvd_sampler(self, resolution: int = 64, vocab_size: int = 2):
        self.resolution = resolution
        self.dvd_sampler = USDMAncestralGuidanceFunctionSampler(sigma_min=1e-5, vocab_size=vocab_size)

    def _sample_from_cond(
        self,
        cond: dict[str, torch.Tensor],
        num_samples: int = 1,
        steps: int = 256,
        cfg_strength: CfgSchedule = image_cfg_schedule,
        seed: int | None = None,
        rescale_t: str | float = "cosine",
        verbose: bool = True,
        noise: torch.Tensor | None = None,
        keep_mask: torch.Tensor | None = None,
        classes: torch.Tensor | None = None,
    ) -> DVDVoxelOutput:
        if seed is not None:
            torch.manual_seed(seed)

        device = self.device
        cond = {k: v.to(device) for k, v in cond.items()}
        if noise is None:
            noise = torch.randint(
                0,
                2,
                (num_samples, 1, self.resolution, self.resolution, self.resolution),
                device=device,
            )
        else:
            noise = noise.to(device)

        sample_kwargs = {}
        if keep_mask is not None or classes is not None:
            if keep_mask is None or classes is None:
                raise ValueError("keep_mask and classes must be provided together")
            sample_kwargs["mask"] = keep_mask.to(device).bool()
            sample_kwargs["classes"] = classes.to(device).long()

        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            res = self.dvd_sampler.sample(
                self.denoiser,
                noise=noise,
                **cond,
                steps=steps,
                cfg_strength=normalize_cfg_schedule(cfg_strength),
                verbose=verbose,
                rescale_t=rescale_t,
                **sample_kwargs,
            )

        samples = res.samples.squeeze(1).detach().cpu().long()
        return DVDVoxelOutput(samples=samples, coords=samples_to_coords(samples))

    def edit_voxels_from_cond(
        self,
        cond: dict[str, torch.Tensor],
        voxels: DVDVoxelOutput | torch.Tensor,
        keep_mask: torch.Tensor,
        steps: int = 128,
        cfg_strength: CfgSchedule = bsp_edit_cfg_schedule,
        seed: int | None = None,
        rescale_t: str | float = "cosine",
        verbose: bool = True,
        init: Literal["base", "random"] = "base",
    ) -> DVDVoxelOutput:
        base = as_voxel_output(voxels, resolution=self.resolution).samples
        if keep_mask.shape != base.shape:
            raise ValueError(f"keep_mask must have shape {tuple(base.shape)}, got {tuple(keep_mask.shape)}")

        if seed is not None:
            torch.manual_seed(seed)

        if init == "base":
            init_samples = base.clone()
            init_samples[~keep_mask.cpu()] = torch.randint(0, 2, init_samples[~keep_mask.cpu()].shape)
        elif init == "random":
            init_samples = torch.randint(0, 2, base.shape)
        else:
            raise ValueError(f"Unsupported init mode: {init}")

        return self._sample_from_cond(
            cond,
            num_samples=base.shape[0],
            steps=steps,
            cfg_strength=cfg_strength,
            seed=None,
            rescale_t=rescale_t,
            verbose=verbose,
            noise=init_samples.unsqueeze(1),
            keep_mask=keep_mask,
            classes=base,
        )


class DVDImageToVoxelPipeline(DVDVoxelPipelineMixin, TrellisImageTo3DPipeline):
    def __init__(self, denoiser: torch.nn.Module, image_cond_model: str = "dinov2_vitl14_reg", resolution: int = 64):
        TrellisImageTo3DPipeline.__init__(self, models={"denoiser": denoiser}, image_cond_model=image_cond_model)
        self._init_dvd_sampler(resolution=resolution)

    def preprocess_image(self, image):
        """
        Preprocess image conditions for DVD voxel generation.

        This follows the TRELLIS foreground crop geometry, but pads and
        composites with a white background to match DVD's training setup.
        """
        import numpy as np
        import rembg
        from PIL import Image, ImageOps

        has_alpha = False
        if image.mode == "RGBA":
            alpha = np.array(image.getchannel("A"))
            if not np.all(alpha == 255):
                has_alpha = True

        if has_alpha:
            image = image.convert("RGBA")
        else:
            image = image.convert("RGB")
            max_size = max(image.size)
            scale = min(1, 1024 / max_size)
            if scale < 1:
                image = image.resize((int(image.width * scale), int(image.height * scale)), Image.Resampling.LANCZOS)
            if getattr(self, "rembg_session", None) is None:
                self.rembg_session = rembg.new_session("u2net")
            image = rembg.remove(image, session=self.rembg_session).convert("RGBA")

        mask_np = np.array(image.getchannel("A"))
        ys, xs = np.nonzero(mask_np > 0.8*255)
        if len(xs) == 0 or len(ys) == 0:
            raise ValueError("No foreground pixels found in image alpha channel.")

        x0, y0, x1, y1 = xs.min(), ys.min(), xs.max(), ys.max()
        center_x = (x0 + x1) / 2.0
        center_y = (y0 + y1) / 2.0
        hsize = max(x1 - x0, y1 - y0) / 2.0
        aug_hsize = hsize * 1.1

        ax0 = int(center_x - aug_hsize)
        ay0 = int(center_y - aug_hsize)
        ax1 = int(center_x + aug_hsize)
        ay1 = int(center_y + aug_hsize)

        width, height = image.size
        pad_l = max(0, -ax0)
        pad_t = max(0, -ay0)
        pad_r = max(0, ax1 - width)
        pad_b = max(0, ay1 - height)
        if pad_l or pad_t or pad_r or pad_b:
            image = ImageOps.expand(
                image,
                border=(pad_l, pad_t, pad_r, pad_b),
                fill=(255, 255, 255, 255),
            )
            ax0 += pad_l
            ay0 += pad_t
            ax1 += pad_l
            ay1 += pad_t
            width, height = image.size

        ax0 = max(0, ax0)
        ay0 = max(0, ay0)
        ax1 = min(width, ax1)
        ay1 = min(height, ay1)
        image = image.crop((ax0, ay0, ax1, ay1))
        image = image.resize((518, 518), Image.Resampling.LANCZOS)

        if image.mode == "RGBA":
            white = Image.new("RGBA", image.size, (255, 255, 255, 255))
            image = Image.alpha_composite(white, image).convert("RGB")
        else:
            image = image.convert("RGB")
        return image

    @classmethod
    def from_files(
        cls,
        config_path: str,
        checkpoint_path: str,
        image_cond_model: str = "dinov2_vitl14_reg",
        resolution: int = 64,
        device: torch.device | str | None = None,
    ) -> "DVDImageToVoxelPipeline":
        denoiser = load_dvd_denoiser(config_path, checkpoint_path)
        pipeline = cls(denoiser, image_cond_model=image_cond_model, resolution=resolution)
        if device is not None:
            pipeline.to(torch.device(device))
        return pipeline

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str,
        variant: str = "base",
        config_name: str | None = None,
        checkpoint_name: str | None = None,
        image_cond_model: str = "dinov2_vitl14_reg",
        resolution: int = 64,
        device: torch.device | str | None = None,
        subfolder: str | None = None,
        revision: str | None = None,
        repo_type: str | None = None,
        cache_dir: str | None = None,
        local_files_only: bool = False,
        token: bool | str | None = None,
    ) -> "DVDImageToVoxelPipeline":
        config_name, checkpoint_name = _dvd_pretrained_filenames("image", variant, config_name, checkpoint_name)
        config_path, checkpoint_path = _resolve_dvd_pretrained_files(
            pretrained_model_name_or_path,
            config_name,
            checkpoint_name,
            subfolder=subfolder,
            revision=revision,
            repo_type=repo_type,
            cache_dir=cache_dir,
            local_files_only=local_files_only,
            token=token,
        )
        return cls.from_files(
            config_path,
            checkpoint_path,
            image_cond_model=image_cond_model,
            resolution=resolution,
            device=device,
        )

    @torch.no_grad()
    def sample_voxels(
        self,
        image,
        num_samples: int = 1,
        seed: int | None = None,
        steps: int = 256,
        cfg_strength: CfgSchedule = image_cfg_schedule,
        preprocess_image: bool = True,
        rescale_t: str | float = "cosine",
        verbose: bool = True,
    ) -> DVDVoxelOutput:
        if preprocess_image:
            image = self.preprocess_image(image)
        cond = self.get_cond([image])
        return self._sample_from_cond(cond, num_samples, steps, cfg_strength, seed, rescale_t, verbose)

    @torch.no_grad()
    def edit_voxels(
        self,
        image,
        voxels: DVDVoxelOutput | torch.Tensor,
        keep_mask: torch.Tensor,
        seed: int | None = None,
        steps: int = 128,
        cfg_strength: CfgSchedule = bsp_edit_cfg_schedule,
        preprocess_image: bool = True,
        rescale_t: str | float = "cosine",
        verbose: bool = True,
        init: Literal["base", "random"] = "base",
    ) -> DVDVoxelOutput:
        if preprocess_image:
            image = self.preprocess_image(image)
        cond = self.get_cond([image])
        return self.edit_voxels_from_cond(cond, voxels, keep_mask, steps, cfg_strength, seed, rescale_t, verbose, init)


class DVDTextToVoxelPipeline(DVDVoxelPipelineMixin, TrellisTextTo3DPipeline):
    def __init__(self, denoiser: torch.nn.Module, text_cond_model: str, resolution: int = 64):
        TrellisTextTo3DPipeline.__init__(self, models={"denoiser": denoiser}, text_cond_model=text_cond_model)
        self._init_dvd_sampler(resolution=resolution)

    @classmethod
    def from_files(
        cls,
        config_path: str,
        checkpoint_path: str,
        text_cond_model: str = "openai/clip-vit-large-patch14",
        resolution: int = 64,
        device: torch.device | str | None = None,
    ) -> "DVDTextToVoxelPipeline":
        denoiser = load_dvd_denoiser(config_path, checkpoint_path)
        pipeline = cls(denoiser, text_cond_model=text_cond_model, resolution=resolution)
        if device is not None:
            pipeline.to(torch.device(device))
        return pipeline

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str,
        variant: str = "base",
        config_name: str | None = None,
        checkpoint_name: str | None = None,
        text_cond_model: str = "openai/clip-vit-large-patch14",
        resolution: int = 64,
        device: torch.device | str | None = None,
        subfolder: str | None = None,
        revision: str | None = None,
        repo_type: str | None = None,
        cache_dir: str | None = None,
        local_files_only: bool = False,
        token: bool | str | None = None,
    ) -> "DVDTextToVoxelPipeline":
        config_name, checkpoint_name = _dvd_pretrained_filenames("text", variant, config_name, checkpoint_name)
        config_path, checkpoint_path = _resolve_dvd_pretrained_files(
            pretrained_model_name_or_path,
            config_name,
            checkpoint_name,
            subfolder=subfolder,
            revision=revision,
            repo_type=repo_type,
            cache_dir=cache_dir,
            local_files_only=local_files_only,
            token=token,
        )
        return cls.from_files(
            config_path,
            checkpoint_path,
            text_cond_model=text_cond_model,
            resolution=resolution,
            device=device,
        )

    @torch.no_grad()
    def sample_voxels(
        self,
        prompt: str,
        num_samples: int = 1,
        seed: int | None = None,
        steps: int = 128,
        cfg_strength: CfgSchedule = text_cfg_schedule,
        rescale_t: str | float = "cosine",
        verbose: bool = True,
    ) -> DVDVoxelOutput:
        cond = self.get_cond([prompt])
        return self._sample_from_cond(cond, num_samples, steps, cfg_strength, seed, rescale_t, verbose)

    @torch.no_grad()
    def edit_voxels(
        self,
        prompt: str,
        voxels: DVDVoxelOutput | torch.Tensor,
        keep_mask: torch.Tensor,
        seed: int | None = None,
        steps: int = 128,
        cfg_strength: CfgSchedule = bsp_edit_cfg_schedule,
        rescale_t: str | float = "cosine",
        verbose: bool = True,
        init: Literal["base", "random"] = "base",
    ) -> DVDVoxelOutput:
        cond = self.get_cond([prompt])
        return self.edit_voxels_from_cond(cond, voxels, keep_mask, steps, cfg_strength, seed, rescale_t, verbose, init)


def run_image_stage2_from_dvd_voxels(
    trellis_pipeline: TrellisImageTo3DPipeline,
    image,
    voxels: DVDVoxelOutput | torch.Tensor,
    seed: int = 42,
    slat_sampler_params: dict | None = None,
    formats: list[str] | None = None,
    preprocess_image: bool = True,
    resolution: int = 64,
) -> dict:
    output = as_voxel_output(voxels, resolution=resolution)
    coords = output.to_trellis_coords(resolution=resolution, device=trellis_pipeline.device)
    return trellis_pipeline.run(
        image,
        seed=seed,
        slat_sampler_params=slat_sampler_params or {},
        formats=formats or ["mesh", "gaussian", "radiance_field"],
        preprocess_image=preprocess_image,
        coords=coords,
    )


def run_text_stage2_from_dvd_voxels(
    trellis_pipeline: TrellisTextTo3DPipeline,
    prompt: str,
    voxels: DVDVoxelOutput | torch.Tensor,
    seed: int = 42,
    slat_sampler_params: dict | None = None,
    formats: list[str] | None = None,
    resolution: int = 64,
) -> dict:
    output = as_voxel_output(voxels, resolution=resolution)
    coords = output.to_trellis_coords(resolution=resolution, device=trellis_pipeline.device)
    return trellis_pipeline.run(
        prompt=prompt,
        seed=seed,
        slat_sampler_params=slat_sampler_params or {},
        formats=formats or ["mesh", "gaussian", "radiance_field"],
        coords=coords,
    )