Create app.py

app.py

@@ -0,0 +1,226 @@
+import os
+import gc
+import logging
+from typing import Any, Dict
+
+import torch
+import yt_dlp
+import gradio as gr
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from huggingface_hub import login, InferenceClient
+
+# Set up basic logging.
+logging.basicConfig(level=logging.INFO)
+
+# -------------------------------
+# Download Audio from Video URL
+# -------------------------------
+def download_audio(url: str) -> str:
+    """
+    Download audio from a video URL and convert it to MP3 format.
+    """
+    ydl_opts = {
+        'format': 'bestaudio/best',
+        'postprocessors': [{
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'mp3',
+            'preferredquality': '192',
+        }],
+    }
+    try:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info = ydl.extract_info(url, download=True)
+            audio_file = ydl.prepare_filename(info)
+            if not audio_file.endswith('.mp3'):
+                audio_file = audio_file.rsplit('.', 1)[0] + '.mp3'
+        logging.info("Audio downloaded successfully: %s", audio_file)
+        return audio_file
+    except Exception as e:
+        logging.error("Error downloading audio: %s", e)
+        raise RuntimeError("Audio download failed") from e
+
+# ---------------------------------------
+# Set Up Speech Recognition Model & Pipe
+# ---------------------------------------
+if torch.cuda.is_available():
+    model_device = "cuda"
+    pipeline_device = 0  # GPU device index for Hugging Face pipeline.
+    torch_dtype = torch.float16
+    speech_model_id = "openai/whisper-large-v3-turbo"
+    batch_size = 16
+else:
+    model_device = "cpu"
+    pipeline_device = -1  # CPU for pipeline.
+    torch_dtype = torch.float32
+    speech_model_id = "openai/whisper-tiny"
+    batch_size = 2
+
+try:
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        speech_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+    )
+    model.to(model_device)
+    processor = AutoProcessor.from_pretrained(speech_model_id)
+except Exception as e:
+    logging.error("Error loading the speech model: %s", e)
+    raise
+
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model=model,
+    tokenizer=processor.tokenizer,
+    feature_extractor=processor.feature_extractor,
+    torch_dtype=torch_dtype,
+    device=pipeline_device,
+)
+
+# --------------------------------------
+# Transcription and SRT Conversion
+# --------------------------------------
+def transcribe_audio(audio_path: str, batch_size: int) -> Dict[str, Any]:
+    """
+    Transcribe the audio file using the configured pipeline.
+    """
+    try:
+        result = pipe(
+            audio_path,
+            chunk_length_s=10,
+            stride_length_s=(4, 2),
+            batch_size=batch_size,
+            return_timestamps=True,
+        )
+        return result
+    except Exception as e:
+        logging.error("Error during transcription: %s", e)
+        raise
+
+def seconds_to_srt_time(seconds: float) -> str:
+    """
+    Convert seconds to SRT time format (HH:MM:SS,mmm).
+    """
+    if seconds is None or not isinstance(seconds, (int, float)):
+        return "00:00:00,000"
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millis = int((seconds - int(seconds)) * 1000)
+    return f"{hours:02}:{minutes:02}:{secs:02},{millis:03}"
+
+def convert_to_srt(transcribed: Dict[str, Any]) -> str:
+    """
+    Convert transcription chunks into SRT format.
+    """
+    srt_output = []
+    if "chunks" in transcribed:
+        for i, chunk in enumerate(transcribed["chunks"], start=1):
+            if chunk.get("timestamp") is not None:
+                start_time = seconds_to_srt_time(chunk["timestamp"][0])
+                end_time = seconds_to_srt_time(chunk["timestamp"][1])
+                srt_output.append(f"{i}\n{start_time} --> {end_time}\n{chunk['text']}\n")
+            else:
+                srt_output.append(f"{i}\n{chunk['text']}\n")
+        return "\n".join(srt_output)
+    else:
+        logging.warning("No chunks found; returning plain text.")
+        return transcribed.get("text", "")
+
+# ------------------------------
+# Hugging Face Login Adjustment
+# ------------------------------
+def hf_login() -> None:
+    """
+    Log in to Hugging Face using the token from environment variables.
+    """
+    huggingface_api_token = os.environ.get('HF_TOKEN')
+    if not huggingface_api_token:
+        raise ValueError("HF_TOKEN not set in environment variables.")
+    login(token=huggingface_api_token)
+    logging.info("Logged in to Hugging Face successfully.")
+
+# Log in once (this can be done at startup)
+hf_login()
+
+# -------------------------------------------
+# Generate Video Chapters from the Transcript
+# -------------------------------------------
+def generate_chapters(srt_text: str) -> str:
+    """
+    Generate video chapters from the SRT transcript using a text generation model.
+    """
+    chapter_model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"  # or another model if desired
+    client = InferenceClient(model=chapter_model_id)
+
+    prompt = (
+        "Based on the following video transcript, generate a numbered list of concise, SEO-friendly video chapters with timestamps. "
+        "Keep related parts together to limit the number of chapters (up to 5-10 chapters). "
+        "Each chapter should be in the format '<timestamp> <chapter title>', where the first chapter starts at 0:00. "
+        "Timestamps should be in the format 'm:ss' as needed. For example:\n\n"
+        "0:00 Intro\n"
+        "1:34 Why the GPT wrapper is bad\n"
+        "2:14 Smart users workflow\n\n"
+        "Only output the chapters list in the provided format. Stop after one list.\n"
+        "Transcript:\n"
+        f"{srt_text}\n\n"
+        "Chapters:"
+    )
+
+    generation_parameters = {
+        "max_new_tokens": 300,
+        "temperature": 0.5,
+        "top_p": 0.95,
+        "do_sample": True,
+    }
+
+    try:
+        generated_text = client.text_generation(prompt, **generation_parameters)
+        return generated_text
+    except Exception as e:
+        logging.error("Error generating chapters: %s", e)
+        raise
+
+# -------------------------------------------
+# Main Processing Function for Gradio UI
+# -------------------------------------------
+def process_video(video_url: str):
+    # Download audio from the provided URL.
+    audio_file = download_audio(video_url)
+    logging.info("Audio file saved as: %s", audio_file)
+    
+    # Transcribe the audio.
+    transcribed_text = transcribe_audio(audio_file, batch_size)
+    
+    # Clean up memory.
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    # Convert transcription to SRT format.
+    srt_text = convert_to_srt(transcribed_text)
+    
+    # Generate chapters from the SRT.
+    chapters = generate_chapters(srt_text)
+    
+    return srt_text, chapters
+
+# -------------------------------------------
+# Gradio Interface Definition
+# -------------------------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("# Video Chapter Generator")
+    
+    with gr.Row():
+        video_url_input = gr.Textbox(label="Video URL", placeholder="Enter video URL here", lines=1)
+    
+    with gr.Row():
+        process_button = gr.Button("Process Video")
+    
+    with gr.Row():
+        srt_output = gr.Textbox(label="SRT Transcript", interactive=False, lines=15, show_copy_button=True)
+    
+    with gr.Row():
+        chapters_output = gr.Textbox(label="Generated Chapters", interactive=False, lines=10, show_copy_button=True)
+    
+    process_button.click(fn=process_video, inputs=video_url_input, outputs=[srt_output, chapters_output])
+
+# Launch the Gradio app
+demo.launch()