Update app.py

app.py

@@ -1,25 +1,95 @@
-import streamlit as st
 import os
-import subprocess
+import torch
+import streamlit as st
+from diffusers import AutoencoderKL, DDIMScheduler
+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
+
+# 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
+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()
+    device = accelerator.device
+
+    # Set the seed
+    set_seed(seed)
+
+    # Model and Tokenizer loading (use pre-trained from Hugging Face)
+    model_name = "stabilityai/stable-diffusion-2-1-base"  # Use appropriate model name
+
+    # Load scheduler, tokenizer, and models
+    val_scheduler = DDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
+    val_scheduler.set_timesteps(50, device=device)
+
+    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")
 
-# Title
-st.title("Multimodal Garment Designer")
+    # 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,
+    )
 
-# Input fields
-dataset_path = st.text_input("Dataset Path:", "./assets/data/dresscode")
-output_dir = st.text_input("Output Directory:", "./output")
-test_order = st.selectbox("Test Order:", ["paired", "unpaired"])
-save_name = st.text_input("Save Name:", "model_output")
+    # Freeze VAE and text encoder
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
 
-# Button to run the eval script
-if st.button("Run Evaluation"):
-    # Construct the command
-    command = f"python eval.py --output_dir {output_dir} --dataset_path {dataset_path} --test_order {test_order} --save_name {save_name}"
+    # Select pipeline (use disentangled option if needed)
+    val_pipe = MGDPipe(
+        text_encoder=text_encoder,
+        vae=vae,
+        unet=unet.to(vae.dtype),
+        tokenizer=tokenizer,
+        scheduler=val_scheduler,
+    ).to(device)
 
-    # Run the command
-    with st.spinner("Running evaluation..."):
-        result = subprocess.run(command, shell=True, capture_output=True, text=True)
+    # 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 result
-    st.code(result.stdout)
-    st.code(result.stderr)
+        # Display the generated image
+        st.image(images[0], caption="Generated Fashion Image", use_column_width=True)