From 62cf5aa51cbdd1a1af5e095c351148f7c6728971 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 2 Dec 2025 05:36:01 -0800
Subject: [PATCH] Don't assume CUDA GPU

---
 modules/ui_image_generation.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/modules/ui_image_generation.py b/modules/ui_image_generation.py
index b8a0653e..484678f4 100644
--- a/modules/ui_image_generation.py
+++ b/modules/ui_image_generation.py
@@ -17,6 +17,7 @@ from modules.image_models import (
     unload_image_model
 )
 from modules.logging_colors import logger
+from modules.torch_utils import get_device
 from modules.utils import gradio
 
 ASPECT_RATIOS = {
@@ -657,7 +658,11 @@ def generate(state):
     if seed == -1:
         seed = np.random.randint(0, 2**32 - 1)
 
-    generator = torch.Generator("cuda").manual_seed(int(seed))
+    device = get_device()
+    if device is None:
+        device = "cpu"
+    generator = torch.Generator(device).manual_seed(int(seed))
+
     all_images = []
 
     # Get pipeline type for parameter adjustment