From abd6f00001710347d56e14e1928876c166e2c9df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 14 May 2021 18:22:27 +0200
Subject: [PATCH] python3Packages.torchvision-bin: init at 0.9.1

This derivation can be used in conjunction with pytorch-bin to avoid
the long build time of PyTorch/TorchVision CUDA builds.
---
 .../python-modules/torchvision/bin.nix        | 59 +++++++++++++++++++
 .../torchvision/binary-hashes.nix             | 22 +++++++
 pkgs/top-level/python-packages.nix            |  2 +
 3 files changed, 83 insertions(+)
 create mode 100644 pkgs/development/python-modules/torchvision/bin.nix
 create mode 100644 pkgs/development/python-modules/torchvision/binary-hashes.nix

diff --git a/pkgs/development/python-modules/torchvision/bin.nix b/pkgs/development/python-modules/torchvision/bin.nix
new file mode 100644
index 00000000000..62980dcca8c
--- /dev/null
+++ b/pkgs/development/python-modules/torchvision/bin.nix
@@ -0,0 +1,59 @@
+{ lib
+, stdenv
+, buildPythonPackage
+, fetchurl
+, isPy37
+, isPy38
+, isPy39
+, patchelf
+, pillow
+, python
+, pytorch-bin
+}:
+
+let
+  pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
+  srcs = import ./binary-hashes.nix version;
+  unsupported = throw "Unsupported system";
+  version = "0.9.1";
+in buildPythonPackage {
+  inherit version;
+
+  pname = "torchvision";
+
+  format = "wheel";
+
+  src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
+
+  disabled = !(isPy37 || isPy38 || isPy39);
+
+  nativeBuildInputs = [
+    patchelf
+  ];
+
+  propagatedBuildInputs = [
+    pillow
+    pytorch-bin
+  ];
+
+  pythonImportsCheck = [ "torchvision" ];
+
+  postFixup = let
+    rpath = lib.makeLibraryPath [ stdenv.cc.cc.lib ];
+  in ''
+    # Note: after patchelf'ing, libcudart can still not be found. However, this should
+    #       not be an issue, because PyTorch is loaded before torchvision and brings
+    #       in the necessary symbols.
+    patchelf --set-rpath "${rpath}:${pytorch-bin}/${python.sitePackages}/torch/lib:" \
+      "$out/${python.sitePackages}/torchvision/_C.so"
+  '';
+
+  meta = with lib; {
+    description = "PyTorch vision library";
+    homepage = "https://pytorch.org/";
+    changelog = "https://github.com/pytorch/vision/releases/tag/v${version}";
+    license = licenses.bsd3;
+    platforms = platforms.linux;
+    maintainers = with maintainers; [ danieldk ];
+  };
+}
diff --git a/pkgs/development/python-modules/torchvision/binary-hashes.nix b/pkgs/development/python-modules/torchvision/binary-hashes.nix
new file mode 100644
index 00000000000..b9a991e24ab
--- /dev/null
+++ b/pkgs/development/python-modules/torchvision/binary-hashes.nix
@@ -0,0 +1,22 @@
+# Warning: use the same CUDA version as pytorch-bin.
+#
+# Precompiled wheels can be found at:
+# https://download.pytorch.org/whl/torch_stable.html
+
+version: {
+  x86_64-linux-37 = {
+    name = "torchvision-${version}-cp37-cp37m-linux_x86_64.whl";
+    url = "https://download.pytorch.org/whl/cu111/torchvision-${version}%2Bcu111-cp37-cp37m-linux_x86_64.whl";
+    hash = "sha256-7EMVB8KZg2I3P4RqnIVk/7OOAPA1OWOipns58cSCUrw=";
+  };
+  x86_64-linux-38 = {
+    name = "torchvision-${version}-cp38-cp38-linux_x86_64.whl";
+    url = "https://download.pytorch.org/whl/cu111/torchvision-${version}%2Bcu111-cp38-cp38-linux_x86_64.whl";
+    hash = "sha256-VjsCBW9Lusr4aDQLqaFh5dpV/5ZJ5PDs7nY4CbCHDTA=";
+  };
+  x86_64-linux-39 = {
+    name = "torchvision-${version}-cp39-cp39-linux_x86_64.whl";
+    url = "https://download.pytorch.org/whl/cu111/torchvision-${version}%2Bcu111-cp39-cp39-linux_x86_64.whl";
+    hash = "sha256-pzR7TBE+WcAmozskoeOVBuMkGJf9tvsaXsUkTcu86N8=";
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 857c01e617e..21eacfb16c2 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -8225,6 +8225,8 @@ in {
 
   torchvision = callPackage ../development/python-modules/torchvision { };
 
+  torchvision-bin = callPackage ../development/python-modules/torchvision/bin.nix { };
+
   tornado = callPackage ../development/python-modules/tornado { };
 
   # Used by circus and grab-site, 2020-08-29