Refactor install_croco function in app.py to compile the curope CUDA extension using torch.utils.cpp_extension.load() instead of subprocess calls, improving compatibility with zerogpu. Added error handling and fallback to PyTorch implementation if compilation fails.

app.py

@@ -3,14 +3,56 @@ import sys
 import spaces
 import torch
 
-@spaces.GPU(duration=10)
+@spaces.GPU
 @torch.autocast("cuda")
 @torch.no_grad()
 def install_croco():
-    subprocess.check_call([
-        sys.executable, "-m", "pip", "install", "-e", 
-        "./extern/CUT3R/src/croco/models/curope"
-    ])
+    """
+    Compile curope CUDA extension using torch.utils.cpp_extension.load()
+    This approach works better with zerogpu than subprocess calls.
+    """
+    try:
+        import os
+        from torch.utils.cpp_extension import load
+        
+        # Path to the curope source files
+        curope_path = "./extern/CUT3R/src/croco/models/curope"
+        
+        # Define source files
+        sources = [
+            os.path.join(curope_path, "curope.cpp"),
+            os.path.join(curope_path, "kernels.cu"),
+        ]
+        
+        # Get CUDA architectures
+        if torch.cuda.is_available():
+            all_cuda_archs = torch.cuda.get_gencode_flags().replace("compute=", "arch=").split()
+        else:
+            # Fallback architectures if CUDA info is not available
+            all_cuda_archs = [
+                '-gencode', 'arch=compute_70,code=sm_70',
+                '-gencode', 'arch=compute_75,code=sm_75', 
+                '-gencode', 'arch=compute_80,code=sm_80',
+                '-gencode', 'arch=compute_86,code=sm_86'
+            ]
+        
+        # Compile the extension
+        curope = load(
+            name="curope",
+            sources=sources,
+            extra_cflags=["-O3"],
+            extra_cuda_cflags=["-O3", "--ptxas-options=-v", "--use_fast_math"] + all_cuda_archs,
+            verbose=True,
+            build_directory="./curope_build"  # Specify build directory
+        )
+        
+        print("Successfully compiled curope CUDA extension")
+        return True
+        
+    except Exception as e:
+        print(f"Failed to compile curope CUDA extension: {e}")
+        print("Will fall back to PyTorch implementation")
+        return False
 
 install_croco()