Update app.py

app.py

@@ -27,6 +27,7 @@ from huggingface_hub import hf_hub_download
 model = create_model('./models/cldm_v21_512_latctrl_coltrans.yaml').cpu()
 ckpt = hf_hub_download(repo_id="xywwww/scene_diffusion", filename="checkpoints/epoch=25-step=112553.ckpt")
 model.load_state_dict(load_state_dict(ckpt), strict=False)
+model = model.cuda()
 ddim_sampler = DDIMSampler(model)
 
 
@@ -38,8 +39,8 @@ def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resoluti
         # detected_map = apply_canny(img, low_threshold, high_threshold)
         # detected_map = HWC3(detected_map)
 
-        #control = torch.from_numpy(img.copy()).float().cuda() / 255.0
-        control = torch.from_numpy(img.copy()).float() / 255.0
+        control = torch.from_numpy(img.copy()).float().cuda() / 255.0
+        #control = torch.from_numpy(img.copy()).float() / 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()
         control = control.to(memory_format=torch.contiguous_format).float()