Update app.py

app.py

@@ -11,16 +11,37 @@ from src.mgd_pipelines.mgd_pipe import MGDPipe
 def load_model():
     try:
         # Define your model loading logic
+        print("Initializing model loading...")
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
-        print("VAE model loaded successfully.")
+        print(f"Device selected: {device}")
         
+        # Load the VAE
+        print("Loading VAE...")
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+        print("VAE loaded successfully.")
+
+        # Load the tokenizer
+        print("Loading tokenizer...")
         tokenizer = CLIPTokenizer.from_pretrained("microsoft/xclip-base-patch32", subfolder="tokenizer")
+        print("Tokenizer loaded successfully.")
+
+        # Load the text encoder
+        print("Loading text encoder...")
         text_encoder = CLIPTextModel.from_pretrained("microsoft/xclip-base-patch32", subfolder="text_encoder")
+        print("Text encoder loaded successfully.")
+
+        # Load the UNet model
+        print("Loading UNet...")
         unet = torch.hub.load("aimagelab/multimodal-garment-designer", model="mgd", pretrained=True)
+        print("UNet loaded successfully.")
+
+        # Load the scheduler
+        print("Loading scheduler...")
         scheduler = DDIMScheduler.from_pretrained("stabilityai/sd-scheduler", subfolder="scheduler")
+        print("Scheduler loaded successfully.")
 
         # Initialize the pipeline
+        print("Initializing pipeline...")
         pipe = MGDPipe(
             text_encoder=text_encoder,
             vae=vae,
@@ -28,7 +49,8 @@ def load_model():
             tokenizer=tokenizer,
             scheduler=scheduler,
         ).to(device)
-        pipe.enable_attention_slicing()  # Enable memory-efficient inference
+        pipe.enable_attention_slicing()
+        print("Pipeline initialized successfully.")
         return pipe
     except Exception as e:
         print(f"Error loading the model: {e}")