Update app.py

app.py

@@ -1,5 +1,18 @@
-# app.py — Íslensk talgreining + talnaraðgreining (works 100 %)
-import os, threading, time, requests
+# 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"
+
+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)
@@ -8,11 +21,7 @@ def keep_awake():
         except: pass
 threading.Thread(target=keep_awake, daemon=True).start()
 
-import gradio as gr
-from transformers import pipeline
-from pyannote.audio import Pipeline
-
-# Your Whisper-small
+# Load your Whisper-small
 asr = pipeline(
     "automatic-speech-recognition",
     model="palli23/whisper-small-sam_spjall",
@@ -22,35 +31,75 @@ asr = pipeline(
     batch_size=8,
 )
 
-# pyannote 3.1 diarization
-diarization = Pipeline.from_pretrained(
-    "pyannote/speaker-diarization-3.1",
-    use_auth_token=True
+# WhisperX setup (diarization + timestamps)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+batch_size = 16
+compute_type = "float16"
+
+# Load WhisperX model (your Whisper-small)
+model = whisperx.load_model("palli23/whisper-small-sam_spjall", device, compute_type=compute_type)
+
+# Load diarization model (pyannote internal to WhisperX)
+diarize_model = whisperx.DiarizationPipeline(
+    use_auth_token=True,
+    device=device,
+    min_speakers=1,
+    max_speakers=6,
 )
 
-def transcribe(audio, diarize=True):
-    if not audio: return "Hladdu upp hljóð"
+# Load alignment model (for timestamps)
+align_model, metadata = whisperx.load_align_model(language_code="is", device=device)
+
+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)
     
-    # Raw transcription
-    text = asr(audio)["text"]
+    # Align for word-level timestamps
+    result = whisperx.align(
+        result["segments"], align_model, metadata, audio, device, return_char_alignments=False
+    )
     
-    if not diarize:
-        return text
+    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 + speaker labels
-    result = diarization(audio)
+    # Diarization
+    diarize_segments = diarize_model(audio)
+    result = whisperx.assign_word_speakers(diarize_segments, result)
+    
+    # Return with speakers + timestamps
     lines = []
-    for turn, _, speaker in result.itertracks(yield_label=True):
-        lines.append(f"[{speaker}] {turn.start:.1f}–{turn.end:.1f}s: {text}")
+    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)
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Íslensk talgreining + talnarar")
-    gr.Markdown("**palli23/whisper-small + pyannote 3.1** • pallinr1@protonmail.com")
+# 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)")
-    chk = gr.Checkbox(label="Virkja talnaraðgreiningu", value=True)
+    diarize = gr.Checkbox(label="Virkja diarization + timestamps", value=True)
+    
     btn = gr.Button("Transcribe", variant="primary")
-    out = gr.Textbox(lines=30, label="Útskrift")
-    btn.click(transcribe, inputs=[audio, chk], outputs=out)
+    out = gr.Textbox(lines=25, label="Útskrift", show_copy_button=True)
+    
+    btn.click(transcribe_with_whisperx, inputs=[audio, diarize], outputs=out)
 
 demo.launch(auth=None, share=True)