"""ZeroGPU VoiceDesign entry point; CPU execution remains in the CLI worker."""
import spaces  # Initialize ZeroGPU before importing torch or loading the model.
import os
import numpy as np
import torch
from huggingface_hub import snapshot_download
from qwen_tts import Qwen3TTSModel

MODEL_ID = 'Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign'
REVISION = '5ecdb67327fd37bb2e042aab12ff7391903235d3'
path = os.getenv('QWEN_MODEL_PATH') or snapshot_download(MODEL_ID, revision=REVISION)
# ZeroGPU records startup CUDA placements and restores them when a GPU is allocated.
model = Qwen3TTSModel.from_pretrained(
    path, device_map='cuda', dtype=torch.bfloat16, attn_implementation='sdpa')


@spaces.GPU(duration=120)
def generate_speech(text, language, description):
    torch.manual_seed(42)
    np.random.seed(42)
    with torch.inference_mode():
        wavs, rate = model.generate_voice_design(
            text=text, language=language, instruct=description,
            non_streaming_mode=True, max_new_tokens=512)
    return np.asarray(wavs[0], dtype=np.float32), rate