Fix device selection on Apple Silicon (MPS) (#205)

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.

This commit is contained in:

Mohamed Firas

2026-08-21 15:51:39 -07:00

committed by

GitHub

parent bb0531a1ae

commit 83ae073127

8 changed files with 25 additions and 7 deletions

									
										py/AILab_BiRefNet.py
									
		+2
		-1
	
												View File
												
				@@ -23,7 +23,8 @@ import importlib.util

				from safetensors.torch import load_file

				import cv2

				device = "cuda" if torch.cuda.is_available() else "cpu"

				from AILab_utils import get_device

				device = get_device()

				# Add model path

				folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))