python3Packages.vqgan-jax: init at unstable-2022-04-20

pkgs/development/python-modules/vqgan-jax/default.nix

@@ -0,0 +1,46 @@
+{ lib
+, buildPythonPackage
+, fetchFromGitHub
+, flax
+, jax
+, jaxlib
+, transformers
+}:
+
+buildPythonPackage rec {
+  pname = "vqgan-jax";
+  version = "unstable-2022-04-20";
+
+  src = fetchFromGitHub {
+    owner = "patil-suraj";
+    repo = "vqgan-jax";
+    rev = "1be20eee476e5d35c30e4ec3ed12222018af8ce4";
+    sha256 = "sha256-OZihAXpE0UsgauQ38XDmAF+lrIgz05uK0ro8SCdVsPc=";
+  };
+
+  format = "setuptools";
+
+  buildInputs = [
+    jaxlib
+  ];
+
+  propagatedBuildInputs = [
+    flax
+    jax
+    transformers
+  ];
+
+  doCheck = false;
+
+  pythonImportsCheck = [
+    "vqgan_jax"
+  ];
+
+  meta = with lib; {
+    description = "JAX implementation of VQGAN";
+    homepage = "https://github.com/patil-suraj/vqgan-jax";
+    # license unknown: https://github.com/patil-suraj/vqgan-jax/issues/9
+    license = lib.licenses.unfree;
+    maintainers = with maintainers; [ r-burns ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -11050,6 +11050,8 @@ in {
 
   vpk = callPackage ../development/python-modules/vpk { };
 
+  vqgan-jax = callPackage ../development/python-modules/vqgan-jax { };
+
   vsts = callPackage ../development/python-modules/vsts { };
 
   vsts-cd-manager = callPackage ../development/python-modules/vsts-cd-manager { };