Enhance install_croco function in app.py by adding checks for source file existence, improved CUDA architecture handling, and a cleanup process for the build directory. Implement an alternative compilation method using setup.py if the primary method fails, ensuring robust error handling and better user feedback during the installation of the curope CUDA extension.

app.py

@@ -13,46 +13,115 @@ def install_croco():
     """
     try:
         import os
+        import shutil
         from torch.utils.cpp_extension import load
         
         # Path to the curope source files
         curope_path = "./extern/CUT3R/src/croco/models/curope"
         
+        # Check if source files exist
+        cpp_file = os.path.join(curope_path, "curope.cpp")
+        cu_file = os.path.join(curope_path, "kernels.cu")
+        
+        if not os.path.exists(cpp_file) or not os.path.exists(cu_file):
+            print(f"Source files not found: {cpp_file}, {cu_file}")
+            return False
+        
+        print(f"Found source files: {cpp_file}, {cu_file}")
+        
         # Define source files
-        sources = [
-            os.path.join(curope_path, "curope.cpp"),
-            os.path.join(curope_path, "kernels.cu"),
-        ]
+        sources = [cpp_file, cu_file]
         
-        # Get CUDA architectures
+        # Get CUDA architectures - be more conservative
         if torch.cuda.is_available():
-            all_cuda_archs = torch.cuda.get_gencode_flags().replace("compute=", "arch=").split()
+            try:
+                # Get the current GPU architecture
+                capability = torch.cuda.get_device_capability()
+                arch_code = f"compute_{capability[0]}{capability[1]}"
+                sm_code = f"sm_{capability[0]}{capability[1]}"
+                all_cuda_archs = ['-gencode', f'arch={arch_code},code={sm_code}']
+                print(f"Using CUDA architecture: {arch_code}")
+            except:
+                # Fallback to common architectures
+                all_cuda_archs = [
+                    '-gencode', 'arch=compute_70,code=sm_70',
+                    '-gencode', 'arch=compute_80,code=sm_80',
+                ]
         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'
-            ]
+            print("CUDA not available!")
+            return False
+        
+        # Create build directory in current working directory
+        build_dir = os.path.abspath("./curope_build_temp")
+        if os.path.exists(build_dir):
+            shutil.rmtree(build_dir)
+        os.makedirs(build_dir, exist_ok=True)
+        print(f"Using build directory: {build_dir}")
         
         # Compile the extension
+        print("Starting CUDA compilation...")
         curope = load(
             name="curope",
             sources=sources,
             extra_cflags=["-O3"],
-            extra_cuda_cflags=["-O3", "--ptxas-options=-v", "--use_fast_math"] + all_cuda_archs,
+            extra_cuda_cflags=["-O3", "--use_fast_math"] + all_cuda_archs,
             verbose=True,
-            build_directory="./curope_build"  # Specify build directory
+            build_directory=build_dir,
+            with_cuda=True
         )
         
         print("Successfully compiled curope CUDA extension")
+        
+        # Clean up build directory
+        try:
+            shutil.rmtree(build_dir)
+        except:
+            pass
+            
         return True
         
     except Exception as e:
-        print(f"Failed to compile curope CUDA extension: {e}")
-        print("Will fall back to PyTorch implementation")
-        return False
+        print(f"Method 1 failed: {e}")
+        print("Trying alternative compilation method...")
+        
+        # Alternative method: use setup.py build_ext --inplace
+        try:
+            import subprocess
+            import sys
+            import os
+            
+            # Change to curope directory
+            original_cwd = os.getcwd()
+            curope_path = "./extern/CUT3R/src/croco/models/curope"
+            os.chdir(curope_path)
+            
+            # Set CUDA environment variables to ensure they're available
+            env = os.environ.copy()
+            if torch.cuda.is_available():
+                cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+                if cuda_home:
+                    env['CUDA_HOME'] = cuda_home
+                    env['CUDA_PATH'] = cuda_home
+            
+            # Run setup.py build_ext --inplace within the GPU context
+            result = subprocess.run([
+                sys.executable, "setup.py", "build_ext", "--inplace"
+            ], env=env, capture_output=True, text=True)
+            
+            os.chdir(original_cwd)
+            
+            if result.returncode == 0:
+                print("Successfully compiled curope using setup.py")
+                return True
+            else:
+                print(f"Setup.py compilation failed: {result.stderr}")
+                return False
+                
+        except Exception as e2:
+            print(f"Alternative method also failed: {e2}")
+            import traceback
+            traceback.print_exc()
+            return False
 
 install_croco()