init

README.md

@@ -1,8 +1,8 @@
 ---
 title: Whisper Transcribe
-emoji: 🌖
-colorFrom: yellow
-colorTo: green
+emoji: 🎙️
+colorFrom: blue
+colorTo: purple
 sdk: gradio
 sdk_version: 5.35.0
 app_file: app.py

app.py

@@ -1,7 +1,461 @@
 import gradio as gr
+import torch
+import torchaudio
+import numpy as np
+import pandas as pd
+import time
+import datetime
+import re
+import subprocess
+import os
+import tempfile
+import spaces
+from transformers import pipeline
+from pyannote.audio import Pipeline
+import requests
+import base64
+from typing import List, Dict, Any, Optional, Tuple
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Install flash attention for acceleration
+try:
+    subprocess.run(
+        "pip install flash-attn --no-build-isolation",
+        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+        shell=True,
+        check=True
+    )
+except subprocess.CalledProcessError:
+    print("Warning: Could not install flash-attn, falling back to default attention")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+class WhisperTranscriber:
+    def __init__(self):
+        self.pipe = None
+        self.diarization_model = None
+        
+    @spaces.GPU
+    def setup_models(self):
+        """Initialize models with GPU acceleration"""
+        if self.pipe is None:
+            print("Loading Whisper model...")
+            self.pipe = pipeline(
+                "automatic-speech-recognition",
+                model="openai/whisper-large-v3-turbo",
+                torch_dtype=torch.float16,
+                device="cuda:0",
+                model_kwargs={"attn_implementation": "flash_attention_2"},
+                return_timestamps=True,
+            )
+            
+        if self.diarization_model is None:
+            print("Loading diarization model...")
+            # Note: You'll need to set up authentication for pyannote models
+            # For demo purposes, we'll handle the case where it's not available
+            try:
+                self.diarization_model = Pipeline.from_pretrained(
+                    "pyannote/speaker-diarization-3.1",
+                    use_auth_token=os.getenv("HF_TOKEN")
+                ).to(torch.device("cuda"))
+            except Exception as e:
+                print(f"Could not load diarization model: {e}")
+                self.diarization_model = None
+    
+    def convert_audio_format(self, audio_path: str) -> str:
+        """Convert audio to 16kHz mono WAV format"""
+        temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        temp_wav_path = temp_wav.name
+        temp_wav.close()
+        
+        try:
+            subprocess.run([
+                "ffmpeg", "-i", audio_path,
+                "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le",
+                temp_wav_path, "-y"
+            ], check=True, capture_output=True)
+            return temp_wav_path
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(f"Audio conversion failed: {e}")
+    
+    @spaces.GPU
+    def transcribe_audio(
+        self,
+        audio_path: str,
+        language: Optional[str] = None,
+        translate: bool = False,
+        prompt: Optional[str] = None
+    ) -> Tuple[List[Dict], str]:
+        """Transcribe audio using Whisper with flash attention"""
+        if self.pipe is None:
+            self.setup_models()
+            
+        print("Starting transcription...")
+        start_time = time.time()
+        
+        # Prepare generation kwargs
+        generate_kwargs = {}
+        if language:
+            generate_kwargs["language"] = language
+        if translate:
+            generate_kwargs["task"] = "translate"
+        if prompt:
+            generate_kwargs["prompt_ids"] = self.pipe.tokenizer.encode(prompt)
+            
+        # Transcribe with timestamps
+        result = self.pipe(
+            audio_path,
+            return_timestamps=True,
+            generate_kwargs=generate_kwargs
+        )
+        
+        # Extract segments and detected language
+        segments = []
+        if "chunks" in result:
+            for chunk in result["chunks"]:
+                segment = {
+                    "start": float(chunk["timestamp"][0] or 0),
+                    "end": float(chunk["timestamp"][1] or 0),
+                    "text": chunk["text"].strip(),
+                }
+                segments.append(segment)
+        else:
+            # Fallback for different result format
+            segments = [{
+                "start": 0.0,
+                "end": 0.0,
+                "text": result["text"]
+            }]
+        
+        detected_language = getattr(result, 'language', language or 'unknown')
+        
+        transcription_time = time.time() - start_time
+        print(f"Transcription completed in {transcription_time:.2f} seconds")
+        
+        return segments, detected_language
+    
+    @spaces.GPU  
+    def perform_diarization(
+        self,
+        audio_path: str,
+        num_speakers: Optional[int] = None
+    ) -> Tuple[List[Dict], int]:
+        """Perform speaker diarization"""
+        if self.diarization_model is None:
+            print("Diarization model not available, assigning single speaker")
+            return [], 1
+            
+        print("Starting diarization...")
+        start_time = time.time()
+        
+        # Load audio for diarization
+        waveform, sample_rate = torchaudio.load(audio_path)
+        
+        # Perform diarization
+        diarization = self.diarization_model(
+            {"waveform": waveform, "sample_rate": sample_rate},
+            num_speakers=num_speakers,
+        )
+        
+        # Convert to list format
+        diarize_segments = []
+        diarization_list = list(diarization.itertracks(yield_label=True))
+        
+        for turn, _, speaker in diarization_list:
+            diarize_segments.append({
+                "start": turn.start,
+                "end": turn.end, 
+                "speaker": speaker
+            })
+            
+        unique_speakers = {speaker for _, _, speaker in diarization_list}
+        detected_num_speakers = len(unique_speakers)
+        
+        diarization_time = time.time() - start_time
+        print(f"Diarization completed in {diarization_time:.2f} seconds")
+        
+        return diarize_segments, detected_num_speakers
+    
+    def merge_transcription_and_diarization(
+        self,
+        transcription_segments: List[Dict],
+        diarization_segments: List[Dict]
+    ) -> List[Dict]:
+        """Merge transcription segments with speaker information"""
+        if not diarization_segments:
+            # No diarization available, assign single speaker
+            for segment in transcription_segments:
+                segment["speaker"] = "SPEAKER_00"
+            return transcription_segments
+            
+        print("Merging transcription and diarization...")
+        diarize_df = pd.DataFrame(diarization_segments)
+        
+        final_segments = []
+        for segment in transcription_segments:
+            # Calculate intersection with diarization segments
+            diarize_df["intersection"] = np.maximum(0, 
+                np.minimum(diarize_df["end"], segment["end"]) - 
+                np.maximum(diarize_df["start"], segment["start"])
+            )
+            
+            # Find speaker with maximum intersection
+            dia_tmp = diarize_df[diarize_df["intersection"] > 0]
+            if len(dia_tmp) > 0:
+                speaker = (
+                    dia_tmp.groupby("speaker")["intersection"]
+                    .sum()
+                    .sort_values(ascending=False)
+                    .index[0]
+                )
+            else:
+                speaker = "SPEAKER_00"
+                
+            segment["speaker"] = speaker
+            segment["duration"] = segment["end"] - segment["start"]
+            final_segments.append(segment)
+            
+        return final_segments
+    
+    def group_segments_by_speaker(
+        self,
+        segments: List[Dict],
+        max_gap: float = 1.0,
+        max_duration: float = 30.0
+    ) -> List[Dict]:
+        """Group consecutive segments from the same speaker"""
+        if not segments:
+            return segments
+            
+        grouped_segments = []
+        current_group = segments[0].copy()
+        sentence_end_pattern = r"[.!?]+\s*$"
+        
+        for segment in segments[1:]:
+            time_gap = segment["start"] - current_group["end"]
+            current_duration = current_group["end"] - current_group["start"]
+            
+            # Conditions for combining segments
+            can_combine = (
+                segment["speaker"] == current_group["speaker"] and
+                time_gap <= max_gap and
+                current_duration < max_duration and
+                not re.search(sentence_end_pattern, current_group["text"])
+            )
+            
+            if can_combine:
+                # Merge segments
+                current_group["end"] = segment["end"]
+                current_group["text"] += " " + segment["text"]
+                current_group["duration"] = current_group["end"] - current_group["start"]
+            else:
+                # Start new group
+                grouped_segments.append(current_group)
+                current_group = segment.copy()
+                
+        grouped_segments.append(current_group)
+        
+        # Clean up text
+        for segment in grouped_segments:
+            segment["text"] = re.sub(r"\s+", " ", segment["text"]).strip()
+            segment["text"] = re.sub(r"\s+([.,!?])", r"\1", segment["text"])
+            
+        return grouped_segments
+    
+    @spaces.GPU
+    def process_audio(
+        self,
+        audio_file,
+        num_speakers: Optional[int] = None,
+        language: Optional[str] = None,
+        translate: bool = False,
+        prompt: Optional[str] = None,
+        group_segments: bool = True
+    ) -> Dict[str, Any]:
+        """Main processing function"""
+        if audio_file is None:
+            return {"error": "No audio file provided"}
+            
+        try:
+            # Setup models if not already done
+            self.setup_models()
+            
+            # Convert audio format
+            wav_path = self.convert_audio_format(audio_file)
+            
+            try:
+                # Transcribe audio
+                transcription_segments, detected_language = self.transcribe_audio(
+                    wav_path, language, translate, prompt
+                )
+                
+                # Perform diarization
+                diarization_segments, detected_num_speakers = self.perform_diarization(
+                    wav_path, num_speakers
+                )
+                
+                # Merge transcription and diarization
+                final_segments = self.merge_transcription_and_diarization(
+                    transcription_segments, diarization_segments
+                )
+                
+                # Group segments if requested
+                if group_segments:
+                    final_segments = self.group_segments_by_speaker(final_segments)
+                
+                return {
+                    "segments": final_segments,
+                    "language": detected_language,
+                    "num_speakers": detected_num_speakers or 1,
+                    "total_segments": len(final_segments)
+                }
+                
+            finally:
+                # Clean up temporary file
+                if os.path.exists(wav_path):
+                    os.unlink(wav_path)
+                    
+        except Exception as e:
+            return {"error": f"Processing failed: {str(e)}"}
+
+# Initialize transcriber
+transcriber = WhisperTranscriber()
+
+def format_segments_for_display(result: Dict[str, Any]) -> str:
+    """Format segments for display in Gradio"""
+    if "error" in result:
+        return f"❌ Error: {result['error']}"
+        
+    segments = result.get("segments", [])
+    language = result.get("language", "unknown")
+    num_speakers = result.get("num_speakers", 1)
+    
+    output = f"🎯 **Detection Results:**\n"
+    output += f"- Language: {language}\n"
+    output += f"- Speakers: {num_speakers}\n"
+    output += f"- Segments: {len(segments)}\n\n"
+    
+    output += "📝 **Transcription:**\n\n"
+    
+    for i, segment in enumerate(segments, 1):
+        start_time = str(datetime.timedelta(seconds=int(segment["start"])))
+        end_time = str(datetime.timedelta(seconds=int(segment["end"])))
+        speaker = segment.get("speaker", "SPEAKER_00")
+        text = segment["text"]
+        
+        output += f"**{speaker}** ({start_time} → {end_time})\n"
+        output += f"{text}\n\n"
+        
+    return output
+
+def process_audio_gradio(
+    audio_file,
+    num_speakers,
+    language,
+    translate,
+    prompt,
+    group_segments
+):
+    """Gradio interface function"""
+    result = transcriber.process_audio(
+        audio_file=audio_file,
+        num_speakers=num_speakers if num_speakers > 0 else None,
+        language=language if language != "auto" else None,
+        translate=translate,
+        prompt=prompt if prompt.strip() else None,
+        group_segments=group_segments
+    )
+    
+    formatted_output = format_segments_for_display(result)
+    return formatted_output, result
+
+# Create Gradio interface
+with gr.Blocks(
+    title="🎙️ Whisper Transcription with Speaker Diarization",
+    theme=gr.themes.Soft()
+) as demo:
+    gr.Markdown("""
+    # 🎙️ Advanced Audio Transcription & Speaker Diarization
+    
+    Upload an audio file to get accurate transcription with speaker identification, powered by:
+    - **Whisper Large V3 Turbo** with Flash Attention for fast transcription
+    - **Pyannote 3.1** for speaker diarization
+    - **ZeroGPU** acceleration for optimal performance
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            audio_input = gr.Audio(
+                label="🎵 Upload Audio File",
+                type="filepath",
+                sources=["upload", "microphone"]
+            )
+            
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
+                num_speakers = gr.Slider(
+                    minimum=0,
+                    maximum=20,
+                    value=0,
+                    step=1,
+                    label="Number of Speakers (0 = auto-detect)"
+                )
+                
+                language = gr.Dropdown(
+                    choices=["auto", "en", "es", "fr", "de", "it", "pt", "ru", "ja", "ko", "zh"],
+                    value="auto",
+                    label="Language"
+                )
+                
+                translate = gr.Checkbox(
+                    label="Translate to English",
+                    value=False
+                )
+                
+                prompt = gr.Textbox(
+                    label="Vocabulary Prompt (names, acronyms, etc.)",
+                    placeholder="Enter names, technical terms, or context...",
+                    lines=2
+                )
+                
+                group_segments = gr.Checkbox(
+                    label="Group segments by speaker",
+                    value=True
+                )
+            
+            process_btn = gr.Button("🚀 Transcribe Audio", variant="primary", size="lg")
+        
+        with gr.Column():
+            output_text = gr.Markdown(
+                label="📝 Transcription Results",
+                value="Upload an audio file and click 'Transcribe Audio' to get started!"
+            )
+            
+            output_json = gr.JSON(
+                label="🔧 Raw Output (JSON)",
+                visible=False
+            )
+    
+    # Event handlers
+    process_btn.click(
+        fn=process_audio_gradio,
+        inputs=[
+            audio_input,
+            num_speakers,
+            language,
+            translate,
+            prompt,
+            group_segments
+        ],
+        outputs=[output_text, output_json],
+        show_progress=True
+    )
+    
+    # Examples
+    gr.Markdown("### 📋 Usage Tips:")
+    gr.Markdown("""
+    - **Supported formats**: MP3, WAV, M4A, FLAC, OGG, and more
+    - **Max duration**: Recommended under 10 minutes for optimal performance
+    - **Speaker detection**: Works best with clear, distinct voices
+    - **Languages**: Supports 100+ languages with auto-detection
+    - **Vocabulary**: Add names and technical terms in the prompt for better accuracy
+    """)
+
+if __name__ == "__main__":
+    demo.launch(debug=True)

config.py

@@ -0,0 +1,69 @@
+# Configuration settings for the Whisper Transcription Space
+
+# Model configurations
+WHISPER_MODEL = "openai/whisper-large-v3-turbo"
+DIARIZATION_MODEL = "pyannote/speaker-diarization-3.1"
+
+# Audio processing settings
+AUDIO_SAMPLE_RATE = 16000
+AUDIO_CHANNELS = 1
+MAX_AUDIO_DURATION = 600  # 10 minutes in seconds
+
+# Transcription settings
+DEFAULT_BEAM_SIZE = 5
+DEFAULT_LANGUAGE = None  # Auto-detect
+DEFAULT_TRANSLATE = False
+
+# Diarization settings
+MAX_SPEAKERS = 20
+DEFAULT_NUM_SPEAKERS = None  # Auto-detect
+
+# Segment grouping settings
+MAX_SEGMENT_GAP = 1.0  # seconds
+MAX_SEGMENT_DURATION = 30.0  # seconds
+
+# Flash attention settings
+FLASH_ATTENTION_ENABLED = True
+TORCH_DTYPE = "float16"
+
+# ZeroGPU settings
+GPU_MEMORY_FRACTION = 0.8
+CUDA_DEVICE = "cuda:0"
+
+# Gradio interface settings
+GRADIO_THEME = "soft"
+GRADIO_DEBUG = False
+GRADIO_SHARE = False
+
+# Environment variables
+HF_TOKEN_ENV_VAR = "HF_TOKEN"
+
+# Supported audio formats
+SUPPORTED_AUDIO_FORMATS = [
+    ".mp3", ".wav", ".m4a", ".flac", ".ogg", 
+    ".aac", ".wma", ".opus", ".webm"
+]
+
+# Language codes
+SUPPORTED_LANGUAGES = {
+    "auto": "Auto-detect",
+    "en": "English",
+    "es": "Spanish", 
+    "fr": "French",
+    "de": "German",
+    "it": "Italian",
+    "pt": "Portuguese",
+    "ru": "Russian",
+    "ja": "Japanese",
+    "ko": "Korean",
+    "zh": "Chinese",
+    "ar": "Arabic",
+    "hi": "Hindi",
+    "tr": "Turkish",
+    "pl": "Polish",
+    "nl": "Dutch",
+    "sv": "Swedish",
+    "da": "Danish",
+    "no": "Norwegian",
+    "fi": "Finnish"
+} 

requirements.txt

@@ -0,0 +1,34 @@
+# Core ML libraries
+torch>=2.0.0
+torchaudio>=2.0.0
+transformers>=4.35.0
+accelerate>=0.24.0
+
+# Gradio and Spaces
+gradio>=4.0.0
+spaces>=0.19.0
+
+# Audio processing and transcription
+ffmpeg-python>=0.2.0
+librosa>=0.10.0
+soundfile>=0.12.0
+
+# Speaker diarization
+pyannote.audio>=3.1.0
+pyannote.core>=5.0.0
+pyannote.database>=5.0.0
+pyannote.metrics>=3.2.0
+
+# Data processing
+pandas>=1.5.0
+numpy>=1.24.0
+
+# Utility libraries
+requests>=2.28.0
+typing-extensions>=4.5.0
+
+# Flash attention (will be installed at runtime)
+# flash-attn (installed dynamically in app.py)
+
+# Additional dependencies for ZeroGPU
+huggingface_hub>=0.16.0