update

stable_diffusion/ldm/models/diffusion/ddim.py

@@ -4,10 +4,12 @@ import torch
 import numpy as np
 from tqdm import tqdm
 from functools import partial
+import devicetorch
 
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, \
     extract_into_tensor
 
+device = devicetorch.get(torch)
 
 class DDIMSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -18,8 +20,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(device):
+                attr = attr.to(torch.device(device))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
@@ -238,4 +240,4 @@ class DDIMSampler(object):
             x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
                                           unconditional_guidance_scale=unconditional_guidance_scale,
                                           unconditional_conditioning=unconditional_conditioning)
-        return x_dec
+        return x_dec

stable_diffusion/ldm/models/diffusion/plms.py

@@ -4,9 +4,11 @@ import torch
 import numpy as np
 from tqdm import tqdm
 from functools import partial
+import devicetorch
 
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
 
+device = devicetorch.get(torch)
 
 class PLMSSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -17,8 +19,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(device):
+                attr = attr.to(torch.device(device))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

stable_diffusion/notebook_helpers.py

@@ -14,7 +14,9 @@ from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.util import ismap
 import time
 from omegaconf import OmegaConf
+import devicetorch
 
+device = devicetorch.get(torch)
 
 def download_models(mode):
 
@@ -44,7 +46,7 @@ def load_model_from_config(config, ckpt):
     sd = pl_sd["state_dict"]
     model = instantiate_from_config(config.model)
     m, u = model.load_state_dict(sd, strict=False)
-    model.cuda()
+    model.to(device)
     model.eval()
     return {"model": model}, global_step
 
@@ -117,7 +119,7 @@ def get_cond(mode, selected_path):
         c = rearrange(c, '1 c h w -> 1 h w c')
         c = 2. * c - 1.
 
-        c = c.to(torch.device("cuda"))
+        c = c.to(torch.device(device))
         example["LR_image"] = c
         example["image"] = c_up
 
@@ -267,4 +269,4 @@ def make_convolutional_sample(batch, model, mode="vanilla", custom_steps=None, e
     log["sample"] = x_sample
     log["time"] = t1 - t0
 
-    return log
+    return log