Update app.py

app.py

@@ -6,6 +6,7 @@ Automatically generated by Colab.
 Original file is located at
     https://colab.research.google.com/drive/1_wYfP0IRdb9fpc2zvbg8IqdXGx1dTo7X
 """
+
 from datasets import load_dataset
 from PIL import Image, ImageChops
 from transformers import CLIPProcessor, CLIPModel
@@ -18,17 +19,17 @@ from diffusers import StableDiffusionImg2ImgPipeline
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load model and processor
+# Load CLIP model
 model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
 # Load dataset
 dataset = load_dataset("lirus18/deepfashion", split="train")
 
-# Precompute image vectors
+# Embed a subset of images
 image_vectors = []
 image_indices = []
-N = 500  # use a smaller subset to avoid long loading
+N = 500
 
 for i in range(N):
     img = dataset[i]['image'].convert("RGB")
@@ -48,24 +49,30 @@ def find_similar(user_image, top_k=3, exclude_index=None):
 
     sims = cosine_similarity(query_vec, image_vectors)[0]
     if exclude_index is not None:
-        sims[exclude_index] = -1  # Exclude identical
+        sims[exclude_index] = -1
 
     top_idx = sims.argsort()[-top_k:][::-1]
     return [dataset[image_indices[i]]['image'] for i in top_idx]
 
-# Load Stable Diffusion pipeline
+# 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()
 
-def generate_outfit_from_image(input_image):
+def generate_outfits(input_image, n=10):
     prompt = "fashion outfit design inspired by the clothing item"
     init_image = input_image.resize((512, 512))
-    result = pipe(prompt=prompt, image=init_image, strength=0.7, guidance_scale=7.5)
-    return result.images[0]
+    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 recommendation function
 def recommend_from_upload(uploaded_image):
     uploaded_image = uploaded_image.convert("RGB")
     closest_idx = None
@@ -76,8 +83,18 @@ def recommend_from_upload(uploaded_image):
             break
 
     similar_imgs = find_similar(uploaded_image, top_k=3, exclude_index=closest_idx)
-    generated_img = generate_outfit_from_image(uploaded_image)
-    return [uploaded_image] + similar_imgs + [generated_img]
+    generated_imgs = generate_outfits(uploaded_image, n=10)
+
+    return [uploaded_image] + similar_imgs + generated_imgs
+
+# 5 clickable example images (must be uploaded to the repo)
+example_paths = [
+    ["fashion_examples/example1.jpg"],
+    ["fashion_examples/example2.jpg"],
+    ["fashion_examples/example3.jpg"],
+    ["fashion_examples/example4.jpg"],
+    ["fashion_examples/example5.jpg"]
+]
 
 # Gradio Interface
 demo = gr.Interface(
@@ -88,12 +105,13 @@ demo = gr.Interface(
         gr.Image(label="Similar Item 1"),
         gr.Image(label="Similar Item 2"),
         gr.Image(label="Similar Item 3"),
-        gr.Image(label="Generated New Outfit"),
+        gr.Gallery(label="AI-Generated Outfits (x10)").style(grid=(5, 2), height="auto"),
     ],
     title="👗 Fashion Outfit Recommender",
-    description="Upload a clothing image to see 3 similar outfits and 1 AI-generated one!"
+    description="Upload a clothing image to get 3 similar items from the dataset and 10 AI-generated outfit designs.",
+    examples=example_paths
 )
 
-# Only launch if main
 if __name__ == "__main__":
     demo.launch()
+