From dfdc51246c678b585e1bdfdb7d2f202b0ca0e362 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Wed, 3 Jan 2024 22:38:13 +0200
Subject: SwinIR: use prefer_half

---
 extensions-builtin/SwinIR/scripts/swinir_model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'extensions-builtin')

diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py
index bc427fea..6a8e21b0 100644
--- a/extensions-builtin/SwinIR/scripts/swinir_model.py
+++ b/extensions-builtin/SwinIR/scripts/swinir_model.py
@@ -1,6 +1,7 @@
 import logging
 import sys
 
+import torch
 from PIL import Image
 
 from modules import devices, modelloader, script_callbacks, shared, upscaler_utils
@@ -69,7 +70,7 @@ class UpscalerSwinIR(Upscaler):
         model_descriptor = modelloader.load_spandrel_model(
             filename,
             device=self._get_device(),
-            dtype=devices.dtype,
+            prefer_half=(devices.dtype == torch.float16),
             expected_architecture="SwinIR",
         )
         if getattr(shared.opts, 'SWIN_torch_compile', False):
-- 
cgit v1.2.3