Fix for macs (#7)

* Fix uv_unwrapper for osx. Mesh gen times at ~5s on m1 max

* [uv_unwrapper] Fix setup.py regression

* Update README.md

README.md

@@ -51,6 +51,8 @@ Stable Fast 3D can also run on Macs via the MPS backend, with the texture baker
 
 Note that support is **experimental** and not guaranteed to give the same performance and/or quality as the CUDA backend.
 
+You will need to install OpenMP runtime to enable clang support for `-fopenmp`. Follow the tutorial here https://mac.r-project.org/openmp/ 
+
 MPS backend support was tested on M1 max 64GB with the latest PyTorch nightly release. We recommend you install the latest PyTorch (2.4.0 as of writing) and/or the nightly version to avoid any issues that my arise with older PyTorch versions.
 
 You also need to run the code with `PYTORCH_ENABLE_MPS_FALLBACK=1`.

uv_unwrapper/setup.py

@@ -1,3 +1,4 @@
+import torch
 import glob
 import os
 
@@ -15,16 +16,16 @@ def get_extensions():
     if debug_mode:
         print("Compiling in debug mode")
 
+    is_mac = True if torch.backends.mps.is_available() else False
     extension = CppExtension
 
     extra_link_args = []
     extra_compile_args = {
         "cxx": [
             "-O3" if not debug_mode else "-O0",
             "-fdiagnostics-color=always",
-            "-fopenmp",
-            "-march=native",
-        ],
+            ("-Xclang " if is_mac else "") + "-fopenmp",
+        ] + ["-march=native"] if not is_mac else [],
     }
     if debug_mode:
         extra_compile_args["cxx"].append("-g")
@@ -54,8 +55,8 @@ def get_extensions():
                 "c10",
                 "torch",
                 "torch_cpu",
-                "torch_python",
-            ],
+                "torch_python"
+            ] + ["omp"] if is_mac else [],
         )
     )
 

uv_unwrapper/uv_unwrapper/csrc/bvh.h

@@ -2,7 +2,9 @@
 
 #include <cfloat>
 #include <cmath>
+#ifndef __ARM_ARCH_ISA_A64
 #include <immintrin.h>
+#endif
 #include <limits>
 #include <vector>
 

uv_unwrapper/uv_unwrapper/unwrap.py

@@ -235,8 +235,12 @@ def _calculate_tangents(
             # t_nrm_idx is always the same as t_pos_idx
             vn_idx[i] = triangle_idxs[:, i]
 
-        tangents = torch.zeros_like(vertex_normals)
-        tansum = torch.zeros_like(vertex_normals)
+        if(torch.backends.mps.is_available()):
+            tangents = torch.zeros_like(vertex_normals).contiguous()
+            tansum = torch.zeros_like(vertex_normals).contiguous()
+        else:
+            tangents = torch.zeros_like(vertex_normals)
+            tansum = torch.zeros_like(vertex_normals)
 
         # Compute tangent space for each triangle
         duv1 = tex[1] - tex[0]