update

app.py

@@ -20,7 +20,7 @@ from torch import autocast
 import cv2
 import imageio
 import devicetorch
-device = devicetorch.get(torch)
+DEVICE = devicetorch.get(torch)
 
 sys.path.append("./stable_diffusion")
 
@@ -85,7 +85,7 @@ class CompVisDenoiser(K.external.CompVisDenoiser):
     def forward(self, input_0, input_1, sigma, **kwargs):
         c_out, c_in = [append_dims(x, input_0.ndim) for x in self.get_scalings(sigma)]
         # 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.cpu().float()).to(device), **kwargs)
+        eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma.cpu().float()).to(DEVICE), **kwargs)
         
         return input_0 + eps_0 * c_out, eps_1
 
@@ -122,7 +122,6 @@ def sample_euler_ancestral(model, x_0, x_1, sigmas, height, width, extra_args=No
 
     mask_list = []
     image_list = []
-    model.to(device)
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised_0, denoised_1 = model(x_0, x_1, sigmas[i] * s_in, **extra_args)
         image_list.append(denoised_0)
@@ -153,7 +152,7 @@ args = parser.parse_args()
 
 config = OmegaConf.load(args.config)
 model = load_model_from_config(config, args.ckpt, args.vae_ckpt)
-model.eval().to(device)
+model.eval().to(DEVICE)
 model_wrap = CompVisDenoiser(model)
 model_wrap_cfg = CFGDenoiser(model_wrap)
 null_token = model.get_learned_conditioning([""])
@@ -186,8 +185,8 @@ def generate(
     if instruction == "":
         return [input_image, seed]
     
-    model.to(device)
-    if device == "cuda":
+    model.to(DEVICE)
+    if DEVICE == "cuda":
         with torch.no_grad(), autocast("cuda"), model.ema_scope():
             return run(model, input_image, input_image_copy, width, height, instruction, steps, seed)
     else: