Release v1.0 including v2_fast improvements and dynamic compute_type selection

app.py

@@ -22,7 +22,8 @@ with gr.Blocks() as demo:
     gr.Markdown(description_string)
     with gr.Row():
         with gr.Column(scale=1):
-            model_version = gr.Dropdown(label="Model Version", choices=["v2_fast", "v2.0"], value="v2_fast")
+            model_version = gr.Dropdown(label="Model Version", choices=["v2_fast", "v1.0"], value="v2_fast")
+
             input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
 
         with gr.Column(scale=1):

settings.py

@@ -1,6 +1,8 @@
 DEBUG_MODE = True
-MODEL_PATH_V2 = "langtech-veu/whisper-timestamped-cs"
+MODEL_PATH_V1 = "projecte-aina/whisper-large-v3-tiny-caesar"
 MODEL_PATH_V2_FAST = "langtech-veu/faster-whisper-timestamped-cs"
 LEFT_CHANNEL_TEMP_PATH = "temp_mono_speaker2.wav"
 RIGHT_CHANNEL_TEMP_PATH = "temp_mono_speaker1.wav"
-RESAMPLING_FREQ = 16000
+RESAMPLING_FREQ = 16000
+BATCH_SIZE = 1
+TASK = "transcribe"

whisper_cs_dev.py

@@ -1,5 +1,5 @@
 from faster_whisper import WhisperModel
-import whisper_timestamped as whisper_ts
+from transformers import pipeline
 from pydub import AudioSegment
 import os
 import torchaudio
@@ -10,8 +10,9 @@ import sys
 from pathlib import Path
 import glob
 import ctypes
+import numpy as np
 
-from settings import DEBUG_MODE, MODEL_PATH_V2_FAST, MODEL_PATH_V2, LEFT_CHANNEL_TEMP_PATH, RIGHT_CHANNEL_TEMP_PATH, RESAMPLING_FREQ
+from settings import DEBUG_MODE, MODEL_PATH_V2_FAST, MODEL_PATH_V1, LEFT_CHANNEL_TEMP_PATH, RIGHT_CHANNEL_TEMP_PATH, RESAMPLING_FREQ, BATCH_SIZE, TASK
 
 def load_cudnn():
 
@@ -54,44 +55,45 @@ def load_cudnn():
 
 def get_settings():
 
-    if DEBUG_MODE: print(f"Entering get_settings function...")
+    if DEBUG_MODE: 
+        print(f"Entering get_settings function...")
 
     is_cuda_available = torch.cuda.is_available()
     if is_cuda_available:
         device = "cuda"
-        compute_type = "float16"
+        compute_type = "default" 
+
     else:
         device = "cpu"
-        compute_type = "int8"
-    if DEBUG_MODE: print(f"is_cuda_available: {is_cuda_available}")
-    if DEBUG_MODE: print(f"device: {device}")
-    if DEBUG_MODE: print(f"compute_type: {compute_type}")
+        compute_type = "default" 
 
+    if DEBUG_MODE: print(f"[SETTINGS] Device: {device}")
     if DEBUG_MODE: print(f"Exited get_settings function.")
 
     return device, compute_type
 
 
+
 def load_model(use_v2_fast, device, compute_type):
 
-    if DEBUG_MODE: print(f"Entering load_model function...")
-    
-    if DEBUG_MODE: print(f"use_v2_fast: {use_v2_fast}")
+    if DEBUG_MODE: 
+        print(f"Entering load_model function...")
+        print(f"[MODEL LOADING] use_v2_fast: {use_v2_fast}")
 
     if use_v2_fast:    
-        if DEBUG_MODE: print(f"Loading {MODEL_PATH_V2_FAST} using {device} with {compute_type}...")
         model = WhisperModel(
             MODEL_PATH_V2_FAST,
             device = device,
             compute_type = compute_type,
         )
     else:
-        if DEBUG_MODE: print(f"Loading {MODEL_PATH_V2} using {device} with {compute_type}...")
-        # TODO add compute_type to load model
-        model = whisper_ts.load_model(
-            MODEL_PATH_V2, 
-            device = device,
-        )
+        model = pipeline(
+            task="automatic-speech-recognition",
+            model=MODEL_PATH_V1,
+            chunk_length_s=30,
+            device=device,
+            token=os.getenv("HF_TOKEN")
+            ) 
 
     if DEBUG_MODE: print(f"Exiting load_model function...")
     
@@ -109,21 +111,36 @@ def split_input_stereo_channels(audio_path):
     elif ext == ".mp3":
         audio = AudioSegment.from_file(audio_path, format="mp3")
     else:
-        raise ValueError(f"Unsupported file format for: {audio_path}")
+        raise ValueError(f"[FORMAT AUDIO] Unsupported file format for: {audio_path}")
 
     channels = audio.split_to_mono()
 
     if len(channels) != 2:
-        raise ValueError(f"Audio {audio_path} has {len(channels)} channels (instead of 2).")
+        raise ValueError(f"[FORMAT AUDIO] Audio {audio_path} has {len(channels)} channels (instead of 2).")
 
     channels[0].export(RIGHT_CHANNEL_TEMP_PATH, format="wav")  # Right
     channels[1].export(LEFT_CHANNEL_TEMP_PATH, format="wav")  # Left
 
     if DEBUG_MODE: print(f"Exited split_input_stereo_channels function.")
 
+def compute_type_to_audio_dtype(compute_type: str, device: str) -> np.dtype:
+    if DEBUG_MODE: print(f"Entering compute_type_to_audio_dtype function.")
+
+    compute_type = compute_type.lower()
+
+    if device.startswith("cuda"):
+        if "float16" in compute_type or "int8" in compute_type:
+            audio_np_dtype = np.float16
+        else:
+            audio_np_dtype = np.float32
+    else:
+        audio_np_dtype = np.float32
+
+    if DEBUG_MODE: print(f"Exited compute_type_to_audio_dtype function.")
+    return audio_np_dtype
 
-def format_audio(audio_path):
 
+def format_audio(audio_path: str, compute_type: str, device: str) -> np.ndarray:
     if DEBUG_MODE: print(f"Entering format_audio function...")
 
     input_audio, sample_rate = torchaudio.load(audio_path)
@@ -135,81 +152,47 @@ def format_audio(audio_path):
     input_audio = resampler(input_audio)
     input_audio = input_audio.squeeze()
 
-    if DEBUG_MODE: print(f"Exited format_audio function.")
-    
-    return input_audio, RESAMPLING_FREQ
+    np_dtype = compute_type_to_audio_dtype(compute_type, device)
 
+    input_audio = input_audio.numpy().astype(np_dtype)
 
-def process_waveforms():
+    if DEBUG_MODE: 
+        print(f"[FORMAT AUDIO] Audio dtype for actual_compute_type: {input_audio.dtype}")
+        print(f"Exited format_audio function.")
+    return input_audio
 
-    if DEBUG_MODE: print(f"Entering process_waveforms function...")
 
-    left_waveform, _ = format_audio(LEFT_CHANNEL_TEMP_PATH)
-    right_waveform, _ = format_audio(RIGHT_CHANNEL_TEMP_PATH)
 
-    # TODO should this be equal to compute_type?
-    left_waveform = left_waveform.numpy().astype("float16")
-    right_waveform = right_waveform.numpy().astype("float16")
+def process_waveforms(device: str, compute_type: str):
 
-    if DEBUG_MODE: print(f"Exited process_waveforms function.")
+    if DEBUG_MODE: print(f"Entering process_waveforms function...")
+
+    left_waveform  = format_audio(LEFT_CHANNEL_TEMP_PATH, compute_type, device)
+    right_waveform = format_audio(RIGHT_CHANNEL_TEMP_PATH, compute_type, device)
 
+    if DEBUG_MODE: print(f"Exited process_waveforms function.")
     return left_waveform, right_waveform
 
 
-def transcribe_audio_no_fast_model(model, audio_path):
+def transcribe_pipeline(audio, model):
+    if DEBUG_MODE: print(f"Entering transcribe_pipeline function.")
 
-    if DEBUG_MODE: print(f"Entering transcribe_audio_no_fast_model function...")
-    
-    result = whisper_ts.transcribe(
-        model,
-        audio_path,
-        beam_size=5,
-        best_of=5,
-        temperature=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
-        vad=False,  
-        detect_disfluencies=True,
-    )
-    
-    words = []
-    for segment in result.get('segments', []):
-        for word in segment.get('words', []):
-            word_text = word.get('word', '').strip()
-            if word_text.startswith(' '):
-                word_text = word_text[1:]
-            
-            words.append({
-                'word': word_text,
-                'start': word.get('start', 0),
-                'end': word.get('end', 0),
-                'confidence': word.get('confidence', 0)
-            })
-    
-    return {
-        'audio_path': audio_path,
-        'text': result['text'].strip(),
-        'segments': result.get('segments', []),
-        'words': words,
-        'duration': result.get('duration', 0),
-        'success': True
-    }
+    text = model(audio, batch_size=BATCH_SIZE, generate_kwargs={"task": TASK}, return_timestamps=True)["text"]
 
-    if DEBUG_MODE: print(f"Exited transcribe_audio_no_fast_model function.")
+    if DEBUG_MODE: print(f"Exited transcribe_pipeline function.")
 
+    return text
 
-def transcribe_channels(left_waveform, right_waveform, model, use_v2_fast):
+
+def transcribe_channels(left_waveform, right_waveform, model):
 
     if DEBUG_MODE: print(f"Entering transcribe_channels function...")
 
-    if DEBUG_MODE: print(f"Preparing to transcribe...")
+    left_result, _ = model.transcribe(left_waveform, beam_size=5, task="transcribe")
+    right_result, _ = model.transcribe(right_waveform, beam_size=5, task="transcribe")
 
-    if use_v2_fast:
-        left_result, _ = model.transcribe(left_waveform, beam_size=5, task="transcribe")
-        right_result, _ = model.transcribe(right_waveform, beam_size=5, task="transcribe")
-        left_result = list(left_result)
-        right_result = list(right_result)
-    else:
-        left_result = transcribe_audio_no_fast_model(model, left_waveform)
-        right_result = transcribe_audio_no_fast_model(model, right_waveform)
+    left_result = list(left_result)
+    right_result = list(right_result)
 
     if DEBUG_MODE: print(f"Exited transcribe_channels function.")
 
@@ -270,37 +253,27 @@ def post_merge_consecutive_segments_from_text(transcription_text: str) -> str:
     return merged_transcription.strip()
 
 
-def get_segments(result, speaker_label, use_v2_fast):
+def get_segments(result, speaker_label):
 
     if DEBUG_MODE: print(f"Entering get_segments function...")
 
-    if use_v2_fast:
-        segments = result
-        final_segments = [
-            (seg.start, seg.end, speaker_label, post_process_transcription(seg.text.strip()))
-            for seg in segments if seg.text
-        ]
-    else:
-        segments = result.get("segments", [])
-        if not segments:
-            final_segments = []
-        final_segments = [
-            (seg.get("start", 0.0), seg.get("end", 0.0), speaker_label,
-                post_process_transcription(seg.get("text", "").strip()))
-            for seg in segments if seg.get("text")
-        ]
+    segments = result
+    final_segments = [
+        (seg.start, seg.end, speaker_label, post_process_transcription(seg.text.strip()))
+        for seg in segments if seg.text
+    ]
 
     if DEBUG_MODE: print(f"EXited get_segments function.")
 
     return final_segments
         
 
-def post_process_transcripts(left_result, right_result, use_v2_fast):
+def post_process_transcripts(left_result, right_result):
 
     if DEBUG_MODE: print(f"Entering post_process_transcripts function...")
 
-    left_segs = get_segments(left_result, "Speaker 1", use_v2_fast)
-    right_segs = get_segments(right_result, "Speaker 2", use_v2_fast)
+    left_segs = get_segments(left_result, "Speaker 1")
+    right_segs = get_segments(right_result, "Speaker 2")
 
     merged_transcript = sorted(
         left_segs + right_segs,
@@ -320,8 +293,6 @@ def post_process_transcripts(left_result, right_result, use_v2_fast):
 def cleanup_temp_files(*file_paths):
 
     if DEBUG_MODE: print(f"Entered cleanup_temp_files function...")
-
-    if DEBUG_MODE: print(f"File paths to remove: {file_paths}")
         
     for path in file_paths:
         if path and os.path.exists(path):
@@ -331,28 +302,44 @@ def cleanup_temp_files(*file_paths):
     if DEBUG_MODE: print(f"Exited cleanup_temp_files function.")
 
 
-def generate(audio_path, use_v2_fast):
 
+
+def generate(audio_path, use_v2_fast):
     if DEBUG_MODE: print(f"Entering generate function...")
 
     start = time.time()
 
     load_cudnn()
-    device, compute_type = get_settings()
-    model = load_model(use_v2_fast, device, compute_type)
-    split_input_stereo_channels(audio_path)
-    left_waveform, right_waveform = process_waveforms()
-    left_result, right_result = transcribe_channels(left_waveform, right_waveform, model, use_v2_fast)
-    output = post_process_transcripts(left_result, right_result, use_v2_fast)
-    cleanup_temp_files(LEFT_CHANNEL_TEMP_PATH, RIGHT_CHANNEL_TEMP_PATH)
-
-    end = time.time()
-    elapsed_secs = end - start
+    device, requested_compute_type = get_settings()
+    model = load_model(use_v2_fast, device, requested_compute_type)
 
-    if DEBUG_MODE: print(f"elapsed_secs: {elapsed_secs}")
+    if use_v2_fast:
+        actual_compute_type = model.model.compute_type
+    else:
+        actual_compute_type = "float32" #HF pipeline safe default
+    
+    if DEBUG_MODE:
+        print(f"[SETTINGS] Requested compute_type: {requested_compute_type}")
+        print(f"[SETTINGS] Actual compute_type: {actual_compute_type}")
+    
+    if use_v2_fast:
+        split_input_stereo_channels(audio_path)
+        left_waveform, right_waveform = process_waveforms(device, actual_compute_type)
+        left_result, right_result = transcribe_channels(left_waveform, right_waveform, model)
+        output = post_process_transcripts(left_result, right_result)
+        cleanup_temp_files(LEFT_CHANNEL_TEMP_PATH, RIGHT_CHANNEL_TEMP_PATH)
+    else:
+        audio = format_audio(audio_path, actual_compute_type, device)
+        merged_results = transcribe_pipeline(audio, model)
+        output = post_process_transcription(merged_results)
 
-    if DEBUG_MODE: print(f"Exited generate function.")
+    end = time.time()
     
-    return output
+    audio_duration = torchaudio.info(audio_path).num_frames / torchaudio.info(audio_path).sample_rate
+    rtf = (end - start) / audio_duration
 
+    if DEBUG_MODE: print(f"[LATENCY]: {end - start}")
+    if DEBUG_MODE: print(f"[RTF]: {rtf:.2f}")
+    if DEBUG_MODE: print(f"Exited generate function.")
 
+    return output