JaehaerysNL/StableDiffusion
1
0
Fork 0
mirror of https://github.com/Stability-AI/stablediffusion.git synced 2024-12-22 23:55:00 +00:00
Code Issues Projects Releases Packages Wiki Activity

Don't hardcode the device id to 'cuda', instead use what the model is configured to

Browse source
This commit is contained in:
cmdr2 2022-12-16 16:24:09 +05:30
parent cc77f2300d
commit 3aad790367
3 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
ldm/models/diffusion/ddim.py
View file

@ -16,8 +16,8 @@ class DDIMSampler(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 != torch.device(self.model.device):
attr = attr.to(torch.device(self.model.device))
setattr(self, name, attr)
def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

4
ldm/models/diffusion/dpm_solver/sampler.py
View file

@ -19,8 +19,8 @@ class DPMSolverSampler(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 != torch.device(self.model.device):
attr = attr.to(torch.device(self.model.device))
setattr(self, name, attr)
@torch.no_grad()

4
ldm/models/diffusion/plms.py
View file

@ -18,8 +18,8 @@ class PLMSSampler(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 != torch.device(self.model.device):
attr = attr.to(torch.device(self.model.device))
setattr(self, name, attr)
def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):