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)