Update app.py

app.py

@@ -6,38 +6,41 @@ 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
+from PIL import Image, ImageChops
 from transformers import CLIPProcessor, CLIPModel
 from sklearn.metrics.pairwise import cosine_similarity
 import torch
 import numpy as np
 import gradio as gr
+from diffusers import StableDiffusionImg2ImgPipeline
 
-# Load dataset and model
-dataset = load_dataset("lirus18/deepfashion", split="train")
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+# Device setup
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model and processor
+model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# Load dataset
+dataset = load_dataset("lirus18/deepfashion", split="train")
 
-# Precompute vectors
+# Precompute image vectors
 image_vectors = []
 image_indices = []
-N = 500
+N = 500  # use a smaller subset to avoid long loading
 
 for i in range(N):
-    image = dataset[i]['image'].convert("RGB")
-    inputs = processor(images=image, return_tensors="pt").to(device)
+    img = dataset[i]['image'].convert("RGB")
+    inputs = processor(images=img, return_tensors="pt").to(device)
     with torch.no_grad():
-        embedding = model.get_image_features(**inputs)
-    image_vectors.append(embedding.cpu().numpy().squeeze())
+        emb = model.get_image_features(**inputs)
+    image_vectors.append(emb.cpu().numpy().squeeze())
     image_indices.append(i)
 
 image_vectors = np.array(image_vectors)
 
+# Find similar images
 def find_similar(user_image, top_k=3, exclude_index=None):
     inputs = processor(images=user_image.convert("RGB"), return_tensors="pt").to(device)
     with torch.no_grad():
@@ -45,14 +48,12 @@ 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
+        sims[exclude_index] = -1  # Exclude identical
 
     top_idx = sims.argsort()[-top_k:][::-1]
     return [dataset[image_indices[i]]['image'] for i in top_idx]
 
-# Placeholder for Stable Diffusion (optional)
-from diffusers import StableDiffusionImg2ImgPipeline
-
+# Load Stable Diffusion pipeline
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
@@ -61,11 +62,10 @@ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
 def generate_outfit_from_image(input_image):
     prompt = "fashion outfit design inspired by the clothing item"
     init_image = input_image.resize((512, 512))
-    generated = pipe(prompt=prompt, image=init_image, strength=0.7, guidance_scale=7.5)
-    return generated.images[0]
-
-from PIL import ImageChops
+    result = pipe(prompt=prompt, image=init_image, strength=0.7, guidance_scale=7.5)
+    return result.images[0]
 
+# Main recommendation function
 def recommend_from_upload(uploaded_image):
     uploaded_image = uploaded_image.convert("RGB")
     closest_idx = None
@@ -77,7 +77,6 @@ def recommend_from_upload(uploaded_image):
 
     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]
 
 # Gradio Interface
@@ -95,4 +94,6 @@ demo = gr.Interface(
     description="Upload a clothing image to see 3 similar outfits and 1 AI-generated one!"
 )
 
-demo.launch()
+# Only launch if main
+if __name__ == "__main__":
+    demo.launch()