Update server.py

server.py

@@ -1,11 +1,6 @@
 #!/usr/bin/env python3
 """
-Faster Qwen3-TTS Demo Server (CPU Optimizado + Parches Anti-CUDA y Anti-None)
-
-Usage:
-    python demo/server.py
-    python demo/server.py --model Qwen/Qwen3-TTS-12Hz-1.7B-Base --port 7860
-    python demo/server.py --no-preload  # skip startup model load
+Qwen3-TTS Demo Server (Librería Oficial - CPU Nativo)
 """
 
 import argparse
@@ -35,73 +30,12 @@ from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 torch.set_num_threads(4)
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-# ==============================================================================
-# 🛡️ ESCUDO TOTAL ANTI-CUDA Y ANTI-NONE
-# ==============================================================================
-import site
-
-def _apply_anti_cuda_shield():
-    # 1. Parche físico para el ValueError de la librería
-    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)
-    except Exception: pass
-
-    # 2. Neutralizar validaciones internas de CUDA
-    if hasattr(torch.cuda, '_lazy_init'):
-        torch.cuda._lazy_init = lambda *args, **kwargs: None
-    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"
-
-    # 3. Interceptar .cuda()
-    torch.Tensor.cuda = lambda self, *args, **kwargs: self
-    torch.nn.Module.cuda = lambda self, *args, **kwargs: self
-
-    # 4. Interceptar y redirigir .to('cuda') hacia .to('cpu')
-    _orig_tensor_to = torch.Tensor.to
-    def _tensor_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_tensor_to(self, *new_args, **kwargs)
-    torch.Tensor.to = _tensor_to_mock
-
-    _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
-
-_apply_anti_cuda_shield()
-
+# Importamos la librería OFICIAL de Alibaba
 try:
-    from faster_qwen3_tts import FasterQwen3TTS
-    import faster_qwen3_tts.model as fq_model
-    
-    # Clon del PredictorGraph para CPU
-    class CPU_PredictorGraph:
-        def __init__(self, model, *args, **kwargs):
-            self.model = model
-            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
-        def __getattr__(self, name): return getattr(self.model, name)
-
-    fq_model.PredictorGraph = CPU_PredictorGraph
+    from qwen_tts import Qwen3TTSModel
 except ImportError:
-    print("Error: faster_qwen3_tts not found.")
+    print("Error: qwen-tts no está instalado. Revisa tu requirements.txt")
     sys.exit(1)
-# ==============================================================================
 
 from nano_parakeet import from_pretrained as _parakeet_from_pretrained
 
@@ -116,12 +50,12 @@ _ALL_MODELS =[
 _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]
+    AVAILABLE_MODELS =[m for m in _ALL_MODELS if m in _allowed]
 else:
     AVAILABLE_MODELS = list(_ALL_MODELS)
 
 BASE_DIR = Path(__file__).resolve().parent
-_ASSET_DIR = Path(os.environ.get("ASSET_DIR", "/tmp/faster-qwen3-tts-assets"))
+_ASSET_DIR = Path(os.environ.get("ASSET_DIR", "/tmp/qwen3-tts-assets"))
 PRESET_TRANSCRIPTS = _ASSET_DIR / "samples" / "parity" / "icl_transcripts.txt"
 PRESET_REFS =[
     ("ref_audio_3", _ASSET_DIR / "ref_audio_3.wav", "Clone 1"),
@@ -176,20 +110,10 @@ def _load_preset_refs() -> None:
             "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
-
-app = FastAPI(title="Faster Qwen3-TTS Demo")
+app = FastAPI(title="Qwen3-TTS Demo Oficial")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-_model_cache: OrderedDict[str, FasterQwen3TTS] = OrderedDict()
+_model_cache: OrderedDict[str, Qwen3TTSModel] = OrderedDict()
 _model_cache_max: int = int(os.environ.get("MODEL_CACHE_SIZE", "1"))
 _active_model_name: str | None = None
 _loading = False
@@ -220,7 +144,7 @@ def _get_cached_ref_path(content: bytes) -> str:
     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"
+        path = Path(tempfile.gettempdir()) / f"qwen3_tts_ref_{digest}.wav"
         if not path.exists(): path.write_bytes(content)
         _ref_cache[digest] = str(path)
         return str(path)
@@ -251,8 +175,8 @@ async def get_status():
     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[]
+            model_type = "official"
+            speakers = active.get_supported_speakers() or[]
         except Exception: pass
     return {
         "loaded": active is not None, "model": _active_model_name, "loading": _loading,
@@ -280,11 +204,16 @@ async def load_model(model_id: str = Form(...)):
         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)
+            
+            # Carga NATIVA de la librería oficial
+            new_model = Qwen3TTSModel.from_pretrained(
+                model_id, 
+                device_map="cpu", 
+                dtype=torch.float32
+            )
             _model_cache[model_id] = new_model
             _model_cache.move_to_end(model_id)
             _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}
@@ -319,59 +248,46 @@ async def generate_stream(
         try:
             model = _model_cache.get(_active_model_name)
             t0 = time.perf_counter()
-            total_audio_s = 0.0
-            voice_clone_ms = 0.0
             
+            # Generación estándar empaquetada en un solo bloque para evitar crasheos de chunks
             if mode == "voice_clone":
-                gen = model.generate_voice_clone_streaming(
+                audio_list, sr = model.generate_voice_clone(
                     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
+                    x_vector_only_mode=xvec_only, temperature=temperature, top_k=top_k, 
+                    repetition_penalty=repetition_penalty, max_new_tokens=360
                 )
             elif mode == "custom":
-                gen = model.generate_custom_voice_streaming(
+                audio_list, sr = model.generate_custom_voice(
                     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
+                    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,
+                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
                 )
 
-            ttfa_ms, total_gen_ms = None, 0.0
-
-            for chunk, sr, timing in gen:
-                # 🛡️ PROTECCIÓN ANTI-NONE APLICADA
-                timing = timing or {}
-                prefill = timing.get('prefill_ms')
-                decode = timing.get('decode_ms')
-                
-                # Convertimos a float de forma segura (0.0 si es None)
-                prefill_val = float(prefill) if prefill is not None else 0.0
-                decode_val = float(decode) if decode is not None else 0.0
-
-                total_gen_ms += (prefill_val + decode_val)
-                if ttfa_ms is None: ttfa_ms = total_gen_ms
-                
-                chunk_audio = _concat_audio(chunk)
-                total_audio_s += len(chunk_audio) / sr
-                rtf = total_audio_s / (total_gen_ms / 1000) if total_gen_ms > 0 else 0.0
-                
-                payload = {
-                    "type": "chunk", "audio_b64": _to_wav_b64(chunk_audio, sr), "sample_rate": sr,
-                    "ttfa_ms": round(ttfa_ms), "voice_clone_ms": round(voice_clone_ms),
-                    "rtf": round(rtf, 3), "total_audio_s": round(total_audio_s, 3),
-                    "elapsed_ms": round((time.perf_counter() - t0) * 1000, 3)
-                }
-                loop.call_soon_threadsafe(queue.put_nowait, json.dumps(payload))
-
-            loop.call_soon_threadsafe(queue.put_nowait, json.dumps({
-                "type": "done", "ttfa_ms": round(ttfa_ms or 0), "voice_clone_ms": round(voice_clone_ms),
-                "rtf": round(rtf, 3) if 'rtf' in locals() else 0.0,
-                "total_audio_s": round(total_audio_s, 3), "total_ms": round((time.perf_counter() - t0) * 1000)
-            }))
+            elapsed = time.perf_counter() - t0
+            chunk_audio = _concat_audio(audio_list)
+            dur = len(chunk_audio) / sr
+            rtf = dur / elapsed if elapsed > 0 else 0.0
+            ttfa_ms = round(elapsed * 1000)
+            
+            # Enviamos el audio completo como un único "Chunk"
+            payload = {
+                "type": "chunk", "audio_b64": _to_wav_b64(chunk_audio, sr), "sample_rate": sr,
+                "ttfa_ms": ttfa_ms, "voice_clone_ms": 0, "rtf": round(rtf, 3), 
+                "total_audio_s": round(dur, 3), "elapsed_ms": ttfa_ms
+            }
+            loop.call_soon_threadsafe(queue.put_nowait, json.dumps(payload))
+            
+            # Enviamos señal de "Done"
+            done_payload = {
+                "type": "done", "ttfa_ms": ttfa_ms, "voice_clone_ms": 0,
+                "rtf": round(rtf, 3), "total_audio_s": round(dur, 3), "total_ms": ttfa_ms
+            }
+            loop.call_soon_threadsafe(queue.put_nowait, json.dumps(done_payload))
+
         except Exception as e:
             loop.call_soon_threadsafe(queue.put_nowait, json.dumps({"type": "error", "message": str(e)}))
         finally:
@@ -414,11 +330,18 @@ async def generate_non_streaming(
     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)
+            audio_list, sr = model.generate_voice_clone(
+                text=text, language=language, ref_audio=tmp_path, ref_text=ref_text, 
+                x_vector_only_mode=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)
+            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)
+            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
@@ -429,7 +352,7 @@ async def generate_non_streaming(
         return JSONResponse({"audio_b64": _to_wav_b64(audio, sr), "sample_rate": sr, "metrics": {"total_ms": round(elapsed * 1000), "audio_duration_s": round(dur, 3), "rtf": round(rtf, 3)}})
 
 def main():
-    parser = argparse.ArgumentParser(description="Faster Qwen3-TTS Demo Server")
+    parser = argparse.ArgumentParser(description="Qwen3-TTS Demo Server")
     parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-Base", help="Model to preload at startup")
     parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", 7860)))
     parser.add_argument("--host", default="0.0.0.0")
@@ -438,11 +361,16 @@ def main():
 
     if not args.no_preload:
         global _active_model_name, _parakeet
-        print(f"Loading model: {args.model}")
-        _startup_model = FasterQwen3TTS.from_pretrained(args.model, device="cpu", dtype=torch.float32)
+        print(f"Loading official model: {args.model}")
+        
+        _startup_model = Qwen3TTSModel.from_pretrained(
+            args.model, 
+            device_map="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")