From d822573429d8cdecc40441caea51771b57a777a0 Mon Sep 17 00:00:00 2001
From: Oleksandr Chychkan <chychkan@gmail.com>
Date: Mon, 31 Jan 2022 11:08:39 -0800
Subject: [PATCH] Add support for Apple M1 laptops

---
 requirements_3.9_arm64.txt | 14 ++++++++++++++
 scripts/0_setup.sh         | 33 +++++++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 requirements_3.9_arm64.txt

diff --git a/requirements_3.9_arm64.txt b/requirements_3.9_arm64.txt
new file mode 100644
index 0000000..056f8c6
--- /dev/null
+++ b/requirements_3.9_arm64.txt
@@ -0,0 +1,14 @@
+numpy==1.21.4
+opencv-python==4.5.3.56
+numexpr==2.8.1
+h5py==3.1.0
+tqdm==4.62.3
+colorama==0.4.4
+cython==0.29.26
+ffmpeg-python==0.2.0
+Pillow==8.4.0
+scikit-image==0.19.0
+numpy==1.21.4
+scipy==1.7.3
+tensorflow-macos==2.7.0
+PyQt6==6.0.3
\ No newline at end of file
diff --git a/scripts/0_setup.sh b/scripts/0_setup.sh
index 1908632..a4e6f94 100755
--- a/scripts/0_setup.sh
+++ b/scripts/0_setup.sh
@@ -20,11 +20,40 @@ source .dfl/env/bin/activate
 version=$(python -V | cut -f 2 -d ' ' | cut -f 1,2 -d .)
 reqs_file='requirements.txt'
 
+version_suffix=''
 if [[ ! -z "$version" && -f "requirements_$version.txt" ]]; then
-  reqs_file="requirements_$version.txt"
+  version_suffix="_$version"
 fi
 
+architecture=''
+if [[ ! -z "$(uname -a | grep -oE ' arm64$')" && -f "requirements${version_suffix}_arm64.txt" ]]; then
+  architecture='arm64'
+fi
+
+architecture_suffix=''
+if [[ ! -z "$architecture" ]]; then
+  architecture_suffix="_arm64"
+fi
+
+reqs_file="requirements${version_suffix}${architecture_suffix}.txt"
+
 echo "Using $reqs_file for $(python -V)"
-pip install -r $reqs_file
+
+if [[ ! -z "$architecture" ]]; then
+  if [[ -z "$(brew ls --versions hdf5)" ]]; then
+    echo "ERROR: HDF5 needs to be installed to run DeepFaceLab on M1 chip."
+    echo "You can install it with 'brew install hdf5'. For more details, see https://formulae.brew.sh/formula/hdf5"
+    echo "Once it is installed, run ./scripts/0_setup.sh again"
+    exit 1
+  fi
+
+  cython_pkg="$(cat $reqs_file | grep -E 'cython==.+')"
+  pip --no-cache-dir install $cython_pkg
+
+  h5py_pkg="$(cat $reqs_file | grep -E 'h5py==.+')"
+  HDF5_DIR="$(brew --prefix hdf5)" pip --no-cache-dir install --no-build-isolation $h5py_pkg
+fi
+
+pip --no-cache-dir install -r $reqs_file
 
 echo "Done."