zlei6/OpenMontage
1
0
Fork 0
mirror of https://github.com/calesthio/OpenMontage.git synced 2026-08-14 04:37:13 +08:00
Code Issues 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

49
tools/enhancement/face_restore.py
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}"

50
tools/enhancement/upscale.py
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

70
tools/video/_shared.py
View File

@@ -99,7 +99,7 @@ LTX_LOCAL_VARIANTS = {
"default_width": 768,
"default_height": 512,
"default_num_frames": 121,
"fps": 24,
"fps": 30,
},
}
@@ -147,6 +147,39 @@ LTX2_FRAME_COUNTS = {
}
def get_torch_device() -> str:
"""Return best available torch device: cuda > mps (Apple Silicon Metal) > cpu.
Priority order:
1. cuda — NVIDIA GPU (fastest for most diffusion models)
2. mps — Apple Silicon Metal (M1/M2/M3/M4/M5, macOS >= 12.3)
3. cpu — fallback, always available but slow
MPS detection is guarded for torch builds that lack ``torch.backends.mps``
(e.g. older pip wheels or Linux builds). We check both build-time support
(``is_built()``) and runtime availability (``is_available()``).
"""
try:
import torch as _torch # noqa: PLC0415
except ImportError:
return "cpu"
if _torch.cuda.is_available():
return "cuda"
# Guard: torch.backends.mps may not exist on older/non-macOS builds
try:
mps_backend = getattr(_torch, "backends", None)
mps_backend = getattr(mps_backend, "mps", None) if mps_backend else None
if mps_backend is not None:
# Check build-time support first, then runtime availability
is_built = getattr(mps_backend, "is_built", lambda: True)()
is_available = getattr(mps_backend, "is_available", lambda: False)()
if is_built and is_available:
return "mps"
except Exception:
pass
return "cpu"
def local_generation_enabled() -> bool:
return os.environ.get("VIDEO_GEN_LOCAL_ENABLED", "").lower() in {"true", "1", "yes"}
@@ -165,9 +198,15 @@ def local_generation_status() -> ToolStatus:
def local_install_instructions() -> str:
return (
"Enable local video generation and install the diffusers stack:\n"
" set VIDEO_GEN_LOCAL_ENABLED=true\n"
" pip install diffusers transformers accelerate torch pillow requests\n"
"Use a GPU with the VRAM profile listed on the selected tool."
" export VIDEO_GEN_LOCAL_ENABLED=true\n"
" uv pip install diffusers transformers accelerate torch pillow requests\n"
"\n"
"GPU support — pick what matches your hardware:\n"
" NVIDIA CUDA — works out of the box with the above\n"
" Apple Silicon (MPS, macOS >= 12.3) — works out of the box; no extra build\n"
" CPU fallback — slow but functional on any machine\n"
"\n"
"VRAM profile: see the selected tool's resource_profile for minimum VRAM."
)
@@ -201,13 +240,30 @@ def load_diffusers_pipeline(pipeline_class: str, model_id: str, enable_offload:
}
pipeline_name = pipeline_map.get(pipeline_class, pipeline_class)
pipeline_class_obj = getattr(diffusers, pipeline_name)
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
device = get_torch_device()
# bfloat16 is only reliable on CUDA; MPS uses float16 for inference,
# CPU must use float32 (float16 is emulated and unreliable on CPU)
if device == "cuda" and torch.cuda.is_bf16_supported():
dtype = torch.bfloat16
elif device == "cpu":
dtype = torch.float32
else:
dtype = torch.float16
pipeline = pipeline_class_obj.from_pretrained(model_id, torch_dtype=dtype)
if enable_offload:
pipeline.enable_model_cpu_offload()
if device == "cuda":
pipeline.enable_model_cpu_offload()
else:
# enable_model_cpu_offload() is CUDA-only; fall back to direct device placement
pipeline = pipeline.to(device)
else:
pipeline = pipeline.to("cuda")
pipeline = pipeline.to(device)
if hasattr(pipeline, "enable_attention_slicing"):
pipeline.enable_attention_slicing()
if hasattr(pipeline, "vae") and pipeline.vae is not None:
if hasattr(pipeline.vae, "enable_tiling"):