From 2cd8656fbd5b78a04df1762e78fc30e31c311112 Mon Sep 17 00:00:00 2001
From: Tianyuan <tianyuan.wang@outlook.com>
Date: Mon, 13 May 2024 15:21:42 +0100
Subject: [PATCH] Update CPU-only support for building the project.

Setup.py adds CUDA extensions based on whether there is a GPU or not. If not presented, add CPU extensions only. Also added exception handling for module importing where error may occur.
---
 alphapose/utils/roi_align/roi_align.py |  6 ++-
 detector/nms/nms_wrapper.py            |  6 ++-
 setup.py                               | 60 +++++++++++++++-----------
 3 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/alphapose/utils/roi_align/roi_align.py b/alphapose/utils/roi_align/roi_align.py
index b7610758..555982e6 100644
--- a/alphapose/utils/roi_align/roi_align.py
+++ b/alphapose/utils/roi_align/roi_align.py
@@ -3,8 +3,10 @@
 from torch.autograd.function import once_differentiable
 from torch.nn.modules.utils import _pair
 
-from . import roi_align_cuda
-
+try:
+    from . import roi_align_cuda
+except ImportError:
+    print("Import roi_align_cuda module failed, skipping...")
 
 class RoIAlignFunction(Function):
 
diff --git a/detector/nms/nms_wrapper.py b/detector/nms/nms_wrapper.py
index 9a051e72..5e5b85c6 100644
--- a/detector/nms/nms_wrapper.py
+++ b/detector/nms/nms_wrapper.py
@@ -1,7 +1,11 @@
 import numpy as np
 import torch
 
-from . import nms_cpu, nms_cuda
+try:
+    from . import nms_cpu, nms_cuda
+except ImportError:
+    print("Import nms_cpu and nms_cuda failed, importing only nms_cpu...")
+    from . import nms_cpu
 from .soft_nms_cpu import soft_nms_cpu
 
 
diff --git a/setup.py b/setup.py
index 25d14cd4..dcf7037b 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,8 @@
 import numpy as np
 from Cython.Build import cythonize
 from setuptools import Extension, find_packages, setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
+import torch
 
 MAJOR = 0
 MINOR = 5
@@ -116,6 +117,14 @@ def make_cuda_ext(name, module, sources):
         })
 
 
+def make_cpp_ext(name, module, sources):
+
+    return CppExtension(
+        name='{}.{}'.format(module, name),
+        sources=[os.path.join(*module.split('.'), p) for p in sources]
+    )
+
+
 def get_ext_modules():
     ext_modules = []
     # only windows visual studio 2013+ support compile c/cuda extensions
@@ -128,33 +137,36 @@ def get_ext_modules():
                 name='soft_nms_cpu',
                 module='detector.nms',
                 sources=['src/soft_nms_cpu.pyx']),
-            make_cuda_ext(
+            make_cpp_ext(
                 name='nms_cpu',
                 module='detector.nms',
                 sources=['src/nms_cpu.cpp']),
-            make_cuda_ext(
-                name='nms_cuda',
-                module='detector.nms',
-                sources=['src/nms_cuda.cpp', 'src/nms_kernel.cu']),
-            make_cuda_ext(
-                name='roi_align_cuda',
-                module='alphapose.utils.roi_align',
-                sources=['src/roi_align_cuda.cpp', 'src/roi_align_kernel.cu']),
-            make_cuda_ext(
-                name='deform_conv_cuda',
-                module='alphapose.models.layers.dcn',
-                sources=[
-                    'src/deform_conv_cuda.cpp',
-                    'src/deform_conv_cuda_kernel.cu'
-                ]),
-            make_cuda_ext(
-                name='deform_pool_cuda',
-                module='alphapose.models.layers.dcn',
-                sources=[
-                    'src/deform_pool_cuda.cpp',
-                    'src/deform_pool_cuda_kernel.cu'
-                ]),
         ]
+        if torch.cuda.is_available():  # Skip cuda extensions if no GPU available
+            ext_modules.extend([
+                make_cuda_ext(
+                    name='nms_cuda',
+                    module='detector.nms',
+                    sources=['src/nms_cuda.cpp', 'src/nms_kernel.cu']),
+                make_cuda_ext(
+                    name='roi_align_cuda',
+                    module='alphapose.utils.roi_align',
+                    sources=['src/roi_align_cuda.cpp', 'src/roi_align_kernel.cu']),
+                make_cuda_ext(
+                    name='deform_conv_cuda',
+                    module='alphapose.models.layers.dcn',
+                    sources=[
+                        'src/deform_conv_cuda.cpp',
+                        'src/deform_conv_cuda_kernel.cu'
+                    ]),
+                make_cuda_ext(
+                    name='deform_pool_cuda',
+                    module='alphapose.models.layers.dcn',
+                    sources=[
+                        'src/deform_pool_cuda.cpp',
+                        'src/deform_pool_cuda_kernel.cu'
+                    ]),
+            ])
     return ext_modules