import logging import math import os import tempfile import time from typing import Dict, Any import yt_dlp as youtube_dl from fastapi import FastAPI, UploadFile, Form, HTTPException from fastapi.responses import HTMLResponse import jax.numpy as jnp import numpy as np from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE from transformers.pipelines.audio_utils import ffmpeg_read from whisper_jax import FlaxWhisperPipline app = FastAPI(title="Whisper JAX: The Fastest Whisper API ⚡️") logger = logging.getLogger("whisper-jax-app") logger.setLevel(logging.DEBUG) ch = logging.StreamHandler() ch.setLevel(logging.DEBUG) formatter = logging.Formatter("%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d %H:%M:%S") ch.setFormatter(formatter) logger.addHandler(ch) checkpoint = "openai/whisper-large-v3" BATCH_SIZE = 32 CHUNK_LENGTH_S = 30 NUM_PROC = 32 FILE_LIMIT_MB = 10000 YT_LENGTH_LIMIT_S = 15000 # limit to 2 hour YouTube files pipeline = FlaxWhisperPipline(checkpoint, dtype=jnp.bfloat16, batch_size=BATCH_SIZE) stride_length_s = CHUNK_LENGTH_S / 6 chunk_len = round(CHUNK_LENGTH_S * pipeline.feature_extractor.sampling_rate) stride_left = stride_right = round(stride_length_s * pipeline.feature_extractor.sampling_rate) step = chunk_len - stride_left - stride_right # do a pre-compile step so that the first user to use the demo isn't hit with a long transcription time logger.debug("Compiling forward call...") start = time.time() random_inputs = { "input_features": np.ones( (BATCH_SIZE, pipeline.model.config.num_mel_bins, 2 * pipeline.model.config.max_source_positions) ) } random_timestamps = pipeline.forward(random_inputs, batch_size=BATCH_SIZE, return_timestamps=True) compile_time = time.time() - start logger.debug(f"Compiled in {compile_time}s") def timeit(func): async def wrapper(*args, **kwargs): start_time = time.time() result = await func(*args, **kwargs) end_time = time.time() execution_time = end_time - start_time if isinstance(result, dict): result['total_execution_time'] = execution_time else: result = {'result': result, 'total_execution_time': execution_time} return result return wrapper @app.post("/transcribe_audio") @timeit async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcribe", return_timestamps: bool = False) -> Dict[str, Any]: logger.debug("Starting transcribe_chunked_audio function") logger.debug(f"Received parameters - task: {task}, return_timestamps: {return_timestamps}") logger.debug("Checking for audio file...") if not audio_file: logger.warning("No audio file") raise HTTPException(status_code=400, detail="No audio file submitted!") logger.debug(f"Audio file received: {audio_file.filename}") try: # Read the file content file_content = await audio_file.read() file_size = len(file_content) file_size_mb = file_size / (1024 * 1024) logger.debug(f"File size: {file_size} bytes ({file_size_mb:.2f}MB)") except Exception as e: logger.error(f"Error reading file: {str(e)}", exc_info=True) raise HTTPException(status_code=500, detail=f"Error reading file: {str(e)}") if file_size_mb > FILE_LIMIT_MB: logger.warning(f"Max file size exceeded: {file_size_mb:.2f}MB > {FILE_LIMIT_MB}MB") raise HTTPException(status_code=400, detail=f"File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB.") try: logger.debug("Performing ffmpeg read on audio file") inputs = ffmpeg_read(file_content, pipeline.feature_extractor.sampling_rate) inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate} logger.debug("ffmpeg read completed successfully") except Exception as e: logger.error(f"Error in ffmpeg read: {str(e)}", exc_info=True) raise HTTPException(status_code=500, detail=f"Error processing audio file: {str(e)}") logger.debug("Calling tqdm_generate to transcribe audio") try: text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps) logger.debug(f"Transcription completed. Runtime: {runtime:.2f}s") except Exception as e: logger.error(f"Error in tqdm_generate: {str(e)}", exc_info=True) raise HTTPException(status_code=500, detail=f"Error transcribing audio: {str(e)}") logger.debug("Transcribe_chunked_audio function completed successfully") return {"text": text, "runtime": runtime, "timing_info": timing_info} @app.post("/transcribe_youtube") @timeit async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe", return_timestamps: bool = False) -> Dict[str, Any]: logger.debug("Loading YouTube file...") try: html_embed_str = _return_yt_html_embed(yt_url) except Exception as e: logger.error("Error generating YouTube HTML embed:", exc_info=True) raise HTTPException(status_code=500, detail="Error generating YouTube HTML embed") with tempfile.TemporaryDirectory() as tmpdirname: filepath = os.path.join(tmpdirname, "video.mp4") try: logger.debug("Downloading YouTube audio...") download_yt_audio(yt_url, filepath) except Exception as e: logger.error("Error downloading YouTube audio:", exc_info=True) raise HTTPException(status_code=500, detail="Error downloading YouTube audio") try: logger.debug(f"Opening downloaded audio file: {filepath}") with open(filepath, "rb") as f: inputs = f.read() except Exception as e: logger.error("Error reading downloaded audio file:", exc_info=True) raise HTTPException(status_code=500, detail="Error reading downloaded audio file") inputs = ffmpeg_read(inputs, pipeline.feature_extractor.sampling_rate) inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate} logger.debug("Done loading YouTube file") try: logger.debug("Calling tqdm_generate to transcribe YouTube audio") text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps) except Exception as e: logger.error("Error transcribing YouTube audio:", exc_info=True) raise HTTPException(status_code=500, detail="Error transcribing YouTube audio") return {"html_embed": html_embed_str, "text": text, "runtime": runtime, "timing_info": timing_info} def tqdm_generate(inputs: dict, task: str, return_timestamps: bool): start_time = time.time() logger.debug(f"Starting tqdm_generate - task: {task}, return_timestamps: {return_timestamps}") inputs_len = inputs["array"].shape[0] logger.debug(f"Input array length: {inputs_len}") all_chunk_start_idx = np.arange(0, inputs_len, step) num_samples = len(all_chunk_start_idx) num_batches = math.ceil(num_samples / BATCH_SIZE) logger.debug(f"Number of samples: {num_samples}, Number of batches: {num_batches}") logger.debug("Preprocessing audio for inference") try: dataloader = pipeline.preprocess_batch(inputs, chunk_length_s=CHUNK_LENGTH_S, batch_size=BATCH_SIZE) logger.debug("Preprocessing completed successfully") except Exception as e: logger.error(f"Error in preprocessing: {str(e)}", exc_info=True) raise model_outputs = [] transcription_start_time = time.time() logger.debug("Starting transcription...") try: for i, batch in enumerate(dataloader): logger.debug(f"Processing batch {i+1}/{num_batches} with {len(batch)} samples") batch_output = pipeline.forward(batch, batch_size=BATCH_SIZE, task=task, return_timestamps=True) model_outputs.append(batch_output) logger.debug(f"Batch {i+1} processed successfully") except Exception as e: logger.error(f"Error during batch processing: {str(e)}", exc_info=True) raise transcription_runtime = time.time() - transcription_start_time logger.debug(f"Transcription completed in {transcription_runtime:.2f}s") logger.debug("Post-processing transcription results") try: post_processed = pipeline.postprocess(model_outputs, return_timestamps=True) logger.debug("Post-processing completed successfully") except Exception as e: logger.error(f"Error in post-processing: {str(e)}", exc_info=True) raise text = post_processed["text"] if return_timestamps: timestamps = post_processed.get("chunks") timestamps = [ f"[{format_timestamp(chunk['timestamp'][0])} -> {format_timestamp(chunk['timestamp'][1])}] {chunk['text']}" for chunk in timestamps ] text = "\n".join(str(feature) for feature in timestamps) total_processing_time = time.time() - start_time logger.debug("tqdm_generate function completed successfully") return text, transcription_runtime, { "transcription_time": transcription_runtime, "total_processing_time": total_processing_time } def _return_yt_html_embed(yt_url): video_id = yt_url.split("?v=")[-1] HTML_str = ( f'
' "
" ) return HTML_str def download_yt_audio(yt_url, filename): info_loader = youtube_dl.YoutubeDL() try: logger.debug(f"Extracting info for YouTube URL: {yt_url}") info = info_loader.extract_info(yt_url, download=False) except youtube_dl.utils.DownloadError as err: logger.error("Error extracting YouTube info:", exc_info=True) raise HTTPException(status_code=400, detail=str(err)) file_length = info["duration_string"] file_h_m_s = file_length.split(":") file_h_m_s = [int(sub_length) for sub_length in file_h_m_s] if len(file_h_m_s) == 1: file_h_m_s.insert(0, 0) if len(file_h_m_s) == 2: file_h_m_s.insert(0, 0) file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2] if file_length_s > YT_LENGTH_LIMIT_S: yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S)) file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s)) raise HTTPException(status_code=400, detail=f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.") ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"} with youtube_dl.YoutubeDL(ydl_opts) as ydl: try: logger.debug(f"Downloading YouTube audio to {filename}") ydl.download([yt_url]) except youtube_dl.utils.ExtractorError as err: logger.error("Error downloading YouTube audio:", exc_info=True) raise HTTPException(status_code=400, detail=str(err)) def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."): if seconds is not None: milliseconds = round(seconds * 1000.0) hours = milliseconds // 3_600_000 milliseconds -= hours * 3_600_000 minutes = milliseconds // 60_000 milliseconds -= minutes * 60_000 seconds = milliseconds // 1_000 milliseconds -= seconds * 1_000 hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else "" return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}" else: # we have a malformed timestamp so just return it as is return seconds