Update app.py

app.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """
-HuggingFace Space entry point for OmniVoice demo - OPTIMIZED FOR CPU.
+HuggingFace Space entry point for OmniVoice demo.
+
 """
 
 import logging
@@ -14,40 +15,31 @@ logging.basicConfig(
 logging.getLogger("omnivoice").setLevel(logging.DEBUG)
 
 import numpy as np
+import spaces
 import torch
 from omnivoice import OmniVoice, OmniVoiceGenerationConfig
 from omnivoice.cli.demo import build_demo
 
-# --- [TỐI ƯU CPU] Cấu hình Threading cho PyTorch ---
-# Thay số 4 bằng số nhân thực (Physical Cores) của CPU bạn để đạt hiệu năng tốt nhất
-num_cores = os.cpu_count() or 4
-torch.set_num_threads(max(1, num_cores // 2)) 
-torch.set_num_interop_threads(1)
-
 # ---------------------------------------------------------------------------
 # Model loading
 # ---------------------------------------------------------------------------
 CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
 
 print(f"Loading model from {CHECKPOINT} to cpu ...")
-
-# --- [TỐI ƯU CPU] Chuyển đổi dtype sang bfloat16 (hoặc float32 nếu CPU quá cũ) ---
-# bfloat16 chạy rất nhanh trên các CPU hiện đại hỗ trợ AVX-512/AMX
-chosen_dtype = torch.bfloat16 if torch.cuda.is_available() or hasattr(torch, 'bfloat16') else torch.float32
-
 model = OmniVoice.from_pretrained(
     CHECKPOINT,
     device_map="cpu",
-    dtype=chosen_dtype, 
+    dtype=torch.float16,
     load_asr=True,
 )
 sampling_rate = model.sampling_rate
-print(f"Model loaded successfully with {chosen_dtype}!")
+print("Model loaded successfully!")
 
 # ---------------------------------------------------------------------------
 # Generation logic
 # ---------------------------------------------------------------------------
 
+
 def _gen_core(
     text,
     language,
@@ -66,11 +58,8 @@ def _gen_core(
     if not text or not text.strip():
         return None, "Please enter the text to synthesize."
 
-    # --- [TỐI ƯU CPU] Mặc định giảm num_step xuống 16 hoặc 20 để chạy nhanh hơn ---
-    steps = int(num_step) if num_step else 16 
-
     gen_config = OmniVoiceGenerationConfig(
-        num_step=steps,
+        num_step=int(num_step or 32),
         guidance_scale=float(guidance_scale) if guidance_scale is not None else 2.0,
         denoise=bool(denoise) if denoise is not None else True,
         preprocess_prompt=bool(preprocess_prompt),
@@ -100,9 +89,7 @@ def _gen_core(
         kw["instruct"] = instruct.strip()
 
     try:
-        # Tăng tốc tính toán bằng cách tắt autograd
-        with torch.no_grad():
-            audio = model.generate(**kw)
+        audio = model.generate(**kw)
     except Exception as e:
         return None, f"Error: {type(e).__name__}: {e}"
 
@@ -111,8 +98,11 @@ def _gen_core(
 
 
 # ---------------------------------------------------------------------------
-# Wrapper (Bỏ ZeroGPU decorator vì chạy trên CPU local)
+# ZeroGPU wrapper
 # ---------------------------------------------------------------------------
+
+
+@spaces.GPU(duration=60)
 def generate_fn(*args, **kwargs):
     return _gen_core(*args, **kwargs)
 
@@ -123,4 +113,4 @@ def generate_fn(*args, **kwargs):
 demo = build_demo(model, CHECKPOINT, generate_fn=generate_fn)
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()