Update app.py

app.py

@@ -27,16 +27,22 @@ def load_models(pretrained_model_path, device):
     # Handle torch.hub checkpoint loading for CPU-only environments
     map_location = torch.device("cpu") if device.type == "cpu" else None
 
-    # Load the UNet model and map to CPU if necessary
+    # Load the UNet model and force map_location for state_dict loading
     unet = torch.hub.load(
         repo_or_dir="aimagelab/multimodal-garment-designer",
         source="github",
         model="mgd",
         pretrained=True,
         dataset="dresscode",  # Change to "vitonhd" if needed
-        map_location=map_location,  # Ensure the model loads on CPU if needed
     )
 
+    # Ensure the model state dict is mapped correctly to the CPU if needed
+    if device.type == "cpu":
+        checkpoint_url = unet.config.get("checkpoint")
+        if checkpoint_url:
+            state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
+            unet.load_state_dict(state_dict)
+
     # Move UNet to the appropriate device
     unet = unet.to(device)