[Feature] Install mmhuman3d without torch (#253)

open-mmlab · Sep 14, 2022 · 24dba46 · 24dba46
1 parent fe9cc4c
commit 24dba46
Showing 1 changed file with 20 additions and 12 deletions.
diff --git a/setup.py b/setup.py
@@ -1,7 +1,12 @@
 from setuptools import find_packages, setup
 
-import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+try:
+    import torch
+    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+    cmd_class = {'build_ext': BuildExtension}
+except ModuleNotFoundError:
+    cmd_class = {}
+    print('Skip building ext ops due to the absence of torch.')
 
 
 def readme():
@@ -19,14 +24,17 @@ def get_version():
 
 def get_extensions():
     extensions = []
-    if torch.cuda.is_available():
-        ext_ops = CUDAExtension(
-            'mmhuman3d.core.renderer.mpr_renderer.cuda.rasterizer',  # noqa: E501
-            [
-                'mmhuman3d/core/renderer/mpr_renderer/cuda/rasterizer.cpp',  # noqa: E501
-                'mmhuman3d/core/renderer/mpr_renderer/cuda/rasterizer_kernel.cu',  # noqa: E501
-            ])
-        extensions.append(ext_ops)
+    try:
+        if torch.cuda.is_available():
+            ext_ops = CUDAExtension(
+                'mmhuman3d.core.renderer.mpr_renderer.cuda.rasterizer',  # noqa: E501
+                [
+                    'mmhuman3d/core/renderer/mpr_renderer/cuda/rasterizer.cpp',  # noqa: E501
+                    'mmhuman3d/core/renderer/mpr_renderer/cuda/rasterizer_kernel.cu',  # noqa: E501
+                ])
+            extensions.append(ext_ops)
+    except Exception as e:
+        print(f'Skip building ext ops: {e}')
     return extensions
 
 
@@ -109,7 +117,7 @@ def gen_packages_items():
 setup(
     name='mmhuman3d',
     version=get_version(),
-    description='OpenMMLab 3D Human Toolbox and Benchmark',
+    description='OpenMMLab 3D Human Parametric Model Toolbox and Benchmark',
     long_description=readme(),
     long_description_content_type='text/markdown',
     author='OpenMMLab',
@@ -118,7 +126,7 @@ def gen_packages_items():
     url='https://github.com/open-mmlab/mmhuman3d',
     packages=find_packages(exclude=('configs', 'tools', 'demo')),
     ext_modules=get_extensions(),
-    cmdclass={'build_ext': BuildExtension},
+    cmdclass=cmd_class,
     include_package_data=True,
     classifiers=[
         'Development Status :: 4 - Beta',