app.py

import os

# Must be set before torch's OpenMP/MKL thread pools initialize on first use.
# This is the most reliable way to make sure BLAS/OpenMP actually uses all
# cores instead of a conservative default.
def _detect_cpu_threads() -> int:
    # Respect a value the platform/container already set -- on cgroup-limited
    # containers (like HF Spaces) this is often the *correct* real quota,
    # whereas os.cpu_count() reports the host machine's full core count and
    # will cause thread oversubscription if trusted blindly.
    for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS"):
        val = os.environ.get(var)
        if val and val.isdigit() and int(val) > 0:
            return int(val)
    # sched_getaffinity reflects CPU-affinity/cgroup restrictions more
    # accurately than os.cpu_count() on Linux; fall back to cpu_count if
    # unavailable (e.g. non-Linux).
    try:
        return len(os.sched_getaffinity(0))
    except AttributeError:
        return os.cpu_count() or 4


_CPU_THREADS = str(_detect_cpu_threads())
os.environ.setdefault("OMP_NUM_THREADS", _CPU_THREADS)
os.environ.setdefault("MKL_NUM_THREADS", _CPU_THREADS)

import tempfile
import time
import traceback
from typing import Optional

import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

from qwen_tts import Qwen3TTSModel

app = FastAPI()

MODEL = None
SUPPORTED_SPEAKERS = []
SUPPORTED_LANGUAGES = []

MODEL_NAME = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
DEFAULT_SPEAKER = "Ryan"       # English male voice; pick any from SUPPORTED_SPEAKERS
DEFAULT_LANGUAGE = "English"   # or "Auto" to let the model detect it

# bf16 has no native hardware acceleration on most x86 CPUs (it gets
# emulated via upcasting), which is often slower than plain fp32. fp32 is
# the reliable baseline for generic CPU inference.
CPU_DTYPE = torch.float32


@app.on_event("startup")
def startup_event():
    global MODEL, SUPPORTED_SPEAKERS, SUPPORTED_LANGUAGES
    try:
        print("=" * 80)
        print("Starting application...")
        print(f"PyTorch version: {torch.__version__}")
        cuda_available = torch.cuda.is_available()
        print(f"CUDA available: {cuda_available}")

        if not cuda_available:
            n_threads = int(_CPU_THREADS)
            # Intra-op parallelism (parallelizes ops like matmul/conv internally).
            torch.set_num_threads(n_threads)
            # Inter-op parallelism (runs independent ops concurrently). Can
            # only be set once, before any parallel work starts, so this
            # must stay early and wrapped defensively.
            try:
                torch.set_num_interop_threads(max(1, n_threads // 2))
            except RuntimeError as e:
                print(f"Could not set interop threads (already initialized): {e}")
            # Denormal floats are handled by a slow FP path on most CPUs;
            # flushing them to zero avoids random latency spikes during
            # generation. Negligible effect on audio quality.
            torch.set_flush_denormal(True)
            print(f"No GPU detected. CPU threads: {n_threads} "
                  f"(OMP_NUM_THREADS={os.environ.get('OMP_NUM_THREADS')})")

        print(f"Loading Qwen3-TTS model: {MODEL_NAME}")
        load_start = time.perf_counter()

        device_map = "cuda:0" if cuda_available else "cpu"
        dtype = torch.bfloat16 if cuda_available else CPU_DTYPE
        load_kwargs = dict(device_map=device_map, dtype=dtype)

        if cuda_available:
            try:
                MODEL = Qwen3TTSModel.from_pretrained(
                    MODEL_NAME,
                    attn_implementation="flash_attention_2",
                    **load_kwargs,
                )
            except Exception as flash_err:
                print(f"flash_attention_2 unavailable ({flash_err}); "
                      f"falling back to default attention implementation.")
                MODEL = Qwen3TTSModel.from_pretrained(MODEL_NAME, **load_kwargs)
        else:
            MODEL = Qwen3TTSModel.from_pretrained(MODEL_NAME, **load_kwargs)

        SUPPORTED_SPEAKERS = MODEL.get_supported_speakers()
        SUPPORTED_LANGUAGES = MODEL.get_supported_languages()

        load_elapsed = time.perf_counter() - load_start
        print(f"Supported speakers: {SUPPORTED_SPEAKERS}")
        print(f"Supported languages: {SUPPORTED_LANGUAGES}")
        print(f"Model loaded successfully in {load_elapsed:.1f}s.")
        print("=" * 80)
    except Exception as exc:
        print(traceback.format_exc())
        MODEL = None
        print(f"Startup failed: {exc}")


@app.get("/")
def health():
    return {
        "status": "running",
        "model_loaded": MODEL is not None,
        "cpu_threads": _CPU_THREADS,
        "supported_speakers": SUPPORTED_SPEAKERS,
        "supported_languages": SUPPORTED_LANGUAGES,
    }


@app.get("/tts")
def generate(
    text: str,
    speaker: str = Query(DEFAULT_SPEAKER, description="Voice to use, e.g. Ryan, Vivian, Aiden"),
    language: str = Query(DEFAULT_LANGUAGE, description="Target language, or 'Auto' to detect"),
    instruct: Optional[str] = Query(
        None, description="Optional natural-language style instruction, e.g. 'Speak happily'"
    ),
    max_new_tokens: Optional[int] = Query(
        None,
        description="Optional cap on generated audio tokens. Lower = faster on CPU, "
                    "but can truncate longer sentences. Omit to use the model default.",
    ),
):
    output_file = None
    try:
        print("=" * 80)
        print(f"Incoming text: {text}")
        print(f"speaker={speaker} language={language} instruct={instruct!r} "
              f"max_new_tokens={max_new_tokens}")

        if MODEL is None:
            raise HTTPException(status_code=500, detail="Model was not loaded.")

        if SUPPORTED_SPEAKERS and speaker not in SUPPORTED_SPEAKERS:
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported speaker '{speaker}'. Choose from: {SUPPORTED_SPEAKERS}",
            )

        generate_kwargs = {}
        if max_new_tokens is not None:
            generate_kwargs["max_new_tokens"] = max_new_tokens

        gen_start = time.perf_counter()
        # inference_mode disables autograd bookkeeping entirely (faster and
        # lower memory than no_grad) -- pure inference, no training ever
        # happens on this path.
        with torch.inference_mode():
            wavs, sample_rate = MODEL.generate_custom_voice(
                text=text,
                language=language,
                speaker=speaker,
                instruct=instruct or "",
                **generate_kwargs,
            )
        gen_elapsed = time.perf_counter() - gen_start
        print(f"Generation took {gen_elapsed:.2f}s")

        output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
        sf.write(output_file, wavs[0], sample_rate)

        size = os.path.getsize(output_file)
        print(f"Output file: {output_file}")
        print(f"File size: {size} bytes")
        if size == 0:
            raise RuntimeError("Generated WAV file is empty.")
        print("=" * 80)

        return FileResponse(
            output_file,
            media_type="audio/wav",
            filename="speech.wav",
            headers={"X-Generation-Seconds": f"{gen_elapsed:.2f}"},
            # Clean up the temp file once the response has been sent.
            background=BackgroundTask(lambda: os.remove(output_file) if os.path.exists(output_file) else None),
        )
    except HTTPException:
        raise
    except Exception:
        print(traceback.format_exc())
        if output_file and os.path.exists(output_file):
            os.remove(output_file)
        raise