Update app.py

app.py

@@ -1,95 +1,109 @@
 import os
+import pandas as np
 import torch
 import streamlit as st
-from diffusers import AutoencoderKL, DDIMScheduler
+from PIL import Image
+from accelerate import Accelerator
+from diffusers import DDIMScheduler, AutoencoderKL
 from transformers import CLIPTextModel, CLIPTokenizer
 from src.mgd_pipelines.mgd_pipe import MGDPipe
 from src.mgd_pipelines.mgd_pipe_disentangled import MGDPipeDisentangled
-from src.utils.image_from_pipe import generate_images_from_mgd_pipe
-from accelerate import Accelerator
-from diffusers.utils import check_min_version
 from src.utils.set_seeds import set_seed
+from src.utils.image_from_pipe import generate_images_from_mgd_pipe
+from datasets.dresscode import DressCodeDataset
 
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.10.0.dev0")
-
-# Set the environment variables for Hugging Face Spaces
+# Set environment variables
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 os.environ["WANDB_START_METHOD"] = "thread"
 
-# Streamlit interface components
-st.title("Fashion Image Generation with Multimodal Garment Designer")
-
-# Streamlit Input Parameters
-category = st.selectbox("Select Category", ["dresses", "upper_body", "lower_body", "all"])
-guidance_scale = st.slider("Guidance Scale", min_value=0.1, max_value=20.0, value=7.5, step=0.1)
-guidance_scale_pose = st.slider("Guidance Scale (Pose)", min_value=0.1, max_value=20.0, value=7.5, step=0.1)
-guidance_scale_sketch = st.slider("Guidance Scale (Sketch)", min_value=0.1, max_value=20.0, value=7.5, step=0.1)
-sketch_cond_rate = st.slider("Sketch Conditioning Rate", min_value=0.1, max_value=1.0, value=0.5, step=0.05)
-start_cond_rate = st.slider("Start Conditioning Rate", min_value=0.1, max_value=1.0, value=0.5, step=0.05)
-seed = st.number_input("Seed", value=42, min_value=1)
-
-# Button to run the image generation
-if st.button("Generate Image"):
-    # Initialize Accelerator (for mixed precision, etc.)
-    accelerator = Accelerator()
+# Function to process inputs and run inference
+def run_inference(prompt, sketch_image=None, category="dresses", seed=None, mixed_precision="fp16"):
+    # Initialize accelerator
+    accelerator = Accelerator(mixed_precision=mixed_precision)
     device = accelerator.device
 
-    # Set the seed
-    set_seed(seed)
+    # Load models and datasets
+    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("ptx0/pseudo-journey-v2", subfolder="scheduler")
 
-    # Model and Tokenizer loading (use pre-trained from Hugging Face)
-    model_name = "stabilityai/stable-diffusion-2-1-base"  # Use appropriate model name
+    # Load UNet (assumed pretrained)
+    unet = torch.hub.load("aimagelab/multimodal-garment-designer", "mgd", pretrained=True)
 
-    # Load scheduler, tokenizer, and models
-    val_scheduler = DDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
-    val_scheduler.set_timesteps(50, device=device)
+    # Freeze VAE and text encoder
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
 
-    tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")
-    text_encoder = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder")
-    vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae")
+    # Set seed for reproducibility
+    if seed is not None:
+        set_seed(seed)
 
-    # Load UNet model (you can use your own model)
-    unet = torch.hub.load(
-        dataset="aimagelab/multimodal-garment-designer",
-        repo_or_dir="aimagelab/multimodal-garment-designer",
-        source="github",
-        model="mgd",
-        pretrained=True,
+    # Load appropriate dataset
+    category = [category]
+    test_dataset = DressCodeDataset(
+        dataroot_path="path_to_dataset", phase="test", category=category, size=(512, 384)
     )
 
-    # Freeze VAE and text encoder
-    vae.requires_grad_(False)
-    text_encoder.requires_grad_(False)
+    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
+
+    # Move models to the device
+    text_encoder.to(device)
+    vae.to(device)
+    unet.to(device).eval()
 
-    # Select pipeline (use disentangled option if needed)
-    val_pipe = MGDPipe(
+    # Handle sketch and text inputs
+    if sketch_image is not None:
+        # Process the sketch (resize, normalize, etc.)
+        sketch_image = sketch_image.resize((512, 384))
+        sketch_tensor = torch.tensor(np.array(sketch_image)).unsqueeze(0).float().to(device)
+
+    # Select pipeline (disentangled if required)
+    val_pipe = MGDPipeDisentangled(
         text_encoder=text_encoder,
         vae=vae,
-        unet=unet.to(vae.dtype),
+        unet=unet,
         tokenizer=tokenizer,
         scheduler=val_scheduler,
     ).to(device)
 
-    # Run image generation using your pipeline
-    with torch.no_grad():
-        # Generate the image
-        images = generate_images_from_mgd_pipe(
-            test_order="test",  # or some predefined order
-            pipe=val_pipe,
-            test_dataloader=None,  # Adjust accordingly, or use pre-existing dataset
-            save_name="generated_image",
-            dataset="dresscode",  # Adjust if needed
-            output_dir=".",  # Save location
-            guidance_scale=guidance_scale,
-            guidance_scale_pose=guidance_scale_pose,
-            guidance_scale_sketch=guidance_scale_sketch,
-            sketch_cond_rate=sketch_cond_rate,
-            start_cond_rate=start_cond_rate,
-            no_pose=False,
-            disentagle=False,  # Adjust if needed
-            seed=seed,
-        )
-
-        # Display the generated image
-        st.image(images[0], caption="Generated Fashion Image", use_column_width=True)
+    val_pipe.enable_attention_slicing()
+
+    # Generate image
+    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 generated_images[0]  # Assuming single image output
+
+# Streamlit UI
+st.title("Fashion Image Generator")
+st.write("Generate colorful fashion images based on a rough sketch and/or a text prompt.")
+
+# Upload a sketch image
+uploaded_sketch = st.file_uploader("Upload a rough sketch (optional)", type=["png", "jpg", "jpeg"])
+
+# Text input for prompt
+prompt = st.text_input("Enter a prompt (optional)", "A red dress with floral patterns")
+
+# Input options
+category = st.text_input("Enter category (optional):", "dresses")
+seed = st.slider("Seed", min_value=1, max_value=100, step=1, value=None)
+precision = st.selectbox("Select precision:", ["fp16", "fp32"])
+
+# Show uploaded sketch image
+if uploaded_sketch is not None:
+    sketch_image = Image.open(uploaded_sketch)
+    st.image(sketch_image, caption="Uploaded Sketch", use_column_width=True)
+
+# Button to generate image
+if st.button("Generate Image"):
+    with st.spinner("Generating image..."):
+        # Run inference with sketch or prompt (or both)
+        result_image = run_inference(prompt, sketch_image, category, seed, precision)
+        st.image(result_image, caption="Generated Image", use_column_width=True)