Update app.py

app.py

@@ -1,73 +1,82 @@
 import os
 import gradio as gr
-import torch
 import whisperx
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN")  # MUST be set in HF Spaces secrets
 
-CT2_MODEL = "palli23/whisper-small-sam_spjall-ct2"
-DIAR_MODEL = "pyannote/speaker-diarization-3.1"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+ASR_MODEL = "palli23/whisper-small-sam_spjall-ct2"
+DIARIZATION_MODEL = "pyannote/speaker-diarization-3.1"
+ALIGN_MODEL = "WAV2VEC2_ASR_LARGE_LV60K_960H"
 
-
-def load_all_models():
-    print("Loading ASR model...")
-    asr_model = whisperx.load_model(
-        CT2_MODEL,
-        device=DEVICE,
-        compute_type="float16" if DEVICE == "cuda" else "int8",
-        token=HF_TOKEN,
+def load_models():
+    print("Loading WhisperX ASR...")
+    asr = whisperx.load_model(
+        model_name=ASR_MODEL,
+        device="cuda" if whisperx.is_cuda_available() else "cpu",
+        compute_type="int8",  # Safe for Spaces
+        hf_token=HF_TOKEN
     )
 
     print("Loading alignment model...")
     align_model, metadata = whisperx.load_align_model(
-        language_code="is",
-        device=DEVICE,
-        token=HF_TOKEN,
+        language_code="is", 
+        model_name=ALIGN_MODEL,
+        hf_token=HF_TOKEN
     )
 
     print("Loading diarization model...")
-    diar_model = whisperx.DiarizationPipeline(
-        DIAR_MODEL,
-        device=DEVICE,
-        token=HF_TOKEN,
+    diar = whisperx.DiarizationPipeline(
+        DIARIZATION_MODEL,
+        hf_token=HF_TOKEN,
+        use_auth_token=True
     )
 
-    return asr_model, align_model, metadata, diar_model
+    return asr, align_model, metadata, diar
+
 
+asr_model, align_model, align_metadata, diar_pipeline = load_models()
 
-print("Initializing...")
-asr_model, align_model, align_metadata, diar_model = load_all_models()
 
+def transcribe(audio):
+    if audio is None:
+        return "No audio provided."
 
-def transcribe(audio_file):
-    audio = whisperx.load_audio(audio_file)
-    result = asr_model.transcribe(audio, batch_size=16)
+    print("Running ASR...")
+    result = asr_model.transcribe(audio)
 
+    print("Running alignment...")
     aligned = whisperx.align(
         result["segments"],
         align_model,
         align_metadata,
         audio,
-        DEVICE,
+        "is"
     )
 
-    diarization = diar_model(audio)
-    final_segments = whisperx.assign_speakers(aligned["segments"], diarization)
+    print("Running diarization...")
+    diarization = diar_pipeline(audio)
+
+    print("Assigning speaker labels...")
+    final_result = whisperx.assign_word_speakers(
+        diarization,
+        aligned
+    )
 
-    output_text = ""
-    for seg in final_segments:
+    text_out = ""
+    for seg in final_result["segments"]:
         speaker = seg.get("speaker", "Unknown")
-        output_text += f"[{speaker}] {seg['text']}\n"
+        text_out += f"[{speaker}] {seg['text']}\n"
 
-    return output_text
+    return text_out
 
 
 ui = gr.Interface(
     fn=transcribe,
     inputs=gr.Audio(type="filepath"),
-    outputs=gr.Textarea(),
-    title="WhisperX Icelandic + Diarization",
+    outputs=gr.Textbox(label="Transcription + Speakers", lines=20),
+    title="WhisperX Icelandic CT2 + Diarization",
+    description="Uses your private CT2 Whisper Small model + alignment + pyannote diarization."
 )
 
-ui.launch()
+if __name__ == "__main__":
+    ui.launch()