Update app.py

app.py

@@ -1,54 +1,58 @@
-import io
 import os
-import shlex
-import subprocess
-import tempfile
-import zipfile
-from functools import partial
 import sys
-
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException
-from fastapi.responses import StreamingResponse
-from fastapi.middleware.cors import CORSMiddleware
-
-import numpy as np
-import rembg
-import torch
-from PIL import Image
-
-from tsr.system import TSR
-from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
+import subprocess
 
 # ----------------------------------------------------------------
 # 0. Bootstrap install of torchmcubes wheel if needed
 # ----------------------------------------------------------------
 
-# Compute the path to this file and to your wheels folder
+# 1) Locate the wheel relative to this file
 BASE_DIR   = os.path.dirname(os.path.abspath(__file__))
 WHEELS_DIR = os.path.join(BASE_DIR, "wheels")
 WHL_NAME   = "torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl"
 WHL_PATH   = os.path.join(WHEELS_DIR, WHL_NAME)
 
-def install_torchmcubes():
-    if not os.path.exists(WHL_PATH):
-        raise FileNotFoundError(f"Cannot find wheel at {WHL_PATH!r}")
-    print(f"[bootstrap] Installing torchmcubes from {WHL_PATH!r} …")
+def install_torchmcubes_wheel(whl_path: str):
+    """
+    Installs the given wheel via pip. Exits on failure.
+    """
+    if not os.path.exists(whl_path):
+        raise FileNotFoundError(f"Cannot find torchmcubes wheel at {whl_path!r}")
+    print(f"[bootstrap] Installing torchmcubes from {whl_path!r} …")
     try:
         subprocess.check_call([
             sys.executable, "-m", "pip", "install",
-            "--upgrade", WHL_PATH
+            "--upgrade", whl_path
         ])
     except subprocess.CalledProcessError as e:
         print(f"[bootstrap][ERROR] pip install failed: {e}")
         sys.exit(1)
 
-# Try to import, and if missing, install then re-import
+# 2) Try import, install if missing
 try:
     import torchmcubes
 except ImportError:
-    install_torchmcubes()
+    install_torchmcubes_wheel(WHL_PATH)
     import torchmcubes
 
+# ----------------------------------------------------------------
+# 1. Now safe to import the rest of your dependencies
+# ----------------------------------------------------------------
+import io
+import zipfile
+import tempfile
+import numpy as np
+import rembg
+import torch
+from PIL import Image
+
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+
+from tsr.system import TSR
+from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
+
 # ------------------------------------------------------------
 # 1. Model & utils initialization (runs at startup)
 # ------------------------------------------------------------