dotfiles/hm/extra/whisperx.nix

{
  pkgs ? import <nixpkgs> { },
}:
pkgs.python3Packages.buildPythonPackage rec {
  pname = "whisperx";
  version = "3.2.0";
  # pypi doesn't have the requirements.txt file, and it's required
  src = pkgs.fetchFromGitHub {
    owner = "m-bain";
    repo = "whisperX";
    rev = "v${version}"; # git doesn't have tags
    hash = "sha256-JQvyR9JW8OiSRI0eywTyWB4VMXmu6cTpBekBWonoJa4=";
  };
  pyproject = true;
  dependencies = [
    pkgs.python3Packages.torch
    pkgs.python3Packages.torchaudio
    (pkgs.python3Packages.faster-whisper.overrideAttrs (old: {
      # 1.0.2 is actually breaking APIs (requires hotwords argument)
      src = pkgs.fetchFromGitHub {
        owner = "SYSTRAN";
        repo = "faster-whisper";
        rev = "v1.0.0";
        hash = "sha256-0fE8X1d6CgDrrHtRudksN/tIGRtBKMvoNwkSVyFNda4=";
      };
    }))
    pkgs.python3Packages.transformers
    pkgs.python3Packages.pyannote-audio # Not in the requirements.txt for some reason
    pkgs.python3Packages.pandas
    pkgs.python3Packages.nltk
  ];
  pythonRelaxDeps = true;
  # torchaudio: 2.5.1a0 is >=2, despite dependency check saying otherwise
  # ctranslate2: was pinned to 4.4.0 to fix some nvidia pinning issue or whatnot
  # pyannote-audio: not sure what pins this to 3.1.1, but 3.3.1 works fine
  # For some reason pyannote-audio doesn't get relaxed when listed in a pythonRelaxDeps array,
  # hence why using true
  build-system = [
    pkgs.python3Packages.setuptools
  ];
  pythonImportsCheck = [
    "whisperx"
  ];
}