From f261a4a53c153c630a506bc5282e9955c36b3ef2 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 11 Mar 2023 11:56:05 +0300
Subject: use selected device instead of always cuda for UniPC sampler

---
 modules/models/diffusion/uni_pc/sampler.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

(limited to 'modules/models/diffusion')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index 6bb3bb21..bf346ff4 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -3,7 +3,8 @@
 import torch
 
 from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
-from modules import shared
+from modules import shared, devices
+
 
 class UniPCSampler(object):
     def __init__(self, model, **kwargs):
@@ -16,8 +17,8 @@ class UniPCSampler(object):
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
+            if attr.device != devices.device:
+                attr = attr.to(devices.device)
         setattr(self, name, attr)
 
     def set_hooks(self, before_sample, after_sample, after_update):
-- 
cgit v1.2.1