Update trellis/pipelines/trellis_image_to_3d.py

trellis/pipelines/trellis_image_to_3d.py

@@ -231,14 +231,17 @@ class TrellisImageTo3DPipeline(Pipeline):
             if scale < 1:
                 input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
                         
-            # Get mask using BiRefNet
-            mask = self._get_birefnet_mask(input)
+            # # Get mask using BiRefNet
+            # mask = self._get_birefnet_mask(input)
             
-            # Convert input to RGBA and apply mask
-            input_rgba = input.convert('RGBA')
-            input_array = np.array(input_rgba)
-            input_array[:, :, 3] = mask * 255  # Apply mask to alpha channel
-            output = Image.fromarray(input_array)
+            # # Convert input to RGBA and apply mask
+            # input_rgba = input.convert('RGBA')
+            # input_array = np.array(input_rgba)
+            # input_array[:, :, 3] = mask * 255  # Apply mask to alpha channel
+            # output = Image.fromarray(input_array)
+            if getattr(self, 'rembg_session', None) is None:
+                self.rembg_session = rembg.new_session('u2net')
+            output = rembg.remove(input, session=self.rembg_session)
 
         # Process the output image
         output_np = np.array(output)
@@ -789,11 +792,11 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
         del new_pipeline.VGGT_model.point_head
         new_pipeline.VGGT_model.eval()
 
-        new_pipeline.birefnet_model = AutoModelForImageSegmentation.from_pretrained(
-            'ZhengPeng7/BiRefNet',
-            trust_remote_code=True
-        ).cpu()
-        new_pipeline.birefnet_model.eval()
+        # new_pipeline.birefnet_model = AutoModelForImageSegmentation.from_pretrained(
+        #     'ZhengPeng7/BiRefNet',
+        #     trust_remote_code=True
+        # ).cpu()
+        # new_pipeline.birefnet_model.eval()
         
         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']