BricksDisplay/OuteTTS-Speaker-Creator

1"""2Alias module to redirect whisper imports to whisperx.3This allows OuteTTS to use whisperx instead of the standard whisper package.4"""5 6import sys7import importlib.util8 9def setup_whisper_alias():10    """Setup alias so that 'import whisper' uses whisperx instead."""11    try:12        # Check if whisperx is available13        whisperx_spec = importlib.util.find_spec("whisperx")14        if whisperx_spec is None:15            print("Warning: whisperx not found, falling back to regular whisper")16            return17        18        # Import whisperx19        import whisperx20        21        # Create a module wrapper that provides whisper-like interface22        class WhisperAlias:23            def __init__(self):24                self.model = whisperx.WhisperModel if hasattr(whisperx, 'WhisperModel') else None25                self.load_model = self._load_model26                27            def _load_model(self, name, **kwargs):28                """Load model with whisperx compatible interface."""29                # Create WhisperX model instance30                device = "cuda" if kwargs.get("device", "auto") == "cuda" else "cpu"31                compute_type = "float16" if device == "cuda" else "int8"32                33                model = whisperx.load_model(34                    name,35                    device=device,36                    compute_type=compute_type37                )38                39                return WhisperXModelWrapper(model, device)40        41        class WhisperXModelWrapper:42            """Wrapper to make whisperx compatible with whisper interface."""43            44            def __init__(self, model, device):45                self.model = model46                self.device = device47                48            def transcribe(self, audio, **kwargs):49                """Transcribe audio with whisper-compatible interface."""50                # Store original word_timestamps setting51                original_word_timestamps = kwargs.get('word_timestamps', False)52                53                # Load audio if it's a file path54                if isinstance(audio, str):55                    audio_data = whisperx.load_audio(audio)56                else:57                    audio_data = audio58                59                # Use whisperx's transcribe method60                batch_size = kwargs.get('batch_size', 16)61                result = self.model.transcribe(audio_data, batch_size=batch_size)62                63                # If word timestamps are requested, perform alignment64                if original_word_timestamps and result.get("segments"):65                    try:66                        # Load alignment model67                        model_a, metadata = whisperx.load_align_model(68                            language_code=result.get("language", "en"),69                            device=self.device70                        )71                        72                        # Align the segments73                        result = whisperx.align(74                            result["segments"],75                            model_a,76                            metadata,77                            audio_data,78                            self.device,79                            return_char_alignments=False80                        )81                    except Exception as e:82                        print(f"Warning: Could not perform alignment: {e}")83                        # Continue without alignment84                85                # Ensure result format is compatible with whisper format86                if "segments" not in result:87                    result["segments"] = []88                89                # Ensure 'text' field exists - concatenate all segment texts90                if "text" not in result:91                    result["text"] = " ".join([segment.get("text", "") for segment in result.get("segments", [])])92                93                # Add words field to segments if word timestamps were requested94                for segment in result.get("segments", []):95                    if original_word_timestamps and "words" not in segment:96                        # If we don't have words but they were requested, create empty words list97                        segment["words"] = []98                99                return result100        101        # Create the alias module102        whisper_alias = WhisperAlias()103        104        # Add to sys.modules so 'import whisper' uses our alias105        sys.modules['whisper'] = whisper_alias106        107        print("✅ Successfully aliased whisper to whisperx")108        109    except ImportError as e:110        print(f"Warning: Could not setup whisper alias: {e}")111        print("Falling back to regular whisper (if available)")112    except Exception as e:113        print(f"Warning: Error setting up whisper alias: {e}")114 115# Auto-setup when module is imported116setup_whisper_alias()