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

{ lib
, python3Packages
, fetchFromGitHub
, python3
}:

# 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
#
# 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";
  # until https://github.com/mozilla/TTS/issues/424 is resolved
  # we treat released models as released versions:
  # https://github.com/mozilla/TTS/wiki/Released-Models
  version = "0.0.9";

  src = fetchFromGitHub {
    owner = "mozilla";
    repo = "TTS";
    rev = "df5899daf4ba4ec89544edf94f9c2e105c544461";
    sha256 = "sha256-lklG8DqG04LKJY93z2axeYhW8gtpbRG41o9ow2gJjuA=";
  };

  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
    # Not required for building/installation but for their development/ci workflow
    sed -i -e '/pylint/d' requirements.txt
    sed -i -e '/cardboardlint/d' requirements.txt setup.py
  '';

  nativeBuildInputs = [ python3Packages.cython ];

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

  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}
  '';

  pytestFlagsArray = [
    # requires tensorflow
    "--ignore=tests/test_tacotron2_tf_model.py"
    "--ignore=tests/test_vocoder_tf_melgan_generator.py"
    "--ignore=tests/test_vocoder_tf_pqmf.py"
  ];

  meta = with lib; {
    homepage = "https://github.com/mozilla/TTS";
    description = "Deep learning for Text to Speech";
    license = licenses.mpl20;
    maintainers = with maintainers; [ hexa mic92 ];
  };
}