Create app.py

app.py

@@ -0,0 +1,309 @@
+import os
+import time
+import json
+import random
+import string
+import pathlib
+import tempfile
+import logging
+
+import torch
+import whisperx
+import librosa
+import numpy as np
+import requests
+
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import JSONResponse
+
+app = FastAPI(title="WhisperX API")
+
+# -------------------------------
+# Logging and Model Setup
+# -------------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("whisperx_api")
+
+device = "cpu"
+compute_type = "int8"
+torch.set_num_threads(os.cpu_count())
+
+# Pre-load models for different sizes
+models = {
+    "tiny": whisperx.load_model("tiny", device, compute_type=compute_type, vad_method='silero'),
+    "base": whisperx.load_model("base", device, compute_type=compute_type, vad_method='silero'),
+    "small": whisperx.load_model("small", device, compute_type=compute_type, vad_method='silero'),
+    "large": whisperx.load_model("large", device, compute_type=compute_type, vad_method='silero'),
+    "large-v2": whisperx.load_model("large-v2", device, compute_type=compute_type, vad_method='silero'),
+    "large-v3": whisperx.load_model("large-v3", device, compute_type=compute_type, vad_method='silero'),
+}
+
+def seconds_to_srt_time(seconds: float) -> str:
+    """Convert seconds (float) into SRT timestamp format (HH:MM:SS,mmm)."""
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millis = int((seconds - int(seconds)) * 1000)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
+
+# -------------------------------
+# Vocal Extraction Function
+# -------------------------------
+def get_vocals(input_file):
+    try:
+        session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(11))
+        file_id = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(11))
+        file_content = pathlib.Path(input_file).read_bytes()
+        file_len = len(file_content)
+        r = requests.post(
+            f'https://politrees-audio-separator-uvr.hf.space/gradio_api/upload?upload_id={file_id}', 
+            files={'files': open(input_file, 'rb')}
+        )
+        json_data = r.json()
+
+        headers = {
+            'accept': '*/*',
+            'accept-language': 'en-US,en;q=0.5',
+            'content-type': 'application/json',
+            'origin': 'https://politrees-audio-separator-uvr.hf.space',
+            'priority': 'u=1, i',
+            'referer': 'https://politrees-audio-separator-uvr.hf.space/?__theme=system',
+            'sec-ch-ua': '"Not(A:Brand";v="99", "Brave";v="133", "Chromium";v="133"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Windows"',
+            'sec-fetch-dest': 'empty',
+            'sec-fetch-mode': 'cors',
+            'sec-fetch-site': 'same-origin',
+            'sec-fetch-storage-access': 'none',
+            'sec-gpc': '1',
+            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36',
+        }
+
+        params = {
+            '__theme': 'system',
+        }
+
+        json_payload = {
+            'data': [
+                {
+                    'path': json_data[0],
+                    'url': 'https://politrees-audio-separator-uvr.hf.space/gradio_api/file=' + json_data[0],
+                    'orig_name': pathlib.Path(input_file).name,
+                    'size': file_len,
+                    'mime_type': 'audio/wav',
+                    'meta': {'_type': 'gradio.FileData'},
+                },
+                'MelBand Roformer | Vocals by Kimberley Jensen',
+                256,
+                False,
+                5,
+                0,
+                '/tmp/audio-separator-models/',
+                'output',
+                'wav',
+                0.9,
+                0,
+                1,
+                'NAME_(STEM)_MODEL',
+                'NAME_(STEM)_MODEL',
+                'NAME_(STEM)_MODEL',
+                'NAME_(STEM)_MODEL',
+                'NAME_(STEM)_MODEL',
+                'NAME_(STEM)_MODEL',
+                'NAME_(STEM)_MODEL',
+            ],
+            'event_data': None,
+            'fn_index': 5,
+            'trigger_id': 28,
+            'session_hash': session_hash,
+        }
+
+        response = requests.post(
+            'https://politrees-audio-separator-uvr.hf.space/gradio_api/queue/join',
+            params=params,
+            headers=headers,
+            json=json_payload,
+        )
+
+        max_retries = 5
+        retry_delay = 5
+        retry_count = 0
+        while retry_count < max_retries:
+            try:
+                logger.info(f"Connecting to stream... Attempt {retry_count + 1}")
+                r = requests.get(
+                    f'https://politrees-audio-separator-uvr.hf.space/gradio_api/queue/data?session_hash={session_hash}',
+                    stream=True
+                )
+                if r.status_code != 200:
+                    raise Exception(f"Failed to connect: HTTP {r.status_code}")
+                logger.info("Connected successfully.")
+                for line in r.iter_lines():
+                    if line:
+                        json_resp = json.loads(line.decode('utf-8').replace('data: ', ''))
+                        logger.info(json_resp)
+                        if 'process_completed' in json_resp['msg']:
+                            logger.info("Process completed.")
+                            output_url = json_resp['output']['data'][1]['url']
+                            logger.info(f"Output URL: {output_url}")
+                            return output_url
+                logger.info("Stream ended prematurely. Reconnecting...")
+            except Exception as e:
+                logger.error(f"Error occurred: {e}. Retrying...")
+            retry_count += 1
+            time.sleep(retry_delay)
+        logger.error("Max retries reached. Exiting.")
+        return None
+    except Exception as ex:
+        logger.error(f"Unexpected error in get_vocals: {ex}")
+        return None
+
+def split_audio_by_pause(audio, sr, pause_threshold, top_db=30, energy_threshold=0.03):
+    intervals = librosa.effects.split(audio, top_db=top_db)
+    merged_intervals = []
+    current_start, current_end = intervals[0]
+    for start, end in intervals[1:]:
+        gap_duration = (start - current_end) / sr
+        if gap_duration < pause_threshold:
+            current_end = end
+        else:
+            merged_intervals.append((current_start, current_end))
+            current_start, current_end = start, end
+    merged_intervals.append((current_start, current_end))
+    # Filter out segments with low average RMS energy
+    filtered_intervals = []
+    for start, end in merged_intervals:
+        segment = audio[start:end]
+        rms = np.mean(librosa.feature.rms(y=segment))
+        if rms >= energy_threshold:
+            filtered_intervals.append((start, end))
+    return filtered_intervals
+
+# -------------------------------
+# Main Transcription Function
+# -------------------------------
+def transcribe(audio_file, model_size="base", debug=False, pause_threshold=0.0, vocal_extraction=False, language="en"):
+    start_time = time.time()
+    srt_output = ""
+    debug_log = []
+    subtitle_index = 1
+
+    try:
+        # Optionally extract vocals first
+        if vocal_extraction:
+            debug_log.append("Vocal extraction enabled; processing input file for vocals...")
+            extracted_url = get_vocals(audio_file)
+            if extracted_url is not None:
+                debug_log.append("Vocal extraction succeeded; downloading extracted audio...")
+                response = requests.get(extracted_url)
+                if response.status_code == 200:
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
+                        tmp.write(response.content)
+                        audio_file = tmp.name
+                    debug_log.append("Extracted audio downloaded and saved for transcription.")
+                else:
+                    debug_log.append("Failed to download extracted audio; proceeding with original file.")
+            else:
+                debug_log.append("Vocal extraction failed; proceeding with original audio.")
+
+        # Load audio file (resampled to 16kHz)
+        audio, sr = librosa.load(audio_file, sr=16000)
+        debug_log.append(f"Audio loaded: {len(audio)/sr:.2f} seconds at {sr} Hz")
+
+        # Select model and set batch size
+        model = models[model_size]
+        batch_size = 8 if model_size == "tiny" else 4
+
+        # Transcribe using specified language (or auto-detect)
+        if language:
+            transcript = model.transcribe(audio, batch_size=batch_size, language=language)
+        else:
+            transcript = model.transcribe(audio, batch_size=batch_size)
+            language = transcript.get("language", "unknown")
+
+        # Load alignment model for the given language
+        model_a, metadata = whisperx.load_align_model(language_code=language, device=device)
+
+        if pause_threshold > 0:
+            segments = split_audio_by_pause(audio, sr, pause_threshold)
+            debug_log.append(f"Audio split into {len(segments)} segment(s) using pause threshold of {pause_threshold}s")
+            for seg_idx, (seg_start, seg_end) in enumerate(segments):
+                audio_segment = audio[seg_start:seg_end]
+                seg_duration = (seg_end - seg_start) / sr
+                debug_log.append(f"Segment {seg_idx+1}: start={seg_start/sr:.2f}s, duration={seg_duration:.2f}s")
+                seg_transcript = model.transcribe(audio_segment, batch_size=batch_size, language=language)
+                seg_aligned = whisperx.align(
+                    seg_transcript["segments"], model_a, metadata, audio_segment, device
+                )
+                for segment in seg_aligned["segments"]:
+                    for word in segment["words"]:
+                        adjusted_start = word['start'] + seg_start/sr
+                        adjusted_end = word['end'] + seg_start/sr
+                        start_timestamp = seconds_to_srt_time(adjusted_start)
+                        end_timestamp = seconds_to_srt_time(adjusted_end)
+                        srt_output += f"{subtitle_index}\n{start_timestamp} --> {end_timestamp}\n{word['word']}\n\n"
+                        subtitle_index += 1
+        else:
+            # Process the entire audio without splitting
+            transcript = model.transcribe(audio, batch_size=batch_size, language=language)
+            aligned = whisperx.align(
+                transcript["segments"], model_a, metadata, audio, device
+            )
+            for segment in aligned["segments"]:
+                for word in segment["words"]:
+                    start_timestamp = seconds_to_srt_time(word['start'])
+                    end_timestamp = seconds_to_srt_time(word['end'])
+                    srt_output += f"{subtitle_index}\n{start_timestamp} --> {end_timestamp}\n{word['word']}\n\n"
+                    subtitle_index += 1
+
+        debug_log.append(f"Language used: {language}")
+        debug_log.append(f"Batch size: {batch_size}")
+        debug_log.append(f"Processed in {time.time()-start_time:.2f}s")
+
+    except Exception as e:
+        logger.error("Error during transcription:", exc_info=True)
+        srt_output = "Error occurred during transcription"
+        debug_log.append(f"ERROR: {str(e)}")
+
+    if debug:
+        return srt_output, "\n".join(debug_log)
+    return srt_output
+
+# -------------------------------
+# FastAPI Endpoints
+# -------------------------------
+@app.post("/transcribe")
+async def transcribe_endpoint(
+    audio_file: UploadFile = File(...),
+    model_size: str = Form("base"),
+    debug: bool = Form(False),
+    pause_threshold: float = Form(0.0),
+    vocal_extraction: bool = Form(False),
+    language: str = Form("en")
+):
+    try:
+        # Save the uploaded file to a temporary location
+        suffix = pathlib.Path(audio_file.filename).suffix
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+            tmp.write(await audio_file.read())
+            tmp_path = tmp.name
+
+        result = transcribe(tmp_path, model_size=model_size, debug=debug,
+                            pause_threshold=pause_threshold,
+                            vocal_extraction=vocal_extraction,
+                            language=language)
+
+        os.remove(tmp_path)
+
+        if debug:
+            srt_text, debug_info = result
+            return JSONResponse(content={"srt": srt_text, "debug": debug_info})
+        else:
+            return JSONResponse(content={"srt": result})
+    except Exception as e:
+        logger.error(f"Error in transcribe_endpoint: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+@app.get("/")
+async def root():
+    return {"message": "WhisperX API is running."}