From dbd0369271df21d2a6ff0a1991e1608fc151dd78 Mon Sep 17 00:00:00 2001 From: matthewcroughan Date: Fri, 24 Feb 2023 16:49:47 +0000 Subject: [PATCH] projects/koboldai: init --- flake.lock | 18 ++++ flake.nix | 5 + modules/aipython3/overlays.nix | 1 + packages/apispec-webframeworks/default.nix | 36 +++++++ projects/koboldai/default.nix | 20 ++++ projects/koboldai/package.nix | 109 +++++++++++++++++++++ 6 files changed, 189 insertions(+) create mode 100644 packages/apispec-webframeworks/default.nix create mode 100644 projects/koboldai/default.nix create mode 100644 projects/koboldai/package.nix diff --git a/flake.lock b/flake.lock index e632c08..2d87dbe 100644 --- a/flake.lock +++ b/flake.lock @@ -37,6 +37,23 @@ "type": "github" } }, + "koboldai-src": { + "flake": false, + "locked": { + "lastModified": 1668957963, + "narHash": "sha256-fKQ/6LiMmrfSWczC5kcf6M9cpuF9dDYl2gJ4+6ZLSdY=", + "owner": "koboldai", + "repo": "koboldai-client", + "rev": "f2077b8e58db6bd47a62bf9ed2649bb0711f9678", + "type": "github" + }, + "original": { + "owner": "koboldai", + "ref": "1.19.2", + "repo": "koboldai-client", + "type": "github" + } + }, "nixpkgs": { "locked": { "lastModified": 1675763311, @@ -57,6 +74,7 @@ "inputs": { "flake-parts": "flake-parts", "invokeai-src": "invokeai-src", + "koboldai-src": "koboldai-src", "nixpkgs": "nixpkgs" } } diff --git a/flake.nix b/flake.nix index ced9020..331505e 100644 --- a/flake.nix +++ b/flake.nix @@ -14,6 +14,10 @@ url = "github:invoke-ai/InvokeAI/v2.2.5"; flake = false; }; + koboldai-src = { + url = "github:koboldai/koboldai-client/1.19.2"; + flake = false; + }; flake-parts = { url = "github:hercules-ci/flake-parts"; inputs.nixpkgs-lib.follows = "nixpkgs"; @@ -28,6 +32,7 @@ ./modules/dependency-sets ./modules/aipython3 ./projects/invokeai + ./projects/koboldai ]; }; } diff --git a/modules/aipython3/overlays.nix b/modules/aipython3/overlays.nix index dec97d3..2c9912a 100644 --- a/modules/aipython3/overlays.nix +++ b/modules/aipython3/overlays.nix @@ -21,6 +21,7 @@ pkgs: { callPackage = final.callPackage; rmCallPackage = path: args: rm (callPackage path args); in { + apispec-webframeworks = callPackage ../../packages/apispec-webframeworks { }; pydeprecate = callPackage ../../packages/pydeprecate { }; taming-transformers-rom1504 = callPackage ../../packages/taming-transformers-rom1504 { }; diff --git a/packages/apispec-webframeworks/default.nix b/packages/apispec-webframeworks/default.nix new file mode 100644 index 0000000..1ce4320 --- /dev/null +++ b/packages/apispec-webframeworks/default.nix @@ -0,0 +1,36 @@ +{ lib, python3Packages }: + +python3Packages.buildPythonPackage rec { + pname = "apispec-webframeworks"; + version = "0.5.2"; + disabled = python3Packages.pythonOlder "3.6"; + + src = python3Packages.fetchPypi { + inherit pname version; + hash = "sha256-DbNbJnkUs/jFYqygJhlX28tBdvJV6swiUgJ3AQgY3PM="; + }; + + propagatedBuildInputs = with python3Packages; [ + apispec + packaging + ]; + + nativeCheckInputs = with python3Packages; [ + pytestCheckHook + mock + flask + tornado + bottle + ]; + + doCheck = false; + + pythonImportsCheck = [ "apispec_webframeworks" ]; + + meta = with lib; { + description = "apispec plugin for integrating with various web frameworks"; + homepage = "https://github.com/marshmallow-code/apispec-webframeworks"; + license = licenses.mit; + maintainers = [ maintainers.sikmir ]; + }; +} diff --git a/projects/koboldai/default.nix b/projects/koboldai/default.nix new file mode 100644 index 0000000..97a683c --- /dev/null +++ b/projects/koboldai/default.nix @@ -0,0 +1,20 @@ +{ inputs, lib, ... }: + +{ + perSystem = { config, pkgs, ... }: let + inherit (config.dependencySets) aipython3-amd aipython3-nvidia; + + src = inputs.koboldai-src; + + mkKoboldAIVariant = args: pkgs.callPackage ./package.nix ({ inherit src; } // args); + in { + packages = { + koboldai-nvidia = mkKoboldAIVariant { + aipython3 = aipython3-nvidia; + }; + koboldai-amd = mkKoboldAIVariant { + aipython3 = aipython3-amd; + }; + }; + }; +} diff --git a/projects/koboldai/package.nix b/projects/koboldai/package.nix new file mode 100644 index 0000000..283c92d --- /dev/null +++ b/projects/koboldai/package.nix @@ -0,0 +1,109 @@ +{ aipython3 +, lib +, src +, wsl ? false +, fetchFromGitHub +, writeShellScriptBin +, runCommand +, tmpDir ? "/tmp/nix-koboldai" +, stateDir ? "$HOME/.koboldai/state" +}: +let + overrides = { + transformers = aipython3.transformers.overrideAttrs (old: rec { + propagatedBuildInputs = old.propagatedBuildInputs ++ [ aipython3.huggingface-hub ]; + pname = "transformers"; + version = "4.24.0"; + src = fetchFromGitHub { + owner = "huggingface"; + repo = pname; + rev = "refs/tags/v${version}"; + hash = "sha256-aGtTey+QK12URZcGNaRAlcaOphON4ViZOGdigtXU1g0="; + }; + }); + bleach = aipython3.bleach.overrideAttrs (old: rec { + pname = "bleach"; + version = "4.1.0"; + src = fetchFromGitHub { + owner = "mozilla"; + repo = pname; + rev = "refs/tags/v${version}"; + hash = "sha256-YuvH8FvZBqSYRt7ScKfuTZMsljJQlhFR+3tg7kABF0Y="; + }; + }); + }; + # The original kobold-ai program wants to write models settings and user + # scripts to the current working directory, but tries to write to the + # /nix/store erroneously due to mismanagement of the current working + # directory in its source code. The patching below replicates the original + # functionality of the program by making symlinks in the source code + # directory that point to ${tmpDir} + # + # The wrapper script we have made for the program will then create another + # symlink that points to ${stateDir}, ultimately the default symlink trail + # looks like the following + # + # /nix/store/kobold-ai/models -> /tmp/nix-koboldai -> ~/.koboldai/state + patchedSrc = runCommand "koboldAi-patchedSrc" {} '' + cp -r --no-preserve=mode ${src} ./src + cd src + rm -rf models settings userscripts + cd - + substituteInPlace ./src/aiserver.py --replace 'os.system("")' 'STATE_DIR = os.path.expandvars("${stateDir}")' + substituteInPlace ./src/aiserver.py --replace 'cache_dir="cache"' "cache_dir=os.path.join(STATE_DIR, 'cache')" + substituteInPlace ./src/aiserver.py --replace 'shutil.rmtree("cache/")' 'shutil.rmtree(os.path.join(STATE_DIR, "cache"))' + substituteInPlace ./src/aiserver.py --replace "app.config['SESSION_TYPE'] = 'filesystem'" "app.config['SESSION_TYPE'] = 'memcached'" + mv ./src $out + ln -s ${tmpDir}/models/ $out/models + ln -s ${tmpDir}/settings/ $out/settings + ln -s ${tmpDir}/userscripts/ $out/userscripts + ''; + koboldPython = aipython3.python.withPackages (_: with aipython3; [ + overrides.bleach + overrides.transformers + colorama + flask + flask-socketio + flask-session + eventlet + dnspython + markdown + sentencepiece + protobuf + marshmallow + loguru + termcolor + psutil + torch-bin + torchvision-bin + apispec + apispec-webframeworks + lupa + memcached + ]); +in +(writeShellScriptBin "koboldai" '' + if [ -d "/usr/lib/wsl/lib" ] + then + echo "Running via WSL (Windows Subsystem for Linux), setting LD_LIBRARY_PATH" + set -x + export LD_LIBRARY_PATH="/usr/lib/wsl/lib" + set +x + fi + rm -rf ${tmpDir} + mkdir -p ${tmpDir} + mkdir -p ${stateDir}/models ${stateDir}/cache ${stateDir}/settings ${stateDir}/userscripts + ln -s ${stateDir}/models/ ${tmpDir}/models + ln -s ${stateDir}/settings/ ${tmpDir}/settings + ln -s ${stateDir}/userscripts/ ${tmpDir}/userscripts + ${koboldPython}/bin/python ${patchedSrc}/aiserver.py $@ +'').overrideAttrs + (_: { + meta = { + maintainers = [ lib.maintainers.matthewcroughan ]; + license = lib.licenses.agpl3; + description = "browser-based front-end for AI-assisted writing with multiple local & remote AI models"; + homepage = "https://github.com/KoboldAI/KoboldAI-Client"; + mainProgram = "koboldai"; + }; + })