Mohamed Firas
83ae073127
Several nodes selected their device with a module-level `device = "cuda" if torch.cuda.is_available() else "cpu"`. On Apple Silicon `torch.cuda.is_available()` is False, so they silently ran on CPU even when ComfyUI had selected MPS. The model still runs, just far slower, so there is no error -- only a job that sits at 0% for minutes (reads as a hang). A single BiRefNet image took ~18 min on CPU; on MPS a 21s video processed in ~9 min (~1s/frame). Defer to ComfyUI's own device management via a shared helper `get_device()` in AILab_utils (wraps `comfy.model_management.get_torch_device()`), so nodes follow whatever backend ComfyUI chose -- matching what the SAM2/SAM3/ Florence2/Segment/SDMatte nodes already do. Centralizing it keeps the check from drifting again. Applied to BiRefNet, RMBG, ClothSegment, BodySegment, FashionSegment, FaceSegment, and SegmentV2. LamaRemover already uses get_torch_device() and is left unchanged. Fixes #200. Related to #91, #135.
130 lines
3.7 KiB
Python
130 lines
3.7 KiB
Python
# Shared utility helpers
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from comfy.utils import common_upscale
|
|
import comfy.model_management as mm
|
|
|
|
|
|
def get_device():
|
|
"""Return the torch device ComfyUI selected (CUDA, MPS, or CPU).
|
|
|
|
Defers to ComfyUI's device management instead of assuming CUDA-or-CPU, so nodes
|
|
follow whatever backend ComfyUI chose. Fixes models silently running on CPU on
|
|
Apple Silicon (MPS). See issues #200 and #135.
|
|
"""
|
|
return mm.get_torch_device()
|
|
|
|
|
|
def tensor2pil(image: torch.Tensor) -> Image.Image:
|
|
return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
|
|
def pil2tensor(image: Image.Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def pil2mask(image: Image.Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def resize_image(img: Image.Image, width: int, height: int) -> Image.Image:
|
|
return img.resize((width, height), resample=Image.LANCZOS)
|
|
|
|
|
|
def blend_overlay(img_1: Image.Image, img_2: Image.Image) -> Image.Image:
|
|
arr1 = np.array(img_1).astype(float) / 255.0
|
|
arr2 = np.array(img_2).astype(float) / 255.0
|
|
mask = arr2 < 0.5
|
|
result = np.zeros_like(arr1)
|
|
result[mask] = 2 * arr1[mask] * arr2[mask]
|
|
result[~mask] = 1 - 2 * (1 - arr1[~mask]) * (1 - arr2[~mask])
|
|
return Image.fromarray(np.clip(result * 255, 0, 255).astype(np.uint8))
|
|
|
|
|
|
def fill_mask(width: int, height: int, mask: Image.Image, box=(0, 0), color=0) -> Image.Image:
|
|
bg = Image.new("L", (width, height), color)
|
|
bg.paste(mask, box, mask)
|
|
return bg
|
|
|
|
|
|
def empty_image(width: int, height: int, batch_size: int = 1) -> torch.Tensor:
|
|
return torch.zeros([batch_size, height, width, 3])
|
|
|
|
|
|
def upscale_mask(mask: torch.Tensor, width: int, height: int) -> torch.Tensor:
|
|
if mask.ndim == 3:
|
|
mask = mask.unsqueeze(1)
|
|
mask = common_upscale(mask, width, height, "bicubic", "disabled")
|
|
mask = mask.squeeze(1)
|
|
return mask
|
|
|
|
|
|
def extract_alpha_mask(image: torch.Tensor) -> torch.Tensor:
|
|
alpha = image[..., 3]
|
|
if alpha.max() > 1.0:
|
|
alpha = alpha / 255.0
|
|
if len(alpha.shape) == 4:
|
|
alpha = alpha[:, :, :, 0]
|
|
return alpha.unsqueeze(1) if alpha.ndim == 3 else alpha
|
|
|
|
|
|
def ensure_mask_shape(mask: torch.Tensor | None) -> torch.Tensor | None:
|
|
if mask is None:
|
|
return None
|
|
if mask.ndim == 2:
|
|
return mask.unsqueeze(0)
|
|
if mask.ndim == 4 and mask.shape[1] == 1:
|
|
return mask.squeeze(1)
|
|
return mask
|
|
|
|
|
|
COLOR_PRESETS = {
|
|
"black": "#000000",
|
|
"white": "#FFFFFF",
|
|
"red": "#FF0000",
|
|
"green": "#00FF00",
|
|
"blue": "#0000FF",
|
|
"yellow": "#FFFF00",
|
|
"cyan": "#00FFFF",
|
|
"magenta": "#FF00FF",
|
|
"gray": "#808080",
|
|
"silver": "#C0C0C0",
|
|
"maroon": "#800000",
|
|
"olive": "#808000",
|
|
"purple": "#800080",
|
|
"teal": "#008080",
|
|
"navy": "#000080",
|
|
"orange": "#FFA500",
|
|
"pink": "#FFC0CB",
|
|
"brown": "#A52A2A",
|
|
"violet": "#EE82EE",
|
|
"indigo": "#4B0082",
|
|
"light_gray": "#D3D3D3",
|
|
"dark_gray": "#A9A9A9",
|
|
"light_blue": "#ADD8E6",
|
|
"dark_blue": "#00008B",
|
|
"light_green": "#90EE90",
|
|
"dark_green": "#006400",
|
|
}
|
|
|
|
|
|
def color_format(color: str) -> str:
|
|
if not color:
|
|
return ""
|
|
|
|
color = color.strip().upper()
|
|
if not color.startswith("#"):
|
|
color = f"#{color}"
|
|
|
|
color = color[1:]
|
|
if len(color) == 3:
|
|
r, g, b = color[0], color[1], color[2]
|
|
return f"#{r}{r}{g}{g}{b}{b}"
|
|
if len(color) < 6:
|
|
raise ValueError(f"Invalid color format: {color}")
|
|
if len(color) > 6:
|
|
color = color[:6]
|
|
|
|
return f"#{color}"
|