#!/usr/bin/env bash
# S2S server with Qwen3-TTS — the named voices the Reachy app GUI expects.
# Strictly CPU: no CUDA package installed, none used at runtime.
set -euo pipefail

VLLM_PORT=50000                # must match vLLM --port
VLLM_URL="http://127.0.0.1:${VLLM_PORT}/v1"
MODEL="qwen3.5-9b"             # must match vLLM --served-model-name
VENV="$HOME/.venvs/s2s-qwen"   # separate venv, keeps the kokoro one intact
WS_PORT=8765                   # where the robot connects
API_KEY=""                     # vLLM doesn't check it; set if your endpoint does

# CPU real-time trade-off (see README): quantized talker, codec untouched.
# The 1.7B talker (even Q4_K_M) starved the audio stream on this CPU.
TTS_MODEL="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
TTS_TALKER_QUANT="Q8_0"
TTS_CODEC_QUANT="BF16"

[ -d "$VENV" ] || python3 -m venv "$VENV"
source "$VENV/bin/activate"
pip install -qU pip

# Qwen3-TTS GGML engine, CPU build (pure C++, no CUDA)
pip install -q "qwentts-cpp-python==0.3.0+cpu" \
  -f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/tree/main/whl/cpu

# The GGML backend is not in a PyPI release yet — install from git main.
# TODO: replace with `pip install speech-to-speech` once PyPI has > 0.2.10.
# Uninstall first: the git build reuses the 0.2.10 version number, so pip
# would otherwise keep whatever is already installed instead of today's main.
pip uninstall -yq speech-to-speech 2>/dev/null || true
pip install -q "speech-to-speech @ git+https://github.com/huggingface/speech-to-speech.git"

# CPU torch last (the installs above can drag the CUDA build back in),
# then drop any CUDA leftovers and prove the venv is clean
pip install -q --force-reinstall torch torchaudio --index-url https://download.pytorch.org/whl/cpu
pip list --format=freeze | { grep -i '^nvidia-' || true; } | cut -d= -f1 | xargs -r pip uninstall -yq
python -c 'import torch; assert "+cpu" in torch.__version__, torch.__version__; print("torch", torch.__version__, "- CPU only, OK")'

curl -fsS "$VLLM_URL/models" >/dev/null || { echo "vLLM not reachable at $VLLM_URL" >&2; exit 1; }

# Mixed GGUF quants aren't a CLI option upstream, so this launcher fetches the
# two files itself and hands their paths to the normal main()
cat > "$VENV/s2s_launch.py" <<'PY'
import os, sys
from faster_qwen3_tts import FasterQwen3TTS

def _gguf_paths(model_name):
    from huggingface_hub import hf_hub_download
    from qwentts_cpp.models import GGUF_REPO, _MODEL_TO_TALKER_STEM, _normalize_quant
    talker_q = _normalize_quant(os.environ["QWEN3_TTS_TALKER_QUANT"])
    codec_q = _normalize_quant(os.environ["QWEN3_TTS_CODEC_QUANT"])
    stem = _MODEL_TO_TALKER_STEM[model_name]
    talker = hf_hub_download(GGUF_REPO, f"{stem}-{talker_q}.gguf")
    codec = hf_hub_download(GGUF_REPO, f"qwen-tokenizer-12hz-{codec_q}.gguf")
    print(f"TTS GGUFs: talker {talker_q}, codec {codec_q}", file=sys.stderr)
    return talker, codec

_orig = FasterQwen3TTS.from_pretrained.__func__
def _patched(cls, *args, **kwargs):
    if kwargs.get("backend") == "ggml":
        talker, codec = _gguf_paths(args[0] if args else kwargs["model_name"])
        kwargs["gguf_talker_path"] = talker
        kwargs["gguf_codec_path"] = codec
    return _orig(cls, *args, **kwargs)
FasterQwen3TTS.from_pretrained = classmethod(_patched)

# No voice barge-in: the speaker sits centimeters from the mic, so the robot's
# own echo can outweigh a voice across the room and cut it mid-sentence.
# The app's stop button still works (client-side cancel, different path).
from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig
RuntimeConfig.interrupt_response_enabled = property(lambda self: False)

from speech_to_speech.s2s_pipeline import main
sys.exit(main())
PY

export CUDA_VISIBLE_DEVICES=""
export QWEN3_TTS_TALKER_QUANT="$TTS_TALKER_QUANT" QWEN3_TTS_CODEC_QUANT="$TTS_CODEC_QUANT"

# VAD loosened for the robot's quiet mic (barge-in is disabled in the
# launcher, so a sensitive threshold can't make the robot cut itself)
exec python "$VENV/s2s_launch.py" \
  --mode realtime \
  --ws_port "$WS_PORT" \
  --thresh 0.3 --min_speech_ms 256 --min_speech_continuation_ms 128 \
  --device cpu \
  --stt parakeet-tdt \
  --llm_backend responses-api \
  --tts qwen3 \
  --qwen3_tts_backend ggml \
  --qwen3_tts_model_name "$TTS_MODEL" \
  --model_name "$MODEL" \
  --responses_api_base_url "$VLLM_URL" \
  --responses_api_api_key "$API_KEY" \
  --responses_api_stream