Update aurasr_upscaler.py

AUTOMATIC1111/stable-diffusion-webui-extensions#356 (comment)
AndreyRGW · Oct 27, 2024 · 71b6821 · 71b6821
1 parent da26c5a
commit 71b6821
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/scripts/aurasr_upscaler.py b/scripts/aurasr_upscaler.py
@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from modules.upscaler import Upscaler, UpscalerData
+from modules import devices
 from scripts.aurasr_tools import runAuraSR
 
 @dataclass
@@ -22,7 +23,9 @@ def __init__(self, dirname, fields: Fields = None):
         super().__init__()
 
     def do_upscale(self, img, selected_model):
-        return runAuraSR(img)
+        # Get the appropriate device for AuraSR
+        device = devices.get_device_for('aurasr')
+        return runAuraSR(img, device)
 
 class Class0(BaseClass, Upscaler):
     def __init__(self, dirname):