diarization1Mæló

app.py

@@ -1,78 +1,99 @@
-# app.py for HF Spaces (ZeroGPU safe pyannote)
+# ============================================================
+# app.py – Whisper-small + Pyannote 3.1 (ZeroGPU örugg útgáfa)
+# ============================================================
+
 import os
 import gradio as gr
 import spaces
 import tempfile
 import torch
-
-from torch.serialization import safe_globals
-from pyannote.audio.core.model import Model
-from pyannote.audio.core.task import Task, Specifications
-from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization
-from typing import OrderedDict
-
 from transformers import pipeline
 from pyannote.audio import Pipeline
+from torch.serialization import safe_globals
 
-# Required patches for ZeroGPU
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-torch.serialization.add_safe_globals({
-    "OrderedDict": OrderedDict,
-})
-
-MODEL_NAME = "palli23/whisper-small-sam_spjall"
+# ------------------------------------------------------------
+# STILLT MODELNÖFN
+# ------------------------------------------------------------
+ASR_MODEL = "palli23/whisper-small-sam_spjall"
+DIAR_MODEL = "pyannote/speaker-diarization-3.1"
 
+# ------------------------------------------------------------
+# Aðalfallið – keyrir á ZeroGPU (120s GPU max)
+# ------------------------------------------------------------
 @spaces.GPU(duration=120)
 def transcribe_with_diarization(audio_path):
+
     if not audio_path:
-        return "Hladdu upp hljóðskrá"
-    
-    # Fix strict unpickling in torch 2.6 (ZeroGPU)
+        return "Hladdu upp hljóðskrá."
+
+    # ----------------------------
+    # 1. PYTORCH SAFE GLOBALS FIX
+    # ----------------------------
+    # PyTorch 2.6+ ZeroGPU unpickling patch – MANDATORY
     with safe_globals([
         torch.torch_version.TorchVersion,
-        Model,
-        Task,
-        Specifications,
-        SpeakerDiarization,
-        OrderedDict,
+        "pyannote.audio.core.task.Specifications",
+        "pyannote.audio.core.model.Model",
+        "pyannote.audio.pipelines.speaker_diarization.SpeakerDiarization"
     ]):
+
+        # ----------------------------
+        # 2. Load diarization pipeline
+        # ----------------------------
         diarization = Pipeline.from_pretrained(
-            "pyannote/speaker-diarization-3.1",
-            use_auth_token=os.getenv("HF_TOKEN")
+            DIAR_MODEL,
+            token=os.getenv("HF_TOKEN")   # <--- RÉTT FYRIR PYANNOTE 3.1
         ).to("cuda")
-    
-    # Run diarization
-    dia = diarization(audio_path)
 
-    # Whisper model
+    # Keyra diarization
+    diar = diarization(audio_path)
+
+    # ----------------------------
+    # 3. Whisper ASR
+    # ----------------------------
     asr = pipeline(
-        "automatic-speech-recognition",
-        model=MODEL_NAME,
+        task="automatic-speech-recognition",
+        model=ASR_MODEL,
         device=0,
-        use_auth_token=os.getenv("HF_TOKEN"),
+        token=os.getenv("HF_TOKEN")
     )
 
-    # segment-by-segment ASR
-    result = []
-    for turn, _, speaker in dia.itertracks(yield_label=True):
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            diarization.crop(audio_path, turn).export(f.name, format="wav")
-            chunk = f.name
+    # ----------------------------
+    # 4. Skera út segment + Greina texta
+    # ----------------------------
+    final_output = []
+
+    for turn, _, speaker in diar.itertracks(yield_label=True):
+
+        # Vista tímabundna WAV fyrir hvert segment
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            diar.crop(audio_path, turn).export(tmp.name, format="wav")
+            seg_path = tmp.name
 
-        text = asr(chunk)["text"].strip()
-        os.unlink(chunk)
-        result.append(f"[MÆLENDI {speaker}] {text}")
+        # ASR texti
+        text = asr(seg_path)["text"].strip()
 
-    return "\n".join(result) or "Enginn texti heyrðist."
+        # Vista niðurstöðu
+        final_output.append(f"[MÆLENDI {speaker}] {text}")
 
+        # Hreinsa
+        os.unlink(seg_path)
+
+    return "\n".join(final_output) if final_output else "Ekkert heyrt í hljóðinu."
+
+
+# ------------------------------------------------------------
+# GRADIO UI
+# ------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# Íslenskt ASR + Mælendagreining")
-    gr.Markdown("Whisper-small + pyannote 3.1 (ZeroGPU örugg útgáfa)")
-    
-    audio = gr.Audio(type="filepath", label="Hljóðskrá")
-    btn = gr.Button("Transcribe með mælendum")
-    out = gr.Textbox(lines=35, label="Úttak")
+    gr.Markdown("# 🎙️ Íslenskt tal → texti + mælendagreining")
+    gr.Markdown("Whisper-small + pyannote 3.1 • Virkar á ZeroGPU • 5 mín hljóð max")
+
+    audio_input = gr.Audio(type="filepath", label="Hladdu upp hljóðskrá (.wav / .mp3)")
+    out_box = gr.Textbox(lines=30, label="Útskrift + mælendur")
 
-    btn.click(transcribe_with_diarization, inputs=audio, outputs=out)
+    run_button = gr.Button("Transcribe með mælendum", variant="primary")
+    run_button.click(transcribe_with_diarization, inputs=audio_input, outputs=out_box)
 
+# Spaces auth
 demo.launch(auth=("beta", "beta2025"))