Update app.py

app.py

@@ -11,7 +11,10 @@ from src.mgd_pipelines.mgd_pipe import MGDPipe
 def load_model():
     # Define your model loading logic
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", subfolder="vae")
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+    print("VAE model loaded successfully.")
+except OSError as e:
+    print(f"Error loading the model: {e}")
     tokenizer = CLIPTokenizer.from_pretrained("microsoft/xclip-base-patch32", subfolder="tokenizer")
     text_encoder = CLIPTextModel.from_pretrained("microsoft/xclip-base-patch32", subfolder="text_encoder")
     unet = torch.hub.load("aimagelab/multimodal-garment-designer", model="mgd", pretrained=True)