fix: add MultiScaleDeformableAttention CUDA extension and enhance CUDA toolkit setup

.gitattributes

@@ -95,3 +95,5 @@ pretrained_weights/insightface/models/buffalo_l/2d106det.onnx filter=lfs diff=lf
 pretrained_weights/insightface/models/buffalo_l/det_10g.onnx filter=lfs diff=lfs merge=lfs -text
 pretrained_weights/liveportrait/landmark.onnx filter=lfs diff=lfs merge=lfs -text
 pretrained_weights/liveportrait_animals/**/*.pth filter=lfs diff=lfs merge=lfs -text
+
+MultiScaleDeformableAttention.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text

app.py

@@ -12,6 +12,10 @@ from src.config.argument_config import ArgumentConfig
 from src.config.inference_config import InferenceConfig
 import spaces
 
+ROOT = pathlib.Path(__file__).resolve().parent
+OPS_DIR = ROOT / "src" / "utils" / "dependencies" / "XPose" / "models" / "UniPose" / "ops"
+sys.path.insert(0, str(OPS_DIR))
+
 CUDA_RUN_URL = (
     "https://developer.download.nvidia.com/compute/cuda/11.8.0/"
     "local_installers/cuda_11.8.0_520.61.05_linux.run"
@@ -22,27 +26,34 @@ TORCH_WHL_INDEX = "https://download.pytorch.org/whl/torch_stable.html"
 
 def ensure_cuda_toolkit():
     """Download & install the CUDA toolkit *silently* if it is not present."""
+    print("Checking for CUDA toolkit...")
     if pathlib.Path(f"{CUDA_HOME_PATH}/bin/nvcc").exists():
+        print(f"CUDA toolkit already installed at {CUDA_HOME_PATH}")
         return  # toolkit already installed
 
+    print(f"CUDA toolkit not found. Downloading from {CUDA_RUN_URL}...")
     run_file = f"/tmp/{pathlib.Path(CUDA_RUN_URL).name}"
     subprocess.run(["wget", "-q", CUDA_RUN_URL, "-O", run_file], check=True)
+    print(f"Download complete. Making installer executable...")
     subprocess.run(["chmod", "+x", run_file], check=True)
+    print("Installing CUDA toolkit (this may take a while)...")
     subprocess.run([run_file, "--silent", "--toolkit"], check=True)
+    print("CUDA toolkit installation complete.")
 
     # --- environment variables expected by CUDA extensions -------------------
+    print("Setting up CUDA environment variables...")
     os.environ["CUDA_HOME"] = CUDA_HOME_PATH
     os.environ["PATH"] = f"{CUDA_HOME_PATH}/bin:" + os.environ.get("PATH", "")
     os.environ["LD_LIBRARY_PATH"] = (
         f"{CUDA_HOME_PATH}/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")
     )
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+    print("CUDA environment setup complete.")
 
 
 def build_xpose_ops():
     """Build the MultiScaleDeformableAttention CUDA extension with enhanced error handling."""
     try:
-        # Check if already built to avoid rebuilding
         import MultiScaleDeformableAttention
         print("MultiScaleDeformableAttention already installed")
         return True
@@ -53,21 +64,17 @@ def build_xpose_ops():
             current_dir, "src/utils/dependencies/XPose/models/UniPose/ops")
         try:
             os.chdir(ops_dir)
-            # Set environment variables to help with compilation
             os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
-            os.environ["MAX_JOBS"] = "1"  # Limit to avoid memory issues
 
             try:
-                # Try building with fewer optimizations
                 subprocess.run(
-                    [sys.executable, "setup.py", "build", "install", "--user"],
+                    [sys.executable, "setup.py", "build", "install"],
                     check=True,
                     env={**os.environ, "CFLAGS": "-O0", "CXXFLAGS": "-O0"}
                 )
                 print("MultiScaleDeformableAttention built successfully")
 
                 try:
-                    # Verify it was properly built
                     import MultiScaleDeformableAttention
                     built_success = True
                 except ImportError:
@@ -77,17 +84,14 @@ def build_xpose_ops():
             except subprocess.CalledProcessError as e:
                 print(f"Build error: {e}")
 
-                # If the standard build fails, try with minimal compilation
                 try:
                     print("Attempting simplified build...")
                     subprocess.run(
-                        [sys.executable, "setup.py", "build", "install", "--user",
-                         "--global-option=--disable-cuda"],
+                        [sys.executable, "setup.py", "build", "install"],
                         check=True
                     )
                     print("Simplified build completed")
 
-                    # Verify import works
                     try:
                         import MultiScaleDeformableAttention
                         print("MultiScaleDeformableAttention imported after simplified build")

src/utils/dependencies/XPose/models/UniPose/ops/MultiScaleDeformableAttention.cpython-310-x86_64-linux-gnu.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e522d249550dbec7577971415180c6685de6e5e864f89e780398ce8f32e71ad
+size 9901144