Update app.py

app.py

@@ -9,37 +9,41 @@ from src.mgd_pipelines.mgd_pipe import MGDPipe
 # Initialize the model and other components
 @st.cache_resource
 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")
-    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)
-    scheduler = DDIMScheduler.from_pretrained("stabilityai/sd-scheduler", subfolder="scheduler")
+    try:
+        # 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")
+        print("VAE model loaded successfully.")
+        
+        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)
+        scheduler = DDIMScheduler.from_pretrained("stabilityai/sd-scheduler", subfolder="scheduler")
 
-    pipe = MGDPipe(
-        text_encoder=text_encoder,
-        vae=vae,
-        unet=unet.to(vae.dtype),
-        tokenizer=tokenizer,
-        scheduler=scheduler,
-    ).to(device)
-    return pipe
+        pipe = MGDPipe(
+            text_encoder=text_encoder,
+            vae=vae,
+            unet=unet.to(vae.dtype),
+            tokenizer=tokenizer,
+            scheduler=scheduler,
+        ).to(device)
+        return pipe
+    except OSError as e:
+        print(f"Error loading the model: {e}")
+        return None
 
 pipe = load_model()
 
 def generate_images(pipe, text_input=None, sketch=None):
     # Generate images from text or sketch or both
     images = []
-    if text_input:
-        prompt = [text_input]
-        images.extend(pipe(prompt=prompt))
-    if sketch:
-        sketch_image = Image.open(sketch).convert("RGB")
-        images.extend(pipe(sketch=sketch_image))
+    if pipe:
+        if text_input:
+            prompt = [text_input]
+            images.extend(pipe(prompt=prompt))
+        if sketch:
+            sketch_image = Image.open(sketch).convert("RGB")
+            images.extend(pipe(sketch=sketch_image))
     return images
 
 # Streamlit UI
@@ -65,6 +69,9 @@ if st.button("Generate"):
             sketches = BytesIO(sketch_file.read()) if sketch_file else None
             images = generate_images(pipe, text_input=text_input, sketch=sketches)
 
-        # Display results
-        for i, img in enumerate(images):
-            st.image(img, caption=f"Generated Image {i+1}")
+        if images:
+            # Display results
+            for i, img in enumerate(images):
+                st.image(img, caption=f"Generated Image {i+1}")
+        else:
+            st.error("Failed to generate images. Please check the model or inputs.")