Fixes for ZeroGPU

gradio_app.py

@@ -1,6 +1,5 @@
 import os
-
-os.system("pip install ./texture_baker/ ./uv_unwrapper/")
+import spaces
 
 import random
 import tempfile
@@ -18,48 +17,55 @@ from gradio_pointcloudeditor import PointCloudEditor
 from PIL import Image
 from transparent_background import Remover
 
-import spar3d.utils as spar3d_utils
-from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
-from spar3d.system import SPAR3D
 
-os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.environ.get("TMPDIR", "/tmp"), "gradio")
+@spaces.GPU
+def setup():
+    os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.environ.get("TMPDIR", "/tmp"), "gradio")
+    
+    os.system("USE_CUDA=1 pip install -vv ./texture_baker ./uv_unwrapper")
 
-bg_remover = Remover()  # default setting
+    bg_remover = Remover()  # default setting
 
-COND_WIDTH = 512
-COND_HEIGHT = 512
-COND_DISTANCE = 2.2
-COND_FOVY = 0.591627
-BACKGROUND_COLOR = [0.5, 0.5, 0.5]
+    COND_WIDTH = 512
+    COND_HEIGHT = 512
+    COND_DISTANCE = 2.2
+    COND_FOVY = 0.591627
+    BACKGROUND_COLOR = [0.5, 0.5, 0.5]
 
-# Cached. Doesn't change
-c2w_cond = spar3d_utils.default_cond_c2w(COND_DISTANCE)
-intrinsic, intrinsic_normed_cond = spar3d_utils.create_intrinsic_from_fov_rad(
-    COND_FOVY, COND_HEIGHT, COND_WIDTH
-)
+    # Cached. Doesn't change
+    c2w_cond = spar3d_utils.default_cond_c2w(COND_DISTANCE)
+    intrinsic, intrinsic_normed_cond = spar3d_utils.create_intrinsic_from_fov_rad(
+        COND_FOVY, COND_HEIGHT, COND_WIDTH
+    )
 
-generated_files = []
+    generated_files = []
 
-# Delete previous gradio temp dir folder
-if os.path.exists(os.environ["GRADIO_TEMP_DIR"]):
-    print(f"Deleting {os.environ['GRADIO_TEMP_DIR']}")
-    import shutil
+    # Delete previous gradio temp dir folder
+    if os.path.exists(os.environ["GRADIO_TEMP_DIR"]):
+        print(f"Deleting {os.environ['GRADIO_TEMP_DIR']}")
+        import shutil
 
-    shutil.rmtree(os.environ["GRADIO_TEMP_DIR"])
+        shutil.rmtree(os.environ["GRADIO_TEMP_DIR"])
 
-device = spar3d_utils.get_device()
+    device = spar3d_utils.get_device()
 
-model = SPAR3D.from_pretrained(
-    "stabilityai/stable-point-aware-3d",
-    config_name="config.yaml",
-    weight_name="model.safetensors",
-)
-model.eval()
-model = model.to(device)
+    model = SPAR3D.from_pretrained(
+        "stabilityai/stable-point-aware-3d",
+        config_name="config.yaml",
+        weight_name="model.safetensors",
+    )
+    model.eval()
+    model = model.to(device)
 
-example_files = [
-    os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
-]
+    example_files = [
+        os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
+    ]
+
+setup()
+
+import spar3d.utils as spar3d_utils
+from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
+from spar3d.system import SPAR3D
 
 
 def forward_model(
@@ -332,6 +338,7 @@ def regenerate_run(
     )
 
 
+@spaces.GPU
 def run_button(
     run_btn,
     input_image,
@@ -348,6 +355,7 @@ def run_button(
     texture_resolution,
 ):
     if run_btn == "Run":
+
         if torch.cuda.is_available():
             torch.cuda.reset_peak_memory_stats()
 

texture_baker/setup.py

@@ -52,6 +52,9 @@ def get_extensions():
         extra_compile_args["nvcc"].append("-g")
         extra_link_args.extend(["-O0", "-g"])
 
+    if use_cuda:
+        extra_compile_args["nvcc"].append("--gpu-architecture=sm_80")
+
     define_macros = []
     extensions = []
     libraries = []