Update app.py

app.py

@@ -1,46 +1,52 @@
 import os, io
-import gradio as gr
 import numpy as np
 import soundfile as sf
+import requests
 import torch
+import gradio as gr
+import spaces  # <-- needed for GPU Zero
+
 from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
 
-# ---------- EDIT THESE TO YOUR ACTUAL MODEL IDS ----------
+# --------------------- CONFIG: EDIT THESE ---------------------
 MODEL_CATALOG = {
     "Iban (ASR)": {
-        "repo_id": "mds04/iban_transcription",   # <- change to exact name
+        "repo_id": "mds04/iban_transcription",          # <-- exact repo id
         "language": "iban",
-        "task": "transcribe",
     },
     "Bukar Sadong (ASR)": {
-        "repo_id": "mds04/bukar_sadong_transcription",  # <- change to exact name
+        "repo_id": "mds04/bukar_sadong_transcription",   # <-- exact repo id
         "language": "bukar-sadong",
-        "task": "transcribe",
     },
 }
-# ---------------------------------------------------------
+DEFAULT_MODEL = "Iban (ASR)"
+DEFAULT_TASK = "transcribe"   # or "translate" if your model supports it
+DEFAULT_FORCE_LANG = True
+DEFAULT_MAX_TOKENS = 256
+# --------------------------------------------------------------
 
-# Private model(s)? Add Space secret HF_TOKEN (Settings → Secrets)
+# Private model(s)? Add Space Secret: HF_TOKEN
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Lazy loaders so we don't load both models at startup
-_MODEL_CACHE = {}
+# Lazy cache to avoid loading both models at startup
+_MODEL_CACHE: dict[str, tuple[AutoProcessor, AutoModelForSpeechSeq2Seq]] = {}
 
-def _load_bundle(key):
-    if key in _MODEL_CACHE:
-        return _MODEL_CACHE[key]
-    info = MODEL_CATALOG[key]
+def _load_bundle(model_key: str):
+    if model_key in _MODEL_CACHE:
+        return _MODEL_CACHE[model_key]
+    info = MODEL_CATALOG[model_key]
     proc = AutoProcessor.from_pretrained(info["repo_id"], token=HF_TOKEN)
-    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+    mdl = AutoModelForSpeechSeq2Seq.from_pretrained(
         info["repo_id"], token=HF_TOKEN, torch_dtype=dtype
     ).to(device).eval()
-    _MODEL_CACHE[key] = (proc, model)
-    return _MODEL_CACHE[key]
+    _MODEL_CACHE[model_key] = (proc, mdl)
+    return _MODEL_CACHE[model_key]
 
-def _resample_to_16k(x, sr):
+def _resample_to_16k(x: np.ndarray, sr: int) -> np.ndarray:
+    """Naive linear resampler to 16k (no librosa)."""
     if sr == 16000:
         return x.astype(np.float32)
     duration = x.shape[0] / sr
@@ -48,62 +54,103 @@ def _resample_to_16k(x, sr):
     t_new = np.linspace(0.0, duration, num=int(duration * 16000), endpoint=False)
     return np.interp(t_new, t_old, x).astype(np.float32)
 
-def _load_audio_16k(path: str) -> np.ndarray:
-    with open(path, "rb") as f:
-        data, sr = sf.read(io.BytesIO(f.read()))
+def _read_audio_bytes(path_or_url: str) -> bytes:
+    if path_or_url.startswith("http://") or path_or_url.startswith("https://"):
+        r = requests.get(path_or_url, timeout=30)
+        r.raise_for_status()
+        return r.content
+    with open(path_or_url, "rb") as f:
+        return f.read()
+
+def _load_audio_16k(input_obj) -> np.ndarray:
+    """
+    Accepts:
+      - str filepath,
+      - dict from Gradio v4 with {'path': <url or filepath>, 'meta': {...}}
+    Returns mono float32 @ 16k
+    """
+    if isinstance(input_obj, dict) and "path" in input_obj:
+        path_or_url = input_obj["path"]
+    elif isinstance(input_obj, str):
+        path_or_url = input_obj
+    else:
+        raise ValueError("Unsupported audio input format")
+
+    raw = _read_audio_bytes(path_or_url)
+    data, sr = sf.read(io.BytesIO(raw))
     if data.ndim == 2:
-        data = data.mean(axis=1)  # stereo -> mono
+        data = data.mean(axis=1)
     return _resample_to_16k(data, sr)
 
-def transcribe(model_choice, audio_path, force_lang, task_choice, max_tokens):
-    if not audio_path:
+# --------- IMPORTANT FOR GPU ZERO: decorate the main handler ----------
+@spaces.GPU  # <- tells Space to allocate GPU for this function
+def transcribe(model_choice, audio_input, task_choice, force_lang, max_tokens):
+    """
+    model_choice: str (dropdown)
+    audio_input: filepath or dict with 'path'
+    task_choice: "transcribe" | "translate"
+    force_lang: bool
+    max_tokens: int (slider)
+    """
+    if not audio_input:
         return "Please upload or record audio."
+
     processor, model = _load_bundle(model_choice)
-    audio = _load_audio_16k(audio_path)
+    audio = _load_audio_16k(audio_input)
     inputs = processor(audio=audio, sampling_rate=16000, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
     gen_kwargs = dict(max_new_tokens=int(max_tokens), do_sample=False)
 
-    # Try to force language/task if supported (Whisper-style)
-    if force_lang or task_choice:
-        if hasattr(processor, "get_decoder_prompt_ids"):
-            lang = MODEL_CATALOG[model_choice]["language"] if force_lang else None
-            try:
-                gen_kwargs["forced_decoder_ids"] = processor.get_decoder_prompt_ids(
-                    language=lang,
-                    task=task_choice
-                )
-            except Exception:
-                pass
+    # Force language (Whisper-style) if available and requested
+    if force_lang and hasattr(processor, "get_decoder_prompt_ids"):
+        try:
+            lang = MODEL_CATALOG[model_choice]["language"]
+            gen_kwargs["forced_decoder_ids"] = processor.get_decoder_prompt_ids(
+                language=lang, task=task_choice
+            )
+        except Exception:
+            pass
 
     with torch.no_grad():
         ids = model.generate(**inputs, **gen_kwargs)
 
-    text = processor.batch_decode(ids, skip_special_tokens=True)[0]
-    return text
+    return processor.batch_decode(ids, skip_special_tokens=True)[0]
+# ---------------------------------------------------------------------
 
 with gr.Blocks(title="Iban & Bukar Sadong ASR") as demo:
     gr.Markdown("## Iban & Bukar Sadong Transcription\nSelect a model, then upload or record audio.")
+
     with gr.Row():
         model_choice = gr.Dropdown(
             choices=list(MODEL_CATALOG.keys()),
-            value=list(MODEL_CATALOG.keys())[0],
+            value=DEFAULT_MODEL,
             label="Model"
         )
+
     with gr.Row():
-        audio_in = gr.Audio(sources=["microphone","upload"], type="filepath", label="Audio")
+        # Use type="filepath" so we get a path; code also supports remote URLs
+        audio_in = gr.Audio(sources=["microphone", "upload"], type="filepath", label="Audio")
+
     with gr.Row():
-        force_lang = gr.Checkbox(value=True, label="Force model’s language prompt")
-        task_choice = gr.Dropdown(choices=["transcribe","translate"], value="transcribe", label="Task")
-        max_tokens = gr.Slider(64, 512, value=256, step=16, label="Max new tokens")
+        task_choice = gr.Dropdown(
+            choices=["transcribe", "translate"],
+            value=DEFAULT_TASK,
+            label="Task"
+        )
+        force_lang = gr.Checkbox(
+            value=DEFAULT_FORCE_LANG,
+            label="Force model’s language prompt"
+        )
+        max_tokens = gr.Slider(64, 512, value=DEFAULT_MAX_TOKENS, step=16, label="Max new tokens")
 
-    btn = gr.Button("Transcribe")
     out = gr.Textbox(label="Transcription", lines=4)
+    btn = gr.Button("Transcribe")
 
+    # IMPORTANT: inputs here must match the function signature order
     btn.click(
-        transcribe,
-        inputs=[model_choice, audio_in, force_lang, task_choice, max_tokens],
+        fn=transcribe,
+        inputs=[model_choice, audio_in, task_choice, force_lang, max_tokens],
         outputs=out
     )