Trellis loading Bug

README.md

@@ -5,7 +5,7 @@ colorFrom: yellow
 colorTo: purple
 sdk: gradio
 python_version: 3.10.13
-sdk_version: 5.21.0
+sdk_version: 5.22.0
 app_file: app.py
 pinned: true
 short_description: Transform Your Images into Mesmerizing Hexagon Grids

app.py

@@ -754,6 +754,7 @@ def replace_input_with_sketch_image(sketch_image):
 
 @spaces.GPU(progress=gr.Progress(track_tqdm=True))
 def load_trellis_model():
+    gr.Info("TRELLIS_PIPELINE load start", 60)
     global TRELLIS_PIPELINE
     loaded = False
     if TRELLIS_PIPELINE == None:
@@ -772,7 +773,9 @@ def load_trellis_model():
             print(f"Error preloading TRELLIS_PIPELINE: {e}")
             gr.Error(f"Failed to load TRELLIS_PIPELINE: {e}")
             TRELLIS_PIPELINE = None
-    
+    else:
+        loaded = True
+        print("TRELLIS_PIPELINE already loaded\n")
 
 def load_3d_models(is_open: bool = True) -> bool:
     if is_open:
@@ -781,21 +784,20 @@ def load_3d_models(is_open: bool = True) -> bool:
         image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
         depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", ignore_mismatched_sizes=True)
         print("DPT models loaded\n")
-        # is_open = load_trellis_model()
-        # if is_open:
+        if not constants.IS_SHARED_SPACE:
+            load_trellis_model()
         print("3D models loaded")
         gr.Info("3D models loaded.")
-        #else:
-        #     gr.Error("Failed to load TRELLIS_PIPELINE.")
     return gr.update(interactive = is_open)
 
 def unload_3d_models(is_open: bool = False) -> bool:
     if not is_open:
         gr.Info("Unloading 3D models...")
         global image_processor, depth_model, TRELLIS_PIPELINE
-        if TRELLIS_PIPELINE:
-            TRELLIS_PIPELINE.to("cpu")
-            TRELLIS_PIPELINE = None
+        if not constants.IS_SHARED_SPACE:
+            if TRELLIS_PIPELINE:
+                TRELLIS_PIPELINE.cpu()
+                TRELLIS_PIPELINE = None
         if depth_model:
             del image_processor
             del depth_model
@@ -946,74 +948,79 @@ 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")
-    # 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)
-        outputs = TRELLIS_PIPELINE.run(
-            processed_image,
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": steps,
-                "cfg_strength": 7.5,
-            },
-            slat_sampler_params={
-                "steps": steps,
-                "cfg_strength": 3.0,
-            },
-        )
+    if TRELLIS_PIPELINE is None:
+        gr.Warning(f"Trellis Pipeline is not initialized: {TRELLIS_PIPELINE.device()}")
+        return [None, None, depth_img]
+    else:
+        # 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)
+            outputs = TRELLIS_PIPELINE.run(
+                processed_image,
+                seed=seed,
+                formats=["gaussian", "mesh"],
+                preprocess_image=False,
+                sparse_structure_sampler_params={
+                    "steps": steps,
+                    "cfg_strength": 7.5,
+                },
+                slat_sampler_params={
+                    "steps": steps,
+                    "cfg_strength": 3.0,
+                },
+            )
 
-        # Validate the mesh
-        mesh = outputs['mesh'][0]
-        meshisdict = isinstance(mesh, dict)
-        if meshisdict:
-            vertices = mesh['vertices']
-            faces = mesh['faces']
-        else:
-            vertices = mesh.vertices
-            faces = mesh.faces
+            # Validate the mesh
+            mesh = outputs['mesh'][0]
+            meshisdict = isinstance(mesh, dict)
+            if meshisdict:
+                vertices = mesh['vertices']
+                faces = mesh['faces']
+            else:
+                vertices = mesh.vertices
+                faces = mesh.faces
 
-        print(f"Mesh vertices: {vertices.shape}, faces: {faces.shape}")
-        if faces.max() >= vertices.shape[0]:
-            raise ValueError(f"Invalid mesh: face index {faces.max()} exceeds vertex count {vertices.shape[0]}")
-    except Exception as e:
-        gr.Warning(f"Error generating 3D asset: {e}")
-        print(f"Error generating 3D asset: {e}")
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-        return None,None, depth_img
-
-    # Ensure data is on GPU and has correct type
-    if not vertices.is_cuda or not faces.is_cuda:
-        raise ValueError("Mesh data must be on GPU")
-    if vertices.dtype != torch.float32 or faces.dtype != torch.int32:
-        if meshisdict:
-            mesh['faces'] = faces.to(torch.int32)
-            mesh['vertices'] = vertices.to(torch.float32)
-        else:
-            mesh.faces = faces.to(torch.int32)
-            mesh.vertices = vertices.to(torch.float32)
+            print(f"Mesh vertices: {vertices.shape}, faces: {faces.shape}")
+            if faces.max() >= vertices.shape[0]:
+                raise ValueError(f"Invalid mesh: face index {faces.max()} exceeds vertex count {vertices.shape[0]}")
+        except Exception as e:
+            gr.Warning(f"Error generating 3D asset: {e}")
+            print(f"Error generating 3D asset: {e}")
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+            return None,None, depth_img
+
+        # Ensure data is on GPU and has correct type
+        if not vertices.is_cuda or not faces.is_cuda:
+            raise ValueError("Mesh data must be on GPU")
+        if vertices.dtype != torch.float32 or faces.dtype != torch.int32:
+            if meshisdict:
+                mesh['faces'] = faces.to(torch.int32)
+                mesh['vertices'] = vertices.to(torch.float32)
+            else:
+                mesh.faces = faces.to(torch.int32)
+                mesh.vertices = vertices.to(torch.float32)
 
-    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
+        user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
+        os.makedirs(user_dir, exist_ok=True)
 
-    video = render_utils.render_video(outputs['gaussian'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['color']
-    try:
-        video_geo = render_utils.render_video(outputs['mesh'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['normal']
-        video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-    except Exception as e:
-        gr.Info(f"Error rendering video: {e}")
-        print(f"Error rendering video: {e}")
-    video_path = os.path.join(user_dir, f'{output_name}.mp4')
-    imageio.mimsave(video_path, video, fps=8)
+        video = render_utils.render_video(outputs['gaussian'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['color']
+        try:
+            video_geo = render_utils.render_video(outputs['mesh'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['normal']
+            video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+        except Exception as e:
+            gr.Info(f"Error rendering video: {e}")
+            print(f"Error rendering video: {e}")
+        video_path = os.path.join(user_dir, f'{output_name}.mp4')
+        imageio.mimsave(video_path, video, fps=8)
 
-    #snapshot_results = render_utils.render_snapshot_depth(outputs['mesh'][0], resolution=1280, r=1, fov=80)
-    #depth_snapshot = Image.fromarray(snapshot_results['normal'][0]).convert("L")
-    depth_snapshot = depth_img
+        #snapshot_results = render_utils.render_snapshot_depth(outputs['mesh'][0], resolution=1280, r=1, fov=80)
+        #depth_snapshot = Image.fromarray(snapshot_results['normal'][0]).convert("L")
+        depth_snapshot = depth_img
+
+        state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], output_name)
 
-    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], output_name)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
@@ -1595,9 +1602,9 @@ with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',
     #     outputs=[depth_map_output, model_output, model_file], scroll_to_output=True
     # )
     accordian_3d.expand(
-        fn=load_trellis_model,
-        trigger_mode="always_last"
-    ).then(
+    #     fn=load_trellis_model,
+    #     trigger_mode="always_last"
+    # ).then(
         fn=load_3d_models,
         trigger_mode="always_last",
         outputs=[generate_3d_asset_button],
@@ -1675,13 +1682,13 @@ if __name__ == "__main__":
 
     # image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
     # depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", ignore_mismatched_sizes=True)
-    # if constants.IS_SHARED_SPACE:
-    #     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
-    #     except:
-    #         pass
+    if constants.IS_SHARED_SPACE:
+        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
+        except:
+            pass
     hexaGrid.queue(default_concurrency_limit=1,max_size=12,api_open=False)
     hexaGrid.launch(allowed_paths=["assets","/","./assets","images","./images", "./images/prerendered", 'e:/TMP'], favicon_path="./assets/favicon.ico", max_file_size="10mb")