Update app.py

app.py

@@ -1,62 +1,77 @@
-import streamlit as st
 import os
+import torch
+import streamlit as st
 from PIL import Image
-from src.eval import main  # Import the modified main function from evl.py
-
-# Title and Description
-st.title("Fashion Image Generator")
-st.write("Upload a rough sketch, set parameters, and generate realistic garment images.")
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers import AutoencoderKL, DDIMScheduler
+from src.mgd_pipelines.mgd_pipe import MGDPipe
+from src.mgd_pipelines.mgd_pipe_disentangled import MGDPipeDisentangled
 
-# File Upload Section
-uploaded_file = st.file_uploader("Upload your rough sketch (PNG, JPG, JPEG):", type=["png", "jpg", "jpeg"])
+# Function to load models
+def load_models(pretrained_model_name_or_path, device):
+    tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder")
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+    scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+    scheduler.set_timesteps(50, device=device)
+    
+    unet = torch.hub.load(
+        repo_or_dir="aimagelab/multimodal-garment-designer",
+        model="mgd",
+        pretrained=True,
+        source="github",
+    )
+    return tokenizer, text_encoder, vae, scheduler, unet
 
-# Sidebar for Parameters
-st.sidebar.title("Model Configuration")
-pretrained_model_path = st.sidebar.text_input("Pretrained Model Path", "runwayml/stable-diffusion-inpainting")
-dataset_path = st.sidebar.text_input("Dataset Path", "./datasets/dresscode")
-output_dir = st.sidebar.text_input("Output Directory", "./outputs")
-guidance_scale_sketch = st.sidebar.slider("Sketch Guidance Scale", 1.0, 10.0, 7.5)
-batch_size = st.sidebar.number_input("Batch Size", min_value=1, max_value=16, value=1)
-mixed_precision = st.sidebar.selectbox("Mixed Precision Mode", ["fp16", "fp32"], index=0)
-seed = st.sidebar.number_input("Random Seed", value=42, step=1)
+# Function to generate images
+def generate_image(sketch, prompt, tokenizer, text_encoder, vae, scheduler, unet, device):
+    # Preprocess inputs
+    sketch = sketch.resize((512, 384)).convert("RGB")
+    sketch_tensor = torch.tensor([torch.tensor(sketch, dtype=torch.float32).permute(2, 0, 1) / 255.0]).to(device)
+    
+    # Tokenize prompt
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Initialize pipeline
+    pipeline = MGDPipe(
+        text_encoder=text_encoder.to(device),
+        vae=vae.to(device),
+        unet=unet.to(device),
+        tokenizer=tokenizer,
+        scheduler=scheduler,
+    ).to(device)
 
-# Run Button
-if st.button("Generate Image"):
-    if uploaded_file:
-        # Save uploaded sketch locally
-        os.makedirs("temp_uploads", exist_ok=True)
-        sketch_path = os.path.join("temp_uploads", uploaded_file.name)
-        with open(sketch_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
+    # Generate image
+    pipeline.enable_attention_slicing()
+    with torch.inference_mode():
+        outputs = pipeline(images=sketch_tensor, text=inputs["input_ids"], guidance_scale=7.5)
+    
+    return outputs[0]
 
-        # Prepare arguments for the backend
-        args = {
-            "pretrained_model_name_or_path": pretrained_model_path,
-            "dataset": "dresscode",
-            "dataset_path": dataset_path,
-            "output_dir": output_dir,
-            "guidance_scale": 7.5,
-            "guidance_scale_sketch": guidance_scale_sketch,
-            "mixed_precision": mixed_precision,
-            "batch_size": batch_size,
-            "seed": seed,
-            "save_name": "generated_image",  # Output file name
-        }
+# Streamlit UI
+st.title("Garment Designer")
+st.write("Upload a sketch and provide a text description to generate garment designs!")
 
-        # Run the backend model
-        st.write("Generating image...")
-        try:
-            output_path = main(args)  # Call your backend main function
-            st.write("Image generation complete!")
+# User Inputs
+uploaded_file = st.file_uploader("Upload your sketch", type=["png", "jpg", "jpeg"])
+text_prompt = st.text_input("Enter a text description for the garment")
 
-            # Display the generated image
-            output_image_path = os.path.join(output_dir, "generated_image.png")  # Update if needed
-            if os.path.exists(output_image_path):
-                output_image = Image.open(output_image_path)
-                st.image(output_image, caption="Generated Image", use_column_width=True)
-            else:
-                st.error("Image generation failed. No output file found.")
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
+# Generate button
+if st.button("Generate"):
+    if uploaded_file and text_prompt:
+        # Load models
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        pretrained_model_path = "your-pretrained-model-path"  # Replace with actual model path
+        tokenizer, text_encoder, vae, scheduler, unet = load_models(pretrained_model_path, device)
+        
+        # Load sketch
+        sketch = Image.open(uploaded_file)
+        
+        # Generate image
+        st.write("Generating the garment design...")
+        output_image = generate_image(sketch, text_prompt, tokenizer, text_encoder, vae, scheduler, unet, device)
+        
+        # Display output
+        st.image(output_image, caption="Generated Garment Design", use_column_width=True)
     else:
-        st.error("Please upload a sketch before generating an image.")
+        st.error("Please upload a sketch and enter a text description.")