Upgrade Trellis to handle landscapes

app.py

@@ -962,7 +962,7 @@ def generate_3d_asset_part1(depth_image_source, randomize_seed, seed, input_imag
     final_seed = np.random.randint(0, constants.MAX_SEED) if randomize_seed else seed
     # Process the image for depth estimation
     depth_img = depth_process_image(image_path, resized_width=1536, z_scale=336)
-    depth_img = resize_image_with_aspect_ratio(depth_img, 1536, 1536)
+    #depth_img = resize_image_with_aspect_ratio(depth_img, 1536, 1536)
     
     user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
     depth_img = save_image_to_temp_png(depth_img, user_dir, f"{output_name}_depth")
@@ -975,6 +975,7 @@ def generate_3d_asset_part2(depth_img, image_path, output_name, seed, steps, mod
     image_raw = Image.open(image_path).convert("RGB")
     resized_image = resize_image_with_aspect_ratio(image_raw, model_resolution, model_resolution)
     depth_img = Image.open(depth_img).convert("RGBA")
+
     if TRELLIS_PIPELINE is None:
         gr.Warning(f"Trellis Pipeline is not initialized: {TRELLIS_PIPELINE.device()}")
         return [None, None, depth_img]
@@ -982,7 +983,7 @@ def generate_3d_asset_part2(depth_img, image_path, output_name, seed, steps, mod
         # Preprocess and run the Trellis pipeline with fixed sampler settings
         try:
             TRELLIS_PIPELINE.cuda()
-            processed_image = TRELLIS_PIPELINE.preprocess_image(resized_image, max_resolution=model_resolution)
+            processed_image = TRELLIS_PIPELINE.preprocess_image(resized_image, max_resolution=model_resolution, remove_bg = False)
             outputs = TRELLIS_PIPELINE.run(
                 processed_image,
                 seed=seed,
@@ -1503,7 +1504,7 @@ with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',
 
     is_multiimage = gr.State(False)
     output_buf = gr.State()
-    ddd_image_path = gr.State("./images/images/Beeuty-1.png")
+    ddd_image_path = gr.State("./images/images/Bee-test-2.png")
     ddd_file_name = gr.State("Hexagon_file")
     with gr.Row():
         gr.Examples(examples=[
@@ -1761,7 +1762,7 @@ if __name__ == "__main__":
         TRELLIS_PIPELINE = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
         TRELLIS_PIPELINE.to(device)
         try:
-           TRELLIS_PIPELINE.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
+           TRELLIS_PIPELINE.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)), 512)    # Preload rembg
         except:
             pass
     hexaGrid.queue(default_concurrency_limit=1,max_size=12,api_open=False)

trellis/pipelines/trellis_image_to_3d.py

@@ -82,7 +82,7 @@ class TrellisImageTo3DPipeline(Pipeline):
         ])
         self.image_cond_model_transform = transform
 
-    def preprocess_image(self, input: Image.Image, max_resolution: int =1024) -> Image.Image:
+    def preprocess_image(self, input: Image.Image, max_resolution: int =1024, remove_bg: bool = True) -> Image.Image:
         """
         Preprocess the input image.
         """
@@ -100,9 +100,12 @@ class TrellisImageTo3DPipeline(Pipeline):
             scale = min(1, max_resolution / 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)
+            if remove_bg:
+                if getattr(self, 'rembg_session', None) is None:
+                    self.rembg_session = rembg.new_session('u2net')
+                output = rembg.remove(input, session=self.rembg_session)
+            else:
+                output = input.convert('RGBA')
         output_np = np.array(output)
         alpha = output_np[:, :, 3]
         bbox = np.argwhere(alpha > 0.8 * 255)
@@ -264,6 +267,7 @@ class TrellisImageTo3DPipeline(Pipeline):
         slat_sampler_params: dict = {},
         formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
         preprocess_image: bool = True,
+        remove_bg: bool = True,
     ) -> dict:
         """
         Run the pipeline.
@@ -276,7 +280,7 @@ class TrellisImageTo3DPipeline(Pipeline):
             preprocess_image (bool): Whether to preprocess the image.
         """
         if preprocess_image:
-            image = self.preprocess_image(image)
+            image = self.preprocess_image(image, remove_bg=remove_bg)
         cond = self.get_cond([image])
         torch.manual_seed(seed)
         coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
@@ -351,6 +355,7 @@ class TrellisImageTo3DPipeline(Pipeline):
         formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
         preprocess_image: bool = True,
         mode: Literal['stochastic', 'multidiffusion'] = 'stochastic',
+        remove_bg: bool = True,
     ) -> dict:
         """
         Run the pipeline with multiple images as condition
@@ -363,7 +368,7 @@ class TrellisImageTo3DPipeline(Pipeline):
             preprocess_image (bool): Whether to preprocess the image.
         """
         if preprocess_image:
-            images = [self.preprocess_image(image) for image in images]
+            images = [self.preprocess_image(image,remove_bg=remove_bg) for image in images]
         cond = self.get_cond(images)
         cond['neg_cond'] = cond['neg_cond'][:1]
         torch.manual_seed(seed)