Update app.py

app.py

@@ -14,7 +14,8 @@ OUTPUT_DIR.mkdir(exist_ok=True)
 
 # Load pre-trained models
 def load_model(model_path, use_cpu=False):
-    model = torch.load(model_path, map_location="cpu" if use_cpu else None)
+    # Explicitly set weights_only=False to load the full model object
+    model = torch.load(model_path, map_location="cpu" if use_cpu else None, weights_only=False)
     if not use_cpu and torch.cuda.is_available():
         model = model.cuda()
     model.eval()