Update app.py

app.py

@@ -14,6 +14,7 @@ from sklearn.metrics.pairwise import cosine_similarity
 import torch
 import numpy as np
 import gradio as gr
+from diffusers import StableDiffusionImg2ImgPipeline
 
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -28,7 +29,7 @@ dataset = load_dataset("lirus18/deepfashion", split="train")
 # Embed a subset of dataset images
 image_vectors = []
 image_indices = []
-N = 500  # You can increase this later if needed
+N = 500
 
 for i in range(N):
     img = dataset[i]['image'].convert("RGB")
@@ -51,13 +52,33 @@ def find_similar(user_image, top_k=3, exclude_index=None):
         sims[exclude_index] = -1
 
     top_idx = sims.argsort()[-top_k:][::-1]
-    return [dataset[image_indices[i]]['image'] for i in top_idx]
+    return [dataset[image_indices[i]]['image'] for i in top_idx], query_vec
 
-# Main function (no generation, just simulation)
+# Load Stable Diffusion
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    low_cpu_mem_usage=True
+).to(device)
+pipe.enable_attention_slicing()
+
+# Generate 10 images
+def generate_outfits(input_image, n=10):
+    prompt = "fashion outfit design inspired by the clothing item"
+    init_image = input_image.resize((512, 512))
+    generated_images = []
+
+    for _ in range(n):
+        result = pipe(prompt=prompt, image=init_image, strength=0.7, guidance_scale=7.5)
+        generated_images.append(result.images[0])
+
+    return generated_images
+
+# Main function
 def recommend_from_upload(uploaded_image):
     uploaded_image = uploaded_image.convert("RGB")
 
-    # Check if uploaded image is already in dataset
+    # Check for duplicates
     closest_idx = None
     for i in range(len(image_indices)):
         dataset_image = dataset[image_indices[i]]['image'].convert("RGB")
@@ -65,15 +86,27 @@ def recommend_from_upload(uploaded_image):
             closest_idx = i
             break
 
-    # Get 3 similar images
-    similar_imgs = find_similar(uploaded_image, top_k=3, exclude_index=closest_idx)
+    # Find similar items
+    similar_imgs, query_vec = find_similar(uploaded_image, top_k=3, exclude_index=closest_idx)
+
+    # Generate 10 new outfits
+    generated_imgs = generate_outfits(uploaded_image, n=10)
 
-    # Use a gray image as the "best generated outfit"
-    placeholder_img = Image.new("RGB", (512, 512), color="gray")
+    # Select best match
+    best_score = -1
+    best_img = None
+    for img in generated_imgs:
+        inputs = processor(images=img, return_tensors="pt").to(device)
+        with torch.no_grad():
+            emb = model.get_image_features(**inputs).cpu().numpy()
+        sim = cosine_similarity(query_vec, emb)[0][0]
+        if sim > best_score:
+            best_score = sim
+            best_img = img
 
-    return [uploaded_image] + similar_imgs + [placeholder_img]
+    return [uploaded_image] + similar_imgs + [best_img]
 
-# Example image paths (must be in root folder)
+# Example paths
 example_paths = [
     ["example1.jpg"],
     ["example2.jpg"],
@@ -82,10 +115,10 @@ example_paths = [
     ["example5.jpg"]
 ]
 
-# Gradio interface with button
+# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 👗 Fashion Outfit Recommender")
-    gr.Markdown("Upload a clothing image to get 3 similar items from the dataset and 1 simulated AI-generated outfit.")
+    gr.Markdown("Upload a clothing image to get 3 similar items from the dataset and 1 AI-generated outfit design.")
 
     with gr.Row():
         image_input = gr.Image(type="pil", label="Upload a clothing item")
@@ -93,11 +126,11 @@ with gr.Blocks() as demo:
     generate_btn = gr.Button("Generate Recommendations")
 
     with gr.Row():
-        output1 = gr.Image(label="Your Input")
-        output2 = gr.Image(label="Similar Item 1")
-        output3 = gr.Image(label="Similar Item 2")
-        output4 = gr.Image(label="Similar Item 3")
-        output5 = gr.Image(label="Simulated AI-Generated Outfit")
+        output1 = gr.Image(label="Your Input", height=512, width=384)
+        output2 = gr.Image(label="Similar Item 1", height=512, width=384)
+        output3 = gr.Image(label="Similar Item 2", height=512, width=384)
+        output4 = gr.Image(label="Similar Item 3", height=512, width=384)
+        output5 = gr.Image(label="AI-Generated Outfit", height=512, width=384)
 
     examples = gr.Examples(
         examples=example_paths,
@@ -113,3 +146,4 @@ if __name__ == "__main__":
     demo.launch()
 
 
+