voicebox/backend/transcribe.py

"""
Whisper ASR module for transcription.
"""

from typing import Optional, List, Dict
import torch
import numpy as np
from pathlib import Path
from .utils.progress import get_progress_manager
from .utils.hf_progress import HFProgressTracker, create_hf_progress_callback


class WhisperModel:
    """Manages Whisper model loading and transcription."""
    
    def __init__(self, model_size: str = "base"):
        self.model = None
        self.processor = None
        self.model_size = model_size
        self.device = self._get_device()
    
    def _get_device(self) -> str:
        """Get the best available device."""
        if torch.cuda.is_available():
            return "cuda"
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            # MPS support for Whisper
            return "cpu"  # Use CPU for stability
        return "cpu"
    
    def is_loaded(self) -> bool:
        """Check if model is loaded."""
        return self.model is not None
    
    def load_model(self, model_size: Optional[str] = None):
        """
        Lazy load the Whisper model.
        
        Args:
            model_size: Model size (tiny, base, small, medium, large)
        """
        if model_size is None:
            model_size = self.model_size
        
        if self.model is not None and self.model_size == model_size:
            return
        
        try:
            from transformers import WhisperProcessor, WhisperForConditionalGeneration
            
            model_name = f"openai/whisper-{model_size}"
            
            # Set up progress tracking
            progress_manager = get_progress_manager()
            progress_model_name = f"whisper-{model_size}"
            
            print(f"Loading Whisper model {model_size} on {self.device}...")
            
            # Set up progress callback
            progress_callback = create_hf_progress_callback(progress_model_name, progress_manager)
            tracker = HFProgressTracker(progress_callback)
            
            # Use progress tracker during download
            with tracker.patch_download():
                self.processor = WhisperProcessor.from_pretrained(model_name)
                self.model = WhisperForConditionalGeneration.from_pretrained(model_name)
            
            self.model.to(self.device)
            self.model_size = model_size
            
            # Mark as complete
            progress_manager.mark_complete(progress_model_name)
            
            print(f"Whisper model {model_size} loaded successfully")
            
        except Exception as e:
            print(f"Error loading Whisper model: {e}")
            progress_manager = get_progress_manager()
            progress_manager.mark_error(f"whisper-{model_size}", str(e))
            raise
    
    def unload_model(self):
        """Unload the model to free memory."""
        if self.model is not None:
            del self.model
            del self.processor
            self.model = None
            self.processor = None
            
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            print("Whisper model unloaded")
    
    async def transcribe(
        self,
        audio_path: str,
        language: Optional[str] = None,
    ) -> str:
        """
        Transcribe audio to text.
        
        Args:
            audio_path: Path to audio file
            language: Optional language hint (en or zh)
            
        Returns:
            Transcribed text
        """
        self.load_model()
        
        from .utils.audio import load_audio
        
        # Load audio
        audio, sr = load_audio(audio_path, sample_rate=16000)
        
        # Process audio
        inputs = self.processor(
            audio,
            sampling_rate=16000,
            return_tensors="pt",
        )
        inputs = inputs.to(self.device)
        
        # Set language if provided
        forced_decoder_ids = None
        if language:
            lang_code = "en" if language == "en" else "zh"
            forced_decoder_ids = self.processor.get_decoder_prompt_ids(
                language=lang_code,
                task="transcribe",
            )
        
        # Generate transcription
        with torch.no_grad():
            predicted_ids = self.model.generate(
                inputs["input_features"],
                forced_decoder_ids=forced_decoder_ids,
            )
        
        # Decode
        transcription = self.processor.batch_decode(
            predicted_ids,
            skip_special_tokens=True,
        )[0]
        
        return transcription.strip()
    
    async def transcribe_with_timestamps(
        self,
        audio_path: str,
        language: Optional[str] = None,
    ) -> List[Dict[str, any]]:
        """
        Transcribe audio with word-level timestamps.
        
        Args:
            audio_path: Path to audio file
            language: Optional language hint
            
        Returns:
            List of word segments with timestamps
        """
        self.load_model()
        
        from .utils.audio import load_audio
        
        # Load audio
        audio, sr = load_audio(audio_path, sample_rate=16000)
        
        # Process audio
        inputs = self.processor(
            audio,
            sampling_rate=16000,
            return_tensors="pt",
        )
        inputs = inputs.to(self.device)
        
        # Set language if provided
        forced_decoder_ids = None
        if language:
            lang_code = "en" if language == "en" else "zh"
            forced_decoder_ids = self.processor.get_decoder_prompt_ids(
                language=lang_code,
                task="transcribe",
            )
        
        # Generate with timestamps
        with torch.no_grad():
            predicted_ids = self.model.generate(
                inputs["input_features"],
                forced_decoder_ids=forced_decoder_ids,
                return_timestamps=True,
            )
        
        # Decode with timestamps
        result = self.processor.batch_decode(
            predicted_ids,
            skip_special_tokens=False,
        )[0]
        
        # Parse timestamps (simplified - would need more robust parsing)
        # For now, return basic transcription
        # TODO: Implement proper timestamp parsing
        transcription = self.processor.batch_decode(
            predicted_ids,
            skip_special_tokens=True,
        )[0]
        
        return [
            {
                "text": transcription,
                "start": 0.0,
                "end": len(audio) / sr,
            }
        ]


# Global model instance
_whisper_model: Optional[WhisperModel] = None


def get_whisper_model() -> WhisperModel:
    """Get or create Whisper model instance."""
    global _whisper_model
    if _whisper_model is None:
        _whisper_model = WhisperModel()
    return _whisper_model


def unload_whisper_model():
    """Unload Whisper model to free memory."""
    global _whisper_model
    if _whisper_model is not None:
        _whisper_model.unload_model()