fixes

app.py

@@ -1,20 +1,20 @@
-from concurrent.futures import ThreadPoolExecutor
-from pathlib import Path
 import csv
 import datetime
 import gc
 import os
 import re
 import shutil
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
 
 import gradio as gr
 import gradio.themes as gr_themes
-from huggingface_hub import hf_hub_download
-from nemo.collections.asr.models import ASRModel
 import numpy as np
-from pydub import AudioSegment
 import spaces
 import torch
+from huggingface_hub import hf_hub_download
+from nemo.collections.asr.models import ASRModel
+from pydub import AudioSegment
 
 try:
     from nemo.collections.asr.models import SortformerEncLabelModel
@@ -108,7 +108,9 @@ def get_audio_segment(audio_path, start_second, end_second):
             return None
         return frame_rate, samples
     except Exception as e:
-        print(f"Error clipping audio {audio_path} from {start_second}s to {end_second}s: {e}")
+        print(
+            f"Error clipping audio {audio_path} from {start_second}s to {end_second}s: {e}"
+        )
         return None
 
 
@@ -160,7 +162,9 @@ def remove_dc_offset(samples: np.ndarray) -> np.ndarray:
     return samples - np.mean(samples, dtype=np.float32)
 
 
-def fft_bandpass(samples: np.ndarray, sr: int, low_hz: float, high_hz: float) -> np.ndarray:
+def fft_bandpass(
+    samples: np.ndarray, sr: int, low_hz: float, high_hz: float
+) -> np.ndarray:
     samples = np.asarray(samples, dtype=np.float32)
     if samples.size == 0:
         return samples
@@ -205,7 +209,9 @@ def spectral_denoise(
     mask = np.clip(mask, min_mask, 1.0)
 
     cleaned_stft = magnitude * mask * np.exp(1j * phase)
-    cleaned = librosa.istft(cleaned_stft, hop_length=hop, win_length=n_fft, length=len(samples))
+    cleaned = librosa.istft(
+        cleaned_stft, hop_length=hop, win_length=n_fft, length=len(samples)
+    )
     return cleaned.astype(np.float32, copy=False)
 
 
@@ -289,7 +295,9 @@ def preprocess_audio_for_transcription(
     samples = np.clip(raw / max_abs, -1.0, 1.0)
 
     samples = remove_dc_offset(samples)
-    samples = spectral_denoise(samples, strength=1.25, noise_percentile=15.0, min_mask=0.06)
+    samples = spectral_denoise(
+        samples, strength=1.25, noise_percentile=15.0, min_mask=0.06
+    )
     samples = fft_bandpass(samples, sr=target_sr, low_hz=120.0, high_hz=3600.0)
     samples = dynamic_rms_normalize(
         samples=samples,
@@ -332,7 +340,11 @@ def _parse_rttm_line(line: str):
     speaker = parts[7]
     if start is None or dur is None or dur <= 0:
         return None
-    return {"start": start, "end": start + dur, "speaker": normalize_speaker_label(speaker)}
+    return {
+        "start": start,
+        "end": start + dur,
+        "speaker": normalize_speaker_label(speaker),
+    }
 
 
 def _parse_simple_segment_line(line: str):
@@ -360,7 +372,9 @@ def parse_diarization_output(raw_output, audio_duration_sec=None) -> list:
         e = _try_float(end)
         if s is None or e is None or e <= s:
             return
-        parsed.append({"start": s, "end": e, "speaker": normalize_speaker_label(speaker)})
+        parsed.append(
+            {"start": s, "end": e, "speaker": normalize_speaker_label(speaker)}
+        )
 
     def walk(obj):
         if obj is None:
@@ -413,7 +427,11 @@ def parse_diarization_output(raw_output, audio_duration_sec=None) -> list:
             return
 
         if isinstance(obj, (list, tuple)):
-            if len(obj) >= 3 and _try_float(obj[0]) is not None and _try_float(obj[1]) is not None:
+            if (
+                len(obj) >= 3
+                and _try_float(obj[0]) is not None
+                and _try_float(obj[1]) is not None
+            ):
                 append_seg(obj[0], obj[1], obj[2])
                 return
             for item in obj:
@@ -421,7 +439,9 @@ def parse_diarization_output(raw_output, audio_duration_sec=None) -> list:
             return
 
         if hasattr(obj, "start") and hasattr(obj, "end"):
-            append_seg(getattr(obj, "start"), getattr(obj, "end"), getattr(obj, "speaker", "0"))
+            append_seg(
+                getattr(obj, "start"), getattr(obj, "end"), getattr(obj, "speaker", "0")
+            )
 
     walk(raw_output)
 
@@ -453,7 +473,10 @@ def merge_adjacent_speaker_segments(segments: list, max_gap_sec: float = 0.15) -
     merged = [segments[0].copy()]
     for seg in segments[1:]:
         last = merged[-1]
-        if seg["speaker"] == last["speaker"] and seg["start"] - last["end"] <= max_gap_sec:
+        if (
+            seg["speaker"] == last["speaker"]
+            and seg["start"] - last["end"] <= max_gap_sec
+        ):
             last["end"] = max(last["end"], seg["end"])
         else:
             merged.append(seg.copy())
@@ -547,7 +570,9 @@ def transcribe_default_with_timestamps(transcribe_path: str):
     return segments
 
 
-def _overlap_seconds(a_start: float, a_end: float, b_start: float, b_end: float) -> float:
+def _overlap_seconds(
+    a_start: float, a_end: float, b_start: float, b_end: float
+) -> float:
     return max(0.0, min(a_end, b_end) - max(a_start, b_start))
 
 
@@ -555,7 +580,9 @@ def _join_tokens(tokens: list) -> str:
     return " ".join(t for t in tokens if t).strip()
 
 
-def split_asr_by_diarization_segments(asr_segments: list, diar_segments: list, asr_words: list = None) -> list:
+def split_asr_by_diarization_segments(
+    asr_segments: list, diar_segments: list, asr_words: list = None
+) -> list:
     if not diar_segments:
         return []
 
@@ -634,7 +661,9 @@ def _clean_token_spacing(text: str) -> str:
 
 
 def _capitalize_first_alpha(text: str) -> str:
-    return re.sub(r"^([^A-Za-z]*)([a-z])", lambda m: m.group(1) + m.group(2).upper(), text)
+    return re.sub(
+        r"^([^A-Za-z]*)([a-z])", lambda m: m.group(1) + m.group(2).upper(), text
+    )
 
 
 def _capitalize_after_full_stop(text: str) -> str:
@@ -702,9 +731,13 @@ UZ_ORDINAL_TO_CARDINAL = {
     "o'ninchi": "o'n",
     "oninchi": "o'n",
 }
-UZ_MONTHS_PATTERN = r"yanvar|fevral|mart|aprel|may|iyun|iyul|avgust|sentabr|oktabr|noyabr|dekabr"
+UZ_MONTHS_PATTERN = (
+    r"yanvar|fevral|mart|aprel|may|iyun|iyul|avgust|sentabr|oktabr|noyabr|dekabr"
+)
 
-_TOKEN_CORE_RE = re.compile(r"^([^A-Za-z0-9'`ʻʼ’‘]*)([A-Za-z0-9'`ʻʼ’‘]+)([^A-Za-z0-9'`ʻʼ’‘]*)$")
+_TOKEN_CORE_RE = re.compile(
+    r"^([^A-Za-z0-9'`ʻʼ’‘]*)([A-Za-z0-9'`ʻʼ’‘]+)([^A-Za-z0-9'`ʻʼ’‘]*)$"
+)
 
 
 def _normalize_uz_word(word: str) -> str:
@@ -729,7 +762,12 @@ def _normalize_uz_word(word: str) -> str:
 def _is_uz_number_like(word: str) -> bool:
     if not word:
         return False
-    if word in UZ_CARDINAL or word in UZ_SCALES or word == "yuz" or word in UZ_ORDINAL_TO_CARDINAL:
+    if (
+        word in UZ_CARDINAL
+        or word in UZ_SCALES
+        or word == "yuz"
+        or word in UZ_ORDINAL_TO_CARDINAL
+    ):
         return True
     return re.match(r"^.+(?:inchi|nchi)$", word) is not None
 
@@ -882,7 +920,9 @@ def normalize_uzbek_date_forms(text: str) -> str:
     return text
 
 
-def postprocess_segment_texts(segment_timestamps: list, diarization_enabled: bool) -> list:
+def postprocess_segment_texts(
+    segment_timestamps: list, diarization_enabled: bool
+) -> list:
     for ts in segment_timestamps:
         txt = str(ts.get("segment", "") or "")
         txt = _clean_token_spacing(txt)
@@ -905,7 +945,9 @@ def resolve_player_audio_path(prepared_path, fallback_path: str) -> str:
 
 
 @spaces.GPU
-def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=True, use_diarization=False):
+def get_transcripts_and_raw_times(
+    audio_path, session_dir, use_preprocessing=True, use_diarization=False
+):
     if not audio_path:
         gr.Error("No audio file path provided for transcription.", duration=None)
         return (
@@ -924,8 +966,12 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
     original_path_name = Path(audio_path).name
     audio_name = Path(audio_path).stem
 
-    csv_button_update = gr.DownloadButton(label="Download Transcript (CSV)", visible=False)
-    srt_button_update = gr.DownloadButton(label="Download Transcript (SRT)", visible=False)
+    csv_button_update = gr.DownloadButton(
+        label="Download Transcript (CSV)", visible=False
+    )
+    srt_button_update = gr.DownloadButton(
+        label="Download Transcript (SRT)", visible=False
+    )
 
     transcribe_path = audio_path
     info_path_name = original_path_name
@@ -952,7 +998,9 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                 processed_audio = preprocess_audio_for_transcription(
                     audio=audio, target_sr=16000, frame_ms=500, target_rms_db=-20.0
                 )
-                processed_audio_path = Path(session_dir, f"{audio_name}_asr_preprocessed.wav")
+                processed_audio_path = Path(
+                    session_dir, f"{audio_name}_asr_preprocessed.wav"
+                )
                 processed_audio.export(processed_audio_path, format="wav")
                 transcribe_path = processed_audio_path.as_posix()
                 info_path_name = f"{original_path_name} (preprocessed)"
@@ -974,12 +1022,17 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
 
             if duration_sec > 480:
                 try:
-                    gr.Info("Audio longer than 8 minutes. Applying long audio settings.", duration=3)
+                    gr.Info(
+                        "Audio longer than 8 minutes. Applying long audio settings.",
+                        duration=3,
+                    )
                     model.change_attention_model("rel_pos_local_attn", [256, 256])
                     model.change_subsampling_conv_chunking_factor(1)
                     long_audio_settings_applied = True
                 except Exception as setting_e:
-                    gr.Warning(f"Could not apply long audio settings: {setting_e}", duration=5)
+                    gr.Warning(
+                        f"Could not apply long audio settings: {setting_e}", duration=5
+                    )
 
             if device == "cuda":
                 model.to(torch.bfloat16)
@@ -990,16 +1043,7 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                 try:
                     gr.Info("Running ASR and diarization in parallel...", duration=3)
 
-                    diar_input_path = transcribe_path
-                    if not use_preprocessing:
-                        diar_audio_path = Path(session_dir, f"{audio_name}_diar_16k_mono.wav")
-                        diar_audio = audio
-                        if diar_audio.channels != 1:
-                            diar_audio = diar_audio.set_channels(1)
-                        if diar_audio.frame_rate != 16000:
-                            diar_audio = diar_audio.set_frame_rate(16000)
-                        diar_audio.export(diar_audio_path, format="wav")
-                        diar_input_path = diar_audio_path.as_posix()
+                    diar_input_path = audio_path
 
                     dmodel = get_diar_model()
                     dmodel.to(device)
@@ -1010,9 +1054,13 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
 
                     def _run_diar():
                         try:
-                            diar_output_local = dmodel.diarize(audio=diar_input_path, batch_size=1)
+                            diar_output_local = dmodel.diarize(
+                                audio=diar_input_path, batch_size=1
+                            )
                         except TypeError:
-                            diar_output_local = dmodel.diarize(audio=[diar_input_path], batch_size=1)
+                            diar_output_local = dmodel.diarize(
+                                audio=[diar_input_path], batch_size=1
+                            )
 
                         diar_segments_local = parse_diarization_output(
                             diar_output_local,
@@ -1042,17 +1090,27 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                             diar_segments=diar_segments,
                             asr_words=asr_words,
                         )
-                        segment_timestamps = merge_consecutive_transcript_rows(segment_timestamps)
+                        segment_timestamps = merge_consecutive_transcript_rows(
+                            segment_timestamps
+                        )
 
                         if not segment_timestamps:
-                            gr.Warning("No aligned diarized rows. Using ASR segmentation.", duration=7)
+                            gr.Warning(
+                                "No aligned diarized rows. Using ASR segmentation.",
+                                duration=7,
+                            )
                             segment_timestamps = asr_segments
 
                     gr.Info("Diarization + ASR complete.", duration=2)
 
                 except Exception as diar_e:
-                    gr.Warning(f"Diarization failed: {diar_e}. Using standard ASR segmentation.", duration=7)
-                    segment_timestamps = transcribe_default_with_timestamps(transcribe_path)
+                    gr.Warning(
+                        f"Diarization failed: {diar_e}. Using standard ASR segmentation.",
+                        duration=7,
+                    )
+                    segment_timestamps = transcribe_default_with_timestamps(
+                        transcribe_path
+                    )
             else:
                 segment_timestamps = transcribe_default_with_timestamps(transcribe_path)
 
@@ -1070,7 +1128,9 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                 ]
                 for ts in segment_timestamps
             ]
-            raw_times_data = [[float(ts["start"]), float(ts["end"])] for ts in segment_timestamps]
+            raw_times_data = [
+                [float(ts["start"]), float(ts["end"])] for ts in segment_timestamps
+            ]
 
             try:
                 csv_file_path = Path(session_dir, f"transcription_{audio_name}.csv")
@@ -1082,7 +1142,9 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                     value=csv_file_path, visible=True, label="Download Transcript (CSV)"
                 )
             except Exception as csv_e:
-                gr.Error(f"Failed to create transcript CSV file: {csv_e}", duration=None)
+                gr.Error(
+                    f"Failed to create transcript CSV file: {csv_e}", duration=None
+                )
 
             if segment_timestamps:
                 try:
@@ -1091,10 +1153,14 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                     with open(srt_file_path, "w", encoding="utf-8") as f:
                         f.write(srt_content)
                     srt_button_update = gr.DownloadButton(
-                        value=srt_file_path, visible=True, label="Download Transcript (SRT)"
+                        value=srt_file_path,
+                        visible=True,
+                        label="Download Transcript (SRT)",
                     )
                 except Exception as srt_e:
-                    gr.Warning(f"Failed to create transcript SRT file: {srt_e}", duration=5)
+                    gr.Warning(
+                        f"Failed to create transcript SRT file: {srt_e}", duration=5
+                    )
 
             gr.Info("Transcription complete.", duration=2)
             return (
@@ -1116,7 +1182,10 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                 srt_button_update,
             )
         except FileNotFoundError:
-            gr.Error(f"Audio file not found for transcription: {Path(transcribe_path).name}", duration=None)
+            gr.Error(
+                f"Audio file not found for transcription: {Path(transcribe_path).name}",
+                duration=None,
+            )
             return (
                 [["Error", "Error", "N/A", "File not found for transcription"]],
                 [[0.0, 0.0]],
@@ -1140,7 +1209,9 @@ def get_transcripts_and_raw_times(audio_path, session_dir, use_preprocessing=Tru
                         model.change_attention_model("rel_pos")
                         model.change_subsampling_conv_chunking_factor(-1)
                     except Exception as revert_e:
-                        gr.Warning(f"Issue reverting model settings: {revert_e}", duration=5)
+                        gr.Warning(
+                            f"Issue reverting model settings: {revert_e}", duration=5
+                        )
 
                 if device == "cuda":
                     model.cpu()
@@ -1230,7 +1301,9 @@ nvidia_theme = gr_themes.Default(
 
 with gr.Blocks(theme=nvidia_theme) as demo:
     model_display_name = MODEL_NAME.split("/")[-1] if "/" in MODEL_NAME else MODEL_NAME
-    gr.Markdown(f"<h1 style='text-align:center;margin:0 auto;'>Speech Transcription with {model_display_name}</h1>")
+    gr.Markdown(
+        f"<h1 style='text-align:center;margin:0 auto;'>Speech Transcription with {model_display_name}</h1>"
+    )
     gr.HTML(article)
 
     current_audio_path_state = gr.State(None)
@@ -1248,18 +1321,32 @@ with gr.Blocks(theme=nvidia_theme) as demo:
 
     with gr.Tabs():
         with gr.TabItem("Audio File"):
-            file_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio File")
-            gr.Examples(examples=examples, inputs=[file_input], label="Example Audio Files")
-            file_transcribe_btn = gr.Button("Transcribe Uploaded File", variant="primary")
+            file_input = gr.Audio(
+                sources=["upload"], type="filepath", label="Upload Audio File"
+            )
+            gr.Examples(
+                examples=examples, inputs=[file_input], label="Example Audio Files"
+            )
+            file_transcribe_btn = gr.Button(
+                "Transcribe Uploaded File", variant="primary"
+            )
 
         with gr.TabItem("Microphone"):
-            mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Record Audio")
-            mic_transcribe_btn = gr.Button("Transcribe Microphone Input", variant="primary")
+            mic_input = gr.Audio(
+                sources=["microphone"], type="filepath", label="Record Audio"
+            )
+            mic_transcribe_btn = gr.Button(
+                "Transcribe Microphone Input", variant="primary"
+            )
 
     gr.Markdown("---")
     with gr.Row():
-        download_btn_csv = gr.DownloadButton(label="Download Transcript (CSV)", visible=False)
-        download_btn_srt = gr.DownloadButton(label="Download Transcript (SRT)", visible=False)
+        download_btn_csv = gr.DownloadButton(
+            label="Download Transcript (CSV)", visible=False
+        )
+        download_btn_srt = gr.DownloadButton(
+            label="Download Transcript (SRT)", visible=False
+        )
 
     vis_timestamps_df = gr.DataFrame(
         headers=["Start (s)", "End (s)", "Speaker", "Segment"],