depot/third_party/nixpkgs/pkgs/tools/audio/tts/default.nix

{ lib
, python3Packages
, fetchFromGitHub
, python3
, fetchpatch
}:

# USAGE:
# $ tts-server --list_models
# # pick your favorite vocoder/tts model
# $ tts-server --model_name tts_models/en/ljspeech/glow-tts --vocoder_name vocoder_models/universal/libri-tts/fullband-melgan
#
# If you upgrade from an old version you may have to delete old models from ~/.local/share/tts
# Also note that your tts version might not support all available models so check:
#   https://github.com/coqui-ai/TTS/releases/tag/v0.0.12
#
# For now, for deployment check the systemd unit in the pull request:
#   https://github.com/NixOS/nixpkgs/pull/103851#issue-521121136

python3Packages.buildPythonApplication rec {
  pname = "tts";
  version = "0.0.12";

  src = fetchFromGitHub {
    owner = "coqui-ai";
    repo = "TTS";
    rev = "v${version}";
    sha256 = "sha256-0M9wcdBmuTK+NvEGsXEdoYiVFjw8G2MRUwmi1PJgmzI=";
  };

  patches = [
    # https://github.com/coqui-ai/TTS/pull/435
    (fetchpatch {
      url = "https://github.com/coqui-ai/TTS/commit/97f98e4c4584ef14ed2f4885aa02c162d9364a00.patch";
      sha256 = "sha256-DAZYOOAe+6TYBF5ukFq5HRwm49askEvNEivuwb/oCWM=";
    })
  ];

  preBuild = ''
    # numba jit tries to write to its cache directory
    export HOME=$TMPDIR
    # we only support pytorch models right now
    sed -i -e '/tensorflow/d' requirements.txt

    sed -i -e 's!librosa==[^"]*!librosa!' requirements.txt setup.py
    sed -i -e 's!unidecode==[^"]*!unidecode!' requirements.txt setup.py
    sed -i -e 's!bokeh==[^"]*!bokeh!' requirements.txt setup.py
    sed -i -e 's!numba==[^"]*!numba!' requirements.txt setup.py
    sed -i -e 's!numpy==[^"]*!numpy!' requirements.txt setup.py
    sed -i -e 's!umap-learn==[^"]*!umap-learn!' requirements.txt setup.py
    # Not required for building/installation but for their development/ci workflow
    sed -i -e '/black/d' requirements.txt
    sed -i -e '/isor/d' requirements.txt
    sed -i -e '/pylint/d' requirements.txt
    sed -i -e '/cardboardlint/d' requirements.txt setup.py
  '';

  nativeBuildInputs = [ python3Packages.cython ];

  propagatedBuildInputs = with python3Packages; [
    attrdict
    bokeh
    flask
    fuzzywuzzy
    gdown
    inflect
    jieba
    librosa
    matplotlib
    phonemizer
    pypinyin
    pysbd
    pytorch
    scipy
    soundfile
    tensorboardx
    tqdm
    umap-learn
    unidecode
  ];

  postInstall = ''
    cp -r TTS/server/templates/ $out/${python3.sitePackages}/TTS/server
    # cython modules are not installed for some reasons
    (
      cd TTS/tts/layers/glow_tts/monotonic_align
      ${python3Packages.python.interpreter} setup.py install --prefix=$out
    )
  '';

  checkInputs = with python3Packages; [ pytestCheckHook ];

  disabledTests = [
    # RuntimeError: fft: ATen not compiled with MKL support
    "test_torch_stft"
    "test_stft_loss"
    "test_multiscale_stft_loss"
    # AssertionErrors that I feel incapable of debugging
    "test_phoneme_to_sequence"
    "test_text2phone"
    "test_parametrized_gan_dataset"
  ];

  preCheck = ''
    # use the installed TTS in $PYTHONPATH instead of the one from source to also have cython modules.
    mv TTS{,.old}
  '';

  disabledTestPaths = [
    # requires tensorflow
    "tests/test_tacotron2_tf_model.py"
    "tests/test_vocoder_tf_melgan_generator.py"
    "tests/test_vocoder_tf_pqmf.py"
  ];

  meta = with lib; {
    homepage = "https://github.com/coqui-ai/TTS";
    changelog = "https://github.com/coqui-ai/TTS/releases/tag/v${version}";
    description = "Deep learning toolkit for Text-to-Speech, battle-tested in research and production";
    license = licenses.mpl20;
    maintainers = with maintainers; [ hexa mic92 ];
  };
}