Update TripoSR/obj_gen.py

TripoSR/obj_gen.py

@@ -13,13 +13,12 @@ from tsr.system import TSR
 from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
 
 import argparse
-# from dotenv import load_dotenv
 
-# load_dotenv()
 
-
-
-device = "cpu"
+if torch.cuda.is_available():
+    device = "cuda:0"
+else:
+    device = "cpu"
 
 model = TSR.from_pretrained(
     "stabilityai/TripoSR",
@@ -73,11 +72,16 @@ def run_example(image_pil):
     return preprocessed, mesh_name_obj, mesh_name_glb
 
 def generate_obj_from_image(image_pil, path="output.obj"):
-    # Preprocess the image without removing the background and with a foreground ratio of 0.9
-    preprocessed = preprocess(image_pil, True, 0.9)
-    
-    # Generate the mesh and get the paths to the .obj and .glb files
-    mesh_paths = generate(preprocessed, 256, ["obj"], path)
+    try:
+        # Preprocess the image without removing the background and with a foreground ratio of 0.9
+        print("Preprocessing image")
+        preprocessed = preprocess(image_pil, True, 0.9)
+        print("Generating mesh")
+        # Generate the mesh and get the paths to the .obj and .glb files
+        mesh_paths = generate(preprocessed, 256, ["obj"], path)
+    except Exception as e:
+        print(f"Error generating mesh: {e}")
+        return None
     
     # Return the path to the .obj file
     return mesh_paths[0]