Update app.py

app.py

@@ -12,6 +12,7 @@ import einops
 import math
 import random
 import pytorch_lightning as pl
+import spaces
 
 def download_file(url, filename):
     response = requests.get(url, stream=True)
@@ -56,7 +57,7 @@ from utils.image import auto_resize
 
 config = OmegaConf.load("configs/model/ccsr_stage2.yaml")
 model = instantiate_from_config(config)
-ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu")
+ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu", weights_only=True)
 load_state_dict(model, ckpt, strict=True)
 model.freeze()
 
@@ -65,6 +66,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 @torch.no_grad()
+@spaces.GPU
 def process(
     control_img: Image.Image,
     num_samples: int,