update

customnet/ddim.py

@@ -9,6 +9,12 @@ from einops import rearrange
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
 from ldm.models.diffusion.sampling_util import renorm_thresholding, norm_thresholding, spatial_norm_thresholding
 
+if torch.cuda.is_available():
+  _device = "cuda"
+elif torch.backends.mps.is_available():
+  _device = "mps"
+else:
+  _device = "cpu"
 
 class DDIMSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -27,9 +33,10 @@ 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 _device == "cuda":
+            if type(attr) == torch.Tensor:
+                if attr.device != torch.device("cuda"):
+                    attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

ldm/models/diffusion/ddim.py

@@ -9,6 +9,12 @@ from einops import rearrange
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
 from ldm.models.diffusion.sampling_util import renorm_thresholding, norm_thresholding, spatial_norm_thresholding
 
+if torch.cuda.is_available():
+  _device = "cuda"
+elif torch.backends.mps.is_available():
+  _device = "mps"
+else:
+  _device = "cpu"
 
 class DDIMSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -27,9 +33,10 @@ 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 _device == "cuda":
+            if type(attr) == torch.Tensor:
+                if attr.device != torch.device("cuda"):
+                    attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
@@ -322,4 +329,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

ldm/models/diffusion/plms.py

@@ -8,6 +8,12 @@ from functools import partial
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
 from ldm.models.diffusion.sampling_util import norm_thresholding
 
+if torch.cuda.is_available():
+  _device = "cuda"
+elif torch.backends.mps.is_available():
+  _device = "mps"
+else:
+  _device = "cpu"
 
 class PLMSSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -17,9 +23,10 @@ class PLMSSampler(object):
         self.schedule = schedule
 
     def register_buffer(self, name, attr):
-        if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
+        if _device == "cuda":
+            if type(attr) == torch.Tensor:
+                if attr.device != torch.device("cuda"):
+                    attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):