Delete trellis/pipelines

trellis/pipelines/__init__.py

@@ -1,25 +0,0 @@
-from . import samplers
-from .trellis_image_to_3d import TrellisImageTo3DPipeline
-from .trellis_text_to_3d import TrellisTextTo3DPipeline
-
-
-def from_pretrained(path: str):
-    """
-    Load a pipeline from a model folder or a Hugging Face model hub.
-
-    Args:
-        path: The path to the model. Can be either local path or a Hugging Face model name.
-    """
-    import os
-    import json
-    is_local = os.path.exists(f"{path}/pipeline.json")
-
-    if is_local:
-        config_file = f"{path}/pipeline.json"
-    else:
-        from huggingface_hub import hf_hub_download
-        config_file = hf_hub_download(path, "pipeline.json")
-
-    with open(config_file, 'r') as f:
-        config = json.load(f)
-    return globals()[config['name']].from_pretrained(path)

trellis/pipelines/base.py

@@ -1,68 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-from .. import models
-
-
-class Pipeline:
-    """
-    A base class for pipelines.
-    """
-    def __init__(
-        self,
-        models: dict[str, nn.Module] = None,
-    ):
-        if models is None:
-            return
-        self.models = models
-        for model in self.models.values():
-            model.eval()
-
-    @staticmethod
-    def from_pretrained(path: str) -> "Pipeline":
-        """
-        Load a pretrained model.
-        """
-        import os
-        import json
-        is_local = os.path.exists(f"{path}/pipeline.json")
-
-        if is_local:
-            config_file = f"{path}/pipeline.json"
-        else:
-            from huggingface_hub import hf_hub_download
-            config_file = hf_hub_download(path, "pipeline.json")
-
-        with open(config_file, 'r') as f:
-            args = json.load(f)['args']
-
-        _models = {}
-        for k, v in args['models'].items():
-            try:
-                _models[k] = models.from_pretrained(f"{path}/{v}")
-            except:
-                _models[k] = models.from_pretrained(v)
-
-        new_pipeline = Pipeline(_models)
-        new_pipeline._pretrained_args = args
-        return new_pipeline
-
-    @property
-    def device(self) -> torch.device:
-        for model in self.models.values():
-            if hasattr(model, 'device'):
-                return model.device
-        for model in self.models.values():
-            if hasattr(model, 'parameters'):
-                return next(model.parameters()).device
-        raise RuntimeError("No device found.")
-
-    def to(self, device: torch.device) -> None:
-        for model in self.models.values():
-            model.to(device)
-
-    def cuda(self) -> None:
-        self.to(torch.device("cuda"))
-
-    def cpu(self) -> None:
-        self.to(torch.device("cpu"))

trellis/pipelines/samplers/__init__.py

@@ -1,2 +0,0 @@
-from .base import Sampler
-from .flow_euler import FlowEulerSampler, FlowEulerCfgSampler, FlowEulerGuidanceIntervalSampler

trellis/pipelines/samplers/base.py

@@ -1,20 +0,0 @@
-from typing import *
-from abc import ABC, abstractmethod
-
-
-class Sampler(ABC):
-    """
-    A base class for samplers.
-    """
-
-    @abstractmethod
-    def sample(
-        self,
-        model,
-        **kwargs
-    ):
-        """
-        Sample from a model.
-        """
-        pass
-    

trellis/pipelines/samplers/classifier_free_guidance_mixin.py

@@ -1,12 +0,0 @@
-from typing import *
-
-
-class ClassifierFreeGuidanceSamplerMixin:
-    """
-    A mixin class for samplers that apply classifier-free guidance.
-    """
-
-    def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, **kwargs):
-        pred = super()._inference_model(model, x_t, t, cond, **kwargs)
-        neg_pred = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
-        return (1 + cfg_strength) * pred - cfg_strength * neg_pred

trellis/pipelines/samplers/flow_euler.py

@@ -1,199 +0,0 @@
-from typing import *
-import torch
-import numpy as np
-from tqdm import tqdm
-from easydict import EasyDict as edict
-from .base import Sampler
-from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin
-from .guidance_interval_mixin import GuidanceIntervalSamplerMixin
-
-
-class FlowEulerSampler(Sampler):
-    """
-    Generate samples from a flow-matching model using Euler sampling.
-
-    Args:
-        sigma_min: The minimum scale of noise in flow.
-    """
-    def __init__(
-        self,
-        sigma_min: float,
-    ):
-        self.sigma_min = sigma_min
-
-    def _eps_to_xstart(self, x_t, t, eps):
-        assert x_t.shape == eps.shape
-        return (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * eps) / (1 - t)
-
-    def _xstart_to_eps(self, x_t, t, x_0):
-        assert x_t.shape == x_0.shape
-        return (x_t - (1 - t) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t)
-
-    def _v_to_xstart_eps(self, x_t, t, v):
-        assert x_t.shape == v.shape
-        eps = (1 - t) * v + x_t
-        x_0 = (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * v
-        return x_0, eps
-
-    def _inference_model(self, model, x_t, t, cond=None, **kwargs):
-        t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
-        return model(x_t, t, cond, **kwargs)
-
-    def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs):
-        pred_v = self._inference_model(model, x_t, t, cond, **kwargs)
-        pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v)
-        return pred_x_0, pred_eps, pred_v
-
-    @torch.no_grad()
-    def sample_once(
-        self,
-        model,
-        x_t,
-        t: float,
-        t_prev: float,
-        cond: Optional[Any] = None,
-        **kwargs
-    ):
-        """
-        Sample x_{t-1} from the model using Euler method.
-        
-        Args:
-            model: The model to sample from.
-            x_t: The [N x C x ...] tensor of noisy inputs at time t.
-            t: The current timestep.
-            t_prev: The previous timestep.
-            cond: conditional information.
-            **kwargs: Additional arguments for model inference.
-
-        Returns:
-            a dict containing the following
-            - 'pred_x_prev': x_{t-1}.
-            - 'pred_x_0': a prediction of x_0.
-        """
-        pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
-        pred_x_prev = x_t - (t - t_prev) * pred_v
-        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0})
-
-    @torch.no_grad()
-    def sample(
-        self,
-        model,
-        noise,
-        cond: Optional[Any] = None,
-        steps: int = 50,
-        rescale_t: float = 1.0,
-        verbose: bool = True,
-        **kwargs
-    ):
-        """
-        Generate samples from the model using Euler method.
-        
-        Args:
-            model: The model to sample from.
-            noise: The initial noise tensor.
-            cond: conditional information.
-            steps: The number of steps to sample.
-            rescale_t: The rescale factor for t.
-            verbose: If True, show a progress bar.
-            **kwargs: Additional arguments for model_inference.
-
-        Returns:
-            a dict containing the following
-            - 'samples': the model samples.
-            - 'pred_x_t': a list of prediction of x_t.
-            - 'pred_x_0': a list of prediction of x_0.
-        """
-        sample = noise
-        t_seq = np.linspace(1, 0, steps + 1)
-        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
-        t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
-        ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
-        for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose):
-            out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
-            sample = out.pred_x_prev
-            ret.pred_x_t.append(out.pred_x_prev)
-            ret.pred_x_0.append(out.pred_x_0)
-        ret.samples = sample
-        return ret
-
-
-class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
-    """
-    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance.
-    """
-    @torch.no_grad()
-    def sample(
-        self,
-        model,
-        noise,
-        cond,
-        neg_cond,
-        steps: int = 50,
-        rescale_t: float = 1.0,
-        cfg_strength: float = 3.0,
-        verbose: bool = True,
-        **kwargs
-    ):
-        """
-        Generate samples from the model using Euler method.
-        
-        Args:
-            model: The model to sample from.
-            noise: The initial noise tensor.
-            cond: conditional information.
-            neg_cond: negative conditional information.
-            steps: The number of steps to sample.
-            rescale_t: The rescale factor for t.
-            cfg_strength: The strength of classifier-free guidance.
-            verbose: If True, show a progress bar.
-            **kwargs: Additional arguments for model_inference.
-
-        Returns:
-            a dict containing the following
-            - 'samples': the model samples.
-            - 'pred_x_t': a list of prediction of x_t.
-            - 'pred_x_0': a list of prediction of x_0.
-        """
-        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, **kwargs)
-
-
-class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSampler):
-    """
-    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.
-    """
-    @torch.no_grad()
-    def sample(
-        self,
-        model,
-        noise,
-        cond,
-        neg_cond,
-        steps: int = 50,
-        rescale_t: float = 1.0,
-        cfg_strength: float = 3.0,
-        cfg_interval: Tuple[float, float] = (0.0, 1.0),
-        verbose: bool = True,
-        **kwargs
-    ):
-        """
-        Generate samples from the model using Euler method.
-        
-        Args:
-            model: The model to sample from.
-            noise: The initial noise tensor.
-            cond: conditional information.
-            neg_cond: negative conditional information.
-            steps: The number of steps to sample.
-            rescale_t: The rescale factor for t.
-            cfg_strength: The strength of classifier-free guidance.
-            cfg_interval: The interval for classifier-free guidance.
-            verbose: If True, show a progress bar.
-            **kwargs: Additional arguments for model_inference.
-
-        Returns:
-            a dict containing the following
-            - 'samples': the model samples.
-            - 'pred_x_t': a list of prediction of x_t.
-            - 'pred_x_0': a list of prediction of x_0.
-        """
-        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)

trellis/pipelines/samplers/guidance_interval_mixin.py

@@ -1,15 +0,0 @@
-from typing import *
-
-
-class GuidanceIntervalSamplerMixin:
-    """
-    A mixin class for samplers that apply classifier-free guidance with interval.
-    """
-
-    def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs):
-        if cfg_interval[0] <= t <= cfg_interval[1]:
-            pred = super()._inference_model(model, x_t, t, cond, **kwargs)
-            neg_pred = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
-            return (1 + cfg_strength) * pred - cfg_strength * neg_pred
-        else:
-            return super()._inference_model(model, x_t, t, cond, **kwargs)

trellis/pipelines/trellis_image_to_3d.py

@@ -1,376 +0,0 @@
-from typing import *
-from contextlib import contextmanager
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from tqdm import tqdm
-from easydict import EasyDict as edict
-from torchvision import transforms
-from PIL import Image
-import rembg
-from .base import Pipeline
-from . import samplers
-from ..modules import sparse as sp
-from ..representations import Gaussian, Strivec, MeshExtractResult
-
-
-class TrellisImageTo3DPipeline(Pipeline):
-    """
-    Pipeline for inferring Trellis image-to-3D models.
-
-    Args:
-        models (dict[str, nn.Module]): The models to use in the pipeline.
-        sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure.
-        slat_sampler (samplers.Sampler): The sampler for the structured latent.
-        slat_normalization (dict): The normalization parameters for the structured latent.
-        image_cond_model (str): The name of the image conditioning model.
-    """
-    def __init__(
-        self,
-        models: dict[str, nn.Module] = None,
-        sparse_structure_sampler: samplers.Sampler = None,
-        slat_sampler: samplers.Sampler = None,
-        slat_normalization: dict = None,
-        image_cond_model: str = None,
-    ):
-        if models is None:
-            return
-        super().__init__(models)
-        self.sparse_structure_sampler = sparse_structure_sampler
-        self.slat_sampler = slat_sampler
-        self.sparse_structure_sampler_params = {}
-        self.slat_sampler_params = {}
-        self.slat_normalization = slat_normalization
-        self.rembg_session = None
-        self._init_image_cond_model(image_cond_model)
-
-    @staticmethod
-    def from_pretrained(path: str) -> "TrellisImageTo3DPipeline":
-        """
-        Load a pretrained model.
-
-        Args:
-            path (str): The path to the model. Can be either local path or a Hugging Face repository.
-        """
-        pipeline = super(TrellisImageTo3DPipeline, TrellisImageTo3DPipeline).from_pretrained(path)
-        new_pipeline = TrellisImageTo3DPipeline()
-        new_pipeline.__dict__ = pipeline.__dict__
-        args = pipeline._pretrained_args
-
-        new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
-        new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
-
-        new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args'])
-        new_pipeline.slat_sampler_params = args['slat_sampler']['params']
-
-        new_pipeline.slat_normalization = args['slat_normalization']
-
-        new_pipeline._init_image_cond_model(args['image_cond_model'])
-
-        return new_pipeline
-    
-    def _init_image_cond_model(self, name: str):
-        """
-        Initialize the image conditioning model.
-        """
-        dinov2_model = torch.hub.load('facebookresearch/dinov2', name, pretrained=True)
-        dinov2_model.eval()
-        self.models['image_cond_model'] = dinov2_model
-        transform = transforms.Compose([
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-        ])
-        self.image_cond_model_transform = transform
-
-    def preprocess_image(self, input: Image.Image) -> Image.Image:
-        """
-        Preprocess the input image.
-        """
-        # if has alpha channel, use it directly; otherwise, remove background
-        has_alpha = False
-        if input.mode == 'RGBA':
-            alpha = np.array(input)[:, :, 3]
-            if not np.all(alpha == 255):
-                has_alpha = True
-        if has_alpha:
-            output = input
-        else:
-            input = input.convert('RGB')
-            max_size = max(input.size)
-            scale = min(1, 1024 / max_size)
-            if scale < 1:
-                input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
-            if getattr(self, 'rembg_session', None) is None:
-                self.rembg_session = rembg.new_session('u2net')
-            output = rembg.remove(input, session=self.rembg_session)
-        output_np = np.array(output)
-        alpha = output_np[:, :, 3]
-        bbox = np.argwhere(alpha > 0.8 * 255)
-        bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
-        center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
-        size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
-        size = int(size * 1.2)
-        bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
-        output = output.crop(bbox)  # type: ignore
-        output = output.resize((518, 518), Image.Resampling.LANCZOS)
-        output = np.array(output).astype(np.float32) / 255
-        output = output[:, :, :3] * output[:, :, 3:4]
-        output = Image.fromarray((output * 255).astype(np.uint8))
-        return output
-
-    @torch.no_grad()
-    def encode_image(self, image: Union[torch.Tensor, list[Image.Image]]) -> torch.Tensor:
-        """
-        Encode the image.
-
-        Args:
-            image (Union[torch.Tensor, list[Image.Image]]): The image to encode
-
-        Returns:
-            torch.Tensor: The encoded features.
-        """
-        if isinstance(image, torch.Tensor):
-            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
-        elif isinstance(image, list):
-            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
-            image = [i.resize((518, 518), Image.LANCZOS) for i in image]
-            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
-            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
-            image = torch.stack(image).to(self.device)
-        else:
-            raise ValueError(f"Unsupported type of image: {type(image)}")
-        
-        image = self.image_cond_model_transform(image).to(self.device)
-        features = self.models['image_cond_model'](image, is_training=True)['x_prenorm']
-        patchtokens = F.layer_norm(features, features.shape[-1:])
-        return patchtokens
-        
-    def get_cond(self, image: Union[torch.Tensor, list[Image.Image]]) -> dict:
-        """
-        Get the conditioning information for the model.
-
-        Args:
-            image (Union[torch.Tensor, list[Image.Image]]): The image prompts.
-
-        Returns:
-            dict: The conditioning information
-        """
-        cond = self.encode_image(image)
-        neg_cond = torch.zeros_like(cond)
-        return {
-            'cond': cond,
-            'neg_cond': neg_cond,
-        }
-
-    def sample_sparse_structure(
-        self,
-        cond: dict,
-        num_samples: int = 1,
-        sampler_params: dict = {},
-    ) -> torch.Tensor:
-        """
-        Sample sparse structures with the given conditioning.
-        
-        Args:
-            cond (dict): The conditioning information.
-            num_samples (int): The number of samples to generate.
-            sampler_params (dict): Additional parameters for the sampler.
-        """
-        # Sample occupancy latent
-        flow_model = self.models['sparse_structure_flow_model']
-        reso = flow_model.resolution
-        noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device)
-        sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
-        z_s = self.sparse_structure_sampler.sample(
-            flow_model,
-            noise,
-            **cond,
-            **sampler_params,
-            verbose=True
-        ).samples
-        
-        # Decode occupancy latent
-        decoder = self.models['sparse_structure_decoder']
-        coords = torch.argwhere(decoder(z_s)>0)[:, [0, 2, 3, 4]].int()
-
-        return coords
-
-    def decode_slat(
-        self,
-        slat: sp.SparseTensor,
-        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
-    ) -> dict:
-        """
-        Decode the structured latent.
-
-        Args:
-            slat (sp.SparseTensor): The structured latent.
-            formats (List[str]): The formats to decode the structured latent to.
-
-        Returns:
-            dict: The decoded structured latent.
-        """
-        ret = {}
-        if 'mesh' in formats:
-            ret['mesh'] = self.models['slat_decoder_mesh'](slat)
-        if 'gaussian' in formats:
-            ret['gaussian'] = self.models['slat_decoder_gs'](slat)
-        if 'radiance_field' in formats:
-            ret['radiance_field'] = self.models['slat_decoder_rf'](slat)
-        return ret
-    
-    def sample_slat(
-        self,
-        cond: dict,
-        coords: torch.Tensor,
-        sampler_params: dict = {},
-    ) -> sp.SparseTensor:
-        """
-        Sample structured latent with the given conditioning.
-        
-        Args:
-            cond (dict): The conditioning information.
-            coords (torch.Tensor): The coordinates of the sparse structure.
-            sampler_params (dict): Additional parameters for the sampler.
-        """
-        # Sample structured latent
-        flow_model = self.models['slat_flow_model']
-        noise = sp.SparseTensor(
-            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
-            coords=coords,
-        )
-        sampler_params = {**self.slat_sampler_params, **sampler_params}
-        slat = self.slat_sampler.sample(
-            flow_model,
-            noise,
-            **cond,
-            **sampler_params,
-            verbose=True
-        ).samples
-
-        std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device)
-        mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device)
-        slat = slat * std + mean
-        
-        return slat
-
-    @torch.no_grad()
-    def run(
-        self,
-        image: Image.Image,
-        num_samples: int = 1,
-        seed: int = 42,
-        sparse_structure_sampler_params: dict = {},
-        slat_sampler_params: dict = {},
-        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
-        preprocess_image: bool = True,
-    ) -> dict:
-        """
-        Run the pipeline.
-
-        Args:
-            image (Image.Image): The image prompt.
-            num_samples (int): The number of samples to generate.
-            sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
-            slat_sampler_params (dict): Additional parameters for the structured latent sampler.
-            preprocess_image (bool): Whether to preprocess the image.
-        """
-        if preprocess_image:
-            image = self.preprocess_image(image)
-        cond = self.get_cond([image])
-        torch.manual_seed(seed)
-        coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
-        slat = self.sample_slat(cond, coords, slat_sampler_params)
-        return self.decode_slat(slat, formats)
-
-    @contextmanager
-    def inject_sampler_multi_image(
-        self,
-        sampler_name: str,
-        num_images: int,
-        num_steps: int,
-        mode: Literal['stochastic', 'multidiffusion'] = 'stochastic',
-    ):
-        """
-        Inject a sampler with multiple images as condition.
-        
-        Args:
-            sampler_name (str): The name of the sampler to inject.
-            num_images (int): The number of images to condition on.
-            num_steps (int): The number of steps to run the sampler for.
-        """
-        sampler = getattr(self, sampler_name)
-        setattr(sampler, f'_old_inference_model', sampler._inference_model)
-
-        if mode == 'stochastic':
-            if num_images > num_steps:
-                print(f"\033[93mWarning: number of conditioning images is greater than number of steps for {sampler_name}. "
-                    "This may lead to performance degradation.\033[0m")
-
-            cond_indices = (np.arange(num_steps) % num_images).tolist()
-            def _new_inference_model(self, model, x_t, t, cond, **kwargs):
-                cond_idx = cond_indices.pop(0)
-                cond_i = cond[cond_idx:cond_idx+1]
-                return self._old_inference_model(model, x_t, t, cond=cond_i, **kwargs)
-        
-        elif mode =='multidiffusion':
-            from .samplers import FlowEulerSampler
-            def _new_inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs):
-                if cfg_interval[0] <= t <= cfg_interval[1]:
-                    preds = []
-                    for i in range(len(cond)):
-                        preds.append(FlowEulerSampler._inference_model(self, model, x_t, t, cond[i:i+1], **kwargs))
-                    pred = sum(preds) / len(preds)
-                    neg_pred = FlowEulerSampler._inference_model(self, model, x_t, t, neg_cond, **kwargs)
-                    return (1 + cfg_strength) * pred - cfg_strength * neg_pred
-                else:
-                    preds = []
-                    for i in range(len(cond)):
-                        preds.append(FlowEulerSampler._inference_model(self, model, x_t, t, cond[i:i+1], **kwargs))
-                    pred = sum(preds) / len(preds)
-                    return pred
-            
-        else:
-            raise ValueError(f"Unsupported mode: {mode}")
-            
-        sampler._inference_model = _new_inference_model.__get__(sampler, type(sampler))
-
-        yield
-
-        sampler._inference_model = sampler._old_inference_model
-        delattr(sampler, f'_old_inference_model')
-
-    @torch.no_grad()
-    def run_multi_image(
-        self,
-        images: List[Image.Image],
-        num_samples: int = 1,
-        seed: int = 42,
-        sparse_structure_sampler_params: dict = {},
-        slat_sampler_params: dict = {},
-        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
-        preprocess_image: bool = True,
-        mode: Literal['stochastic', 'multidiffusion'] = 'stochastic',
-    ) -> dict:
-        """
-        Run the pipeline with multiple images as condition
-
-        Args:
-            images (List[Image.Image]): The multi-view images of the assets
-            num_samples (int): The number of samples to generate.
-            sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
-            slat_sampler_params (dict): Additional parameters for the structured latent sampler.
-            preprocess_image (bool): Whether to preprocess the image.
-        """
-        if preprocess_image:
-            images = [self.preprocess_image(image) for image in images]
-        cond = self.get_cond(images)
-        cond['neg_cond'] = cond['neg_cond'][:1]
-        torch.manual_seed(seed)
-        ss_steps = {**self.sparse_structure_sampler_params, **sparse_structure_sampler_params}.get('steps')
-        with self.inject_sampler_multi_image('sparse_structure_sampler', len(images), ss_steps, mode=mode):
-            coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
-        slat_steps = {**self.slat_sampler_params, **slat_sampler_params}.get('steps')
-        with self.inject_sampler_multi_image('slat_sampler', len(images), slat_steps, mode=mode):
-            slat = self.sample_slat(cond, coords, slat_sampler_params)
-        return self.decode_slat(slat, formats)

trellis/pipelines/trellis_text_to_3d.py

@@ -1,278 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import numpy as np
-from transformers import CLIPTextModel, AutoTokenizer
-import open3d as o3d
-from .base import Pipeline
-from . import samplers
-from ..modules import sparse as sp
-
-
-class TrellisTextTo3DPipeline(Pipeline):
-    """
-    Pipeline for inferring Trellis text-to-3D models.
-
-    Args:
-        models (dict[str, nn.Module]): The models to use in the pipeline.
-        sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure.
-        slat_sampler (samplers.Sampler): The sampler for the structured latent.
-        slat_normalization (dict): The normalization parameters for the structured latent.
-        text_cond_model (str): The name of the text conditioning model.
-    """
-    def __init__(
-        self,
-        models: dict[str, nn.Module] = None,
-        sparse_structure_sampler: samplers.Sampler = None,
-        slat_sampler: samplers.Sampler = None,
-        slat_normalization: dict = None,
-        text_cond_model: str = None,
-    ):
-        if models is None:
-            return
-        super().__init__(models)
-        self.sparse_structure_sampler = sparse_structure_sampler
-        self.slat_sampler = slat_sampler
-        self.sparse_structure_sampler_params = {}
-        self.slat_sampler_params = {}
-        self.slat_normalization = slat_normalization
-        self._init_text_cond_model(text_cond_model)
-
-    @staticmethod
-    def from_pretrained(path: str) -> "TrellisTextTo3DPipeline":
-        """
-        Load a pretrained model.
-
-        Args:
-            path (str): The path to the model. Can be either local path or a Hugging Face repository.
-        """
-        pipeline = super(TrellisTextTo3DPipeline, TrellisTextTo3DPipeline).from_pretrained(path)
-        new_pipeline = TrellisTextTo3DPipeline()
-        new_pipeline.__dict__ = pipeline.__dict__
-        args = pipeline._pretrained_args
-
-        new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
-        new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
-
-        new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args'])
-        new_pipeline.slat_sampler_params = args['slat_sampler']['params']
-
-        new_pipeline.slat_normalization = args['slat_normalization']
-
-        new_pipeline._init_text_cond_model(args['text_cond_model'])
-
-        return new_pipeline
-    
-    def _init_text_cond_model(self, name: str):
-        """
-        Initialize the text conditioning model.
-        """
-        # load model
-        model = CLIPTextModel.from_pretrained(name)
-        tokenizer = AutoTokenizer.from_pretrained(name)
-        model.eval()
-        model = model.cuda()
-        self.text_cond_model = {
-            'model': model,
-            'tokenizer': tokenizer,
-        }
-        self.text_cond_model['null_cond'] = self.encode_text([''])
-
-    @torch.no_grad()
-    def encode_text(self, text: List[str]) -> torch.Tensor:
-        """
-        Encode the text.
-        """
-        assert isinstance(text, list) and all(isinstance(t, str) for t in text), "text must be a list of strings"
-        encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
-        tokens = encoding['input_ids'].cuda()
-        embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
-        
-        return embeddings
-        
-    def get_cond(self, prompt: List[str]) -> dict:
-        """
-        Get the conditioning information for the model.
-
-        Args:
-            prompt (List[str]): The text prompt.
-
-        Returns:
-            dict: The conditioning information
-        """
-        cond = self.encode_text(prompt)
-        neg_cond = self.text_cond_model['null_cond']
-        return {
-            'cond': cond,
-            'neg_cond': neg_cond,
-        }
-
-    def sample_sparse_structure(
-        self,
-        cond: dict,
-        num_samples: int = 1,
-        sampler_params: dict = {},
-    ) -> torch.Tensor:
-        """
-        Sample sparse structures with the given conditioning.
-        
-        Args:
-            cond (dict): The conditioning information.
-            num_samples (int): The number of samples to generate.
-            sampler_params (dict): Additional parameters for the sampler.
-        """
-        # Sample occupancy latent
-        flow_model = self.models['sparse_structure_flow_model']
-        reso = flow_model.resolution
-        noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device)
-        sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
-        z_s = self.sparse_structure_sampler.sample(
-            flow_model,
-            noise,
-            **cond,
-            **sampler_params,
-            verbose=True
-        ).samples
-        
-        # Decode occupancy latent
-        decoder = self.models['sparse_structure_decoder']
-        coords = torch.argwhere(decoder(z_s)>0)[:, [0, 2, 3, 4]].int()
-
-        return coords
-
-    def decode_slat(
-        self,
-        slat: sp.SparseTensor,
-        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
-    ) -> dict:
-        """
-        Decode the structured latent.
-
-        Args:
-            slat (sp.SparseTensor): The structured latent.
-            formats (List[str]): The formats to decode the structured latent to.
-
-        Returns:
-            dict: The decoded structured latent.
-        """
-        ret = {}
-        if 'mesh' in formats:
-            ret['mesh'] = self.models['slat_decoder_mesh'](slat)
-        if 'gaussian' in formats:
-            ret['gaussian'] = self.models['slat_decoder_gs'](slat)
-        if 'radiance_field' in formats:
-            ret['radiance_field'] = self.models['slat_decoder_rf'](slat)
-        return ret
-    
-    def sample_slat(
-        self,
-        cond: dict,
-        coords: torch.Tensor,
-        sampler_params: dict = {},
-    ) -> sp.SparseTensor:
-        """
-        Sample structured latent with the given conditioning.
-        
-        Args:
-            cond (dict): The conditioning information.
-            coords (torch.Tensor): The coordinates of the sparse structure.
-            sampler_params (dict): Additional parameters for the sampler.
-        """
-        # Sample structured latent
-        flow_model = self.models['slat_flow_model']
-        noise = sp.SparseTensor(
-            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
-            coords=coords,
-        )
-        sampler_params = {**self.slat_sampler_params, **sampler_params}
-        slat = self.slat_sampler.sample(
-            flow_model,
-            noise,
-            **cond,
-            **sampler_params,
-            verbose=True
-        ).samples
-
-        std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device)
-        mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device)
-        slat = slat * std + mean
-        
-        return slat
-
-    @torch.no_grad()
-    def run(
-        self,
-        prompt: str,
-        num_samples: int = 1,
-        seed: int = 42,
-        sparse_structure_sampler_params: dict = {},
-        slat_sampler_params: dict = {},
-        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
-    ) -> dict:
-        """
-        Run the pipeline.
-
-        Args:
-            prompt (str): The text prompt.
-            num_samples (int): The number of samples to generate.
-            seed (int): The random seed.
-            sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
-            slat_sampler_params (dict): Additional parameters for the structured latent sampler.
-            formats (List[str]): The formats to decode the structured latent to.
-        """
-        cond = self.get_cond([prompt])
-        torch.manual_seed(seed)
-        coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
-        slat = self.sample_slat(cond, coords, slat_sampler_params)
-        return self.decode_slat(slat, formats)
-    
-    def voxelize(self, mesh: o3d.geometry.TriangleMesh) -> torch.Tensor:
-        """
-        Voxelize a mesh.
-
-        Args:
-            mesh (o3d.geometry.TriangleMesh): The mesh to voxelize.
-            sha256 (str): The SHA256 hash of the mesh.
-            output_dir (str): The output directory.
-        """
-        vertices = np.asarray(mesh.vertices)
-        aabb = np.stack([vertices.min(0), vertices.max(0)])
-        center = (aabb[0] + aabb[1]) / 2
-        scale = (aabb[1] - aabb[0]).max()
-        vertices = (vertices - center) / scale
-        vertices = np.clip(vertices, -0.5 + 1e-6, 0.5 - 1e-6)
-        mesh.vertices = o3d.utility.Vector3dVector(vertices)
-        voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
-        vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])
-        return torch.tensor(vertices).int().cuda()
-
-    @torch.no_grad()
-    def run_variant(
-        self,
-        mesh: o3d.geometry.TriangleMesh,
-        prompt: str,
-        num_samples: int = 1,
-        seed: int = 42,
-        slat_sampler_params: dict = {},
-        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
-    ) -> dict:
-        """
-        Run the pipeline for making variants of an asset.
-
-        Args:
-            mesh (o3d.geometry.TriangleMesh): The base mesh.
-            prompt (str): The text prompt.
-            num_samples (int): The number of samples to generate.
-            seed (int): The random seed
-            slat_sampler_params (dict): Additional parameters for the structured latent sampler.
-            formats (List[str]): The formats to decode the structured latent to.
-        """
-        cond = self.get_cond([prompt])
-        coords = self.voxelize(mesh)
-        coords = torch.cat([
-            torch.arange(num_samples).repeat_interleave(coords.shape[0], 0)[:, None].int().cuda(),
-            coords.repeat(num_samples, 1)
-        ], 1)
-        torch.manual_seed(seed)
-        slat = self.sample_slat(cond, coords, slat_sampler_params)
-        return self.decode_slat(slat, formats)