diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..db1c93c --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +# Devenv +.devenv* +devenv.local.nix +devenv.local.yaml + +# direnv +.direnv + +# pre-commit +.pre-commit-config.yaml +src/protestswap/__pycache__ diff --git a/flake.nix b/flake.nix index 4e6764c..2030bde 100644 --- a/flake.nix +++ b/flake.nix @@ -24,52 +24,82 @@ projectName = project.pyproject.project.name; python = pkgs.python3; - inswapperModel = pkgs.stdenv.mkDerivation { - name = "inswapper"; + inswapperModel = pkgs.stdenv.mkDerivation { + name = "inswapper_128"; src = pkgs.fetchurl { url = "https://huggingface.co/thebiglaskowski/inswapper_128.onnx/resolve/main/inswapper_128.onnx"; hash = "sha256-5KPwjHU8ty0E4Qqg99vj3uu/OVZ9Tq1tzgjpiqSeFq8="; }; - unpackPhase = ''true''; - phases = [ - # "installPhase" - ]; + phases = [ "installPhase" ]; - installPhase = let dest = "$out/var/lib/${projectName}"; in /*shell*/'' - ls -l "$src" - mkdir -p "${dest}" - cp $src ${dest}/inswapper_128.onnx + installPhase = '' + mkdir -p "$out/var/lib/insightface/models/" + cp $src "$out/var/lib/insightface/models/inswapper_128.onnx" ''; }; + buffaloModel = pkgs.stdenv.mkDerivation { + name = "buffalo_l_Model"; + src = pkgs.fetchurl { + url = "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip"; + hash = "sha256-gP/jfYpZQNWac4TCAaKjjUdB8vPFHu9G67KCGKewyi8="; + }; + unpackPhase = /*shell*/'' + ls -l $src + echo $src + + mkdir -p "$out/var/lib/insightface/buffalo_l/" + unar -D -o $out/var/lib/insightface/buffalo_l/ $src + ''; + + nativeBuildInputs = [ pkgs.unar ]; + }; + postInstall = /*shell*/'' + mkdir -p $out/var/lib/insightface/models/buffalo_l + cp ${inswapperModel} $out/var/lib/insightface/models/inswapper_128.onnx + ''; + # ${pkgs.unar}/bin/unar -D -o $out/var/lib/insightface/models/buffalo_l ${buffaloModel} + # ''; + packageAttrs = project.renderers.buildPythonPackage { + inherit python; + }; + packageAttrsLive = project.renderers.mkPythonEditablePackage { + inherit python; + }; + protestswap = python.pkgs.buildPythonPackage ( packageAttrs # // { inherit postInstall; } + ); in { devShells.${system}.default = let - arg = project.renderers.withPackages { inherit python;}; - pythonEnv = python.withPackages arg; + arg = project.renderers.mkPythonEditablePackage { inherit python; }; + pythonEnv = python.pkgs.mkPythonEditablePackage ( arg ); in pkgs.mkShell { - packages = [ pythonEnv inswapperModel]; - shellHook = '' - export FOO=${inswapperModel}/var/lib/${projectName}/inswapper_128.onnx - echo "Binary at '$FOO'" + packages = [ pkgs.python3 pkgs.uv + # buffaloModel + protestswap pkgs.gcc14 ]; + shellHook = /*shell*/ + '' + ls -l ${buffaloModel} + + INSIGHTFACE_ROOT_DIR=$(mktemp -d /tmp/insightface.XXXXXXXX) + export INSIGHTFACE_ROOT_DIR + + mkdir -p "$INSIGHTFACE_ROOT_DIR" + + ln -s ${inswapperModel}/var/lib/insightface/models "$INSIGHTFACE_ROOT_DIR/models" + ln -s ${buffaloModel}/var/lib/insightface/buffalo_l "$INSIGHTFACE_ROOT_DIR/buffalo_l" + + # export PROTESTSWAP_ROOT="/tmp/protestswap_models" + # mkdir -p "$PROTESTSWAP_ROOT/buffalo_l" + # + # cp ${inswapperModel} "$PROTESTSWAP_ROOT/inswapper_128.onnx" + unset PYTHONPATH + uv sync + . .venv/bin/activate ''; }; - packages.${system} = let - attrs = project.renderers.buildPythonPackage { - inherit python; - }; - protestswap = (python.pkgs.buildPythonPackage ( attrs // { - buildInputs = [ inswapperModel ]; - } )); - in { - # inherit inswapperModel; - # "${projectName}" = protestswap; - default = pkgs.symlinkJoin { - name = "${projectName}"; - paths = [inswapperModel protestswap]; - }; - }; + packages.${system}.default = protestswap; }; } diff --git a/pyproject.toml b/pyproject.toml index 78d99e8..27da1e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,11 @@ license-files = ["LICEN[CS]E*"] description = "A minimal pyproject.toml" authors = [ ] requires-python = ">=3.13" -dependencies = [ "insightface", "opencv-python", "matplotlib" ] +dependencies = [ +"insightface", +"onnxruntime", +"opencv-python", "matplotlib" +] [project.scripts] protestswap-cli = "protestswap.cli:main" @@ -13,3 +17,6 @@ protestswap-cli = "protestswap.cli:main" [build-system] requires = ["uv_build >= 0.8.17, <0.9.0"] build-backend = "uv_build" + +[tool.setuptools.packages] +find = {namespaces = false} # Disable implicit namespaces diff --git a/src/protestswap/cli.py b/src/protestswap/cli.py index bcd35ad..67c4726 100644 --- a/src/protestswap/cli.py +++ b/src/protestswap/cli.py @@ -1,7 +1,29 @@ -import protestswap +import os +import logging + +from protestswap.protestswap import ProtestFaceSwapper + +logger = logging.getLogger("protestswap") + +DEFAULT_INSIGHTFACE_DIR = "/var/lib/insightface" +INSIGHTFACE_VAR = "INSIGHTFACE_ROOT_DIR" + def main(): - print("hello protest") + insightface_dir = None + if os.path.exists(DEFAULT_INSIGHTFACE_DIR): + insightface_dir = DEFAULT_INSIGHTFACE_DIR + elif os.environ.get(INSIGHTFACE_VAR): + insightface_dir = os.environ.get(INSIGHTFACE_VAR) + else: + logger.warning( + f"No directory at '{DEFAULT_INSIGHTFACE_DIR}' and '{INSIGHTFACE_VAR}' not set yet." + ) + insightface_dir = f"{os.environ.get('HOME')}/.cache/insightface" + + # swapper = ProtestFaceSwapper(insightface_dir) + # swapper.prepare_model() + if __name__ == "__main__": main() diff --git a/src/protestswap/protestswap.py b/src/protestswap/protestswap.py index d04736c..1b7f584 100644 --- a/src/protestswap/protestswap.py +++ b/src/protestswap/protestswap.py @@ -1,5 +1,8 @@ #!/usr/bin/env python3 +import logging +import os + import cv2 import random # TODO: seed! @@ -9,11 +12,14 @@ from insightface.app import FaceAnalysis from insightface.model_zoo.inswapper import INSwapper from insightface.app.common import Face +logger = logging.getLogger("protestswap") + SWAPPING_MODEL = 'inswapper_128.onnx' -class Faceswapper: - def __init__(self): - self._app : FaceAnalysis = FaceAnalysis(name='buffalo_l') + +class ProtestFaceSwapper: + def __init__(self, root_dir): + self._app : FaceAnalysis = FaceAnalysis(name='buffalo_l', root=insightface_dir) self._model : INSwapper self._source : cv2.typing.MatLike|None = None self._source_faces : list[Face] = [] @@ -49,9 +55,9 @@ class Faceswapper: def main(): - swapper = Faceswapper() - swapper.prepare_model() + swapper = ProtestFaceSwapper(insightface_dir) + swapper.prepare_model() if __name__ == '__main__':