Fix black texture bug: detect degenerate sampler trajectories and retry

trellis2/pipelines/trellis2_image_to_3d.py

@@ -392,20 +392,62 @@ class Trellis2ImageTo3DPipeline(Pipeline):
             self.models['shape_slat_decoder'].low_vram = False
         return ret
     
+    def _detect_degenerate_tex(
+        self,
+        flow_model,
+        noise: SparseTensor,
+        concat_cond: SparseTensor,
+        cond: dict,
+        sampler_params: dict,
+        threshold: float = 1.037,
+    ) -> bool:
+        """
+        Run 2 probe steps to detect degenerate texture flow trajectory.
+
+        Returns True if the trajectory is degenerate (will produce black texture).
+        """
+        steps = sampler_params.get('steps', 12)
+        rescale_t = sampler_params.get('rescale_t', 3.0)
+        t_seq = np.linspace(1, 0, steps + 1)
+        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
+        t_seq = t_seq.tolist()
+
+        sample = noise
+        for step_idx in range(2):
+            t, t_prev = t_seq[step_idx], t_seq[step_idx + 1]
+            out = self.tex_slat_sampler.sample_once(
+                flow_model, sample, t, t_prev,
+                concat_cond=concat_cond,
+                **{k: v for k, v in cond.items()},
+                **{k: v for k, v in sampler_params.items() if k not in ('steps', 'rescale_t')},
+            )
+            sample = out.pred_x_prev
+
+        x0_std = out.pred_x_0.feats.std().item()
+        return x0_std > threshold
+
     def sample_tex_slat(
         self,
         cond: dict,
         flow_model,
         shape_slat: SparseTensor,
         sampler_params: dict = {},
+        max_retries: int = 3,
+        retry_noise_scale: float = 0.5,
     ) -> SparseTensor:
         """
         Sample structured latent with the given conditioning.
-        
+
+        Includes early detection of degenerate texture flow trajectories
+        (black texture bug). If detected after 2 probe steps, retries with
+        scaled noise.
+
         Args:
             cond (dict): The conditioning information.
             shape_slat (SparseTensor): The structured latent for shape
             sampler_params (dict): Additional parameters for the sampler.
+            max_retries (int): Max retries on degenerate detection.
+            retry_noise_scale (float): Noise scale factor on retry.
         """
         # Sample structured latent
         std = torch.tensor(self.shape_slat_normalization['std'])[None].to(shape_slat.device)
@@ -413,10 +455,24 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         shape_slat = (shape_slat - mean) / std
 
         in_channels = flow_model.in_channels if isinstance(flow_model, nn.Module) else flow_model[0].in_channels
-        noise = shape_slat.replace(feats=torch.randn(shape_slat.coords.shape[0], in_channels - shape_slat.feats.shape[1]).to(self.device))
+        n_noise_feats = in_channels - shape_slat.feats.shape[1]
         sampler_params = {**self.tex_slat_sampler_params, **sampler_params}
         if self.low_vram:
             flow_model.to(self.device)
+
+        noise_feats = torch.randn(shape_slat.coords.shape[0], n_noise_feats).to(self.device)
+        for attempt in range(max_retries + 1):
+            noise = shape_slat.replace(feats=noise_feats)
+
+            if self._detect_degenerate_tex(flow_model, noise, shape_slat, cond, sampler_params):
+                if attempt < max_retries:
+                    noise_feats = torch.randn(shape_slat.coords.shape[0], n_noise_feats).to(self.device) * retry_noise_scale
+                    print(f"\033[93m[tex] Degenerate detected, retry {attempt+1}/{max_retries} with noise_scale={retry_noise_scale}\033[0m")
+                    continue
+                else:
+                    print(f"\033[93m[tex] Degenerate detected but retries exhausted, proceeding anyway\033[0m")
+            break
+
         slat = self.tex_slat_sampler.sample(
             flow_model,
             noise,
@@ -432,7 +488,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         std = torch.tensor(self.tex_slat_normalization['std'])[None].to(slat.device)
         mean = torch.tensor(self.tex_slat_normalization['mean'])[None].to(slat.device)
         slat = slat * std + mean
-        
+
         return slat
 
     def decode_tex_slat(