Upload 74 files

trellis/modules/attention/__init__.py

@@ -14,7 +14,7 @@ def __from_env():
     env_attn_backend = os.environ.get('ATTN_BACKEND')
     env_sttn_debug = os.environ.get('ATTN_DEBUG')
     
-    if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'sdpa', 'naive']:
+    if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'sdpa', 'naive']:
         BACKEND = env_attn_backend
     if env_sttn_debug is not None:
         DEBUG = env_sttn_debug == '1'

trellis/modules/attention/full_attn.py

@@ -7,6 +7,8 @@ if BACKEND == 'xformers':
     import xformers.ops as xops
 elif BACKEND == 'flash_attn':
     import flash_attn
+elif BACKEND == 'flash_attn_3':
+    import flash_attn_interface as flash_attn_3
 elif BACKEND == 'sdpa':
     from torch.nn.functional import scaled_dot_product_attention as sdpa
 elif BACKEND == 'naive':
@@ -118,6 +120,14 @@ def scaled_dot_product_attention(*args, **kwargs):
             out = flash_attn.flash_attn_kvpacked_func(q, kv)
         elif num_all_args == 3:
             out = flash_attn.flash_attn_func(q, k, v)
+    elif BACKEND == 'flash_attn_3':
+        if num_all_args == 1:
+            out = flash_attn_3.flash_attn_qkvpacked_func(qkv)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+            out = flash_attn_3.flash_attn_func(q, k, v)
+        elif num_all_args == 3:
+            out = flash_attn_3.flash_attn_func(q, k, v)
     elif BACKEND == 'sdpa':
         if num_all_args == 1:
             q, k, v = qkv.unbind(dim=2)

trellis/modules/sparse/__init__.py

@@ -21,7 +21,7 @@ def __from_env():
         BACKEND = env_sparse_backend
     if env_sparse_debug is not None:
         DEBUG = env_sparse_debug == '1'
-    if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn']:
+    if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn', 'flash_attn_3']:
         ATTN = env_sparse_attn
         
     print(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}")

trellis/modules/sparse/attention/full_attn.py

@@ -7,6 +7,8 @@ if ATTN == 'xformers':
     import xformers.ops as xops
 elif ATTN == 'flash_attn':
     import flash_attn
+elif ATTN == 'flash_attn_3':
+    import flash_attn_interface as flash_attn_3
 else:
     raise ValueError(f"Unknown attention module: {ATTN}")
 
@@ -206,6 +208,22 @@ def sparse_scaled_dot_product_attention(*args, **kwargs):
             out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
         elif num_all_args == 3:
             out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+    elif ATTN == 'flash_attn_3':
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+            cu_seqlens_kv = cu_seqlens_q.clone()
+            max_q_seqlen = max_kv_seqlen = max(q_seqlen)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+            max_q_seqlen = max(q_seqlen)
+            max_kv_seqlen = max(kv_seqlen)
+        elif num_all_args == 3:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+            max_q_seqlen = max(q_seqlen)
+            max_kv_seqlen = max(kv_seqlen)
+        out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen)
     else:
         raise ValueError(f"Unknown attention module: {ATTN}")
     

trellis/modules/sparse/attention/serialized_attn.py

@@ -9,6 +9,8 @@ if ATTN == 'xformers':
     import xformers.ops as xops
 elif ATTN == 'flash_attn':
     import flash_attn
+elif ATTN == 'flash_attn_3':
+    import flash_attn_interface as flash_attn_3
 else:
     raise ValueError(f"Unknown attention module: {ATTN}")
 
@@ -168,6 +170,9 @@ def sparse_serialized_scaled_dot_product_self_attention(
             out = xops.memory_efficient_attention(q, k, v)          # [B, N, H, C]
         elif ATTN == 'flash_attn':
             out = flash_attn.flash_attn_qkvpacked_func(qkv_feats)   # [B, N, H, C]
+        elif ATTN == 'flash_attn_3':
+            q, k, v = qkv_feats.unbind(dim=2)                       # [B, N, H, C]
+            out = flash_attn_3.flash_attn_func(q, k, v)             # [B, N, H, C]
         else:
             raise ValueError(f"Unknown attention module: {ATTN}")
         out = out.reshape(B * N, H, C)                              # [M, H, C]
@@ -183,6 +188,11 @@ def sparse_serialized_scaled_dot_product_self_attention(
             cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
                         .to(qkv.device).int()
             out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+        elif ATTN == 'flash_attn_3':
+            q, k, v = qkv_feats.unbind(dim=1)                       # [M, H, C]
+            cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
+                        .to(qkv.device).int()
+            out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max(seq_lens), max(seq_lens)) # [M, H, C]
 
     out = out[bwd_indices]      # [T, H, C]
 

trellis/modules/sparse/attention/windowed_attn.py

@@ -8,6 +8,8 @@ if ATTN == 'xformers':
     import xformers.ops as xops
 elif ATTN == 'flash_attn':
     import flash_attn
+elif ATTN == 'flash_attn_3':
+    import flash_attn_interface as flash_attn_3
 else:
     raise ValueError(f"Unknown attention module: {ATTN}")
 
@@ -110,6 +112,9 @@ def sparse_windowed_scaled_dot_product_self_attention(
             out = xops.memory_efficient_attention(q, k, v)          # [B, N, H, C]
         elif ATTN == 'flash_attn':
             out = flash_attn.flash_attn_qkvpacked_func(qkv_feats)   # [B, N, H, C]
+        elif ATTN == 'flash_attn_3':
+            q, k, v = qkv_feats.unbind(dim=2)                       # [B, N, H, C]
+            out = flash_attn_3.flash_attn_func(q, k, v)             # [B, N, H, C]
         else:
             raise ValueError(f"Unknown attention module: {ATTN}")
         out = out.reshape(B * N, H, C)                              # [M, H, C]
@@ -125,6 +130,11 @@ def sparse_windowed_scaled_dot_product_self_attention(
             cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
                         .to(qkv.device).int()
             out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+        elif ATTN == 'flash_attn_3':
+            q, k, v = qkv_feats.unbind(dim=1)                       # [M, H, C]
+            cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
+                        .to(qkv.device).int()
+            out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max(seq_lens), max(seq_lens)) # [M, H, C]
 
     out = out[bwd_indices]      # [T, H, C]
 

trellis/modules/sparse/basic.py

@@ -133,6 +133,37 @@ class SparseTensor:
     def dim(self) -> int:
         return len(self.shape)
     
+    @property
+    def ndim(self) -> int:
+        return self.dim()
+
+    @property
+    def dtype(self):
+        return self.feats.dtype
+
+    @property
+    def device(self):
+        return self.feats.device
+
+    @property
+    def seqlen(self) -> torch.LongTensor:
+        seqlen = self.get_spatial_cache('seqlen')
+        if seqlen is None:
+            seqlen = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device)
+            self.register_spatial_cache('seqlen', seqlen)
+        return seqlen
+
+    @property
+    def cum_seqlen(self) -> torch.LongTensor:
+        cum_seqlen = self.get_spatial_cache('cum_seqlen')
+        if cum_seqlen is None:
+            cum_seqlen = torch.cat([
+                torch.tensor([0], dtype=torch.long, device=self.device),
+                self.seqlen.cumsum(dim=0)
+            ], dim=0)
+            self.register_spatial_cache('cum_seqlen', cum_seqlen)
+        return cum_seqlen
+
     @property
     def layout(self) -> List[slice]:
         return self._layout
@@ -367,6 +398,40 @@ class SparseTensor:
         feats = torch.cat(feats, dim=0).contiguous()
         return SparseTensor(feats=feats, coords=coords)
 
+    def reduce(self, op: str, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        if isinstance(dim, int):
+            dim = (dim,)
+        
+        if op =='mean':
+            red = self.feats.mean(dim=dim, keepdim=keepdim)
+        elif op =='sum':
+            red = self.feats.sum(dim=dim, keepdim=keepdim)
+        elif op == 'prod':
+            red = self.feats.prod(dim=dim, keepdim=keepdim)
+        else:
+            raise ValueError(f"Unsupported reduce operation: {op}")
+        
+        if dim is None or 0 in dim:
+            return red
+        
+        red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)
+        return red
+
+    def mean(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='mean', dim=dim, keepdim=keepdim)
+        
+    def sum(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='sum', dim=dim, keepdim=keepdim)
+        
+    def prod(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='prod', dim=dim, keepdim=keepdim)
+    
+    def std(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        mean = self.mean(dim=dim, keepdim=True)
+        mean2 = self.replace(self.feats ** 2).mean(dim=dim, keepdim=True)
+        std = (mean2 - mean ** 2).sqrt()
+        return std
+
     def register_spatial_cache(self, key, value) -> None:
         """
         Register a spatial cache.

trellis/pipelines/__init__.py

@@ -1,5 +1,6 @@
 from . import samplers
 from .trellis_image_to_3d import TrellisImageTo3DPipeline, TrellisVGGTTo3DPipeline
+from .trellis_hybrid_pipeline import TrellisHybridPipeline
 
 def from_pretrained(path: str):
     """

trellis/pipelines/samplers/classifier_free_guidance_mixin.py

@@ -1,12 +1,39 @@
 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
+
 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
+    def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, guidance_rescale=0.0, **kwargs):
+        if cfg_strength == 1:
+            return super()._inference_model(model, x_t, t, cond, **kwargs)
+        elif cfg_strength == 0:
+            return super()._inference_model(model, x_t, t, neg_cond, **kwargs)
+        else:
+            pred_pos = super()._inference_model(model, x_t, t, cond, **kwargs)
+            pred_neg = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
+            pred = cfg_strength * pred_pos + (1 - cfg_strength) * pred_neg
+
+            # CFG rescale
+            if guidance_rescale > 0:
+                x_0_pos = self._pred_to_xstart(x_t, t, pred_pos)
+                x_0_cfg = self._pred_to_xstart(x_t, t, pred)
+                std_pos = x_0_pos.std(dim=list(range(1, x_0_pos.ndim)), keepdim=True)
+                std_cfg = x_0_cfg.std(dim=list(range(1, x_0_cfg.ndim)), keepdim=True)
+                x_0_rescaled = x_0_cfg * (std_pos / std_cfg)
+                x_0 = guidance_rescale * x_0_rescaled + (1 - guidance_rescale) * x_0_cfg
+                pred = self._xstart_to_pred(x_t, t, x_0)
+                
+            return pred

trellis/pipelines/samplers/flow_euler.py

@@ -50,6 +50,11 @@ class FlowEulerSampler(Sampler):
         assert x_0.shape == x_t.shape
         return (x_t - (1 - self.sigma_min) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t)
 
+    def _pred_to_xstart(self, x_t, t, pred):
+        return (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * pred
+    
+    def _xstart_to_pred(self, x_t, t, x_0):
+        return ((1 - self.sigma_min) * x_t - x_0) / (self.sigma_min + (1 - self.sigma_min) * t)
 
     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)
@@ -178,73 +183,6 @@ class FlowEulerSampler(Sampler):
         torch.cuda.empty_cache()
         return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0, "pred_eps": pred_eps})
 
-    def sample_slat_once_opt_delta_v(
-        self,
-        model,
-        slat_decoder_gs,
-        slat_decoder_mesh,
-        std, 
-        mean,
-        dreamsim_model,
-        learning_rate,
-        input_images,
-        extrinsics, 
-        intrinsics,
-        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.
-        """
-        torch.cuda.empty_cache()
-        with torch.no_grad():
-            pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
-        pred_v_opt_feat = torch.nn.Parameter(pred_v.feats.detach().clone())
-        optimizer = torch.optim.Adam([pred_v_opt_feat], betas=(0.5, 0.9), lr=learning_rate)
-        pred_v_opt = sp.SparseTensor(feats=pred_v_opt_feat, coords=pred_v.coords)
-        total_steps = 5
-        input_images = F.interpolate(input_images, size=(259, 259), mode='bilinear', align_corners=False)
-        with tqdm(total=total_steps, disable=True, desc='Appearance (opt): optimizing') as pbar:
-            for step in range(total_steps):
-                optimizer.zero_grad()
-                pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v_opt)
-                pred_gs = slat_decoder_gs(pred_x_0 * std + mean)
-                # pred_mesh = slat_decoder_mesh(pred_x_0 * std + mean)
-                rend_gs = render_utils.render_frames(pred_gs[0], extrinsics, intrinsics, {'resolution': 259, 'bg_color': (0, 0, 0)}, need_depth=True, opt=True)['color']
-                # rend_mesh = render_utils.render_frames_opt(pred_mesh[0], extrinsics, intrinsics, {'resolution': 518, 'bg_color': (0, 0, 0)}, need_depth=True, opt=True)['color']
-                rend_gs = torch.stack(rend_gs, dim=0)
-                loss_gs = loss_utils.l1_loss(rend_gs, input_images, size_average=False).mean(dim=(1,2,3)) + \
-                    (1 - loss_utils.ssim(rend_gs, input_images, size_average=False)) + \
-                        loss_utils.lpips(rend_gs, input_images, size_average=False).mean(dim=(1,2,3)) + \
-                            dreamsim_model(rend_gs, input_images)
-                loss_gs = loss_gs[loss_gs <= 0.8].mean()
-                # loss_gs = (1 - loss_utils.ssim(rend_gs, input_images)) + loss_utils.lpips(rend_gs, input_images) + dreamsim_model(rend_gs, input_images).mean()
-                # loss_mesh = loss_utils.l1_loss(rend_mesh, input_images) + 0.2 * (1 - loss_utils.ssim(rend_mesh, input_images)) + 0.2 * loss_utils.lpips(rend_mesh, input_images)
-                loss = loss_gs + 0.2 * loss_utils.l1_loss(pred_v_opt_feat, pred_v.feats)
-                loss.backward()
-                optimizer.step()
-                pbar.set_postfix({'loss': loss.item()})
-                pbar.update()
-        
-        pred_x_prev = x_t - (t - t_prev) * pred_v_opt.detach()
-        torch.cuda.empty_cache()
-        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0, "pred_eps": pred_eps})
 
     def sample_opt(
         self,
@@ -342,66 +280,6 @@ class FlowEulerSampler(Sampler):
         ret.samples = sample
         return ret
 
-    def sample_slat_opt_delta_v(
-        self,
-        model,
-        slat_decoder_gs,
-        slat_decoder_mesh,
-        std,
-        mean,
-        dreamsim_model,
-        apperance_learning_rate,
-        start_t,
-        input_images,
-        extrinsics, 
-        intrinsics,
-        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": []})
-        # def cosine_anealing(step, total_steps, start_lr, end_lr):
-        #     return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps))
-        for i, (t, t_prev) in enumerate(tqdm(t_pairs, desc="Sampling", disable=not verbose)):
-            if t > start_t:
-                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)
-            else:
-                # learning_rate = cosine_anealing(i - int(np.where(t_seq <= start_t)[0].min()), int(steps - np.where(t_seq <= start_t)[0].min()), apperance_learning_rate, 1e-5)
-                learning_rate = apperance_learning_rate
-                out = self.sample_slat_once_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, std, mean, dreamsim_model, learning_rate, input_images, extrinsics, intrinsics, 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
 
     @torch.no_grad()
     def sample(
@@ -746,7 +624,7 @@ class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
         return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, **kwargs)
 
 
-class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSampler):
+class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
     """
     Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.
     """
@@ -864,53 +742,6 @@ class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSa
         return super().sample_ss_opt_delta_v(model, ss_decoder, ss_learning_rate, ss_start_t, ss, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)
 
 
-    def sample_slat_opt_delta_v(
-        self,
-        model,
-        slat_decoder_gs,
-        slat_decoder_mesh,
-        std,
-        mean,
-        dreamsim_model,
-        apperance_learning_rate,
-        start_t,
-        input_images,
-        extrinsics, 
-        intrinsics,
-        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_slat_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, std, mean, dreamsim_model, apperance_learning_rate, start_t, input_images, extrinsics, intrinsics,noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)
-
-
 class LatentMatchGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, LatentMatchSampler):
     """
     Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.

trellis/pipelines/samplers/guidance_interval_mixin.py

@@ -1,15 +1,26 @@
 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)
+
 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):
+    def _inference_model(self, model, x_t, t, 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
+            return super()._inference_model(model, x_t, t, cond, cfg_strength=cfg_strength, **kwargs)
         else:
-            return super()._inference_model(model, x_t, t, cond, **kwargs)
+            return super()._inference_model(model, x_t, t, cond, cfg_strength=1, **kwargs)

trellis/pipelines/trellis_hybrid_pipeline.py

@@ -0,0 +1,694 @@
+"""
+Hybrid pipeline: ReconViaGen VGGT-based SS stage + TRELLIS.2 shape_slat / tex_slat stages.
+
+Stage 1  – Sparse Structure (SS)  : TrellisVGGTTo3DPipeline  (ReconViaGen, VGGT-conditioned)
+Stage 2  – Shape SLat             : Trellis2ImageTo3DPipeline (TRELLIS.2, DINOv3-conditioned)
+Stage 3  – Texture SLat           : Trellis2ImageTo3DPipeline (TRELLIS.2, DINOv3-conditioned)
+Stage 4  – Decode / GLB export    : Trellis2ImageTo3DPipeline (TRELLIS.2, o_voxel)
+"""
+
+import sys
+import os
+
+# ── Make trellis2 importable from wheels/TRELLIS.2 ───────────────────────────
+_TRELLIS2_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), '..', '..', 'wheels', 'TRELLIS.2'))
+
+if _TRELLIS2_ROOT not in sys.path:
+    sys.path.insert(0, _TRELLIS2_ROOT)
+# o_voxel is installed into the conda env (no extra sys.path needed)
+
+# ── Standard imports ──────────────────────────────────────────────────────────
+from typing import *
+import torch
+import numpy as np
+from PIL import Image
+
+from trellis2.pipelines import Trellis2ImageTo3DPipeline
+from trellis2.modules.sparse import SparseTensor
+
+from .trellis_image_to_3d import TrellisVGGTTo3DPipeline
+
+
+class TrellisHybridPipeline:
+    """
+    Hybrid 3-stage pipeline:
+
+    1. SS      – uses ReconViaGen's TrellisVGGTTo3DPipeline (VGGT + DINOv2 features)
+    2. Shape   – uses TRELLIS.2's Trellis2ImageTo3DPipeline  (shape_slat, DINOv3 features)
+    3. Texture – uses TRELLIS.2's Trellis2ImageTo3DPipeline  (tex_slat,   DINOv3 features)
+
+    The two sub-pipelines are kept separate so each can be moved between CPU/GPU
+    independently when ``low_vram=True``.
+
+    Args:
+        vggt_pipeline   : Loaded TrellisVGGTTo3DPipeline  (provides SS components).
+        trellis2_pipeline: Loaded Trellis2ImageTo3DPipeline (provides shape/tex slat + decode).
+    """
+
+    def __init__(
+        self,
+        vggt_pipeline: TrellisVGGTTo3DPipeline,
+        trellis2_pipeline: Trellis2ImageTo3DPipeline,
+        low_vram: bool = False,
+    ):
+        self.vggt_pipeline    = vggt_pipeline
+        self.trellis2_pipeline = trellis2_pipeline
+        self.low_vram = low_vram
+        if low_vram:
+            trellis2_pipeline.low_vram = True
+
+    # ── Convenience wrappers ──────────────────────────────────────────────────
+
+    def _vggt_models_to(self, device) -> None:
+        """Move vggt_pipeline inference models to *device* (for low-VRAM mode)."""
+        vp = self.vggt_pipeline
+        if hasattr(vp, 'VGGT_model') and vp.VGGT_model is not None:
+            vp.VGGT_model.to(device)
+        for model in vp.models.values():
+            model.to(device)
+
+    @property
+    def device(self):
+        return self.vggt_pipeline.device
+
+    @property
+    def pbr_attr_layout(self):
+        return self.trellis2_pipeline.pbr_attr_layout
+
+    def preprocess_image(self, image: Image.Image) -> Image.Image:
+        """
+        Preprocess the input image — mirrors Trellis2ImageTo3DPipeline.preprocess_image:
+        if the image already has a real alpha channel it is used directly;
+        otherwise the background is removed with the rembg_model (BiRefNet).
+        The foreground is then cropped and composited onto a black background.
+        """
+        t2p = self.trellis2_pipeline
+
+        # Use existing alpha if present, otherwise run background removal
+        has_alpha = False
+        if image.mode == 'RGBA':
+            alpha = np.array(image)[:, :, 3]
+            if not np.all(alpha == 255):
+                has_alpha = True
+
+        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 has_alpha:
+            output = image
+        else:
+            image = image.convert('RGB')
+            if t2p.low_vram:
+                t2p.rembg_model.to(t2p.device)
+            output = t2p.rembg_model(image)
+            if t2p.low_vram:
+                t2p.rembg_model.cpu()
+
+        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)
+        bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
+        output = output.crop(bbox)
+        output = np.array(output).astype(np.float32) / 255
+        output = output[:, :, :3] * output[:, :, 3:4]
+        return Image.fromarray((output * 255).astype(np.uint8))
+
+    # ── Stage 1: SS via mesh voxelisation ────────────────────────────────────
+
+    @staticmethod
+    def _mesh_to_surface_coords(
+        mesh_result,
+        target_res: int,
+        device: torch.device,
+        simplify_ratio: float = 0.95,
+        fill_holes_max_hole_size: float = 0.04,
+        fill_holes_max_hole_nbe: int = 55,
+        fill_holes_resolution: int = 256,
+        fill_holes_num_views: int = 200,
+    ) -> torch.Tensor:
+        """
+        Convert a MeshExtractResult (surface triangle mesh in [-0.5, 0.5]^3) to
+        surface-only voxel coords [N, 4] = [batch_idx, x, y, z] in [0, target_res).
+
+        Pipeline
+        --------
+        1. Decimate + Fill holes – postprocess_mesh (pyvista decimate → GPU _fill_holes),
+           same approach as postprocessing_utils.py lines 426-437
+        2. Voxelize – Open3D surface voxelization directly at target_res
+        """
+        from trellis.utils.postprocessing_utils import postprocess_mesh
+        import open3d as o3d
+
+        verts_np = mesh_result.vertices.cpu().numpy()
+        faces_np = mesh_result.faces.cpu().numpy()
+
+        # 1. ── Decimate + fill holes (mirrors postprocessing_utils.py L426-437) ─
+        verts_np, faces_np = postprocess_mesh(
+            verts_np, faces_np,
+            simplify=True,
+            simplify_ratio=simplify_ratio,
+            fill_holes=True,
+            fill_holes_max_hole_size=fill_holes_max_hole_size,
+            fill_holes_max_hole_nbe=fill_holes_max_hole_nbe,
+            fill_holes_resolution=fill_holes_resolution,
+            fill_holes_num_views=fill_holes_num_views,
+        )
+
+        # 2. ── Voxelize directly at target_res ───────────────────────────────
+        o3d_mesh = o3d.geometry.TriangleMesh()
+        o3d_mesh.vertices  = o3d.utility.Vector3dVector(
+            np.clip(verts_np.astype(np.float64), -0.5 + 1e-6, 0.5 - 1e-6)
+        )
+        o3d_mesh.triangles = o3d.utility.Vector3iVector(faces_np.astype(np.int32))
+        voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(
+            o3d_mesh,
+            voxel_size=1.0 / target_res,
+            min_bound=(-0.5, -0.5, -0.5),
+            max_bound=( 0.5,  0.5,  0.5),
+        )
+        grid_idx = np.array(
+            [voxel.grid_index for voxel in voxel_grid.get_voxels()], dtype=np.int32
+        )                                               # [N, 3]  in [0, target_res)
+
+        if len(grid_idx) == 0:
+            return torch.zeros((0, 4), dtype=torch.int32, device=device)
+
+        idx = torch.tensor(grid_idx, dtype=torch.int32, device=device)
+        batch = torch.zeros(len(idx), 1, dtype=torch.int32, device=device)
+        coords = torch.cat([batch, idx], dim=1)         # [N, 4]: [0, x, y, z]
+        return coords
+
+    @torch.no_grad()
+    def _run_ss_stage(
+        self,
+        images: List[Image.Image],
+        target_ss_res: int,
+        ss_sampler_params: dict,
+        slat_sampler_params: dict,
+    ) -> torch.Tensor:
+        """
+        Generate a rough mesh via vggt_pipeline, then voxelise it into
+        surface-only coords at target_ss_res^3 for the downstream shape/tex stages.
+
+        Returns:
+            coords : (N, 4) int tensor  [batch_idx, x, y, z]  in [0, target_ss_res)
+        """
+        vp = self.vggt_pipeline
+        # vp.device is dynamic (inferred from model params), so when models are on
+        # CPU it returns 'cpu'. Hardcode the target cuda device instead.
+        cuda_device = torch.device('cuda')
+
+        if self.low_vram:
+            self._vggt_models_to(cuda_device)
+
+        outputs, _, _ = vp.run(
+            image=images,
+            formats=["mesh"],
+            preprocess_image=False,
+            sparse_structure_sampler_params=ss_sampler_params,
+            slat_sampler_params=slat_sampler_params,
+        )
+        mesh_result = outputs["mesh"][0]
+        coords = self._mesh_to_surface_coords(mesh_result, target_ss_res, cuda_device)
+
+        if self.low_vram:
+            self._vggt_models_to('cpu')
+            torch.cuda.empty_cache()
+
+        return coords
+
+    @torch.no_grad()
+    def _run_ss_stage_direct(
+        self,
+        images: List[Image.Image],
+        target_ss_res: int,
+        ss_sampler_params: dict,
+    ) -> torch.Tensor:
+        """
+        Run only ReconViaGen's sparse structure diffusion stage to obtain coords
+        directly, without proceeding to the SLAT/mesh stage.
+
+        Returns:
+            coords : (N, 4) int tensor  [batch_idx, x, y, z]  in [0, target_ss_res)
+        """
+        vp = self.vggt_pipeline
+        cuda_device = torch.device('cuda')
+
+        if self.low_vram:
+            self._vggt_models_to(cuda_device)
+
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(dtype=vp.VGGT_dtype):
+                aggregated_tokens_list, _ = vp.vggt_feat(images)
+            b, n, _, _ = aggregated_tokens_list[0].shape
+            image_cond = vp.encode_image(images).reshape(b, n, -1, 1024)
+            ss_cond = vp.get_ss_cond(image_cond[:, :, 5:], aggregated_tokens_list, 1)
+
+        ss_flow_model = vp.models['sparse_structure_flow_model']
+        sampler_params = {**vp.sparse_structure_sampler_params, **ss_sampler_params}
+        reso = ss_flow_model.resolution
+        ss_noise = torch.randn(1, ss_flow_model.in_channels, reso, reso, reso).to(cuda_device)
+
+        ss_latent = vp.sparse_structure_sampler.sample(
+            ss_flow_model,
+            ss_noise,
+            **ss_cond,
+            **sampler_params,
+            verbose=True,
+        ).samples
+
+        decoder = vp.models['sparse_structure_decoder']
+        decoded = decoder(ss_latent) > 0
+        if target_ss_res != decoded.shape[2]:
+            ratio = decoded.shape[2] // target_ss_res
+            decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
+        coords = torch.argwhere(decoded)[:, [0, 2, 3, 4]].int()
+
+        if self.low_vram:
+            self._vggt_models_to('cpu')
+            torch.cuda.empty_cache()
+
+        return coords
+
+    @staticmethod
+    def _translate_ss_params(ss_sampler_params: dict) -> dict:
+        """
+        Translate ReconViaGen SS sampler param keys to TRELLIS.2 format.
+
+        ReconViaGen uses ``cfg_strength`` / ``cfg_interval``;
+        TRELLIS.2 expects ``guidance_strength`` / ``guidance_interval``.
+        All other keys (steps, guidance_rescale, rescale_t) pass through unchanged.
+        """
+        mapping = {'cfg_strength': 'guidance_strength', 'cfg_interval': 'guidance_interval'}
+        return {mapping.get(k, k): v for k, v in ss_sampler_params.items()}
+
+    @torch.no_grad()
+    def _run_ss_stage_trellis2(
+        self,
+        images: List[Image.Image],
+        target_ss_res: int,
+        strategy: Optional[str] = None,
+        sampler_params: dict = {},
+    ) -> torch.Tensor:
+        """
+        Use TRELLIS.2's sparse structure flow model to generate voxel coordinates.
+
+        - Single image or strategy=None : calls t2p.sample_sparse_structure() directly.
+        - Multi-image with strategy     : uses t2p._multi_image_sample() for fusion.
+
+        Args:
+            sampler_params: Override params in TRELLIS.2 key format
+                            (guidance_strength, guidance_rescale, rescale_t, steps).
+                            Use _translate_ss_params() to convert from ReconViaGen format.
+
+        Returns:
+            coords : (N, 4) int tensor  [batch_idx, x, y, z]  in [0, target_ss_res)
+        """
+        t2p = self.trellis2_pipeline
+
+        if len(images) == 1 or strategy is None:
+            cond_512 = t2p.get_cond(images, 512)
+            return t2p.sample_sparse_structure(
+                cond_512, target_ss_res, num_samples=1,
+                sampler_params=sampler_params,
+            )
+
+        # Multi-image: per-image conds + fusion strategy
+        conds_512 = [t2p.get_cond([img], 512) for img in images]
+        flow_model_ss = t2p.models['sparse_structure_flow_model']
+        reso = flow_model_ss.resolution
+        noise = torch.randn(1, flow_model_ss.in_channels, reso, reso, reso).to(t2p.device)
+        if t2p.low_vram:
+            flow_model_ss.to(t2p.device)
+        z_s = t2p._multi_image_sample(
+            t2p.sparse_structure_sampler, flow_model_ss, noise,
+            conds_512, t2p.sparse_structure_sampler_params, sampler_params,
+            strategy=strategy, verbose=True,
+            tqdm_desc="Sampling sparse structure (TRELLIS.2)",
+        )
+        if t2p.low_vram:
+            flow_model_ss.cpu()
+
+        decoder_ss = t2p.models['sparse_structure_decoder']
+        if t2p.low_vram:
+            decoder_ss.to(t2p.device)
+        decoded = decoder_ss(z_s) > 0
+        if t2p.low_vram:
+            decoder_ss.cpu()
+
+        if target_ss_res != decoded.shape[2]:
+            ratio = decoded.shape[2] // target_ss_res
+            decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
+        return torch.argwhere(decoded)[:, [0, 2, 3, 4]].int()
+
+    # ── Main run (single image) ───────────────────────────────────────────────
+
+    @torch.no_grad()
+    def run(
+        self,
+        images: List[Image.Image],
+        seed: int = 42,
+        ss_sampler_params: dict = {},
+        slat_sampler_params: dict = {},
+        shape_slat_sampler_params: dict = {},
+        tex_slat_sampler_params: dict = {},
+        pipeline_type: str = '1024',
+        preprocess_image: bool = True,
+        return_latent: bool = False,
+        max_num_tokens: int = 49152,
+        ss_source: str = 'direct',
+    ):
+        """
+        Run the full hybrid pipeline.
+
+        Args:
+            images                   : List of (preprocessed) PIL images.
+            seed                     : Random seed.
+            ss_sampler_params        : Override params for the SS sampler (ReconViaGen keys:
+                                       cfg_strength, cfg_interval, guidance_rescale, rescale_t, steps).
+            slat_sampler_params      : Override params for the ReconViaGen SLat sampler
+                                       (only used when ss_source='mesh').
+            shape_slat_sampler_params: Override params for the TRELLIS.2 shape-SLat sampler.
+            tex_slat_sampler_params  : Override params for the TRELLIS.2 tex-SLat sampler.
+            pipeline_type            : '512'          – shape/tex slat at 512 resolution.
+                                       '1024'         – coords at res 64, direct shape_slat_flow_model_1024
+                                                        (no cascade, fast).
+                                       '1024_cascade' – LR 512 → upsample → HR 1024 (higher detail).
+                                       '1536_cascade' – LR 512 → upsample → HR 1536 (highest detail).
+            preprocess_image         : Whether to preprocess images first.
+            return_latent            : If True, also return (shape_slat, tex_slat, res).
+            ss_source                : 'direct'     – ReconViaGen SS diffusion → coords (fast).
+                                       'mesh'       – ReconViaGen full pipeline → mesh → voxelize → coords.
+                                       'mvtrellis2' – TRELLIS.2 SS flow model → coords.
+
+        Returns:
+            List[MeshWithVoxel]  (one per sample, usually 1)
+            optionally: (shape_slat, tex_slat, res)
+        """
+        assert pipeline_type in ('512', '1024', '1024_cascade', '1536_cascade'), \
+            f"pipeline_type must be '512', '1024', '1024_cascade', or '1536_cascade', got '{pipeline_type}'"
+        assert ss_source in ('direct', 'mesh', 'mvtrellis2'), \
+            f"ss_source must be 'direct', 'mesh', or 'mvtrellis2', got '{ss_source}'"
+
+        torch.manual_seed(seed)
+
+        if preprocess_image:
+            images = [self.preprocess_image(img) for img in images]
+
+        # Target SS resolution: 32 for 512/cascade pipelines, 64 for direct-1024
+        target_ss_res = {'512': 32, '1024': 64, '1024_cascade': 32, '1536_cascade': 32}[pipeline_type]
+
+        # ── Stage 1: SS ───────────────────────────────────────────────────────
+        if ss_source == 'direct':
+            # ReconViaGen SS diffusion only → coords directly (no mesh)
+            coords = self._run_ss_stage_direct(images, target_ss_res, ss_sampler_params)
+        elif ss_source == 'mvtrellis2':
+            # TRELLIS.2 SS flow model → coords (single image, no fusion strategy)
+            coords = self._run_ss_stage_trellis2(
+                images, target_ss_res, strategy=None,
+                sampler_params=self._translate_ss_params(ss_sampler_params),
+            )
+        else:  # 'mesh'
+            # ReconViaGen full pipeline → mesh → decimate/fill/voxelize → coords
+            coords = self._run_ss_stage(images, target_ss_res, ss_sampler_params, slat_sampler_params)
+
+        # ── Stages 2 & 3: shape_slat + tex_slat (TRELLIS.2) ──────────────────
+        t2p = self.trellis2_pipeline
+
+        cond_512 = t2p.get_cond(images, 512) if pipeline_type in ('512', '1024_cascade', '1536_cascade') else None
+        cond_1024 = t2p.get_cond(images, 1024) if pipeline_type != '512' else None
+
+        if pipeline_type == '512':
+            shape_slat = t2p.sample_shape_slat(
+                cond_512,
+                t2p.models['shape_slat_flow_model_512'],
+                coords,
+                shape_slat_sampler_params,
+            )
+            tex_slat = t2p.sample_tex_slat(
+                cond_512,
+                t2p.models['tex_slat_flow_model_512'],
+                shape_slat,
+                tex_slat_sampler_params,
+            )
+            res = 512
+        elif pipeline_type == '1024':
+            # Direct 1024: coords at res 64 → shape_slat_flow_model_1024 (no cascade)
+            shape_slat = t2p.sample_shape_slat(
+                cond_1024,
+                t2p.models['shape_slat_flow_model_1024'],
+                coords,
+                shape_slat_sampler_params,
+            )
+            tex_slat = t2p.sample_tex_slat(
+                cond_1024,
+                t2p.models['tex_slat_flow_model_1024'],
+                shape_slat,
+                tex_slat_sampler_params,
+            )
+            res = 1024
+        elif pipeline_type == '1024_cascade':
+            shape_slat, res = t2p.sample_shape_slat_cascade(
+                cond_512, cond_1024,
+                t2p.models['shape_slat_flow_model_512'], t2p.models['shape_slat_flow_model_1024'],
+                512, 1024,
+                coords,
+                shape_slat_sampler_params,
+                max_num_tokens
+            )
+            tex_slat = t2p.sample_tex_slat(
+                cond_1024,
+                t2p.models['tex_slat_flow_model_1024'],
+                shape_slat,
+                tex_slat_sampler_params,
+            )
+        elif pipeline_type == '1536_cascade':
+            shape_slat, res = t2p.sample_shape_slat_cascade(
+                cond_512, cond_1024,
+                t2p.models['shape_slat_flow_model_512'], t2p.models['shape_slat_flow_model_1024'],
+                512, 1536,
+                coords, shape_slat_sampler_params,
+                max_num_tokens
+            )
+            tex_slat = t2p.sample_tex_slat(
+                cond_1024,
+                t2p.models['tex_slat_flow_model_1024'],
+                shape_slat,
+                tex_slat_sampler_params,
+            )
+
+        # ── Stage 4: Decode ───────────────────────────────────────────────────
+        torch.cuda.empty_cache()
+        out_mesh = t2p.decode_latent(shape_slat, tex_slat, res)
+
+        if return_latent:
+            return out_mesh, (shape_slat, tex_slat, res)
+        return out_mesh
+
+    # ── Multi-image run ───────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def run_multi_image(
+        self,
+        images: List[Image.Image],
+        strategy: str = 'average_right',
+        seed: int = 42,
+        ss_sampler_params: dict = {},
+        slat_sampler_params: dict = {},
+        shape_slat_sampler_params: dict = {},
+        tex_slat_sampler_params: dict = {},
+        pipeline_type: str = '1024',
+        preprocess_image: bool = True,
+        return_latent: bool = False,
+        max_num_tokens: int = 49152,
+        ss_source: str = 'direct',
+    ):
+        """
+        Multi-image variant.
+
+        The SS stage uses all images jointly (VGGT processes them together, or
+        TRELLIS.2 multi-image fusion when ss_source='mvtrellis2').
+        The shape/tex slat stages use TRELLIS.2's ``_multi_image_sample`` with
+        the chosen fusion ``strategy``.
+
+        Pipeline types:
+          - '512'          : direct shape_slat_flow_model_512, res=512.
+          - '1024'         : coords at res 64, direct shape_slat_flow_model_1024 (no cascade).
+          - '1024_cascade' : LR 512 → upsample → HR 1024 cascade.
+          - '1536_cascade' : LR 512 → upsample → HR 1536 cascade.
+
+        Args:
+            strategy: 'sequential' | 'average' | 'average_right' | 'weighted_average'
+                      | 'adaptive_guidance_weight' | 'fixed_guidance_rescale'
+                      (passed to Trellis2ImageTo3DPipeline._multi_image_sample)
+        """
+        assert pipeline_type in ('512', '1024', '1024_cascade', '1536_cascade'), \
+            f"pipeline_type must be '512', '1024', '1024_cascade', or '1536_cascade', got '{pipeline_type}'"
+
+        # Single image → fall back to run()
+        if len(images) == 1:
+            return self.run(
+                images, seed=seed,
+                ss_sampler_params=ss_sampler_params,
+                slat_sampler_params=slat_sampler_params,
+                shape_slat_sampler_params=shape_slat_sampler_params,
+                tex_slat_sampler_params=tex_slat_sampler_params,
+                pipeline_type=pipeline_type,
+                preprocess_image=preprocess_image,
+                return_latent=return_latent,
+                max_num_tokens=max_num_tokens,
+                ss_source=ss_source,
+            )
+
+        torch.manual_seed(seed)
+
+        if preprocess_image:
+            images = [self.preprocess_image(img) for img in images]
+
+        target_ss_res = {'512': 32, '1024': 64, '1024_cascade': 32, '1536_cascade': 32}[pipeline_type]
+
+        # ── Stage 1: SS (all images fed jointly) ─────────────────────────────
+        if ss_source == 'direct':
+            coords = self._run_ss_stage_direct(images, target_ss_res, ss_sampler_params)
+        elif ss_source == 'mvtrellis2':
+            # TRELLIS.2 SS flow model with multi-image fusion strategy
+            coords = self._run_ss_stage_trellis2(
+                images, target_ss_res, strategy=strategy,
+                sampler_params=self._translate_ss_params(ss_sampler_params),
+            )
+        else:  # 'mesh'
+            coords = self._run_ss_stage(images, target_ss_res, ss_sampler_params, slat_sampler_params)
+
+        # ── Per-image conditioning ────────────────────────────────────────────
+        t2p = self.trellis2_pipeline
+        conds_512  = [t2p.get_cond([img], 512) for img in images] if pipeline_type in ('512', '1024_cascade', '1536_cascade') else None
+        conds_1024 = [t2p.get_cond([img], 1024) for img in images] if pipeline_type != '512' else None
+
+        std_shape  = torch.tensor(t2p.shape_slat_normalization['std'])[None]
+        mean_shape = torch.tensor(t2p.shape_slat_normalization['mean'])[None]
+        std_tex    = torch.tensor(t2p.tex_slat_normalization['std'])[None]
+        mean_tex   = torch.tensor(t2p.tex_slat_normalization['mean'])[None]
+
+        def _sample_shape(fm, conds, tqdm_desc):
+            noise_slat = SparseTensor(
+                feats=torch.randn(coords.shape[0], fm.in_channels).to(t2p.device),
+                coords=coords,
+            )
+            if t2p.low_vram:
+                fm.to(t2p.device)
+            slat = t2p._multi_image_sample(
+                t2p.shape_slat_sampler, fm, noise_slat,
+                conds, t2p.shape_slat_sampler_params, shape_slat_sampler_params,
+                strategy=strategy, verbose=True, tqdm_desc=tqdm_desc,
+            )
+            if t2p.low_vram:
+                fm.cpu()
+            return slat * std_shape.to(slat.device) + mean_shape.to(slat.device)
+
+        def _sample_tex(fm, conds, shape_slat):
+            s_norm = (shape_slat - mean_shape.to(shape_slat.device)) / std_shape.to(shape_slat.device)
+            in_ch  = fm.in_channels if isinstance(fm, torch.nn.Module) else fm[0].in_channels
+            noise_tex = s_norm.replace(
+                feats=torch.randn(s_norm.coords.shape[0], in_ch - s_norm.feats.shape[1]).to(t2p.device)
+            )
+            if t2p.low_vram:
+                fm.to(t2p.device)
+            slat = t2p._multi_image_sample(
+                t2p.tex_slat_sampler, fm, noise_tex,
+                conds, t2p.tex_slat_sampler_params, tex_slat_sampler_params,
+                strategy=strategy, verbose=True, tqdm_desc="Sampling texture SLat",
+                concat_cond=s_norm,
+            )
+            if t2p.low_vram:
+                fm.cpu()
+            return slat * std_tex.to(slat.device) + mean_tex.to(slat.device)
+
+        if pipeline_type == '512':
+            shape_slat = _sample_shape(t2p.models['shape_slat_flow_model_512'], conds_512, "Sampling shape SLat")
+            tex_slat   = _sample_tex(t2p.models['tex_slat_flow_model_512'], conds_512, shape_slat)
+            res = 512
+
+        elif pipeline_type == '1024':
+            # Direct 1024: coords at res 64 → shape_slat_flow_model_1024 (no cascade)
+            shape_slat = _sample_shape(t2p.models['shape_slat_flow_model_1024'], conds_1024, "Sampling shape SLat (1024 direct)")
+            tex_slat   = _sample_tex(t2p.models['tex_slat_flow_model_1024'], conds_1024, shape_slat)
+            res = 1024
+
+        else:  # '1024_cascade' or '1536_cascade' — cascade: LR (512) → upsample → HR (1024)
+            target_res = {'1024_cascade': 1024, '1536_cascade': 1536}[pipeline_type]
+
+            # LR stage: multi-image sample at 512 resolution
+            fm_lr = t2p.models['shape_slat_flow_model_512']
+            noise_lr = SparseTensor(
+                feats=torch.randn(coords.shape[0], fm_lr.in_channels).to(t2p.device),
+                coords=coords,
+            )
+            if t2p.low_vram:
+                fm_lr.to(t2p.device)
+            slat_lr = t2p._multi_image_sample(
+                t2p.shape_slat_sampler, fm_lr, noise_lr,
+                conds_512, t2p.shape_slat_sampler_params, shape_slat_sampler_params,
+                strategy=strategy, verbose=True, tqdm_desc="Sampling shape SLat (LR)",
+            )
+            if t2p.low_vram:
+                fm_lr.cpu()
+            slat_lr = slat_lr * std_shape.to(slat_lr.device) + mean_shape.to(slat_lr.device)
+
+            # Upsample LR slat to HR coords
+            if t2p.low_vram:
+                t2p.models['shape_slat_decoder'].to(t2p.device)
+                t2p.models['shape_slat_decoder'].low_vram = True
+            hr_coords_raw = t2p.models['shape_slat_decoder'].upsample(slat_lr, upsample_times=4)
+            if t2p.low_vram:
+                t2p.models['shape_slat_decoder'].cpu()
+                t2p.models['shape_slat_decoder'].low_vram = False
+
+            hr_resolution = target_res
+            while True:
+                quant_coords = torch.cat([
+                    hr_coords_raw[:, :1],
+                    ((hr_coords_raw[:, 1:] + 0.5) / 512 * (hr_resolution // 16)).int(),
+                ], dim=1)
+                hr_c = quant_coords.unique(dim=0)
+                if hr_c.shape[0] < max_num_tokens or hr_resolution == 1024:
+                    if hr_resolution != target_res:
+                        print(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
+                    break
+                hr_resolution -= 128
+
+            # HR stage: multi-image sample at hr_resolution using conds_1024
+            fm_hr = t2p.models['shape_slat_flow_model_1024']
+            noise_hr = SparseTensor(
+                feats=torch.randn(hr_c.shape[0], fm_hr.in_channels).to(t2p.device),
+                coords=hr_c,
+            )
+            if t2p.low_vram:
+                fm_hr.to(t2p.device)
+            shape_slat = t2p._multi_image_sample(
+                t2p.shape_slat_sampler, fm_hr, noise_hr,
+                conds_1024, t2p.shape_slat_sampler_params, shape_slat_sampler_params,
+                strategy=strategy, verbose=True, tqdm_desc="Sampling shape SLat (HR)",
+            )
+            if t2p.low_vram:
+                fm_hr.cpu()
+            shape_slat = shape_slat * std_shape.to(shape_slat.device) + mean_shape.to(shape_slat.device)
+            res = hr_resolution
+
+            tex_slat = _sample_tex(t2p.models['tex_slat_flow_model_1024'], conds_1024, shape_slat)
+
+        # ── Stage 4: Decode ───────────────────────────────────────────────────
+        torch.cuda.empty_cache()
+        out_mesh = t2p.decode_latent(shape_slat, tex_slat, res)
+
+        if return_latent:
+            return out_mesh, (shape_slat, tex_slat, res)
+        return out_mesh

trellis/pipelines/trellis_image_to_3d.py

@@ -20,8 +20,7 @@ from scipy.spatial.transform import Rotation
 from transformers import AutoModelForImageSegmentation
 import rembg
 # for app_refine.py, please uncomment these lines
-# from dreamsim import dreamsim 
-# from tqdm import tqdm
+from tqdm import tqdm
 
 def export_point_cloud(xyz, color):
     # Convert tensors to numpy arrays if needed
@@ -256,7 +255,7 @@ class TrellisImageTo3DPipeline(Pipeline):
         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.1)
+        size = int(size * 1)
         height, width = alpha.shape
         if not recenter:
             center = [width / 2, height / 2]
@@ -340,8 +339,8 @@ class TrellisImageTo3DPipeline(Pipeline):
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
         ])
         
-        input_images = transform_image(image).unsqueeze(0).to(self.device)
-        
+        input_images = transform_image(image).unsqueeze(0).to(self.device, dtype=torch.float16)
+
         with torch.no_grad():
             preds = self.birefnet_model(input_images)[-1].sigmoid().cpu()
         
@@ -616,68 +615,6 @@ class TrellisImageTo3DPipeline(Pipeline):
         slat = slat * std + mean
         return slat
 
-    def sample_slat_opt(
-        self,
-        apperance_learning_rate,
-        start_t,
-        input_images: torch.Tensor,
-        extrinsics: torch.Tensor,
-        intrinsics: torch.Tensor,
-        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']
-        slat_decoder_gs = self.models['slat_decoder_gs']
-        slat_decoder_mesh = self.models['slat_decoder_mesh']
-        noise = sp.SparseTensor(
-            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
-            coords=coords,
-        )
-        std = torch.tensor(self.slat_normalization['std'])[None].to(self.device)
-        mean = torch.tensor(self.slat_normalization['mean'])[None].to(self.device)
-        sampler_params = {**self.slat_sampler_params, **sampler_params}
-        slat = self.slat_sampler.sample_slat_opt_delta_v(
-            flow_model,
-            slat_decoder_gs,
-            slat_decoder_mesh,
-            std,
-            mean,
-            self.dreamsim_model,
-            apperance_learning_rate,
-            start_t,
-            input_images,
-            extrinsics, 
-            intrinsics,
-            noise,
-            **cond,
-            **sampler_params,
-            verbose=True
-        ).samples
-
-        slat = slat * std + mean
-        # from trellis.utils import render_utils, postprocessing_utils
-        # import imageio
-        # std = torch.tensor(self.slat_normalization['std'])[None].to(noise.device)
-        # mean = torch.tensor(self.slat_normalization['mean'])[None].to(noise.device)
-        # for i in range(sampler_params['steps']):
-        #     latent = slat.pred_x_0[i] * std + mean
-        #     outputs = self.decode_slat(latent, ["mesh", "gaussian"])
-        #     video_geo = render_utils.render_video(outputs['mesh'][0], resolution=512, pitch=0, inverse_direction=True, num_frames=120)['normal']
-        #     video_color = render_utils.render_video(outputs['gaussian'][0], resolution=512, pitch=0, inverse_direction=True, num_frames=120)['color']
-        #     video = [np.concatenate([video_color[i], video_geo[i]], axis=1) for i in range(len(video_color))]
-        #     imageio.mimsave('outputs/slat_iter_{i:02d}.mp4'.format(i=i), video, fps=15)
-        return slat
-
     def get_input(self, batch_data):
         std = torch.tensor(self.slat_normalization['std'])[None].to(self.device)
         mean = torch.tensor(self.slat_normalization['mean'])[None].to(self.device)
@@ -936,13 +873,16 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
     ):
 
         torch.manual_seed(seed)
-        aggregated_tokens_list, _ = self.vggt_feat(image)
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(dtype=self.VGGT_dtype):
+                aggregated_tokens_list, _ = self.vggt_feat(image)
         b, n, _, _ = aggregated_tokens_list[0].shape
         image_cond = self.encode_image(image).reshape(b, n, -1, 1024)
         
         # if coords is None:
         ss_flow_model = self.models['sparse_structure_flow_model']
-        ss_cond = self.get_ss_cond(image_cond[:, :, 5:], aggregated_tokens_list, num_samples)
+        with torch.no_grad():
+            ss_cond = self.get_ss_cond(image_cond[:, :, 5:], aggregated_tokens_list, num_samples)
         # Sample structured latent
         ss_sampler_params = {**self.sparse_structure_sampler_params, **sparse_structure_sampler_params}
         reso = ss_flow_model.resolution
@@ -965,69 +905,11 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
         # slat_steps = {**self.slat_sampler_params, **slat_sampler_params}.get('steps')
         # with self.inject_sampler_multi_image('slat_sampler', len(image), slat_steps, mode=mode):
         #     slat = self.sample_slat(cond, coords, slat_sampler_params)
-
-        slat_cond = self.get_slat_cond(image_cond, aggregated_tokens_list, num_samples)
+        with torch.no_grad():
+            slat_cond = self.get_slat_cond(image_cond, aggregated_tokens_list, num_samples)
         slat = self.sample_slat(slat_cond, coords, slat_sampler_params)
         return self.decode_slat(slat, formats), coords, ss_noise
 
-    def run_refine(
-        self,
-        image: Union[torch.Tensor, list[Image.Image]],
-        ss_learning_rate: float,
-        ss_start_t: float,
-        apperance_learning_rate: float,
-        apperance_start_t: float,
-        extrinsics: torch.Tensor,
-        intrinsics: torch.Tensor,
-        ss_noise: torch.Tensor,
-        input_points: torch.Tensor,
-        ss_refine_type: str = 'No',
-        coords: torch.Tensor = None,
-        num_samples: int = 1,
-        seed: int = 42,
-        sparse_structure_sampler_params: dict = {},
-        slat_sampler_params: dict = {},
-        formats: List[str] = ['mesh'],
-        mode: Literal['stochastic', 'multidiffusion'] = 'stochastic',
-    ):
-
-        torch.manual_seed(seed)
-        aggregated_tokens_list, input_images = self.vggt_feat(image)
-        b, n, _, _ = aggregated_tokens_list[0].shape
-        image_cond = self.encode_image(image).reshape(b, n, -1, 1024)
-        
-        if coords is None:
-            ss_cond = self.get_ss_cond(image_cond[:, :, 5:], aggregated_tokens_list, num_samples)
-            ss = torch.zeros(64, 64, 64, dtype=torch.long, device=image_cond.device)
-            ss = ss.index_put_((input_points[:,0], input_points[:,1], input_points[:,2]), torch.tensor(1, dtype=ss.dtype, device=ss.device))
-            ss = ss[None, None]
-            torch.cuda.empty_cache()
-            # Sample structured latent
-            if ss_refine_type == 'noise':
-                coords = self.sample_sparse_structure_opt_noise(ss_cond, ss, ss_learning_rate, num_samples, sparse_structure_sampler_params, ss_noise)
-            elif ss_refine_type == 'deltav':
-                coords = self.sample_sparse_structure_opt(ss_cond, ss, ss_learning_rate, ss_start_t, num_samples, sparse_structure_sampler_params, ss_noise)
-            torch.cuda.empty_cache()
-
-        # pcd = o3d.geometry.PointCloud()
-        # pcd.points = o3d.utility.Vector3dVector(coords[:,1:].cpu().numpy() / 64 - 0.5)
-        # o3d.io.write_point_cloud('outputs/after_coords.ply', pcd)
-
-        # cond = {
-        #     'cond': image_cond.reshape(n, -1, 1024),
-        #     'neg_cond': torch.zeros_like(image_cond.reshape(n, -1, 1024))[:1],
-        # }
-        
-        # slat_steps = {**self.slat_sampler_params, **slat_sampler_params}.get('steps')
-
-        # with self.inject_sampler_multi_image('slat_sampler', len(image), slat_steps, mode=mode):
-        #     # slat = self.sample_slat(cond, coords, slat_sampler_params)
-        #     slat = self.sample_slat_opt(apperance_learning_rate, apperance_start_t, input_images, extrinsics, intrinsics, cond, coords, slat_sampler_params)
-        
-        slat_cond = self.get_slat_cond(image_cond, aggregated_tokens_list, num_samples)
-        slat = self.sample_slat_opt(apperance_learning_rate, apperance_start_t, input_images, extrinsics, intrinsics, slat_cond, coords, slat_sampler_params)
-        return self.decode_slat(slat, formats)
-
     @staticmethod
     def from_pretrained(path: str) -> "TrellisVGGTTo3DPipeline":
         """
@@ -1040,7 +922,7 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
         new_pipeline = TrellisVGGTTo3DPipeline()
         new_pipeline.__dict__ = pipeline.__dict__
         args = pipeline._pretrained_args
-        new_pipeline.VGGT_dtype = torch.float32
+        new_pipeline.VGGT_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
         VGGT_model = VGGT.from_pretrained("Stable-X/vggt-object-v0-1")
         new_pipeline.VGGT_model = VGGT_model.to(new_pipeline.device)
         del new_pipeline.VGGT_model.depth_head
@@ -1065,9 +947,4 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
 
         new_pipeline._init_image_cond_model(args['image_cond_model'])
 
-        # for app_refine.py, please uncomment these lines
-        # model, _ = dreamsim(pretrained=True, device=new_pipeline.device, dreamsim_type="dino_vitb16", cache_dir="weights/dreamsim")
-        # new_pipeline.dreamsim_model = model
-        # new_pipeline.dreamsim_model.eval()
-
         return new_pipeline