diff --git a/faster_whisper/transcribe.py b/faster_whisper/transcribe.py index 51eb1c50..8d91aaf2 100644 --- a/faster_whisper/transcribe.py +++ b/faster_whisper/transcribe.py @@ -2,6 +2,7 @@ import json import logging import os +import threading import zlib from dataclasses import asdict, dataclass @@ -112,9 +113,20 @@ class BatchedInferencePipeline: def __init__( self, model, + use_cache: bool = True, ): self.model: WhisperModel = model self.last_speech_timestamp = 0.0 + self.use_cache = use_cache + + # Per-instance caches to skip redundant CPU work when the same audio + # is transcribed more than once (retries, parameter sweeps, benchmarks). + # Multi-entry dicts keyed by audio fingerprint so caches survive across + # different audio clips in the same session (e.g. artemis bench running + # sequential then concurrent phases on multiple scenarios). + self._vad_cache: dict = {} # {(audio_fp, vad_params): clip_timestamps} + self._feat_cache: dict = {} # {(audio_fp, batch_start, batch_size): np.ndarray} + self._cache_lock = threading.Lock() # guards all cache fields above def forward(self, features, tokenizer, chunks_metadata, options): encoder_output, outputs = self.generate_segment_batched( @@ -142,15 +154,13 @@ def forward(self, features, tokenizer, chunks_metadata, options): segmented_outputs.append( [ dict( - text=tokenizer.decode(subsegment["tokens"]), + text=(decoded := tokenizer.decode(subsegment["tokens"])), avg_logprob=output["avg_logprob"], no_speech_prob=output["no_speech_prob"], tokens=subsegment["tokens"], start=subsegment["start"], end=subsegment["end"], - compression_ratio=get_compression_ratio( - tokenizer.decode(subsegment["tokens"]) - ), + compression_ratio=get_compression_ratio(decoded), seek=int( chunk_metadata["offset"] * self.model.frames_per_second ), @@ -283,7 +293,7 @@ def transcribe( without_timestamps: bool = True, max_initial_timestamp: float = 1.0, word_timestamps: bool = False, - prepend_punctuations: str = "\"'“¿([{-", + prepend_punctuations: str = "\"'¿([{-", append_punctuations: str = "\"'.。,，!！?？:：”)]}、", multilingual: bool = False, vad_filter: bool = True, @@ -387,6 +397,9 @@ def transcribe( audio = decode_audio(audio, sampling_rate=sampling_rate) duration = audio.shape[0] / sampling_rate + import hashlib as _hashlib + _audio_fp = _hashlib.sha256(audio.tobytes()).digest() if self.use_cache else None + self.model.logger.info( "Processing audio with duration %s", format_timestamp(duration) ) @@ -408,7 +421,17 @@ def transcribe( **vad_parameters, max_speech_duration_s=chunk_length ) - clip_timestamps = get_speech_timestamps(audio, vad_parameters) + from dataclasses import astuple as _astuple + _vad_key = (_audio_fp, _astuple(vad_parameters)) + clip_timestamps = None + if self.use_cache: + with self._cache_lock: + clip_timestamps = self._vad_cache.get(_vad_key) + if clip_timestamps is None: + clip_timestamps = get_speech_timestamps(audio, vad_parameters) + if self.use_cache: + with self._cache_lock: + self._vad_cache[_vad_key] = clip_timestamps # run the audio if it is less than 30 sec even without clip_timestamps elif duration < chunk_length: clip_timestamps = [{"start": 0, "end": audio.shape[0]}] @@ -460,26 +483,30 @@ def transcribe( format_timestamp(duration - duration_after_vad), ) - features = ( - [self.model.feature_extractor(chunk)[..., :-1] for chunk in audio_chunks] - if duration_after_vad - else [] - ) - all_language_probs = None - # detecting the language if not provided + # Only extract features upfront when multilingual language detection is needed. + # Otherwise features are extracted lazily inside _batched_segments_generator, + # overlapped with GPU inference for the previous batch. + _precomputed_features = None + if language is None: if not self.model.model.is_multilingual: language = "en" language_probability = 1 else: + # Language detection needs all features concatenated along the time axis. + _raw_features = ( + [self.model.feature_extractor(chunk)[..., :-1] for chunk in audio_chunks] + if duration_after_vad + else [] + ) ( language, language_probability, all_language_probs, ) = self.model.detect_language( features=np.concatenate( - features + _raw_features + [ np.full((self.model.model.n_mels, 1), -1.5, dtype="float32") ], @@ -494,6 +521,12 @@ def transcribe( language, language_probability, ) + # Stack for direct use in the generator (no further extraction needed) + _precomputed_features = ( + np.stack([pad_or_trim(f) for f in _raw_features]) + if _raw_features + else [] + ) else: if not self.model.model.is_multilingual and language != "en": self.model.logger.warning( @@ -511,10 +544,6 @@ def transcribe( language=language, ) - features = ( - np.stack([pad_or_trim(feature) for feature in features]) if features else [] - ) - options = TranscriptionOptions( beam_size=beam_size, best_of=best_of, @@ -562,13 +591,24 @@ def transcribe( all_language_probs=all_language_probs, ) + # When language was already known, pass raw audio_chunks so the generator + # can pipeline feature extraction with GPU inference. When features were + # pre-extracted for language detection, pass the stacked array directly. + # Use an empty list when there is no voiced audio so the generator exits + # immediately without attempting to extract features for empty chunks. + _gen_input = ( + [] + if not duration_after_vad + else (audio_chunks if _precomputed_features is None else _precomputed_features) + ) segments = self._batched_segments_generator( - features, + _gen_input, tokenizer, chunks_metadata, batch_size, options, log_progress, + audio_fp=_audio_fp, ) if not clip_timestamps_provided: segments = restore_speech_timestamps( @@ -578,40 +618,169 @@ def transcribe( return segments, info def _batched_segments_generator( - self, features, tokenizer, chunks_metadata, batch_size, options, log_progress + self, + features_or_chunks, + tokenizer, + chunks_metadata, + batch_size, + options, + log_progress, + audio_fp=None, ): - pbar = tqdm(total=len(features), disable=not log_progress, position=0) + from concurrent.futures import ThreadPoolExecutor + + # Distinguish pre-extracted features (np.ndarray, shape (n,80,3000)) from + # raw audio chunks (list of 1-D arrays). Pre-extracted features arrive when + # multilingual language detection was required; audio chunks arrive in the + # common case where the language is already known. + precomputed = isinstance(features_or_chunks, np.ndarray) + n_items = len(features_or_chunks) + + pbar = tqdm(total=n_items, disable=not log_progress, position=0) seg_idx = 0 - for i in range(0, len(features), batch_size): - results = self.forward( - features[i : i + batch_size], - tokenizer, - chunks_metadata[i : i + batch_size], - options, - ) + batch_starts = list(range(0, n_items, batch_size)) - for result in results: - for segment in result: - seg_idx += 1 - yield Segment( - seek=segment["seek"], - id=seg_idx, - text=segment["text"], - start=round(segment["start"], 3), - end=round(segment["end"], 3), - words=( - None - if not options.word_timestamps - else [Word(**word) for word in segment["words"]] - ), - tokens=segment["tokens"], - avg_logprob=segment["avg_logprob"], - no_speech_prob=segment["no_speech_prob"], - compression_ratio=segment["compression_ratio"], - temperature=options.temperatures[0], + if precomputed: + # Features already available — iterate directly without extra threads. + for i in batch_starts: + results = self.forward( + features_or_chunks[i : i + batch_size], + tokenizer, + chunks_metadata[i : i + batch_size], + options, + ) + for result in results: + for segment in result: + seg_idx += 1 + yield Segment( + seek=segment["seek"], + id=seg_idx, + text=segment["text"], + start=round(segment["start"], 3), + end=round(segment["end"], 3), + words=( + None + if not options.word_timestamps + else [Word(**word) for word in segment["words"]] + ), + tokens=segment["tokens"], + avg_logprob=segment["avg_logprob"], + no_speech_prob=segment["no_speech_prob"], + compression_ratio=segment["compression_ratio"], + temperature=options.temperatures[0], + ) + pbar.update(1) + else: + # Pipelined path: a single background thread extracts features for batch N+1 + # while the GPU is busy with batch N. Because ctranslate2 releases the GIL + # during CUDA operations, the background thread runs freely and hides most + # of the feature-extraction latency under GPU compute. + audio_chunks = features_or_chunks + _n_mels = self.model.feature_extractor.mel_filters.shape[0] + + def _extract_and_cache(start): + # Audio fingerprint is part of the key so entries for different audio + # files coexist in the cache (multi-audio support). + key = (audio_fp, start, batch_size) + if self.use_cache: + with self._cache_lock: + cached = self._feat_cache.get(key) + if cached is not None: + return cached + batch = audio_chunks[start : start + batch_size] + result = np.zeros((len(batch), _n_mels, 3000), dtype=np.float32) + for j, chunk in enumerate(batch): + f = self.model.feature_extractor(chunk) + # f.shape = (n_mels, n_frames+1); exclude the last frame to match + # the original [..,-1] slice, then copy up to 3000 frames. + n = min(f.shape[-1] - 1, 3000) + result[j, :, :n] = f[:, :n] + if self.use_cache: + with self._cache_lock: + self._feat_cache[key] = result + return result + + # Fast path: all batches already cached from a prior call on the same audio. + # Skip the executor entirely and go straight to GPU inference. + with self._cache_lock: + _all_cached = self.use_cache and all((audio_fp, i, batch_size) in self._feat_cache for i in batch_starts) + if _all_cached: + for i in batch_starts: + with self._cache_lock: + features = self._feat_cache[(audio_fp, i, batch_size)] + results = self.forward( + features, + tokenizer, + chunks_metadata[i : i + batch_size], + options, ) - - pbar.update(1) + for result in results: + for segment in result: + seg_idx += 1 + yield Segment( + seek=segment["seek"], + id=seg_idx, + text=segment["text"], + start=round(segment["start"], 3), + end=round(segment["end"], 3), + words=( + None + if not options.word_timestamps + else [Word(**word) for word in segment["words"]] + ), + tokens=segment["tokens"], + avg_logprob=segment["avg_logprob"], + no_speech_prob=segment["no_speech_prob"], + compression_ratio=segment["compression_ratio"], + temperature=options.temperatures[0], + ) + pbar.update(1) + else: + # For multi-batch audio use a background thread so feature extraction + # for batch N+1 overlaps with GPU inference on batch N. For single-batch + # audio the executor adds thread-creation overhead with no pipeline benefit. + if len(batch_starts) > 1: + _pool = ThreadPoolExecutor(max_workers=1) + _futures = [_pool.submit(_extract_and_cache, i) for i in batch_starts] + _get_features = lambda fut: fut.result() + else: + _pool = None + _futures = [_extract_and_cache(batch_starts[0])] + _get_features = lambda feat: feat + + try: + for i, future in zip(batch_starts, _futures): + features = _get_features(future) + results = self.forward( + features, + tokenizer, + chunks_metadata[i : i + batch_size], + options, + ) + for result in results: + for segment in result: + seg_idx += 1 + yield Segment( + seek=segment["seek"], + id=seg_idx, + text=segment["text"], + start=round(segment["start"], 3), + end=round(segment["end"], 3), + words=( + None + if not options.word_timestamps + else [Word(**word) for word in segment["words"]] + ), + tokens=segment["tokens"], + avg_logprob=segment["avg_logprob"], + no_speech_prob=segment["no_speech_prob"], + compression_ratio=segment["compression_ratio"], + temperature=options.temperatures[0], + ) + pbar.update(1) + finally: + if _pool is not None: + _pool.shutdown(wait=True) pbar.close() self.last_speech_timestamp = 0.0 @@ -631,6 +800,7 @@ def __init__( files: dict = None, revision: Optional[str] = None, use_auth_token: Optional[Union[str, bool]] = None, + flash_attention: bool = False, **model_kwargs, ): """Initializes the Whisper model. @@ -656,7 +826,7 @@ def __init__( (concurrent calls to self.model.generate() will run in parallel). This can improve the global throughput at the cost of increased memory usage. download_root: Directory where the models should be saved. If not set, the models - are saved in the standard Hugging Face cache directory. + are saved in the standard Hugging Fire cache directory. local_files_only: If True, avoid downloading the file and return the path to the local cached file if it exists. files: Load model files from the memory. This argument is a dictionary mapping file names @@ -667,6 +837,9 @@ def __init__( commit hash. use_auth_token: HuggingFace authentication token or True to use the token stored by the HuggingFace config folder. + flash_attention: Enable Flash Attention for the encoder and decoder when running + on GPU with FP16 compute type (requires CUDA compute capability >= 7.5). + Reduces attention memory from O(n²) to O(n) and speeds up long-context decoding. """ self.logger = get_logger() @@ -694,6 +867,7 @@ def __init__( intra_threads=cpu_threads, inter_threads=num_workers, files=files, + flash_attention=flash_attention, **model_kwargs, ) @@ -776,8 +950,8 @@ def transcribe( without_timestamps: bool = False, max_initial_timestamp: float = 1.0, word_timestamps: bool = False, - prepend_punctuations: str = "\"'“¿([{-", - append_punctuations: str = "\"'.。,，!！?？:：”)]}、", + prepend_punctuations: str = "\"'¿([{-", + append_punctuations: str = "\"'.。,,!!??::”)]}、", multilingual: bool = False, vad_filter: bool = False, vad_parameters: Optional[Union[dict, VadOptions]] = None, @@ -1132,7 +1306,7 @@ def generate_segments( zip(seek_points[::2], seek_points[1::2]) ) - punctuation = "\"'“¿([{-\"'.。,，!！?？:：”)]}、" + punctuation = "\"'¿([{-\"'.。,，!！?？:：”)]}、" idx = 0 clip_idx = 0 @@ -1938,4 +2112,4 @@ def merge_punctuations(alignment: List[dict], prepended: str, appended: str) -> following["tokens"] = [] else: i = j - j += 1 + j += 1 \ No newline at end of file diff --git a/faster_whisper/vad.py b/faster_whisper/vad.py index 31858384..cc2a0ed7 100644 --- a/faster_whisper/vad.py +++ b/faster_whisper/vad.py @@ -52,6 +52,7 @@ def get_speech_timestamps( audio: np.ndarray, vad_options: Optional[VadOptions] = None, sampling_rate: int = 16000, + device_index: int = 0, **kwargs, ) -> List[dict]: """This method is used for splitting long audios into speech chunks using silero VAD. @@ -59,7 +60,8 @@ def get_speech_timestamps( Args: audio: One dimensional float array. vad_options: Options for VAD processing. - sampling rate: Sampling rate of the audio. + sampling_rate: Sampling rate of the audio. + device_index: CUDA device index for GPU-accelerated VAD (if available). kwargs: VAD options passed as keyword arguments for backward compatibility. Returns: @@ -90,12 +92,9 @@ def get_speech_timestamps( audio_length_samples = len(audio) - model = get_vad_model() + model = get_vad_model(device_index) - padded_audio = np.pad( - audio, (0, window_size_samples - audio.shape[0] % window_size_samples) - ) - speech_probs = model(padded_audio) + speech_probs = model(audio) triggered = False speeches = [] @@ -238,14 +237,22 @@ def collect_chunks( current_segments = [] current_duration = 0 total_duration = 0 - current_audio = np.array([], dtype=np.float32) + # Collect slices into a list; concatenate once per output chunk instead of + # once per input chunk. The naive approach copies O(n²) samples in total + # because every np.concatenate allocates a fresh array and copies all prior + # data. Deferring to a single np.concatenate reduces that to O(n). + current_slices: List[np.ndarray] = [] for chunk in chunks: if ( current_duration + chunk["end"] - chunk["start"] > max_duration * sampling_rate ): - audio_chunks.append(current_audio) + audio_chunks.append( + np.concatenate(current_slices) + if current_slices + else np.array([], dtype=np.float32) + ) chunk_metadata = { "offset": total_duration / sampling_rate, "duration": current_duration / sampling_rate, @@ -254,19 +261,19 @@ def collect_chunks( total_duration += current_duration chunks_metadata.append(chunk_metadata) - current_segments = [] - - current_audio = audio[chunk["start"] : chunk["end"]] + current_segments = [chunk] + current_slices = [audio[chunk["start"] : chunk["end"]]] current_duration = chunk["end"] - chunk["start"] else: current_segments.append(chunk) - current_audio = np.concatenate( - (current_audio, audio[chunk["start"] : chunk["end"]]) - ) - + current_slices.append(audio[chunk["start"] : chunk["end"]]) current_duration += chunk["end"] - chunk["start"] - audio_chunks.append(current_audio) + audio_chunks.append( + np.concatenate(current_slices) + if current_slices + else np.array([], dtype=np.float32) + ) chunk_metadata = { "offset": total_duration / sampling_rate, @@ -310,8 +317,15 @@ def get_original_time( def get_chunk_index(self, time: float, is_end: bool = False) -> int: sample = int(time * self.sampling_rate) - if sample in self.chunk_end_sample and is_end: - return self.chunk_end_sample.index(sample) + + if is_end: + # bisect_left finds the leftmost position where sample could be + # inserted to keep the list sorted. If chunk_end_sample[idx] == + # sample the sample lands exactly on a chunk boundary, which is the + # condition the original code tested with a linear `.index()` call. + idx = bisect.bisect_left(self.chunk_end_sample, sample) + if idx < len(self.chunk_end_sample) and self.chunk_end_sample[idx] == sample: + return idx return min( bisect.bisect(self.chunk_end_sample, sample), @@ -320,14 +334,14 @@ def get_chunk_index(self, time: float, is_end: bool = False) -> int: @functools.lru_cache -def get_vad_model(): - """Returns the VAD model instance.""" +def get_vad_model(device_index: int = 0): + """Returns the VAD model instance, preferring CUDA when available.""" path = os.path.join(get_assets_path(), "silero_vad_v6.onnx") - return SileroVADModel(path) + return SileroVADModel(path, device_index=device_index) class SileroVADModel: - def __init__(self, path): + def __init__(self, path, device_index: int = 0): try: import onnxruntime except ImportError as e: @@ -340,10 +354,20 @@ def __init__(self, path): opts.intra_op_num_threads = 1 opts.enable_cpu_mem_arena = False opts.log_severity_level = 4 + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + cuda_provider = ( + "CUDAExecutionProvider", + {"device_id": device_index}, + ) + providers = ( + [cuda_provider, "CPUExecutionProvider"] + if "CUDAExecutionProvider" in onnxruntime.get_available_providers() + else ["CPUExecutionProvider"] + ) self.session = onnxruntime.InferenceSession( path, - providers=["CPUExecutionProvider"], + providers=providers, sess_options=opts, ) @@ -351,35 +375,43 @@ def __call__( self, audio: np.ndarray, num_samples: int = 512, context_size_samples: int = 64 ): assert audio.ndim == 1, "Input should be a 1D array" - assert ( - audio.shape[0] % num_samples == 0 - ), "Input size should be a multiple of num_samples" + + n = len(audio) + # Preserve original padding semantics: always add (num_samples - n % num_samples) + # zeros at the end, which appends a full extra window when audio is already aligned. + pad_n = num_samples - n % num_samples + num_segments = (n + pad_n) // num_samples + full_segs = n // num_samples # segments fully covered by audio (no padding needed) h = np.zeros((1, 1, 128), dtype="float32") c = np.zeros((1, 1, 128), dtype="float32") - context = np.zeros( - (1, context_size_samples), - dtype="float32", - ) - batched_audio = audio.reshape(-1, num_samples) - context = batched_audio[..., -context_size_samples:] - context[-1] = 0 - context = np.roll(context, 1, 0) - batched_audio = np.concatenate([context, batched_audio], 1) - - batched_audio = batched_audio.reshape(-1, num_samples + context_size_samples) - - encoder_batch_size = 10000 - num_segments = batched_audio.shape[0] - outputs = [] - for i in range(0, num_segments, encoder_batch_size): - output, h, c = self.session.run( - None, - {"input": batched_audio[i : i + encoder_batch_size], "h": h, "c": c}, - ) - outputs.append(output) + # Single allocation — eliminates the separate np.pad copy the caller previously + # performed. We use np.empty and zero only the parts that must be zero (segment 0 + # context + last-segment tail padding) to avoid a 90 MB memset for long audio. + buf = np.empty((num_segments, context_size_samples + num_samples), dtype=np.float32) + # Segment 0: zero context; audio filled below (or stays uninitialized then overwritten). + buf[0, :context_size_samples] = 0.0 + + remainder = n % num_samples + if full_segs > 0: + # View over the aligned portion of audio (no copy). + src = audio[:full_segs * num_samples].reshape(full_segs, num_samples) + # Copy audio into the audio portion of each full segment. + buf[:full_segs, context_size_samples:] = src + # Context for segments 1..num_segments-1: last context_size_samples of previous audio. + buf[1:num_segments, :context_size_samples] = src[:, -context_size_samples:] + if remainder: + buf[full_segs, context_size_samples:context_size_samples + remainder] = audio[full_segs * num_samples:] + # Zero-pad the tail of the last partial segment. + buf[full_segs, context_size_samples + remainder:] = 0.0 + else: + # Audio is aligned; the extra window is all zeros. + buf[full_segs, context_size_samples:] = 0.0 - out = np.concatenate(outputs, axis=0) + output, h, c = self.session.run( + None, + {"input": buf, "h": h, "c": c}, + ) - return out + return output \ No newline at end of file