Update app.py

app.py

@@ -4,56 +4,32 @@ HuggingFace Space entry point for OmniVoice demo.
 
 """
 
-import logging
 import os
-import tempfile
 from typing import Any, Dict
 
-try:
-    import spaces
-    _USING_ZERO_GPU = True
-except ImportError:
-    _USING_ZERO_GPU = False
-
+import numpy as np
+import spaces
 import torch
-import torchaudio
-
 from omnivoice import OmniVoice, OmniVoiceGenerationConfig
 from omnivoice.cli.demo import build_demo
 
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-# ---------------------------------------------------------------------------
-# Hardware detection
-# ---------------------------------------------------------------------------
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info(f"Using device: {DEVICE}")
-
 # ---------------------------------------------------------------------------
 # Model loading
 # ---------------------------------------------------------------------------
 CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
 
-model = None
-if not _USING_ZERO_GPU:
-    # Non-ZeroGPU: load model at startup on the best available device
-    logger.info(f"Loading model from {CHECKPOINT} on {DEVICE} ...")
-    model = OmniVoice.from_pretrained(
-        CHECKPOINT,
-        device_map=DEVICE,
-        dtype=torch.float16,
-        load_asr=True,
-    )
-    logger.info("Model loaded on %s.", DEVICE)
-else:
-    logger.info("ZeroGPU mode: model will be loaded inside @spaces.GPU() function.")
-
-sampling_rate = 16000  # fallback; will be overwritten after model loads
-
+print(f"Loading model from {CHECKPOINT} to cuda ...")
+model = OmniVoice.from_pretrained(
+    CHECKPOINT,
+    device_map="cuda",
+    dtype=torch.float16,
+    load_asr=True,
+)
+sampling_rate = model.sampling_rate
+print("Model loaded successfully!")
 
 # ---------------------------------------------------------------------------
-# Generation logic (outside build_demo so we can wrap with spaces.GPU)
+# Generation logic
 # ---------------------------------------------------------------------------
 
 
@@ -107,40 +83,27 @@ def _gen_core(
             kw["instruct"] = instruct.strip()
 
     try:
-        out_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
         audio = model.generate(**kw)
-        torchaudio.save(out_path, audio[0], sampling_rate)
     except Exception as e:
         return None, f"Error: {type(e).__name__}: {e}"
 
-    return out_path, "Done."
+    waveform = audio[0].squeeze(0).numpy()
+    waveform = (waveform * 32767).astype(np.int16)
+    return (sampling_rate, waveform), "Done."
 
 
 # ---------------------------------------------------------------------------
 # ZeroGPU wrapper
 # ---------------------------------------------------------------------------
-generate_fn = None
-if _USING_ZERO_GPU:
-    @spaces.GPU()
-    def generate_fn(*args, **kwargs):
-        # Lazy-load model on first call (inside GPU context)
-        global model, sampling_rate
-        if model is None:
-            logger.info(f"Loading model from {CHECKPOINT} on cuda (ZeroGPU) ...")
-            model = OmniVoice.from_pretrained(
-                CHECKPOINT,
-                device_map="cuda",
-                dtype=torch.float16,
-                load_asr=True,
-            )
-            sampling_rate = model.sampling_rate
-            logger.info("Model loaded on cuda (ZeroGPU).")
-        return _gen_core(*args, **kwargs)
-
-    logger.info("Using spaces.GPU() wrapper.")
+
+
+@spaces.GPU(duration=60)
+def generate_fn(*args, **kwargs):
+    return _gen_core(*args, **kwargs)
+
 
 # ---------------------------------------------------------------------------
-# Build and launch demo — reuses the full UI from omnivoice.cli.demo
+# Build and launch demo
 # ---------------------------------------------------------------------------
 demo = build_demo(model, CHECKPOINT, generate_fn=generate_fn)