Update app.py

app.py

@@ -57,8 +57,17 @@ 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", weights_only=True)
-load_state_dict(model, ckpt, strict=True)
+
+# Load the checkpoint without weights_only=True
+ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu")
+
+# Extract only the model state dict
+if "state_dict" in ckpt:
+    state_dict = ckpt["state_dict"]
+else:
+    state_dict = ckpt
+
+load_state_dict(model, state_dict, strict=True)
 model.freeze()
 
 # Check if CUDA is available, otherwise use CPU