feat: enhance WhisperX integration with MPS support and refactor device management

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:

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<String, dynamic> _buildVadOptions(String? language) {
     switch (_normalizeLanguageCode(language)) {
       case 'ja':
-        return const <String, dynamic>{
-          'chunk_size': 12,
-          'vad_onset': 0.4,
-          'vad_offset': 0.25,
-        };
+        return const <String, dynamic>{};
       default:
         return const <String, dynamic>{};
     }
@@ -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<WhisperRuntimeInfo> _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 <String>[
+      '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<void> _restartSidecarForRetry() async {
     await _sidecar.dispose();
   }

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<String, dynamic> payload) {
     final List<String> 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<String, dynamic>? asrOptions,
     Map<String, dynamic>? vadOptions,
     Map<String, dynamic>? 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;
     }