mirror of
https://github.com/TaterTotterson/microWakeWord-Trainer-Nvidia-Docker.git
synced 2026-08-12 07:55:33 -06:00
114 lines
3.4 KiB
Bash
Executable File
114 lines
3.4 KiB
Bash
Executable File
#!/bin/bash
|
|
set -euo pipefail
|
|
|
|
PROGPATH="$(realpath "$0")"
|
|
PROGDIR="$(dirname "${PROGPATH}")"
|
|
|
|
KNOWN_ARGS=( data-dir engine gpu no-gpu )
|
|
# shellcheck source=/dev/null
|
|
source "${PROGDIR}/shell.functions"
|
|
|
|
ENGINE="${ENGINE:-${POSITIONAL_ARGS[0]:-}}"
|
|
case "${ENGINE}" in
|
|
omnivoice|qwen3|moss) ;;
|
|
*)
|
|
echo "Usage: setup_modern_tts_envs --engine=<omnivoice|qwen3|moss> [--data-dir=/data]" >&2
|
|
exit 2
|
|
;;
|
|
esac
|
|
|
|
PYTHON_BIN="${MWW_TTS_PYTHON:-python3.12}"
|
|
command -v "${PYTHON_BIN}" >/dev/null 2>&1 || PYTHON_BIN=python3
|
|
|
|
if [ -z "${GPU:-}" ] ; then
|
|
GPU=false
|
|
if [ -c /dev/nvidiactl ] || { command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1 ; } ; then
|
|
GPU=true
|
|
fi
|
|
fi
|
|
|
|
TTS_ROOT="${DATA_DIR}/tts-envs"
|
|
VENV="${TTS_ROOT}/${ENGINE}"
|
|
STACK_VERSION="modern-tts-v1"
|
|
MARKER="${VENV}/.stack-version"
|
|
mkdir -p "${TTS_ROOT}" "${DATA_DIR}/.cache/huggingface"
|
|
|
|
case "${ENGINE}" in
|
|
omnivoice)
|
|
TORCH_VERSION="2.8.0"
|
|
TORCHAUDIO_VERSION="2.8.0"
|
|
PACKAGE_SPEC="git+https://github.com/k2-fsa/OmniVoice.git@28bc0889d92110491d726a9c79f26a895db5a074"
|
|
IMPORT_NAME="omnivoice"
|
|
STACK_ID="${STACK_VERSION}:omnivoice-28bc088:torch-${TORCH_VERSION}"
|
|
;;
|
|
qwen3)
|
|
TORCH_VERSION="2.9.1"
|
|
TORCHAUDIO_VERSION="2.9.1"
|
|
PACKAGE_SPEC="qwen-tts==0.1.1"
|
|
IMPORT_NAME="qwen_tts"
|
|
STACK_ID="${STACK_VERSION}:qwen-tts-0.1.1:torch-${TORCH_VERSION}"
|
|
;;
|
|
moss)
|
|
TORCH_VERSION="2.7.0"
|
|
TORCHAUDIO_VERSION="2.7.0"
|
|
PACKAGE_SPEC="git+https://github.com/OpenMOSS/MOSS-TTS-Nano.git@cc7bdf19c7639c0870dab22045a33b442760f6be"
|
|
IMPORT_NAME="moss_tts_nano"
|
|
STACK_ID="${STACK_VERSION}:moss-cc7bdf1:torch-${TORCH_VERSION}"
|
|
;;
|
|
esac
|
|
|
|
environment_ready() {
|
|
[ -x "${VENV}/bin/python" ] || return 1
|
|
[ -f "${MARKER}" ] || return 1
|
|
[ "$(cat "${MARKER}")" = "${STACK_ID}" ] || return 1
|
|
"${VENV}/bin/python" - "${IMPORT_NAME}" "${GPU}" <<'PY' >/dev/null 2>&1
|
|
import importlib
|
|
import sys
|
|
import torch
|
|
|
|
importlib.import_module(sys.argv[1])
|
|
expect_cuda = sys.argv[2].lower() == "true"
|
|
if expect_cuda and not torch.cuda.is_available():
|
|
raise SystemExit("NVIDIA GPU was detected but this environment cannot use CUDA")
|
|
if torch.cuda.is_available():
|
|
torch.zeros(1, device="cuda")
|
|
PY
|
|
}
|
|
|
|
if environment_ready ; then
|
|
echo "✅ Reusing ${ENGINE} TTS environment: ${VENV}"
|
|
exit 0
|
|
fi
|
|
|
|
echo "===== Preparing isolated ${ENGINE} TTS environment ====="
|
|
rm -rf "${VENV}"
|
|
"${PYTHON_BIN}" -m venv "${VENV}"
|
|
PY="${VENV}/bin/python"
|
|
"${PY}" -m pip install -U pip setuptools wheel
|
|
|
|
if ${GPU} ; then
|
|
TORCH_INDEX="${MWW_TTS_TORCH_INDEX:-https://download.pytorch.org/whl/cu128}"
|
|
echo "→ Installing CUDA torch ${TORCH_VERSION} from ${TORCH_INDEX}"
|
|
"${PY}" -m pip install \
|
|
"torch==${TORCH_VERSION}" \
|
|
"torchaudio==${TORCHAUDIO_VERSION}" \
|
|
--index-url "${TORCH_INDEX}"
|
|
else
|
|
echo "→ Installing CPU torch ${TORCH_VERSION}"
|
|
"${PY}" -m pip install \
|
|
"torch==${TORCH_VERSION}" \
|
|
"torchaudio==${TORCHAUDIO_VERSION}"
|
|
fi
|
|
|
|
echo "→ Installing ${PACKAGE_SPEC}"
|
|
"${PY}" -m pip install "${PACKAGE_SPEC}" "huggingface_hub[hf_xet]"
|
|
|
|
printf '%s\n' "${STACK_ID}" > "${MARKER}"
|
|
|
|
if ! environment_ready ; then
|
|
echo "❌ ${ENGINE} environment failed its import/CUDA check." >&2
|
|
exit 1
|
|
fi
|
|
|
|
echo "✅ ${ENGINE} TTS environment ready: ${VENV}"
|