From a1d99ac5d9ad94a080f457e89af776776653f863 Mon Sep 17 00:00:00 2001 From: lml Date: Sat, 21 Mar 2026 20:50:46 +0800 Subject: [PATCH] feat: enhance WhisperX integration with MPS support and refactor device management --- assets/sidecar/whisperx_worker.py | 151 ++++++++++++++---- lib/services/whisper/whisper_service.dart | 173 +++++++++++++++------ lib/services/whisper/whisperx_sidecar.dart | 16 ++ 3 files changed, 267 insertions(+), 73 deletions(-) diff --git a/assets/sidecar/whisperx_worker.py b/assets/sidecar/whisperx_worker.py index d2bafcb..fcfe2eb 100644 --- a/assets/sidecar/whisperx_worker.py +++ b/assets/sidecar/whisperx_worker.py @@ -11,7 +11,9 @@ import wave from typing import Any, Dict, List, Optional, Tuple import numpy as np +import torch import whisperx +from whisperx.vads.pyannote import Pyannote SAMPLE_RATE = 16000 @@ -42,6 +44,11 @@ JAPANESE_SEGMENTATION_OPTIONS: Dict[str, Any] = { "max_segment_chars": 24, "min_split_chars": 4, } +DEFAULT_WHISPERX_VAD_OPTIONS: Dict[str, Any] = { + "chunk_size": 30, + "vad_onset": 0.500, + "vad_offset": 0.363, +} def configure_stdio() -> None: @@ -188,6 +195,38 @@ def normalize_options(raw_options: Any) -> Dict[str, Any]: return normalized +def normalize_device(value: Any, default: str = "cpu") -> str: + device = str(value or default).strip().lower() + return device or default + + +def ensure_mps_available() -> None: + mps_backend = getattr(torch.backends, "mps", None) + if mps_backend is None or not mps_backend.is_built(): + raise RuntimeError("This PyTorch build does not include MPS support.") + if not mps_backend.is_available(): + raise RuntimeError("MPS is not available on this machine.") + + +def build_vad_model_options(vad_options: Dict[str, Any]) -> Dict[str, Any]: + model_options: Dict[str, Any] = {} + for key in ("vad_onset", "vad_offset"): + if key in vad_options: + model_options[key] = vad_options[key] + return model_options + + +def build_effective_vad_options( + vad_options: Dict[str, Any], use_custom_vad: bool +) -> Dict[str, Any]: + if not use_custom_vad: + return dict(vad_options) + + effective = dict(DEFAULT_WHISPERX_VAD_OPTIONS) + effective.update(vad_options) + return effective + + def build_segmentation_options( language: Optional[str], override_options: Dict[str, Any] ) -> Dict[str, Any]: @@ -794,16 +833,25 @@ def normalize_transcript_segments( class WhisperXWorker: def __init__(self) -> None: - self.models: Dict[Tuple[str, str, str, Optional[str], str, str], Any] = {} + self.models: Dict[ + Tuple[str, str, str, str, Optional[str], str, str], Any + ] = {} + self.vad_models: Dict[Tuple[str, str], Any] = {} self.align_models: Dict[Tuple[str, str], Tuple[Any, Dict[str, Any]]] = {} def clear_device_resources(self, device: Optional[str] = None) -> None: if device is None: self.models.clear() + self.vad_models.clear() self.align_models.clear() else: self.models = { - key: value for key, value in self.models.items() if key[1] != device + key: value + for key, value in self.models.items() + if key[1] != device and key[2] != device + } + self.vad_models = { + key: value for key, value in self.vad_models.items() if key[0] != device } self.align_models = { key: value @@ -813,52 +861,87 @@ class WhisperXWorker: gc.collect() - if device not in (None, "cuda"): - return - try: - import torch - - if torch.cuda.is_available(): + if device in (None, "cuda") and torch.cuda.is_available(): torch.cuda.empty_cache() + if device in (None, "mps"): + empty_cache = getattr(getattr(torch, "mps", None), "empty_cache", None) + if callable(empty_cache): + empty_cache() except Exception: return + def get_vad_model(self, device: str, vad_options: Dict[str, Any]) -> Any: + device = normalize_device(device) + model_options = build_vad_model_options(vad_options) + key = ( + device, + json.dumps(model_options, sort_keys=True, ensure_ascii=False), + ) + if key in self.vad_models: + return self.vad_models[key] + + if device == "cuda": + self.clear_device_resources(device="cuda") + elif device == "mps": + ensure_mps_available() + + model = Pyannote(torch.device(device), **model_options) + self.vad_models[key] = model + return model + def get_model( self, model_name: str, - device: str, + asr_device: str, compute_type: str, language: Optional[str], asr_options: Dict[str, Any], vad_options: Dict[str, Any], + vad_device: Optional[str] = None, ) -> Any: + resolved_vad_device = normalize_device(vad_device, asr_device) + use_custom_vad = resolved_vad_device != asr_device + effective_vad_options = build_effective_vad_options( + vad_options, use_custom_vad + ) key = ( model_name, - device, + asr_device, + resolved_vad_device, compute_type, language, json.dumps(asr_options, sort_keys=True, ensure_ascii=False), - json.dumps(vad_options, sort_keys=True, ensure_ascii=False), + json.dumps(effective_vad_options, sort_keys=True, ensure_ascii=False), ) if key in self.models: return self.models[key] - if device == "cuda": + if asr_device == "cuda" or resolved_vad_device == "cuda": self.clear_device_resources(device="cuda") + load_kwargs: Dict[str, Any] = { + "compute_type": compute_type, + "language": language, + "asr_options": asr_options, + } + if effective_vad_options: + load_kwargs["vad_options"] = effective_vad_options + if use_custom_vad: + load_kwargs["vad_model"] = self.get_vad_model( + resolved_vad_device, effective_vad_options + ) + model = whisperx.load_model( model_name, - device, - compute_type=compute_type, - language=language, - asr_options=asr_options, - vad_options=vad_options, + asr_device, + **load_kwargs, ) self.models[key] = model return model def get_align_model(self, language: str, device: str) -> Tuple[Any, Dict[str, Any]]: + device = normalize_device(device) key = (language, device) if key in self.align_models: return self.align_models[key] @@ -870,6 +953,8 @@ class WhisperXWorker: if cached_key[1] != "cuda" } self.clear_device_resources(device="cuda") + elif device == "mps": + ensure_mps_available() model, metadata = whisperx.load_align_model( language_code=language, @@ -887,11 +972,11 @@ class WhisperXWorker: "cuda_device_count": 0, "cuda_device_name": None, "cuda_compute_types": [], + "mps_built": False, + "mps_available": False, } try: - import torch - payload["torch_version"] = str(getattr(torch, "__version__", "")) payload["torch_cuda_version"] = getattr(torch.version, "cuda", None) payload["cuda_available"] = bool(torch.cuda.is_available()) @@ -900,6 +985,10 @@ class WhisperXWorker: payload["cuda_device_count"] = device_count if device_count > 0: payload["cuda_device_name"] = str(torch.cuda.get_device_name(0)) + mps_backend = getattr(torch.backends, "mps", None) + if mps_backend is not None: + payload["mps_built"] = bool(mps_backend.is_built()) + payload["mps_available"] = bool(mps_backend.is_available()) except Exception as exc: # pylint: disable=broad-except payload["torch_error"] = str(exc) @@ -929,7 +1018,10 @@ class WhisperXWorker: model_name = str(params.get("model") or "small") language = params.get("language") language = str(language) if language else None - device = str(params.get("device") or "cpu") + device = normalize_device(params.get("device"), "cpu") + asr_device = normalize_device(params.get("asr_device"), device) + vad_device = normalize_device(params.get("vad_device"), asr_device) + align_device = normalize_device(params.get("align_device"), asr_device) compute_type = str(params.get("compute_type") or "int8") batch_size = int(params.get("batch_size") or 4) no_align = to_bool(params.get("no_align"), False) @@ -949,11 +1041,12 @@ class WhisperXWorker: ): model = self.get_model( model_name=model_name, - device=device, + asr_device=asr_device, compute_type=compute_type, language=language, asr_options=asr_options, vad_options=vad_options, + vad_device=vad_device, ) model_logs.flush() @@ -970,7 +1063,7 @@ class WhisperXWorker: verbose=False, ) transcribe_logs.flush() - if device == "cuda": + if asr_device == "cuda" or vad_device == "cuda": del model self.clear_device_resources(device="cuda") @@ -980,7 +1073,9 @@ class WhisperXWorker: if not no_align and normalized_segments: emit_status(request_id, "aligning") emit_progress(request_id, 72) - align_model, align_metadata = self.get_align_model(detected_language, device) + align_model, align_metadata = self.get_align_model( + detected_language, align_device + ) align_logs = ProgressLogStream(request_id) with contextlib.redirect_stdout( align_logs @@ -990,12 +1085,12 @@ class WhisperXWorker: align_model, align_metadata, audio, - device, + align_device, return_char_alignments=False, print_progress=True, ) align_logs.flush() - if device == "cuda": + if align_device == "cuda": del align_model self.clear_device_resources(device="cuda") detected_language = str(aligned.get("language") or detected_language) @@ -1026,7 +1121,11 @@ class WhisperXWorker: "payload": payload, } ) - if device == "cuda": + if ( + asr_device == "cuda" + or vad_device == "cuda" + or align_device == "cuda" + ): self.clear_device_resources(device="cuda") def dispatch(self, message: Dict[str, Any]) -> None: diff --git a/lib/services/whisper/whisper_service.dart b/lib/services/whisper/whisper_service.dart index e33e7d2..30d9bca 100644 --- a/lib/services/whisper/whisper_service.dart +++ b/lib/services/whisper/whisper_service.dart @@ -7,17 +7,27 @@ import '../audio/media_to_wav_converter.dart'; import 'whisperx_sidecar.dart'; class _WhisperExecutionConfig { - final String device; + final String asrDevice; + final String vadDevice; + final String alignDevice; final String computeType; final int batchSize; + final bool usingGpu; final bool usesCuda; + final String modeLabel; + final String? deviceName; final String? statusDetail; const _WhisperExecutionConfig({ - required this.device, + required this.asrDevice, + required this.vadDevice, + required this.alignDevice, required this.computeType, required this.batchSize, + required this.usingGpu, required this.usesCuda, + required this.modeLabel, + required this.deviceName, this.statusDetail, }); } @@ -34,8 +44,9 @@ class WhisperService { }; static const String _cpuDevice = 'cpu'; + static const String _mpsDevice = 'mps'; static const String _cpuComputeType = 'int8'; - static const int _cpuBatchSize = 4; + static const int _cpuBatchSize = 8; final WhisperXSidecar _sidecar = WhisperXSidecar(); final MediaToWavConverter _wavConverter = MediaToWavConverter(); @@ -143,30 +154,49 @@ class WhisperService { onRuntimeInfo: onRuntimeInfo, ); } catch (error) { - if (!primaryConfig.usesCuda || !_looksLikeCudaFailure(error)) { - rethrow; - } + if (primaryConfig.usesCuda) { + if (!_looksLikeCudaFailure(error)) { + rethrow; + } - final _WhisperExecutionConfig? degradedConfig = _buildDegradedCudaConfig( - primaryConfig, - ); - if (degradedConfig != null) { - await _restartSidecarForRetry(); - try { - return await _transcribeWithConfig( - wavPath: wavPath, - modelName: modelName, - language: language, - config: degradedConfig, - onStatus: onStatus, - onLog: onLog, - onRuntimeInfo: onRuntimeInfo, - ); - } catch (retryError) { - if (!_looksLikeCudaFailure(retryError)) { - rethrow; + final _WhisperExecutionConfig? degradedConfig = + _buildDegradedCudaConfig(primaryConfig); + if (degradedConfig != null) { + await _restartSidecarForRetry(); + try { + return await _transcribeWithConfig( + wavPath: wavPath, + modelName: modelName, + language: language, + config: degradedConfig, + onStatus: onStatus, + onLog: onLog, + onRuntimeInfo: onRuntimeInfo, + ); + } catch (retryError) { + if (!_looksLikeCudaFailure(retryError)) { + rethrow; + } } } + + await _restartSidecarForRetry(); + return _transcribeWithConfig( + wavPath: wavPath, + modelName: modelName, + language: language, + config: _cpuConfig( + statusDetail: + 'CUDA failed, falling back to CPU ($_cpuComputeType, batch=$_cpuBatchSize)', + ), + onStatus: onStatus, + onLog: onLog, + onRuntimeInfo: onRuntimeInfo, + ); + } + + if (!primaryConfig.usingGpu || !_looksLikeMpsFailure(error)) { + rethrow; } await _restartSidecarForRetry(); @@ -176,7 +206,7 @@ class WhisperService { language: language, config: _cpuConfig( statusDetail: - 'CUDA failed, falling back to CPU ($_cpuComputeType, batch=$_cpuBatchSize)', + 'MPS failed, falling back to CPU ($_cpuComputeType, batch=$_cpuBatchSize)', ), onStatus: onStatus, onLog: onLog, @@ -205,7 +235,9 @@ class WhisperService { wavPath: wavPath, modelName: modelName, language: language, - device: config.device, + device: config.asrDevice, + vadDevice: config.vadDevice, + alignDevice: config.alignDevice, computeType: config.computeType, batchSize: config.batchSize, noAlign: false, @@ -233,11 +265,7 @@ class WhisperService { Map _buildVadOptions(String? language) { switch (_normalizeLanguageCode(language)) { case 'ja': - return const { - 'chunk_size': 12, - 'vad_onset': 0.4, - 'vad_offset': 0.25, - }; + return const {}; default: return const {}; } @@ -270,6 +298,30 @@ class WhisperService { Future<_WhisperExecutionConfig> _resolveExecutionConfig( String whisperxModel, ) async { + if (Platform.isMacOS && _isArm64Process()) { + final WhisperXRuntimeProbe probe = await _loadRuntimeProbe(); + if (probe.canUseMps) { + return const _WhisperExecutionConfig( + asrDevice: _cpuDevice, + vadDevice: _mpsDevice, + alignDevice: _mpsDevice, + computeType: _cpuComputeType, + batchSize: _cpuBatchSize, + usingGpu: true, + usesCuda: false, + modeLabel: 'Mixed CPU + MPS', + deviceName: 'Apple Metal (MPS)', + statusDetail: + 'Using CPU ASR with MPS-accelerated VAD and alignment (int8, batch=8)', + ); + } + + final String statusDetail = probe.mpsBuilt + ? 'MPS is unavailable on this machine, using CPU ($_cpuComputeType, batch=$_cpuBatchSize)' + : 'Installed PyTorch runtime does not expose MPS, using CPU ($_cpuComputeType, batch=$_cpuBatchSize)'; + return _cpuConfig(statusDetail: statusDetail); + } + if (!Platform.isWindows) { return _cpuConfig(); } @@ -286,10 +338,15 @@ class WhisperService { : 'CUDA GPU'; return _WhisperExecutionConfig( - device: 'cuda', + asrDevice: 'cuda', + vadDevice: 'cuda', + alignDevice: 'cuda', computeType: computeType, batchSize: batchSize, + usingGpu: true, usesCuda: true, + modeLabel: 'CUDA GPU', + deviceName: deviceName, statusDetail: 'Using $deviceName on CUDA ($computeType, batch=$batchSize)', ); @@ -302,29 +359,23 @@ class WhisperService { Future _buildRuntimeInfo( _WhisperExecutionConfig config, ) async { - final WhisperXRuntimeProbe? probe = Platform.isWindows + final WhisperXRuntimeProbe? probe = + (Platform.isWindows || (Platform.isMacOS && _isArm64Process())) ? await _loadRuntimeProbe() : null; final bool cudaAvailable = probe?.cudaAvailable == true; final String? torchCudaVersion = probe?.torchCudaVersion; - final String? deviceName = config.usesCuda - ? (probe?.cudaDeviceName?.trim().isNotEmpty == true - ? probe!.cudaDeviceName!.trim() - : 'CUDA GPU') - : null; - final String modeLabel = switch ((config.usesCuda, cudaAvailable)) { - (true, _) => 'CUDA GPU', - (false, true) => 'CPU fallback', - (false, false) => 'CPU', - }; + final String modeLabel = !config.usingGpu && cudaAvailable + ? 'CPU fallback' + : config.modeLabel; return WhisperRuntimeInfo( modeLabel: modeLabel, - deviceName: deviceName, + deviceName: config.deviceName, computeType: config.computeType, batchSize: config.batchSize, - usingGpu: config.usesCuda, + usingGpu: config.usingGpu, cudaAvailable: cudaAvailable, torchCudaVersion: torchCudaVersion, note: config.statusDetail, @@ -333,10 +384,15 @@ class WhisperService { _WhisperExecutionConfig _cpuConfig({String? statusDetail}) { return _WhisperExecutionConfig( - device: _cpuDevice, + asrDevice: _cpuDevice, + vadDevice: _cpuDevice, + alignDevice: _cpuDevice, computeType: _cpuComputeType, batchSize: _cpuBatchSize, + usingGpu: false, usesCuda: false, + modeLabel: 'CPU', + deviceName: null, statusDetail: statusDetail, ); } @@ -356,10 +412,15 @@ class WhisperService { } return _WhisperExecutionConfig( - device: 'cuda', + asrDevice: 'cuda', + vadDevice: 'cuda', + alignDevice: 'cuda', computeType: 'int8', batchSize: smallerBatchSize, + usingGpu: true, usesCuda: true, + modeLabel: 'CUDA GPU', + deviceName: config.deviceName, statusDetail: 'CUDA init failed or VRAM is low, retrying lighter GPU mode (int8, batch=$smallerBatchSize)', ); @@ -409,6 +470,24 @@ class WhisperService { ].any(lower.contains); } + bool _looksLikeMpsFailure(Object error) { + final String lower = error.toString().toLowerCase(); + return [ + 'mps', + 'metal', + 'not implemented for mps', + 'placeholder storage has not been allocated on mps', + 'mps backend out of memory', + 'does not include mps support', + 'is not available on this machine', + ].any(lower.contains); + } + + bool _isArm64Process() { + final String version = Platform.version.toLowerCase(); + return version.contains('arm64') || version.contains('aarch64'); + } + Future _restartSidecarForRetry() async { await _sidecar.dispose(); } diff --git a/lib/services/whisper/whisperx_sidecar.dart b/lib/services/whisper/whisperx_sidecar.dart index f69f586..a9a5878 100644 --- a/lib/services/whisper/whisperx_sidecar.dart +++ b/lib/services/whisper/whisperx_sidecar.dart @@ -13,6 +13,8 @@ class WhisperXRuntimeProbe { final String whisperxVersion; final String? torchVersion; final String? torchCudaVersion; + final bool mpsBuilt; + final bool mpsAvailable; final bool cudaAvailable; final int cudaDeviceCount; final String? cudaDeviceName; @@ -28,6 +30,8 @@ class WhisperXRuntimeProbe { required this.whisperxVersion, required this.torchVersion, required this.torchCudaVersion, + required this.mpsBuilt, + required this.mpsAvailable, required this.cudaAvailable, required this.cudaDeviceCount, required this.cudaDeviceName, @@ -40,6 +44,7 @@ class WhisperXRuntimeProbe { bool get canUseCuda => cudaAvailable && cudaDeviceCount > 0 && cudaComputeTypes.isNotEmpty; + bool get canUseMps => mpsBuilt && mpsAvailable; factory WhisperXRuntimeProbe.fromPayload(Map payload) { final List computeTypes = @@ -56,6 +61,8 @@ class WhisperXRuntimeProbe { whisperxVersion: (payload['whisperx_version'] as String? ?? '').trim(), torchVersion: (payload['torch_version'] as String?)?.trim(), torchCudaVersion: (payload['torch_cuda_version'] as String?)?.trim(), + mpsBuilt: payload['mps_built'] == true, + mpsAvailable: payload['mps_available'] == true, cudaAvailable: payload['cuda_available'] == true, cudaDeviceCount: (payload['cuda_device_count'] as num?)?.toInt() ?? 0, cudaDeviceName: (payload['cuda_device_name'] as String?)?.trim(), @@ -323,6 +330,8 @@ class WhisperXSidecar { required String computeType, required int batchSize, required bool noAlign, + String? vadDevice, + String? alignDevice, Map? asrOptions, Map? vadOptions, Map? segmentationOptions, @@ -335,10 +344,17 @@ class WhisperXSidecar { 'model': modelName, 'language': language, 'device': device, + 'asr_device': device, 'compute_type': computeType, 'batch_size': batchSize, 'no_align': noAlign, }; + if (vadDevice != null && vadDevice.trim().isNotEmpty) { + params['vad_device'] = vadDevice.trim(); + } + if (alignDevice != null && alignDevice.trim().isNotEmpty) { + params['align_device'] = alignDevice.trim(); + } if (asrOptions != null && asrOptions.isNotEmpty) { params['asr_options'] = asrOptions; }