feat(gpu): add device-aware routing for Apple Silicon MPS support

- Add get_torch_device() helper in _shared.py: cuda > mps > cpu
- Guard MPS detection for torch builds lacking torch.backends.mps
- Check both is_built() and is_available() for MPS
- Route load_diffusers_pipeline() to resolved device instead of hardcoded cuda
- Use float32 on CPU (float16 is emulated/unreliable), float16 on MPS, bfloat16 on CUDA
- Guard enable_model_cpu_offload() to CUDA-only; fall back to .to(device) on MPS
- Enable attention slicing for MPS memory safety
- Add inspect-based signature guard for device= arg on RealESRGANer/GFPGANer
- Update install_instructions on all LOCAL_GPU tools to mention MPS/Apple Silicon
This commit is contained in:
Nik Nichols
2026-06-24 16:29:09 -05:00
parent 28dda8c808
commit 211c3fcba4
3 changed files with 129 additions and 40 deletions
+34 -16
View File
@@ -56,7 +56,11 @@ class Upscale(BaseTool):
runtime = ToolRuntime.LOCAL_GPU
dependencies = ["python:realesrgan", "python:torch", "cmd:ffmpeg"]
install_instructions = "pip install realesrgan # Requires PyTorch with CUDA"
install_instructions = (
"uv pip install realesrgan torch\n"
"Works on: CUDA (NVIDIA), MPS (Apple Silicon M-series, macOS >= 12.3), CPU fallback.\n"
"No separate CUDA build needed on macOS — uv pip install torch includes MPS support."
)
agent_skills = ["ffmpeg"]
capabilities = [
@@ -266,6 +270,8 @@ class Upscale(BaseTool):
face_enhance: bool,
):
"""Build and return a RealESRGANer instance."""
import inspect
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
@@ -281,25 +287,37 @@ class Upscale(BaseTool):
if model_name == "RealESRGAN_x4plus_anime_6B":
model_url = f"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/{model_name}.pth"
half = torch.cuda.is_available()
from tools.video._shared import get_torch_device as _get_device
_device = _get_device()
half = _device == "cuda" # fp16 only safe on CUDA; MPS/CPU use fp32 for realesrgan
upsampler = RealESRGANer(
scale=4,
model_path=model_url,
model=model,
dni_weight=denoise_strength,
half=half,
)
upsampler_kwargs: dict = {
"scale": 4,
"model_path": model_url,
"model": model,
"dni_weight": denoise_strength,
"half": half,
}
# Guard: only pass device= if the installed version accepts it
if "device" in inspect.signature(RealESRGANer.__init__).parameters:
upsampler_kwargs["device"] = torch.device(_device)
upsampler = RealESRGANer(**upsampler_kwargs)
if face_enhance:
from gfpgan import GFPGANer
face_enhancer = GFPGANer(
model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
upscale=scale,
arch="clean",
channel_multiplier=2,
bg_upsampler=upsampler,
)
face_kwargs: dict = {
"model_path": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
"upscale": scale,
"arch": "clean",
"channel_multiplier": 2,
"bg_upsampler": upsampler,
}
# Guard: only pass device= if the installed version accepts it
if "device" in inspect.signature(GFPGANer.__init__).parameters:
face_kwargs["device"] = torch.device(_device)
face_enhancer = GFPGANer(**face_kwargs)
# Monkey-patch so the caller can use the same interface
original_enhance = upsampler.enhance