diff --git a/tools/enhancement/face_restore.py b/tools/enhancement/face_restore.py index 277091d3..cc2408b3 100644 --- a/tools/enhancement/face_restore.py +++ b/tools/enhancement/face_restore.py @@ -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}" diff --git a/tools/enhancement/upscale.py b/tools/enhancement/upscale.py index c87d045b..03648405 100644 --- a/tools/enhancement/upscale.py +++ b/tools/enhancement/upscale.py @@ -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 diff --git a/tools/video/_shared.py b/tools/video/_shared.py index 762d0dbf..d4d679c8 100644 --- a/tools/video/_shared.py +++ b/tools/video/_shared.py @@ -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"):