Mirror/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-06-07 21:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

use selected device instead of always cuda for UniPC sampler

Browse Source
This commit is contained in:
AUTOMATIC 2023-03-11 11:56:05 +03:00
parent a11ce2b96c
commit f261a4a53c
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
modules/models/diffusion/uni_pc/sampler.py
View File

@ -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):