Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 from PIL import Image
 from transformers import CLIPTextModel, CLIPTokenizer
 from diffusers import AutoencoderKL, DDIMScheduler
-from src.mgd_pipelines.mgd_pipe import MGDPipe  # Use your implementation of MGDPipe
+from src.mgd_pipelines.mgd_pipe import MGDPipe  # Your MGDPipe implementation
 
 
 # Load models and pipeline
@@ -24,6 +24,9 @@ def load_models(pretrained_model_path, device):
     scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
     scheduler.set_timesteps(50)
 
+    # Handle torch.hub checkpoint loading for CPU-only environments
+    map_location = "cpu" if device.type == "cpu" else None
+
     # Load the UNet model
     unet = torch.hub.load(
         repo_or_dir="aimagelab/multimodal-garment-designer",
@@ -31,7 +34,11 @@ def load_models(pretrained_model_path, device):
         model="mgd",
         pretrained=True,
         dataset="dresscode",  # Change to "vitonhd" if needed
-    ).to(device)
+        map_location=map_location,
+    )
+
+    # Move UNet to the appropriate device
+    unet = unet.to(device)
 
     # Initialize the pipeline
     pipeline = MGDPipe(
@@ -88,9 +95,11 @@ if st.button("Generate"):
     if uploaded_file and text_prompt:
         st.write("Loading models...")
 
-        # Load the pipeline
+        # Detect device (CPU or GPU)
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         pretrained_model_path = "runwayml/stable-diffusion-inpainting"  # Change as required
+
+        # Load the pipeline
         pipeline = load_models(pretrained_model_path, device)
 
         # Load sketch