freedak/OpenMontage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files
tools/enhancement/face_restore.py
+32 -17
View File
@@ -37,7 +37,9 @@ class FaceRestore(BaseTool):
dependencies = ["python:gfpgan", "python:torch"]
install_instructions = (
"pip install gfpgan # Includes CodeFormer support. Requires PyTorch."
"uv pip install gfpgan torch\n"
"Works on: CUDA (NVIDIA), MPS (Apple Silicon M-series, macOS >= 12.3), CPU fallback.\n"
"No CUDA build needed on macOS — uv pip install torch includes MPS support."
)
agent_skills = ["ffmpeg"]
fallback = None
@@ -121,13 +123,18 @@ class FaceRestore(BaseTool):
try:
import cv2
import inspect
from gfpgan import GFPGANer
import torch
from tools.video._shared import get_torch_device as _get_device
except ImportError as e:
return ToolResult(
success=False,
error=f"Missing dependency: {e}. Run: pip install gfpgan",
error=f"Missing dependency: {e}. Run: uv pip install gfpgan",
)
_device = _get_device()
start = time.time()
# Optional background upsampler
@@ -141,18 +148,22 @@ class FaceRestore(BaseTool):
num_in_ch=3, num_out_ch=3, num_feat=64,
num_block=23, num_grow_ch=32, scale=2,
)
bg_upsampler = RealESRGANer(
scale=2,
model_path=(
bg_kwargs: dict = {
"scale": 2,
"model_path": (
"https://github.com/xinntao/Real-ESRGAN/releases/download/"
"v0.2.1/RealESRGAN_x2plus.pth"
),
model=realesrgan_model,
tile=400,
tile_pad=10,
pre_pad=0,
half=True,
)
"model": realesrgan_model,
"tile": 400,
"tile_pad": 10,
"pre_pad": 0,
"half": (_device == "cuda"),
}
# Guard: only pass device= if the installed version accepts it
if "device" in inspect.signature(RealESRGANer.__init__).parameters:
bg_kwargs["device"] = torch.device(_device)
bg_upsampler = RealESRGANer(**bg_kwargs)
except ImportError:
bg_upsampler = None
@@ -172,12 +183,16 @@ class FaceRestore(BaseTool):
# Instantiate restorer
try:
restorer = GFPGANer(
model_path=model_path,
upscale=upscale,
arch=arch,
bg_upsampler=bg_upsampler,
)
restorer_kwargs: dict = {
"model_path": model_path,
"upscale": upscale,
"arch": arch,
"bg_upsampler": bg_upsampler,
}
# Guard: only pass device= if the installed version accepts it
if "device" in inspect.signature(GFPGANer.__init__).parameters:
restorer_kwargs["device"] = torch.device(_device)
restorer = GFPGANer(**restorer_kwargs)
except Exception as e:
return ToolResult(
success=False, error=f"Failed to load {model_name} model: {e}"
tools/enhancement/upscale.py
+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