test 2

app.py

@@ -3,24 +3,24 @@ from diarization import start_diarization
 from transcribe import start_transcribe
 import ffmpeg
 import gradio as gr
-import os
 
 
-def prepare_input(input_file):
+def prepare_input(input_file, progress=gr.Progress()):
     output_file = "input.wav"
+    progress(0.2, desc="Preparing video")
     ffmpeg.input(input_file).audio.output(
         output_file, format="wav").run()
-
-    progress = gr.Progress()
+    progress(0.4, desc="Acquiring diarization")
     start_diarization(output_file, progress)
-    # return start_transcribe(progress)
+    progress(0.6, desc="Transcribing audio")
+    return start_transcribe(progress)
 
 
 video_interface = gr.Interface(
     fn=prepare_input,
     inputs=gr.Video(type="file"),
     outputs="text",
-    title="Test 1"
+    title="Test 2"
 )
 
 if __name__ == "__main__":

diarization.py

@@ -5,8 +5,7 @@ import os
 import torch
 import json
 
-# hugging_face_token = os.environ["HUGGING_FACE_TOKEN"]
-hugging_face_token = "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"
+hugging_face_token = os.environ["HUGGING_FACE_TOKEN"]
 pipeline = Pipeline.from_pretrained(
     'pyannote/speaker-diarization', use_auth_token=hugging_face_token)
 device = torch.device("cuda")
@@ -14,29 +13,25 @@ pipeline.to(device)
 
 
 def start_diarization(input_file, progress: gr.Progress):
-    print("Starting diarization")
-    progress(0, desc="Starting diarization")
     diarization = pipeline(input_file)
 
     sample_groups = []
     speaker_groups = {}
-    print(str(diarization))
-    # for turn, _, speaker in diarization.itertracks(yield_label=True):
-    # print(diarization)
-    # for step in progress.tqdm(diarization.)
-
-    # if (speaker not in sample_groups):
-    #     sample_groups.append(str(speaker))
-
-    # suffix = 1
-    # file_name = f"{speaker}-{suffix}"
-    # while file_name in speaker_groups:
-    #     suffix += 1
-    #     file_name = f"{speaker}-{suffix}"
-    # speaker_groups[file_name] = [turn.start, turn.end]
-
-    # print(f"speaker_groups {file_name}: {speaker_groups[file_name]}")
-    # print(f"start={turn.start:.3f}s stop={turn.end:.3f}s speaker_{speaker}")
+    iterables = diarization.itertracks(yield_label=True)
+    for turn, _, speaker, _ in zip(iterables, progress.tqdm(iterables, desc="Processing diarization")):
+        if (speaker not in sample_groups):
+            sample_groups.append(str(speaker))
+
+        suffix = 1
+        file_name = f"{speaker}-{suffix}"
+        while file_name in speaker_groups:
+            suffix += 1
+            file_name = f"{speaker}-{suffix}"
+        speaker_groups[file_name] = [turn.start, turn.end]
+
+        print(f"speaker_groups {file_name}: {speaker_groups[file_name]}")
+        print(
+            f"start={turn.start:.3f}s stop={turn.end:.3f}s speaker_{speaker}")
 
     save_groups_json(sample_groups, speaker_groups)
     audio_segmentation(input_file, speaker_groups)

transcribe.py

@@ -10,8 +10,8 @@ model = WhisperModel("medium", device="cuda", compute_type="int8_float16")
 
 
 def start_transcribe(progress):
-    sample_groups, speaker_groups = load_groups_json()
-    for speaker in speaker_groups:
+    _, speaker_groups = load_groups_json()
+    for speaker, _ in zip(speaker_groups, progress.tqdm(speaker_groups, desc="Processing diarization")):
         # Transcribe and save temp file
         audiof = f"{speaker}.wav"
         print(f"Loading {audiof}")