update

app.py

@@ -75,15 +75,20 @@ def append_dims(x, target_dims):
 
 class CompVisDenoiser(K.external.CompVisDenoiser):
     def __init__(self, model, quantize=False, device='cpu'):
-        super().__init__( model, quantize, device)
+        super().__init__(model, quantize, device)
     
     def get_eps(self, *args, **kwargs):
         return self.inner_model.apply_model(*args, **kwargs)
     
     def forward(self, input_0, input_1, sigma, **kwargs):
+        print("input_0.device:", input_0.device)
+        print("input_1.device:", input_1.device)
         c_out, c_in = [append_dims(x, input_0.ndim) for x in self.get_scalings(sigma)]
+        print("c_in.device:", c_in.device)
+        print("c_out.device:", c_out.device)
+        print("sigma.device:", sigma.device)
         # eps_0, eps_1 = self.get_eps(input_0 * c_in, input_1 * c_in, self.sigma_to_t(sigma), **kwargs)
-        eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma), **kwargs)
+        eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma.cpu()).cuda(), **kwargs)
         
         return input_0 + eps_0 * c_out, eps_1
 
@@ -112,7 +117,6 @@ def decode_mask(mask, height = 256, width = 256):
     mask = mask.type(torch.uint8).cpu().numpy()
     return mask
 
-@torch.no_grad()
 def sample_euler_ancestral(model, x_0, x_1, sigmas, height, width, extra_args=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
     """Ancestral sampling with Euler method steps."""
     extra_args = {} if extra_args is None else extra_args
@@ -183,17 +187,24 @@ def generate(
 
     if instruction == "":
         return [input_image, seed]
-
+    
+    model.cuda()
+    print("model.device:", model.device)
     with torch.no_grad(), autocast("cuda"), model.ema_scope():
         cond = {}
-        cond["c_crossattn"] = [model.get_learned_conditioning([instruction])]
+        cond["c_crossattn"] = [model.get_learned_conditioning([instruction]).to(model.device)]
         input_image = 2 * torch.tensor(np.array(input_image)).float() / 255 - 1
         input_image = rearrange(input_image, "h w c -> 1 c h w").to(model.device)
-        cond["c_concat"] = [model.encode_first_stage(input_image).mode()]
+        cond["c_concat"] = [model.encode_first_stage(input_image).mode().to(model.device)]
 
         uncond = {}
-        uncond["c_crossattn"] = [null_token]
+        uncond["c_crossattn"] = [null_token.to(model.device)]
         uncond["c_concat"] = [torch.zeros_like(cond["c_concat"][0])]
+        
+        print("cond['c_crossattn'][0].device:", cond["c_crossattn"][0].device)
+        print("cond['c_concat'][0].device:", cond["c_concat"][0].device)
+        print("uncond['c_crossattn'][0].device:", uncond["c_crossattn"][0].device)
+        print("uncond['c_concat'][0].device:", uncond["c_concat"][0].device)
 
         sigmas = model_wrap.get_sigmas(steps)
 
@@ -204,8 +215,10 @@ def generate(
             "image_cfg_scale": image_cfg_scale,
         }
         torch.manual_seed(seed)
-        z_0 = torch.randn_like(cond["c_concat"][0]) * sigmas[0]
-        z_1 = torch.randn_like(cond["c_concat"][0]) * sigmas[0]
+        z_0 = torch.randn_like(cond["c_concat"][0]).to(model.device) * sigmas[0]
+        z_1 = torch.randn_like(cond["c_concat"][0]).to(model.device) * sigmas[0]
+        print("z_0.device:", z_0.device)
+        print("z_1.device:", z_1.device)
         
         z_0, z_1, image_list, mask_list = sample_euler_ancestral(model_wrap_cfg, z_0, z_1, sigmas, height, width, extra_args=extra_args)