Update app.py

app.py

@@ -1,97 +1,67 @@
-import os
-import numpy as np  # Corrected import
-import torch
 import streamlit as st
+import torch
 from PIL import Image
-from accelerate import Accelerator
+from io import BytesIO
 from diffusers import DDIMScheduler, AutoencoderKL
 from transformers import CLIPTextModel, CLIPTokenizer
-from src.mgd_pipelines.mgd_pipe_disentangled import MGDPipeDisentangled
-from src.utils.set_seeds import set_seed
-from src.utils.image_from_pipe import generate_images_from_mgd_pipe
-from src.datasets.dresscode import DressCodeDataset
-
-# Set environment variables
-os.environ["TOKENIZERS_PARALLELISM"] = "true"
-os.environ["WANDB_START_METHOD"] = "thread"
+from src.mgd_pipelines.mgd_pipe import MGDPipe
 
-# Function to process inputs and run inference
-def run_inference(prompt, sketch_image=None, category="dresses", seed=1, mixed_precision="fp16"):
-    accelerator = Accelerator(mixed_precision=mixed_precision)
-    device = accelerator.device
-
-    # Load models and datasets
+# 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", subfolder="vae")
     tokenizer = CLIPTokenizer.from_pretrained("microsoft/xclip-base-patch32", subfolder="tokenizer")
     text_encoder = CLIPTextModel.from_pretrained("microsoft/xclip-base-patch32", subfolder="text_encoder")
-    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", subfolder="vae")
-    val_scheduler = DDIMScheduler.from_pretrained("stabilityai/sd-scheduler", subfolder="scheduler")
-
-    unet = torch.hub.load("aimagelab/multimodal-garment-designer", "mgd", pretrained=True)
-
-    vae.requires_grad_(False)
-    text_encoder.requires_grad_(False)
-
-    if seed is not None:
-        set_seed(seed)
+    unet = torch.hub.load("aimagelab/multimodal-garment-designer", model="mgd", pretrained=True)
+    scheduler = DDIMScheduler.from_pretrained("stabilityai/sd-scheduler", subfolder="scheduler")
 
-    category = [category]
-    test_dataset = DressCodeDataset(
-        dataroot_path="assets\data\dresscode",  # Replace with actual dataset path
-        phase="test",
-        category=category,
-        size=(512, 384),
-    )
-
-    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
-
-    text_encoder.to(device)
-    vae.to(device)
-    unet.to(device).eval()
-
-    if sketch_image is not None:
-        sketch_tensor = (
-            torch.tensor(np.array(sketch_image)).permute(2, 0, 1).unsqueeze(0).float().to(device) / 255.0
-        )
-
-    val_pipe = MGDPipeDisentangled(
+    pipe = MGDPipe(
         text_encoder=text_encoder,
         vae=vae,
-        unet=unet,
+        unet=unet.to(vae.dtype),
         tokenizer=tokenizer,
-        scheduler=val_scheduler,
+        scheduler=scheduler,
     ).to(device)
+    return pipe
 
-    val_pipe.enable_attention_slicing()
+pipe = load_model()
 
-    generated_images = generate_images_from_mgd_pipe(
-        test_dataloader=test_dataloader,
-        pipe=val_pipe,
-        guidance_scale=7.5,
-        seed=seed,
-        sketch_image=sketch_tensor if sketch_image is not None else None,
-        prompt=prompt,
-    )
-
-    return Image.fromarray((generated_images[0] * 255).astype("uint8"))
+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))
+    return images
 
 # Streamlit UI
-st.title("Fashion Image Generator")
-st.write("Generate colorful fashion images based on a rough sketch and/or a text prompt.")
-
-uploaded_sketch = st.file_uploader("Upload a rough sketch (optional)", type=["png", "jpg", "jpeg"])
-prompt = st.text_input("Enter a prompt (optional)", "A red dress with floral patterns")
-category = st.text_input("Enter category (optional):", "dresses")
-seed = st.slider("Seed", min_value=1, max_value=100, step=1, value=1)
-precision = st.selectbox("Select precision:", ["fp16", "fp32"])
-
-if uploaded_sketch is not None:
-    sketch_image = Image.open(uploaded_sketch)
-    st.image(sketch_image, caption="Uploaded Sketch", use_column_width=True)
-
-if st.button("Generate Image"):
-    with st.spinner("Generating image..."):
-        try:
-            result_image = run_inference(prompt, sketch_image, category, seed, precision)
-            st.image(result_image, caption="Generated Image", use_column_width=True)
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
+st.title("Sketch & Text-based Image Generation")
+st.write("Generate images based on rough sketches, text input, or both.")
+
+option = st.radio("Select Input Type", ("Sketch", "Text", "Both"))
+
+if option in ["Sketch", "Both"]:
+    sketch_file = st.file_uploader("Upload a Sketch", type=["png", "jpg", "jpeg"])
+
+if option in ["Text", "Both"]:
+    text_input = st.text_input("Enter Text Prompt", placeholder="Describe the image you want to generate")
+
+if st.button("Generate"):
+    if option == "Sketch" and not sketch_file:
+        st.error("Please upload a sketch.")
+    elif option == "Text" and not text_input:
+        st.error("Please provide text input.")
+    else:
+        # Generate images based on user input
+        with st.spinner("Generating images..."):
+            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}")