Update app.py

app.py

@@ -59,7 +59,10 @@ model = instantiate_from_config(config)
 ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu")
 load_state_dict(model, ckpt, strict=True)
 model.freeze()
-model.to("cuda")
+
+# Check if CUDA is available, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 @torch.no_grad()
 def process(
@@ -113,7 +116,7 @@ def process(
     control_img = np.array(control_img)
 
     # Convert to tensor (NCHW, [0,1])
-    control = torch.tensor(control_img[None] / 255.0, dtype=torch.float32, device=model.device).clamp_(0, 1)
+    control = torch.tensor(control_img[None] / 255.0, dtype=torch.float32, device=device).clamp_(0, 1)
     control = einops.rearrange(control, "n h w c -> n c h w").contiguous()
     height, width = control.size(-2), control.size(-1)
     model.control_scales = [strength] * 13
@@ -122,7 +125,7 @@ def process(
     preds = []
     for _ in tqdm(range(num_samples)):
         shape = (1, 4, height // 8, width // 8)
-        x_T = torch.randn(shape, device=model.device, dtype=torch.float32)
+        x_T = torch.randn(shape, device=device, dtype=torch.float32)
 
         if not tile_diffusion and not tile_vae:
             samples = sampler.sample_ccsr(