diff --git a/backend/python/diffusers/backend.py b/backend/python/diffusers/backend.py index f32187a25..312754a34 100755 --- a/backend/python/diffusers/backend.py +++ b/backend/python/diffusers/backend.py @@ -122,6 +122,21 @@ from diffusers.schedulers import ( UniPCMultistepScheduler, ) +def select_device(request_cuda, device_option, cuda_available, xpu, mps_available): + """Pick the pipeline device. An explicit `device:` model option wins; + otherwise CUDA is used whenever torch reports it available (ROCm + builds included) or the model config forces it with `cuda: true`, + keeping the pre-existing XPU/MPS overrides. CPU is the fallback, not + the default.""" + if device_option: + return device_option + device = "cuda" if (request_cuda or cuda_available) else "cpu" + if xpu: + device = "xpu" + if mps_available: + device = "mps" + return device + def is_float(s): """Check if a string can be converted to float.""" try: @@ -627,12 +642,13 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): # modify LoraAdapter to be relative to modelFileBase request.LoraAdapter = os.path.join(request.ModelPath, request.LoraAdapter) - device = "cpu" if not request.CUDA else "cuda" - if XPU: - device = "xpu" - mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available() - if mps_available: - device = "mps" + device = select_device( + request.CUDA, + self.options.pop("device", None), + torch.cuda.is_available(), + XPU, + hasattr(torch.backends, "mps") and torch.backends.mps.is_available(), + ) self.device = device if request.LoraAdapter: # Check if its a local file and not a directory ( we load lora differently for a safetensor file ) diff --git a/backend/python/diffusers/test.py b/backend/python/diffusers/test.py index 2922f2a03..f2a07d4e3 100644 --- a/backend/python/diffusers/test.py +++ b/backend/python/diffusers/test.py @@ -7,6 +7,7 @@ import time from unittest.mock import patch, MagicMock # Import dynamic loader for testing (these don't need gRPC) +import backend import diffusers_dynamic_loader as loader from diffusers import DiffusionPipeline, StableDiffusionPipeline @@ -425,3 +426,22 @@ class TestGenerateImageOptionsKwargsMerge(unittest.TestCase): self.assertEqual(pipeline.kwargs["num_inference_steps"], 4) finally: os.unlink(dst_path) + + +class TestDeviceSelection(unittest.TestCase): + """Unit tests for backend.select_device (no GPU required).""" + + def test_autodetect_cuda(self): + self.assertEqual(backend.select_device(False, None, True, False, False), "cuda") + + def test_cpu_fallback(self): + self.assertEqual(backend.select_device(False, None, False, False, False), "cpu") + + def test_forced_cuda(self): + self.assertEqual(backend.select_device(True, None, False, False, False), "cuda") + + def test_device_option_wins(self): + self.assertEqual(backend.select_device(True, "cpu", True, True, True), "cpu") + + def test_mps_overrides(self): + self.assertEqual(backend.select_device(False, None, True, False, True), "mps") diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index b748fca05..d629e1be8 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -739,7 +739,7 @@ For image generation models using the `diffusers` backend: | Field | Type | Description | |-------|------|-------------| -| `diffusers.cuda` | bool | Enable CUDA for diffusers | +| `diffusers.cuda` | bool | Force CUDA. By default the backend auto-detects and uses CUDA when a compatible GPU is present (ROCm builds included). Pin the CPU with `options: ["device:cpu"]` | | `diffusers.pipeline_type` | string | Pipeline type (e.g., `stable-diffusion`, `stable-diffusion-xl`) | | `diffusers.scheduler_type` | string | Scheduler type (e.g., `euler`, `ddpm`) | | `diffusers.enable_parameters` | string | Comma-separated parameters to enable |