Update server.py

server.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Faster Qwen3-TTS Demo Server (Forzado Absoluto a CPU)
+Faster Qwen3-TTS Demo Server (CPU Edición Ultra-Resistente)
 """
 
 import argparse
@@ -26,59 +26,49 @@ from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 
-# OPTIMIZACIÓN CPU: Limita el uso excesivo de hilos
+# OPTIMIZACIÓN CPU
 torch.set_num_threads(4)
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
 # ==============================================================================
-# 🛡️ ESCUDO ANTI-CUDA (FORZAR CPU A NIVEL GLOBAL)
+# 🛡️ ESCUDO TOTAL ANTI-CUDA Y ANTI-NONE
 # ==============================================================================
 import site
 
-def _apply_anti_cuda_shield():
-    # 1. Eliminar bloqueo físico de ValueError en la librería original
+def _apply_shield():
+    # 1. Parche físico
     try:
         for p in site.getsitepackages():
             model_py = os.path.join(p, "faster_qwen3_tts", "model.py")
             if os.path.exists(model_py):
                 with open(model_py, "r") as f: code = f.read()
-                if 'raise ValueError("CUDA graphs require CUDA device")' in code:
-                    code = code.replace('raise ValueError("CUDA graphs require CUDA device")', 'pass')
-                    with open(model_py, "w") as f: f.write(code)
+                code = code.replace('raise ValueError("CUDA graphs require CUDA device")', 'pass')
+                with open(model_py, "w") as f: f.write(code)
     except Exception: pass
 
-    # 2. Neutralizar las alertas internas de compilación CUDA de PyTorch
-    if hasattr(torch.cuda, '_lazy_init'):
-        torch.cuda._lazy_init = lambda *args, **kwargs: None
+    # 2. Mock de CUDA
     torch.cuda.is_available = lambda: False
     torch.cuda.current_device = lambda: 0
     torch.cuda.device_count = lambda: 1
-    torch.cuda.get_device_name = lambda x: "CPU"
+    if hasattr(torch.cuda, '_lazy_init'): torch.cuda._lazy_init = lambda *args, **kwargs: None
 
-    # 3. Interceptar llamadas directas .cuda() en Tensors y Models
+    # 3. Forzado de Tensors y Modules a CPU
     torch.Tensor.cuda = lambda self, *args, **kwargs: self
     torch.nn.Module.cuda = lambda self, *args, **kwargs: self
 
-    # 4. Interceptar y redirigir comandos .to('cuda') hacia .to('cpu')
-    _orig_tensor_to = torch.Tensor.to
-    def _tensor_to_mock(self, *args, **kwargs):
+    def _mock_to(self, *args, **kwargs):
         new_args = tuple('cpu' if isinstance(a, str) and 'cuda' in a else a for a in args)
         if 'device' in kwargs and isinstance(kwargs['device'], str) and 'cuda' in kwargs['device']:
             kwargs['device'] = 'cpu'
-        return _orig_tensor_to(self, *new_args, **kwargs)
-    torch.Tensor.to = _tensor_to_mock
+        return _orig_to(self, *new_args, **kwargs)
 
-    _orig_module_to = torch.nn.Module.to
-    def _module_to_mock(self, *args, **kwargs):
-        new_args = tuple('cpu' if isinstance(a, str) and 'cuda' in a else a for a in args)
-        if 'device' in kwargs and isinstance(kwargs['device'], str) and 'cuda' in kwargs['device']:
-            kwargs['device'] = 'cpu'
-        return _orig_module_to(self, *new_args, **kwargs)
-    torch.nn.Module.to = _module_to_mock
+    _orig_to = torch.Tensor.to
+    torch.Tensor.to = _mock_to
+    _orig_mod_to = torch.nn.Module.to
+    torch.nn.Module.to = _mock_to
 
-_apply_anti_cuda_shield()
+_apply_shield()
 
-# 5. Aplicar clon de PredictorGraph
 try:
     from faster_qwen3_tts import FasterQwen3TTS
     import faster_qwen3_tts.model as fq_model
@@ -86,7 +76,7 @@ try:
     class CPU_PredictorGraph:
         def __init__(self, model, *args, **kwargs):
             self.model = model
-            self.device = "cpu"
+            self.device = torch.device("cpu")
         def __call__(self, *args, **kwargs): return self.model(*args, **kwargs)
         def forward(self, *args, **kwargs): return self.model(*args, **kwargs)
         def warmup(self, *args, **kwargs): pass
@@ -94,13 +84,12 @@ try:
 
     fq_model.PredictorGraph = CPU_PredictorGraph
 except ImportError:
-    print("Error: faster_qwen3_tts not found.")
     sys.exit(1)
 # ==============================================================================
 
 from nano_parakeet import from_pretrained as _parakeet_from_pretrained
 
-_ALL_MODELS =[
+_ALL_MODELS = [
     "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
     "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
     "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
@@ -108,248 +97,117 @@ _ALL_MODELS =[
     "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
 ]
 
-_active_models_env = os.environ.get("ACTIVE_MODELS", "")
-if _active_models_env:
-    _allowed = {m.strip() for m in _active_models_env.split(",") if m.strip()}
-    AVAILABLE_MODELS =[m for m in _ALL_MODELS if m in _allowed]
-else:
-    AVAILABLE_MODELS = list(_ALL_MODELS)
+# Configuración de modelos activos
+_active_env = os.environ.get("ACTIVE_MODELS", "Qwen/Qwen3-TTS-12Hz-0.6B-Base,Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
+AVAILABLE_MODELS = [m.strip() for m in _active_env.split(",") if m.strip()]
 
 _ASSET_DIR = Path(os.environ.get("ASSET_DIR", "/tmp/faster-qwen3-tts-assets"))
-PRESET_TRANSCRIPTS = _ASSET_DIR / "samples" / "parity" / "icl_transcripts.txt"
-PRESET_REFS =[
+PRESET_REFS = [
     ("ref_audio_3", _ASSET_DIR / "ref_audio_3.wav", "Clone 1"),
     ("ref_audio_2", _ASSET_DIR / "ref_audio_2.wav", "Clone 2"),
     ("ref_audio", _ASSET_DIR / "ref_audio.wav", "Clone 3"),
 ]
 
-_GITHUB_RAW = "https://raw.githubusercontent.com/andimarafioti/faster-qwen3-tts/main"
-_PRESET_REMOTE = {
-    "ref_audio":   f"{_GITHUB_RAW}/ref_audio.wav",
-    "ref_audio_2": f"{_GITHUB_RAW}/ref_audio_2.wav",
-    "ref_audio_3": f"{_GITHUB_RAW}/ref_audio_3.wav",
-}
-_TRANSCRIPT_REMOTE = f"{_GITHUB_RAW}/samples/parity/icl_transcripts.txt"
-
-def _fetch_preset_assets() -> None:
-    import urllib.request
-    _ASSET_DIR.mkdir(parents=True, exist_ok=True)
-    PRESET_TRANSCRIPTS.parent.mkdir(parents=True, exist_ok=True)
-    if not PRESET_TRANSCRIPTS.exists():
-        try:
-            urllib.request.urlretrieve(_TRANSCRIPT_REMOTE, PRESET_TRANSCRIPTS)
-        except Exception: pass
-    for key, path, _ in PRESET_REFS:
-        if not path.exists() and key in _PRESET_REMOTE:
-            try:
-                urllib.request.urlretrieve(_PRESET_REMOTE[key], path)
-            except Exception: pass
-
-_preset_refs: dict[str, dict] = {}
-
-def _load_preset_transcripts() -> dict[str, str]:
-    if not PRESET_TRANSCRIPTS.exists(): return {}
-    transcripts = {}
-    for line in PRESET_TRANSCRIPTS.read_text(encoding="utf-8").splitlines():
-        if ":" not in line: continue
-        key_part, text = line.split(":", 1)
-        key = key_part.split("(")[0].strip()
-        transcripts[key] = text.strip()
-    return transcripts
-
-def _load_preset_refs() -> None:
-    transcripts = _load_preset_transcripts()
-    for key, path, label in PRESET_REFS:
-        if not path.exists(): continue
-        content = path.read_bytes()
-        cached_path = _get_cached_ref_path(content)
-        _preset_refs[key] = {
-            "id": key,
-            "label": label,
-            "filename": path.name,
-            "path": cached_path,
-            "ref_text": transcripts.get(key, ""),
-            "audio_b64": base64.b64encode(content).decode(),
-        }
-
-def _prime_preset_voice_cache(model: FasterQwen3TTS) -> None:
-    if not _preset_refs: return
-    for preset in _preset_refs.values():
-        try:
-            model._prepare_generation(
-                text="Hello.",
-                ref_audio=preset["path"],
-                ref_text=preset["ref_text"],
-                language="English",
-                xvec_only=True,
-                non_streaming_mode=True,
-            )
-        except Exception:
-            continue
+_preset_refs: dict = {}
 
-app = FastAPI(title="Faster Qwen3-TTS Demo (CPU)")
+def _get_cached_ref_path(content: bytes) -> str:
+    digest = hashlib.sha1(content).hexdigest()
+    path = Path(tempfile.gettempdir()) / f"fq3_ref_{digest}.wav"
+    if not path.exists(): path.write_bytes(content)
+    return str(path)
+
+app = FastAPI(title="Faster Qwen3-TTS CPU")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-_model_cache: OrderedDict[str, FasterQwen3TTS] = OrderedDict()
-_model_cache_max: int = int(os.environ.get("MODEL_CACHE_SIZE", "1"))
+_model_cache: OrderedDict = OrderedDict()
 _active_model_name: str | None = None
 _loading = False
-_ref_cache: dict[str, str] = {}
-_ref_cache_lock = threading.Lock()
 _parakeet = None
 _generation_lock = asyncio.Lock()
-_generation_waiters: int = 0
-
-MAX_TEXT_CHARS = 1000
-MAX_AUDIO_BYTES = 10 * 1024 * 1024
-
-def _to_wav_b64(audio: np.ndarray, sr: int) -> str:
-    buf = io.BytesIO()
-    sf.write(buf, audio.astype(np.float32).squeeze(), sr, format="WAV", subtype="PCM_16")
-    return base64.b64encode(buf.getvalue()).decode()
-
-def _concat_audio(audio_list) -> np.ndarray:
-    if isinstance(audio_list, np.ndarray): return audio_list.astype(np.float32).squeeze()
-    parts =[np.array(a, dtype=np.float32).squeeze() for a in audio_list if len(a) > 0]
-    return np.concatenate(parts) if parts else np.zeros(0, dtype=np.float32)
-
-def _get_cached_ref_path(content: bytes) -> str:
-    digest = hashlib.sha1(content).hexdigest()
-    with _ref_cache_lock:
-        cached = _ref_cache.get(digest)
-        if cached and os.path.exists(cached): return cached
-        path = Path(tempfile.gettempdir()) / f"faster_qwen3_tts_ref_{digest}.wav"
-        if not path.exists(): path.write_bytes(content)
-        _ref_cache[digest] = str(path)
-        return str(path)
-
-_fetch_preset_assets()
-_load_preset_refs()
+_generation_waiters = 0
 
 @app.get("/")
 async def root(): return FileResponse(Path(__file__).parent / "index.html")
 
-@app.post("/transcribe")
-async def transcribe_audio(audio: UploadFile = File(...)):
-    if _parakeet is None: raise HTTPException(status_code=503, detail="Transcription model not loaded")
-    content = await audio.read()
-    def run():
-        wav, sr = sf.read(io.BytesIO(content), dtype="float32", always_2d=False)
-        if wav.ndim > 1: wav = wav.mean(axis=1)
-        wav_t = torch.from_numpy(wav)
-        if sr != 16000: wav_t = torchaudio.functional.resample(wav_t.unsqueeze(0), sr, 16000).squeeze(0)
-        return _parakeet.transcribe(wav_t)
-    return {"text": await asyncio.to_thread(run)}
-
 @app.get("/status")
 async def get_status():
-    speakers =[]
-    model_type = None
-    active = _model_cache.get(_active_model_name) if _active_model_name else None
-    if active is not None:
-        try:
-            model_type = active.model.model.tts_model_type
-            speakers = active.model.get_supported_speakers() or[]
-        except Exception: pass
+    active = _model_cache.get(_active_model_name)
+    speakers = active.model.get_supported_speakers() if active else []
     return {
-        "loaded": active is not None,
-        "model": _active_model_name,
-        "loading": _loading,
-        "available_models": AVAILABLE_MODELS,
-        "model_type": model_type,
-        "speakers": speakers,
-        "transcription_available": _parakeet is not None,
-        "preset_refs": [{"id": p["id"], "label": p["label"], "ref_text": p["ref_text"]} for p in _preset_refs.values()],
-        "queue_depth": _generation_waiters,
-        "cached_models": list(_model_cache.keys()),
+        "loaded": active is not None, "model": _active_model_name, "loading": _loading,
+        "available_models": AVAILABLE_MODELS, "speakers": speakers,
+        "preset_refs": [{"id": k, "label": v["label"]} for k,v in _preset_refs.items()]
     }
 
-@app.get("/preset_ref/{preset_id}")
-async def get_preset_ref(preset_id: str):
-    preset = _preset_refs.get(preset_id)
-    if not preset: raise HTTPException(status_code=404, detail="Preset not found")
-    return preset
-
 @app.post("/load")
 async def load_model(model_id: str = Form(...)):
     global _active_model_name, _loading
     if model_id in _model_cache:
         _active_model_name = model_id
-        _model_cache.move_to_end(model_id)
-        return {"status": "already_loaded", "model": model_id}
+        return {"status": "ok"}
     _loading = True
-    def _do_load():
+    def _do():
         global _active_model_name, _loading
         try:
-            if len(_model_cache) >= _model_cache_max: _model_cache.popitem(last=False)
-            new_model = FasterQwen3TTS.from_pretrained(model_id, device="cpu", dtype=torch.float32)
-            _model_cache[model_id] = new_model
-            _model_cache.move_to_end(model_id)
+            if len(_model_cache) >= 1: _model_cache.popitem(last=False)
+            m = FasterQwen3TTS.from_pretrained(model_id, device="cpu", dtype=torch.float32)
+            _model_cache[model_id] = m
             _active_model_name = model_id
-            _prime_preset_voice_cache(new_model)
         finally: _loading = False
-    async with _generation_lock: await asyncio.to_thread(_do_load)
-    return {"status": "loaded", "model": model_id}
+    async with _generation_lock: await asyncio.to_thread(_do)
+    return {"status": "loaded"}
 
 @app.post("/generate/stream")
 async def generate_stream(
-    text: str = Form(...), language: str = Form("English"), mode: str = Form("voice_clone"),
-    ref_text: str = Form(""), speaker: str = Form(""), instruct: str = Form(""),
-    xvec_only: bool = Form(True), chunk_size: int = Form(8), temperature: float = Form(0.9),
-    top_k: int = Form(50), repetition_penalty: float = Form(1.05),
+    text: str = Form(...), mode: str = Form("voice_clone"), 
     ref_preset: str = Form(""), ref_audio: UploadFile = File(None),
+    chunk_size: int = Form(8), temperature: float = Form(0.9)
 ):
-    if not _active_model_name or _active_model_name not in _model_cache:
-        raise HTTPException(status_code=400, detail="Model not loaded.")
-    
+    model = _model_cache.get(_active_model_name)
+    if not model: raise HTTPException(status_code=400, detail="Carga el modelo primero")
+
     tmp_path = None
-    tmp_is_cached = False
-    if ref_preset and ref_preset in _preset_refs:
-        preset = _preset_refs[ref_preset]
-        tmp_path, tmp_is_cached = preset["path"], True
-        if not ref_text: ref_text = preset["ref_text"]
-    elif ref_audio and ref_audio.filename:
-        content = await ref_audio.read()
-        tmp_path, tmp_is_cached = _get_cached_ref_path(content), True
+    if ref_preset and ref_preset in _preset_refs: tmp_path = _preset_refs[ref_preset]["path"]
+    elif ref_audio: tmp_path = _get_cached_ref_path(await ref_audio.read())
 
     loop = asyncio.get_event_loop()
     queue = asyncio.Queue()
 
-    def run_generation():
+    def run_gen():
         try:
-            model = _model_cache.get(_active_model_name)
             t0 = time.perf_counter()
             total_audio_s = 0.0
             
-            if mode == "voice_clone":
-                gen = model.generate_voice_clone_streaming(
-                    text=text, language=language, ref_audio=tmp_path, ref_text=ref_text,
-                    xvec_only=xvec_only, chunk_size=chunk_size, temperature=temperature,
-                    top_k=top_k, repetition_penalty=repetition_penalty, max_new_tokens=360
-                )
-            elif mode == "custom":
-                gen = model.generate_custom_voice_streaming(
-                    text=text, speaker=speaker, language=language, instruct=instruct,
-                    chunk_size=chunk_size, temperature=temperature, top_k=top_k,
-                    repetition_penalty=repetition_penalty, max_new_tokens=360
-                )
-            else:
-                gen = model.generate_voice_design_streaming(
-                    text=text, instruct=instruct, language=language, chunk_size=chunk_size,
-                    temperature=temperature, top_k=top_k, repetition_penalty=repetition_penalty, max_new_tokens=360
-                )
+            gen = model.generate_voice_clone_streaming(
+                text=text, ref_audio=tmp_path, chunk_size=chunk_size, 
+                temperature=temperature, max_new_tokens=360
+            ) if mode == "voice_clone" else model.generate_voice_design_streaming(
+                text=text, chunk_size=chunk_size, temperature=temperature, max_new_tokens=360
+            )
 
             ttfa_ms, total_gen_ms = None, 0.0
             for chunk, sr, timing in gen:
-                total_gen_ms += timing.get('prefill_ms', 0) + timing.get('decode_ms', 0)
+                # 🛡️ PROTECCIÓN ANTI-NONE: Si timing es None o faltan keys, usamos 0
+                timing = timing or {}
+                prefill = timing.get('prefill_ms') or 0.0
+                decode = timing.get('decode_ms') or 0.0
+                
+                total_gen_ms += (float(prefill) + float(decode))
                 if ttfa_ms is None: ttfa_ms = total_gen_ms
-                chunk_audio = _concat_audio(chunk)
+                
+                chunk_audio = np.concatenate([np.array(a).squeeze() for a in chunk]) if isinstance(chunk, list) else chunk.squeeze()
                 total_audio_s += len(chunk_audio) / sr
+                
+                # RTF Safe
                 rtf = total_audio_s / (total_gen_ms / 1000) if total_gen_ms > 0 else 0.0
+                
+                buf = io.BytesIO()
+                sf.write(buf, chunk_audio.astype(np.float32), sr, format="WAV", subtype="PCM_16")
+                
                 payload = {
-                    "type": "chunk", "audio_b64": _to_wav_b64(chunk_audio, sr), "sample_rate": sr,
-                    "ttfa_ms": round(ttfa_ms), "rtf": round(rtf, 3), "total_audio_s": round(total_audio_s, 3),
-                    "elapsed_ms": round((time.perf_counter() - t0) * 1000, 3)
+                    "type": "chunk", "audio_b64": base64.b64encode(buf.getvalue()).decode(),
+                    "sample_rate": sr, "ttfa_ms": round(ttfa_ms), "rtf": round(rtf, 3),
+                    "total_audio_s": round(total_audio_s, 3)
                 }
                 loop.call_soon_threadsafe(queue.put_nowait, json.dumps(payload))
 
@@ -358,75 +216,32 @@ async def generate_stream(
             loop.call_soon_threadsafe(queue.put_nowait, json.dumps({"type": "error", "message": str(e)}))
         finally:
             loop.call_soon_threadsafe(queue.put_nowait, None)
-            if tmp_path and os.path.exists(tmp_path) and not tmp_is_cached: os.unlink(tmp_path)
 
     async def sse():
-        global _generation_waiters
-        _generation_waiters += 1
-        try:
-            async with _generation_lock:
-                _generation_waiters -= 1
-                thread = threading.Thread(target=run_generation, daemon=True)
-                thread.start()
-                while True:
-                    msg = await queue.get()
-                    if msg is None: break
-                    yield f"data: {msg}\n\n"
-        finally: pass
+        async with _generation_lock:
+            threading.Thread(target=run_gen, daemon=True).start()
+            while True:
+                msg = await queue.get()
+                if msg is None: break
+                yield f"data: {msg}\n\n"
 
     return StreamingResponse(sse(), media_type="text/event-stream")
 
-@app.post("/generate")
-async def generate_non_streaming(
-    text: str = Form(...), language: str = Form("English"), mode: str = Form("voice_clone"),
-    ref_text: str = Form(""), speaker: str = Form(""), instruct: str = Form(""),
-    xvec_only: bool = Form(True), temperature: float = Form(0.9), top_k: int = Form(50),
-    repetition_penalty: float = Form(1.05), ref_preset: str = Form(""), ref_audio: UploadFile = File(None),
-):
-    model = _model_cache.get(_active_model_name)
-    if not model: raise HTTPException(status_code=400, detail="Model not loaded.")
-    
-    tmp_path = None
-    if ref_preset and ref_preset in _preset_refs: tmp_path = _preset_refs[ref_preset]["path"]
-    elif ref_audio: tmp_path = _get_cached_ref_path(await ref_audio.read())
-
-    def run():
-        t0 = time.perf_counter()
-        if mode == "voice_clone":
-            audio_list, sr = model.generate_voice_clone(text=text, language=language, ref_audio=tmp_path, ref_text=ref_text, xvec_only=xvec_only, temperature=temperature, top_k=top_k, repetition_penalty=repetition_penalty, max_new_tokens=360)
-        elif mode == "custom":
-            audio_list, sr = model.generate_custom_voice(text=text, speaker=speaker, language=language, instruct=instruct, temperature=temperature, top_k=top_k, repetition_penalty=repetition_penalty, max_new_tokens=360)
-        else:
-            audio_list, sr = model.generate_voice_design(text=text, instruct=instruct, language=language, temperature=temperature, top_k=top_k, repetition_penalty=repetition_penalty, max_new_tokens=360)
-        elapsed = time.perf_counter() - t0
-        audio = _concat_audio(audio_list)
-        return audio, sr, elapsed, len(audio)/sr
-
-    async with _generation_lock:
-        audio, sr, elapsed, dur = await asyncio.to_thread(run)
-        return JSONResponse({"audio_b64": _to_wav_b64(audio, sr), "sample_rate": sr, "metrics": {"total_ms": round(elapsed * 1000), "audio_duration_s": round(dur, 3)}})
-
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-Base")
-    parser.add_argument("--port", type=int, default=7860)
-    parser.add_argument("--host", default="0.0.0.0")
-    args = parser.parse_args()
-
+    args = argparse.ArgumentParser()
+    args.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-Base")
+    args.add_argument("--port", type=int, default=7860)
+    args = args.parse_args()
+
+    # Carga inicial
+    print(f"Iniciando en CPU...")
+    m = FasterQwen3TTS.from_pretrained(args.model, device="cpu", dtype=torch.float32)
+    _model_cache[args.model] = m
     global _active_model_name, _parakeet
-    
-    print(f"Loading model: {args.model}")
-    _startup_model = FasterQwen3TTS.from_pretrained(args.model, device="cpu", dtype=torch.float32)
-    _model_cache[args.model] = _startup_model
     _active_model_name = args.model
-    _prime_preset_voice_cache(_startup_model)
-
-    print("Loading transcription model (nano-parakeet)…")
     _parakeet = _parakeet_from_pretrained(device="cpu")
-    print("Transcription model ready on CPU.")
     
-    print(f"Server ready on CPU. Port: {args.port}")
-    uvicorn.run(app, host=args.host, port=args.port)
+    uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 if __name__ == "__main__":
     main()