Update app.py

app.py

@@ -1,6 +1,3 @@
-# app.py — Íslensk talgreining with WhisperX Diarization & Timestamps
-# Public, no login, pallinr1@protonmail.com
-
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
@@ -8,98 +5,106 @@ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 import gradio as gr
 import spaces
 import whisperx
-import torch
-from transformers import pipeline
-
-# Keep Space awake
-import threading, time, requests
-def keep_awake():
-    while True:
-        time.sleep(45*60)
-        try:
-            requests.get(f"https://{os.getenv('SPACE_HOST')}")
-        except: pass
-threading.Thread(target=keep_awake, daemon=True).start()
-
-# Load your Whisper-small
-asr = pipeline(
-    "automatic-speech-recognition",
-    model="palli23/whisper-small-sam_spjall",
-    torch_dtype="float16",
-    device=0,
-    chunk_length_s=30,
-    batch_size=8,
-)
 
-# WhisperX setup (diarization + timestamps)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-batch_size = 16
-compute_type = "float16"
+# -----------------------------
+# MODEL SETTINGS
+# -----------------------------
+MODEL_NAME = "palli23/whisper-small-sam_spjall"
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Load WhisperX model (your Whisper-small)
-model = whisperx.load_model("palli23/whisper-small-sam_spjall", device, compute_type=compute_type)
+# -----------------------------
+# LOAD MODELS ONCE (GPU)
+# -----------------------------
+@spaces.GPU(duration=180)
+def load_all_models():
+    device = "cuda"
 
-# Load diarization model (pyannote internal to WhisperX)
-diarize_model = whisperx.DiarizationPipeline(
-    use_auth_token=True,
-    device=device,
-    min_speakers=1,
-    max_speakers=6,
-)
+    # 1. Whisper-small model
+    asr_model = whisperx.load_model(
+        MODEL_NAME,
+        device=device,
+        compute_type="float16"
+    )
 
-# Load alignment model (for timestamps)
-align_model, metadata = whisperx.load_align_model(language_code="is", device=device)
+    # 2. Alignment model
+    align_model, metadata = whisperx.load_align_model(
+        language_code="is",
+        device=device
+    )
+
+    # 3. Diarization model (pyannote)
+    diar_model = whisperx.DiarizationPipeline(
+        model_name="pyannote/speaker-diarization-3.1",
+        device=device,
+        use_auth_token=HF_TOKEN
+    )
+
+    return asr_model, align_model, metadata, diar_model
+
+asr_model, align_model, align_metadata, diar_model = load_all_models()
+
+
+# -----------------------------
+# TRANSCRIPTION + DIARIZATION
+# -----------------------------
+def transcribe_is_with_diar(audio_path):
 
-def transcribe_with_whisperx(audio_path, use_diarization=True):
     if not audio_path:
         return "Hladdu upp hljóðskrá"
-    
+
     # Load audio
     audio = whisperx.load_audio(audio_path)
-    
-    # Transcribe with your model
-    result = model.transcribe(audio, batch_size=batch_size)
-    
-    # Align for word-level timestamps
-    result = whisperx.align(
-        result["segments"], align_model, metadata, audio, device, return_char_alignments=False
+
+    # --- 1. ASR with Whisper-small
+    asr_result = asr_model.transcribe(
+        audio,
+        batch_size=8
+    )
+
+    # --- 2. Alignment (word timestamps)
+    aligned = whisperx.align(
+        asr_result["segments"],
+        align_model,
+        align_metadata,
+        audio,
+        device="cuda"
+    )
+
+    # --- 3. Diarization
+    diarization = diar_model(audio)
+
+    # --- 4. Merge diarization + words
+    final = whisperx.assign_word_speakers(diarization, aligned)
+
+    # Format output text
+    output_lines = []
+    for seg in final["segments"]:
+        speaker = seg.get("speaker", "SPEAKER_00")
+        text = seg.get("text", "")
+        output_lines.append(f"[{speaker}] {text}")
+
+    return "\n".join(output_lines)
+
+
+# -----------------------------
+# BUILD GRADIO UI
+# -----------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("# 🇮🇸 Íslenskt ASR + Raddgreining (Diarization)")
+    gr.Markdown("**Whisper-small + WhisperX** — Hljóð allt að 5 mínútur")
+
+    audio_in = gr.Audio(
+        type="filepath",
+        label="Hladdu upp hljóði (.mp3 / .wav)"
     )
-    
-    if not use_diarization:
-        # Return with timestamps (no speakers)
-        lines = []
-        for segment in result["segments"]:
-            start = segment["start"]
-            end = segment["end"]
-            text = segment["text"]
-            lines.append(f"{start:.1f}s – {end:.1f}s: {text}")
-        return "\n".join(lines)
-    
-    # Diarization
-    diarize_segments = diarize_model(audio)
-    result = whisperx.assign_word_speakers(diarize_segments, result)
-    
-    # Return with speakers + timestamps
-    lines = []
-    for segment in result["segments"]:
-        speaker = segment.get("speaker", "Unknown")
-        start = segment["start"]
-        end = segment["end"]
-        text = segment["text"]
-        lines.append(f"[{speaker}] {start:.1f}s – {end:.1f}s: {text}")
-    return "\n".join(lines)
-
-# UI — public, no login, your email
-with gr.Blocks(title="Íslensk talgreining + WhisperX") as demo:
-    gr.Markdown("# Íslensk talgreining + WhisperX")
-    gr.Markdown("**Whisper-small + diarization + timestamps • pallinr1@protonmail.com**")
-    
-    audio = gr.Audio(type="filepath", label="Hladdu upp hljóð (max 15 mín)")
-    diarize = gr.Checkbox(label="Virkja diarization + timestamps", value=True)
-    
     btn = gr.Button("Transcribe", variant="primary")
-    out = gr.Textbox(lines=25, label="Útskrift", show_copy_button=True)
-    
-    btn.click(transcribe_with_whisperx, inputs=[audio, diarize], outputs=out)
+    output = gr.Textbox(lines=30, label="Útskrift með raddgreiningu")
+
+    btn.click(fn=transcribe_is_with_diar, inputs=audio_in, outputs=output)
 
-demo.launch(auth=None, share=True)
+demo.launch(
+    auth=None,
+    share=True,
+    server_name="0.0.0.0",
+    server_port=7860
+)