Update app.py

app.py

@@ -1,127 +1,62 @@
 import streamlit as st
-import torch
+import os
 from PIL import Image
-from io import BytesIO
-from diffusers import DDIMScheduler, AutoencoderKL
-from transformers import CLIPTextModel, CLIPTokenizer
-from src.mgd_pipelines.mgd_pipe import MGDPipe
-
-# Initialize the model and other components
-@st.cache_resource
-def load_model():
-    try:
-        # Define your model loading logic
-        print("Initializing model loading...")
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        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,
-            unet=unet.to(vae.dtype),
-            tokenizer=tokenizer,
-            scheduler=scheduler,
-        ).to(device)
-        pipe.enable_attention_slicing()
-        print("Pipeline initialized successfully.")
-        return pipe
-    except Exception 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 = []
-    try:
-        if pipe:
-            # Generate from text
-            if text_input:
-                print(f"Generating image from text: {text_input}")
-                images.append(pipe(prompt=[text_input]))
-
-            # Generate from sketch
-            if sketch:
-                print("Generating image from sketch.")
-                sketch_image = Image.open(sketch).convert("RGB")
-                images.append(pipe(sketch=sketch_image))
-    except Exception as e:
-        print(f"Error during image generation: {e}")
-    return images
-
-# Streamlit UI
-st.title("Sketch & Text-based Image Generation")
-st.write("Generate images based on rough sketches, text input, or both.")
-
-# Input options
-option = st.radio("Select Input Type", ("Sketch", "Text", "Both"))
-
-sketch_file = None
-text_input = None
-
-# Get sketch input
-if option in ["Sketch", "Both"]:
-    sketch_file = st.file_uploader("Upload a Sketch", type=["png", "jpg", "jpeg"])
-
-# Get text input
-if option in ["Text", "Both"]:
-    text_input = st.text_input("Enter Text Prompt", placeholder="Describe the image you want to generate")
-
-# Generate button
-if st.button("Generate"):
-    # Ensure the model is loaded
-    if pipe is None:
-        st.error("Model failed to load. Please restart the application.")
-        st.stop()
-
-    # Validate inputs
-    sketches = BytesIO(sketch_file.read()) if sketch_file else None
-    
-    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.")
-    elif option == "Both" and not (sketch_file or text_input):
-        st.error("Please provide both a sketch and a text prompt.")
+from evl 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.")
+
+# File Upload Section
+uploaded_file = st.file_uploader("Upload your rough sketch (PNG, JPG, JPEG):", type=["png", "jpg", "jpeg"])
+
+# Sidebar for Parameters
+st.sidebar.title("Model Configuration")
+pretrained_model_path = st.sidebar.text_input("Pretrained Model Path", "./models")
+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)
+
+# 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())
+
+        # 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
+        }
+
+        # Run the backend model
+        st.write("Generating image...")
+        try:
+            output_path = main(args)  # Call your backend main function
+            st.write("Image generation complete!")
+
+            # 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}")
     else:
-        # Generate images
-        with st.spinner("Generating images..."):
-            images = generate_images(pipe, text_input=text_input, sketch=sketches)
-
-        # Display results
-        if images:
-            for i, img in enumerate(images):
-                if isinstance(img, torch.Tensor):  # Convert tensor to image
-                    img = img.squeeze().permute(1, 2, 0).cpu().numpy()
-                    img = Image.fromarray((img * 255).astype("uint8"))
-                st.image(img, caption=f"Generated Image {i+1}")
-        else:
-            st.error("Failed to generate images. Please check the inputs or model configuration.")
+        st.error("Please upload a sketch before generating an image.")