Resolution for support of MPS on Apple Silicon. R1, R2, R3, and R4 implemented. Installation detection and speech enhanceres for MPS if available and supported

Author: meizhong986

install.py

@@ -749,6 +749,11 @@ def _get_translate_deps_from_registry() -> list:
     return _get_packages_for_step([Extra.TRANSLATE])
 
 
+def _get_llm_deps_from_registry() -> list:
+    """Get LLM server package specs from registry (uvicorn, fastapi, etc.)."""
+    return _get_packages_for_step([Extra.LLM])
+
+
 def _get_gui_deps_from_registry() -> list:
     """Get GUI package specs from registry."""
     return _get_packages_for_step([Extra.GUI])
@@ -887,7 +892,7 @@ def detect_cuda_version(args) -> str:
         args: Parsed command-line arguments
 
     Returns:
-        "cpu", "cu118", or "cu128"
+        "cpu", "cu118", "cu128", or "metal" (Apple Silicon)
     """
     # Explicit user request
     if args.cpu_only:
@@ -902,7 +907,8 @@ def detect_cuda_version(args) -> str:
         gpu_info = detect_gpu()
         if gpu_info.detected:
             log(f"\n  GPU detected: {gpu_info.name}")
-            log(f"  Driver: {gpu_info.driver_version[0]}.{gpu_info.driver_version[1]}")
+            if gpu_info.driver_version:
+                log(f"  Driver: {gpu_info.driver_version[0]}.{gpu_info.driver_version[1]}")
             log(f"  Selected: {gpu_info.cuda_version or 'CPU'}")
             return gpu_info.cuda_version or "cpu"
         else:
@@ -1229,12 +1235,21 @@ def main():
     # All subsequent packages that depend on torch will see it as satisfied.
     #
     print_header("Installing PyTorch", "Step 2/6")
-    torch_url = get_torch_index_url(cuda_version)
-    run_pip(
-        executor,
-        ["install", "torch", "torchaudio", "--index-url", torch_url],
-        f"Install PyTorch ({cuda_version})"
-    )
+    if cuda_version == "metal":
+        # Apple Silicon: install from default PyPI (has arm64 wheels with MPS support)
+        # Do NOT use --index-url — the CPU index provides x86_64-only builds
+        run_pip(
+            executor,
+            ["install", "torch", "torchaudio"],
+            "Install PyTorch (Metal/MPS)"
+        )
+    else:
+        torch_url = get_torch_index_url(cuda_version)
+        run_pip(
+            executor,
+            ["install", "torch", "torchaudio", "--index-url", torch_url],
+            f"Install PyTorch ({cuda_version})"
+        )
 
     # -------------------------------------------------------------------------
     # Step 3: Install core dependencies
@@ -1309,6 +1324,14 @@ def main():
     if translate_deps:
         run_pip(executor, ["install"] + translate_deps, "Install translation packages")
 
+    # LLM server framework (from registry) - uvicorn, fastapi, etc.
+    # These are needed for the local LLM translation server.
+    # Installed unconditionally (lightweight, pure Python) so the server
+    # is ready if the user installs llama-cpp-python now or later.
+    llm_deps = _get_llm_deps_from_registry()
+    if llm_deps:
+        run_pip(executor, ["install"] + llm_deps, "Install LLM server packages")
+
     # -------------------------------------------------------------------------
     # Local LLM (llama-cpp-python) - Included by Default
     # -------------------------------------------------------------------------
@@ -1472,25 +1495,51 @@ def _install_local_llm(executor: StepExecutor, build_from_source: bool):
     is_intel_mac = (sys.platform == "darwin" and platform_module.machine() != "arm64")
 
     if is_apple_silicon:
-        # Apple Silicon: build from source with Metal
-        log("    Apple Silicon detected - building from source with Metal support.")
-        if get_llama_cpp_source_info:
-            git_url, backend, cmake_args, env_vars = get_llama_cpp_source_info()
-            log(f"    Backend: {backend}")
-            for key, value in env_vars.items():
-                log(f"    Setting {key}={value}")
-                os.environ[key] = value
-            if cmake_args:
-                log(f"    Setting CMAKE_ARGS={cmake_args}")
-                os.environ["CMAKE_ARGS"] = cmake_args
-            run_pip(
-                executor,
-                ["install", git_url],
-                f"Install llama-cpp-python ({backend})",
-                allow_fail=True
-            )
-        else:
-            log("    ERROR: llama_build_utils not available")
+        # Apple Silicon: try prebuilt Metal wheel first, then source build
+        log("    Apple Silicon detected - checking for prebuilt Metal wheel...")
+
+        # Strategy 1: Try prebuilt Metal wheel (fast, ~1 min download)
+        prebuilt_installed = False
+        if get_prebuilt_wheel_url:
+            wheel_url, wheel_backend = get_prebuilt_wheel_url(verbose=True)
+            if wheel_url:
+                log(f"    Found prebuilt wheel: {wheel_backend}")
+                prebuilt_installed = run_pip(
+                    executor,
+                    ["install", wheel_url],
+                    f"Install llama-cpp-python ({wheel_backend})",
+                    allow_fail=True
+                )
+                if prebuilt_installed:
+                    run_pip(
+                        executor,
+                        ["install", "llama-cpp-python[server]"],
+                        "Install llama-cpp-python server extras",
+                        allow_fail=True
+                    )
+            else:
+                log("    No prebuilt Metal wheel found.")
+
+        # Strategy 2: Fall back to source build with Metal (slow, ~10 min)
+        if not prebuilt_installed:
+            log("    Building from source with Metal support...")
+            if get_llama_cpp_source_info:
+                git_url, backend, cmake_args, env_vars = get_llama_cpp_source_info()
+                log(f"    Backend: {backend}")
+                for key, value in env_vars.items():
+                    log(f"    Setting {key}={value}")
+                    os.environ[key] = value
+                if cmake_args:
+                    log(f"    Setting CMAKE_ARGS={cmake_args}")
+                    os.environ["CMAKE_ARGS"] = cmake_args
+                run_pip(
+                    executor,
+                    ["install", git_url],
+                    f"Install llama-cpp-python ({backend})",
+                    allow_fail=True
+                )
+            else:
+                log("    ERROR: llama_build_utils not available")
 
     elif is_intel_mac:
         # Intel Mac: CPU only

whisperjav/installer/core/detector.py

@@ -221,13 +221,13 @@ def detect_gpu() -> GPUInfo:
     plat = detect_platform()
     if plat == DetectedPlatform.MACOS_SILICON:
         return GPUInfo(
-            detected=False,
-            name=None,
+            detected=True,  # Apple Silicon HAS a GPU (Metal/MPS)
+            name="Apple Silicon (Metal)",
             driver_version=None,
-            cuda_version=None,
-            torch_index=CPU_TORCH_INDEX,  # macOS uses CPU/Metal, not CUDA
+            cuda_version="metal",  # Sentinel: install torch from default PyPI (has MPS)
+            torch_index="",  # Empty = no --index-url needed, default PyPI has arm64+MPS wheels
             detection_method="platform",
-            message="Apple Silicon detected - CUDA not supported, using Metal",
+            message="Apple Silicon detected - using Metal/MPS acceleration",
         )
 
     # Try nvidia-smi first (most reliable)
@@ -499,6 +499,9 @@ def get_torch_index_url(cuda_version: Optional[str] = None) -> str:
     if cuda_version == "cpu":
         return CPU_TORCH_INDEX
 
+    if cuda_version == "metal":
+        return ""  # Apple Silicon: use default PyPI (has arm64+MPS wheels)
+
     # Find in CUDA matrix
     for entry in CUDA_DRIVER_MATRIX:
         if entry.cuda_version == cuda_version:

whisperjav/modules/speech_enhancement/backends/bs_roformer.py

@@ -26,6 +26,7 @@
     load_audio_to_array,
     create_failed_result,
     resample_audio,
+    resolve_torch_device,
 )
 
 logger = logging.getLogger("whisperjav")
@@ -107,17 +108,23 @@ def _ensure_initialized(self) -> bool:
 
             logger.info(f"Loading BS-RoFormer model for stem: {self._model_name}")
 
-            # Determine device
-            device = self._device
-            if device == "auto":
-                try:
-                    import torch
-                    device = "cuda" if torch.cuda.is_available() else "cpu"
-                except ImportError:
-                    device = "cpu"
+            # Resolve best available device (cuda > mps > cpu)
+            device = resolve_torch_device(self._device)
 
-            # Initialize separator
-            self._separator = BSRoformer(device=device)
+            # Try initializing with the selected device
+            try:
+                self._separator = BSRoformer(device=device)
+            except Exception as dev_err:
+                # If MPS failed (model may not support all ops), fall back to CPU
+                if device == "mps":
+                    logger.info(
+                        f"BS-RoFormer does not support MPS ({dev_err}), "
+                        "falling back to CPU"
+                    )
+                    device = "cpu"
+                    self._separator = BSRoformer(device=device)
+                else:
+                    raise
 
             self._initialized = True
             logger.info(f"BS-RoFormer loaded successfully on {device}")

whisperjav/modules/speech_enhancement/backends/zipenhancer.py

@@ -34,6 +34,7 @@
     load_audio_to_array,
     create_failed_result,
     resample_audio,
+    resolve_torch_device,
 )
 
 logger = logging.getLogger("whisperjav")
@@ -154,21 +155,44 @@ def _init_torch(self) -> bool:
 
             logger.info("Loading ZipEnhancer via ModelScope (torch)...")
 
-            # Determine device
-            device = self.device
-            if device is None:
-                device = "gpu" if torch.cuda.is_available() else "cpu"
-            elif device == "cuda":
-                device = "gpu"  # ModelScope uses "gpu" not "cuda"
-
-            self._pipeline = pipeline(
-                Tasks.acoustic_noise_suppression,
-                model=MODELSCOPE_MODEL_ID,
-                device=device
-            )
+            # Resolve best available device (cuda > mps > cpu)
+            torch_device = resolve_torch_device(self.device)
+
+            # ModelScope uses "gpu" for CUDA, "cpu" for CPU.
+            # For MPS: try passing "mps" directly — ModelScope >= 1.20 may
+            # forward it to PyTorch. If it fails, fall back to CPU.
+            if torch_device == "cuda":
+                ms_device = "gpu"
+            elif torch_device == "mps":
+                ms_device = "mps"  # Will be validated below with fallback
+            else:
+                ms_device = "cpu"
+
+            # Try to create the pipeline with the selected device
+            try:
+                self._pipeline = pipeline(
+                    Tasks.acoustic_noise_suppression,
+                    model=MODELSCOPE_MODEL_ID,
+                    device=ms_device
+                )
+            except Exception as mps_err:
+                # If MPS failed (ModelScope may not support it), fall back to CPU
+                if ms_device == "mps":
+                    logger.info(
+                        f"ModelScope does not support MPS device ({mps_err}), "
+                        "falling back to CPU"
+                    )
+                    ms_device = "cpu"
+                    self._pipeline = pipeline(
+                        Tasks.acoustic_noise_suppression,
+                        model=MODELSCOPE_MODEL_ID,
+                        device=ms_device
+                    )
+                else:
+                    raise
 
             self._initialized = True
-            logger.info(f"ZipEnhancer (torch) loaded successfully on {device}")
+            logger.info(f"ZipEnhancer (torch) loaded successfully on {ms_device}")
             return True
 
         except ImportError as e:
@@ -213,13 +237,14 @@ def _init_onnx(self) -> bool:
                     f"Directory contents: {contents}"
                 )
 
-            # Configure ONNX providers
+            # Configure ONNX providers (CUDA > CoreML/Apple Silicon > CPU)
+            available_providers = ort.get_available_providers()
             providers = ['CPUExecutionProvider']
             if self.device != 'cpu':
-                # Try CUDA first if available
-                available_providers = ort.get_available_providers()
                 if 'CUDAExecutionProvider' in available_providers:
                     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+                elif 'CoreMLExecutionProvider' in available_providers:
+                    providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider']
 
             self._onnx_session = ort.InferenceSession(
                 onnx_model_path,

whisperjav/modules/speech_enhancement/base.py

@@ -206,6 +206,47 @@ def is_lightweight(self) -> bool:
         ...
 
 
+def resolve_torch_device(requested: Optional[str] = None) -> str:
+    """
+    Resolve the best available torch device for speech enhancement.
+
+    Priority: explicit request > CUDA > MPS (Apple Silicon) > CPU.
+    Validates that the device actually works before returning it.
+
+    Args:
+        requested: Explicitly requested device ("cuda", "mps", "cpu", "auto", or None).
+                   None and "auto" both trigger auto-detection.
+
+    Returns:
+        Device string suitable for torch: "cuda", "mps", or "cpu".
+    """
+    try:
+        import torch
+    except ImportError:
+        return "cpu"
+
+    # Explicit request (not auto)
+    if requested and requested not in (None, "auto"):
+        if requested == "cuda" and torch.cuda.is_available():
+            return "cuda"
+        if requested == "mps":
+            if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                return "mps"
+            logger.debug("MPS requested but not available, falling back to CPU")
+            return "cpu"
+        if requested == "cpu":
+            return "cpu"
+        # Unknown device string — fall through to auto-detect
+        logger.debug(f"Unknown device '{requested}', auto-detecting")
+
+    # Auto-detect: CUDA > MPS > CPU
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
 def load_audio_to_array(
     audio: Union[np.ndarray, Path, str],
     target_sample_rate: Optional[int] = None