Update app.py

app.py

@@ -25,16 +25,16 @@ def load_models(pretrained_model_path, device):
     scheduler.set_timesteps(50)
 
     # Handle torch.hub checkpoint loading for CPU-only environments
-    map_location = "cpu" if device.type == "cpu" else None
+    map_location = torch.device("cpu") if device.type == "cpu" else None
 
-    # Load the UNet model
+    # Load the UNet model and map to CPU if necessary
     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,
+        map_location=map_location,  # Ensure the model loads on CPU if needed
     )
 
     # Move UNet to the appropriate device