recover gpu support

app.py

@@ -16,9 +16,14 @@ from cldm.ddim_hacked import DDIMSampler
 import dlib
 from PIL import Image, ImageDraw
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
 model = create_model('./models/cldm_v15.yaml').cpu()
 model.load_state_dict(load_state_dict('./models/control_sd15_landmarks.pth', location='cpu'))
-model = model
+model = model.to(device)
 ddim_sampler = DDIMSampler(model)
 
 detector = dlib.get_frontal_face_detector()
@@ -56,7 +61,7 @@ def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resoluti
             detected_map = get_68landmarks_img(img)
         detected_map = HWC3(detected_map)
 
-        control = torch.from_numpy(detected_map.copy()).float() / 255.0
+        control = torch.from_numpy(detected_map.copy()).float().to(device) / 255.0
         control = torch.stack([control for _ in range(num_samples)], dim=0)
         control = einops.rearrange(control, 'b h w c -> b c h w').clone()
 

cldm/cldm.py

@@ -18,6 +18,12 @@ from ldm.models.diffusion.ddpm import LatentDiffusion
 from ldm.util import log_txt_as_img, exists, instantiate_from_config
 from ldm.models.diffusion.ddim import DDIMSampler
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    device_type = "cuda"
+else:
+    device = torch.device("cpu")
+    device_type = "cpu"
 
 class ControlledUnetModel(UNetModel):
     def forward(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs):
@@ -424,12 +430,12 @@ class ControlLDM(LatentDiffusion):
 
     def low_vram_shift(self, is_diffusing):
         if is_diffusing:
-            self.model = self.model.cpu()
-            self.control_model = self.control_model.cpu()
+            self.model = self.model.to(device)
+            self.control_model = self.control_model.to(device)
             self.first_stage_model = self.first_stage_model.cpu()
             self.cond_stage_model = self.cond_stage_model.cpu()
         else:
             self.model = self.model.cpu()
             self.control_model = self.control_model.cpu()
-            self.first_stage_model = self.first_stage_model.cpu()
-            self.cond_stage_model = self.cond_stage_model.cpu()
+            self.first_stage_model = self.first_stage_model.to(device)
+            self.cond_stage_model = self.cond_stage_model.to(device)

cldm/ddim_hacked.py

@@ -6,6 +6,10 @@ from tqdm import tqdm
 
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
 
 class DDIMSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -16,8 +20,8 @@ class DDIMSampler(object):
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cpu"):
-                attr = attr.to(torch.device("cpu"))
+            if attr.device != device:
+                attr = attr.to(device)
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

ldm/models/diffusion/ddim.py

@@ -6,6 +6,12 @@ from tqdm import tqdm
 
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    device_type = "cuda"
+else:
+    device = torch.device("cpu")
+    device_type = "cpu"
 
 class DDIMSampler(object):
     def __init__(self, model, schedule="linear", **kwargs):
@@ -16,8 +22,8 @@ class DDIMSampler(object):
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cpu"):
-                attr = attr.to(torch.device("cpu"))
+            if attr.device != torch.device(device_type):
+                attr = attr.to(torch.device(device_type))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

ldm/modules/attention.py

@@ -8,6 +8,12 @@ from typing import Optional, Any
 
 from ldm.modules.diffusionmodules.util import checkpoint
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    device_type = "cuda"
+else:
+    device = torch.device("cpu")
+    device_type = "cpu"
 
 try:
     import xformers
@@ -172,7 +178,7 @@ class CrossAttention(nn.Module):
 
         # force cast to fp32 to avoid overflowing
         if _ATTN_PRECISION =="fp32":
-            with torch.autocast(enabled=False, device_type = 'cpu'):
+            with torch.autocast(enabled=False, device_type = device_type):
                 q, k = q.float(), k.float()
                 sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
         else:

ldm/modules/diffusionmodules/util.py

@@ -133,7 +133,7 @@ class CheckpointFunction(torch.autograd.Function):
     def backward(ctx, *output_grads):
         ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
         with torch.enable_grad(), \
-                torch.cpu.amp.autocast(**ctx.gpu_autocast_kwargs):
+                torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
             # Fixes a bug where the first op in run_function modifies the
             # Tensor storage in place, which is not allowed for detach()'d
             # Tensors.

ldm/modules/encoders/modules.py

@@ -7,6 +7,12 @@ from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextMod
 import open_clip
 from ldm.util import default, count_params
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    device_type = "cuda"
+else:
+    device = torch.device("cpu")
+    device_type = "cpu"
 
 class AbstractEncoder(nn.Module):
     def __init__(self):
@@ -42,7 +48,7 @@ class ClassEmbedder(nn.Module):
         c = self.embedding(c)
         return c
 
-    def get_unconditional_conditioning(self, bs, device="cpu"):
+    def get_unconditional_conditioning(self, bs, device=device_type):
         uc_class = self.n_classes - 1  # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
         uc = torch.ones((bs,), device=device) * uc_class
         uc = {self.key: uc}
@@ -57,7 +63,7 @@ def disabled_train(self, mode=True):
 
 class FrozenT5Embedder(AbstractEncoder):
     """Uses the T5 transformer encoder for text"""
-    def __init__(self, version="google/t5-v1_1-large", device="cpu", max_length=77, freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
+    def __init__(self, version="google/t5-v1_1-large", device=device_type, max_length=77, freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
         super().__init__()
         self.tokenizer = T5Tokenizer.from_pretrained(version)
         self.transformer = T5EncoderModel.from_pretrained(version)
@@ -92,7 +98,7 @@ class FrozenCLIPEmbedder(AbstractEncoder):
         "pooled",
         "hidden"
     ]
-    def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
+    def __init__(self, version="openai/clip-vit-large-patch14", device=device_type, max_length=77,
                  freeze=True, layer="last", layer_idx=None):  # clip-vit-base-patch32
         super().__init__()
         assert layer in self.LAYERS
@@ -140,11 +146,11 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
         "last",
         "penultimate"
     ]
-    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cpu", max_length=77,
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device=device_type, max_length=77,
                  freeze=True, layer="last"):
         super().__init__()
         assert layer in self.LAYERS
-        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version)
+        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device(device_type), pretrained=version)
         del model.visual
         self.model = model
 
@@ -194,7 +200,7 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
 
 
 class FrozenCLIPT5Encoder(AbstractEncoder):
-    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cpu",
+    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device=device_type,
                  clip_max_length=77, t5_max_length=77):
         super().__init__()
         self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length)