Create App.PY

App.PY

@@ -0,0 +1,435 @@
+"""
+AI Fashion Stylist — Hugging Face Space app
+Combines: CLIP recommendation engine (Part 3) + 3 GenAI patterns (Part 4):
+  1. Small Language Model      -> generates text   (Qwen2.5-0.5B-Instruct)
+  2. Small Vision-Language Model -> answers questions (BLIP-VQA-base)
+  3. Small Vision-Language pipeline -> generates an image (clothing segmentation + SD inpainting)
+
+Deploy notes:
+- Upload this file + requirements.txt + final_image_embeddings.npy + catalog_metadata.parquet
+  to your HF Space repo root.
+- Update HF_DATASET_REPO / HF_WINNING_MODEL below if your repo names differ.
+- For Path A "Quick Starters" to actually 1-click work, add 2-3 sample photos to a
+  samples/ folder in your Space and update SAMPLE_PHOTOS below.
+"""
+
+import urllib.parse
+
+import numpy as np
+import pandas as pd
+import torch
+import gradio as gr
+import faiss
+from datasets import load_dataset
+from PIL import Image
+from scipy import ndimage
+from transformers import (
+    CLIPModel, CLIPProcessor, pipeline as hf_pipeline,
+    BlipProcessor, BlipForQuestionAnswering,
+    SegformerImageProcessor, AutoModelForSemanticSegmentation,
+)
+from diffusers import StableDiffusionInpaintPipeline
+
+# ---------------------------------------------------------------------------
+# CONFIG — update these to match your own HF repos
+# ---------------------------------------------------------------------------
+HF_DATASET_REPO = "lihicarmeli/fashion-stylist-multimodal-v2"   # your HF dataset repo
+HF_WINNING_MODEL = "openai/clip-vit-base-patch32"                # winning embedding model (Part 3)
+EMBEDDINGS_FILE = "final_image_embeddings.npy"                   # uploaded next to this app.py
+METADATA_FILE = "catalog_metadata.parquet"                       # uploaded next to this app.py
+SAMPLE_PHOTOS = ["samples/demo_woman.jpg", "samples/demo_man.jpg", "samples/demo_teen.jpg"]
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---------------------------------------------------------------------------
+# LOAD DATA + WINNING EMBEDDING MODEL (runs once, on Space startup)
+# ---------------------------------------------------------------------------
+print("Loading dataset from HF Hub...")
+ds = load_dataset(HF_DATASET_REPO)
+df = ds["train"].to_pandas()
+images = [ds["train"][i]["image_improved"] for i in range(len(ds["train"]))]
+
+print("Loading precomputed embeddings...")
+image_embeddings = np.load(EMBEDDINGS_FILE).astype("float32")
+
+print("Loading winning embedding model (CLIP) from HF Hub...")
+win_model = CLIPModel.from_pretrained(HF_WINNING_MODEL).to(DEVICE).eval()
+win_processor = CLIPProcessor.from_pretrained(HF_WINNING_MODEL)
+
+print("Building FAISS index...")
+dimension = image_embeddings.shape[1]
+faiss_img_index = faiss.IndexFlatL2(dimension)
+faiss.normalize_L2(image_embeddings)
+faiss_img_index.add(image_embeddings)
+
+# ---------------------------------------------------------------------------
+# GENERATION MODELS — the 3 "Good Examples" patterns
+# ---------------------------------------------------------------------------
+print("Loading small language model for text generation (Qwen2.5-0.5B-Instruct)...")
+caption_gen_pipe = hf_pipeline(
+    "text-generation",
+    model="Qwen/Qwen2.5-0.5B-Instruct",
+    device=0 if DEVICE == "cuda" else -1,
+)
+
+print("Loading small vision-language model for VQA (BLIP-VQA-base)...")
+vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
+vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to(DEVICE)
+
+print("Loading clothing segmentation model (helper - finds where the clothes are)...")
+seg_processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+seg_model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes").to(DEVICE)
+
+print("Loading image-generation model for the new outfit (Stable Diffusion Inpainting)...")
+inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+    safety_checker=None,
+).to(DEVICE)
+
+GENDERS = sorted(df["gender"].unique().tolist())
+AGE_GROUPS = sorted(df["age_group"].unique().tolist())
+SKIN_TONES = sorted(df["skin_tone"].unique().tolist())
+UNDERTONES = sorted(df["undertone"].unique().tolist())
+STYLES = sorted(df["style_preference"].unique().tolist())
+
+print("Space ready.")
+
+
+# ---------------------------------------------------------------------------
+# EMBEDDING + SEARCH (Part 3 logic, unchanged)
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def embed_query_image(pil_image):
+    inputs = win_processor(images=pil_image, return_tensors="pt").to(DEVICE)
+    feats = win_model.get_image_features(**inputs)
+    feats = feats.pooler_output if hasattr(feats, "pooler_output") else feats
+    return feats.cpu().numpy().astype("float32")
+
+
+@torch.no_grad()
+def embed_query_text(sentence):
+    inputs = win_processor(text=[sentence], return_tensors="pt", padding=True, truncation=True).to(DEVICE)
+    feats = win_model.get_text_features(**inputs)
+    feats = feats.pooler_output if hasattr(feats, "pooler_output") else feats
+    return feats.cpu().numpy().astype("float32")
+
+
+def build_feature_sentence(skin_tone, undertone, style_preference, gender=None, age_group=None):
+    descriptor = " ".join(p for p in [age_group, gender] if p) or "person"
+    return (
+        f"a {descriptor} with {skin_tone} skin tone and {undertone} undertone, "
+        f"wearing a {style_preference} style outfit"
+    )
+
+
+def faiss_filtered_search(query_emb, top_k=3, exclude_idx=None, gender=None, age_group=None):
+    faiss.normalize_L2(query_emb)
+    k = len(df)
+    distances, indices = faiss_img_index.search(query_emb, k)
+    distances, indices = distances[0], indices[0]
+
+    def collect(require_gender, require_age):
+        kept_i, kept_d = [], []
+        for idx, dist in zip(indices, distances):
+            if idx == -1 or (exclude_idx is not None and idx == exclude_idx):
+                continue
+            row = df.iloc[idx]
+            if require_gender and gender and str(row["gender"]).lower() != str(gender).lower():
+                continue
+            if require_age and age_group and str(row["age_group"]).lower() != str(age_group).lower():
+                continue
+            kept_i.append(idx)
+            kept_d.append(dist)
+            if len(kept_i) == top_k:
+                break
+        return kept_i, kept_d
+
+    kept_i, kept_d = collect(True, True)
+    if len(kept_i) < top_k:
+        kept_i, kept_d = collect(True, False)
+    if len(kept_i) < top_k:
+        kept_i, kept_d = collect(False, False)
+
+    return np.array(kept_i), df.iloc[kept_i], np.array(kept_d)
+
+
+# ---------------------------------------------------------------------------
+# GENERATION — 3 "Good Examples" patterns (Part 4)
+# ---------------------------------------------------------------------------
+
+# --- Pattern 1: small Language Model -> generate text ---
+def generate_stylist_caption(row):
+    user_prompt = (
+        f"Write one short, warm sentence (max 25 words) from a fashion stylist, recommending this look: "
+        f"a {row['style_preference']} style outfit in {row['primary_color']} and {row['secondary_color']}, "
+        f"best colors: {row['recommended_colors']}. Be specific and stylish, no hashtags."
+    )
+    messages = [{"role": "user", "content": user_prompt}]
+    output = caption_gen_pipe(messages, max_new_tokens=40, do_sample=True, temperature=0.7)
+    return output[0]["generated_text"][-1]["content"].strip()
+
+
+# --- Pattern 2: small Vision-Language Model -> answer questions about an image ---
+@torch.no_grad()
+def answer_question_about_image(pil_image, question):
+    inputs = vqa_processor(pil_image.convert("RGB"), question, return_tensors="pt").to(DEVICE)
+    output_ids = vqa_model.generate(**inputs, max_new_tokens=20)
+    return vqa_processor.decode(output_ids[0], skip_special_tokens=True)
+
+
+# --- Pattern 3: small Vision-Language Model -> generate the FULL LOOK (full-body image) ---
+FACE_LABEL_ID = 11
+HAIR_LABEL_ID = 2
+
+
+@torch.no_grad()
+def get_face_protect_mask(pil_image, use_geometric_fallback=True):
+    """Returns a boolean mask (True = protect) covering the face + hair, at pil_image's own
+    resolution. Combines the segmentation model's prediction with a fixed geometric ellipse
+    (centered, where a face statistically sits in a headshot crop) as a safety net - so even
+    if segmentation misclassifies an unusual headwear/turban, the visible face is still
+    guaranteed to be protected."""
+    inputs = seg_processor(images=pil_image, return_tensors="pt").to(DEVICE)
+    logits = seg_model(**inputs).logits
+    upsampled = torch.nn.functional.interpolate(
+        logits, size=pil_image.size[::-1], mode="bilinear", align_corners=False
+    )
+    pred_seg = upsampled.argmax(dim=1)[0].cpu().numpy()
+    seg_protect = np.isin(pred_seg, [FACE_LABEL_ID, HAIR_LABEL_ID])
+
+    if not use_geometric_fallback:
+        return seg_protect
+
+    h, w = seg_protect.shape
+    yy, xx = np.mgrid[0:h, 0:w]
+    cy, cx = h * 0.42, w * 0.5       # face center: slightly above vertical middle of a headshot
+    ry, rx = h * 0.30, w * 0.22      # ellipse radii tuned for a tight headshot/bust crop
+    geometric_protect = (((xx - cx) / rx) ** 2 + ((yy - cy) / ry) ** 2) <= 1.0
+
+    return seg_protect | geometric_protect
+
+
+def build_full_look_prompt(row):
+    gender_word = "man" if str(row["gender"]).lower() in ("man", "male") else "woman"
+    return (
+        f"full body fashion photo of a {row['age_group']} {gender_word}, standing, "
+        f"wearing a {row['style_preference']} style outfit: {row['outfit_top']}, "
+        f"{row['outfit_bottom']}, {row['outfit_shoes']}, {row['outfit_accessory']}, "
+        f"in {row['primary_color']} and {row['secondary_color']}, "
+        f"studio lighting, plain background, head to toe, high quality fashion photography"
+    )
+
+
+def generate_new_outfit_image(pil_image, row, target_gender=None, canvas_size=(512, 1024),
+                               head_width_frac=0.45, steps=40, guidance_scale=8.0):
+    """Generates the FULL LOOK: keeps the original face pixel-identical (scaled down to a
+    realistic head-to-body proportion) and generates the rest of a standing figure wearing
+    the complete recommended outfit (top, bottom, shoes, accessory) around it."""
+    if target_gender is not None and str(row["gender"]).lower() != str(target_gender).lower():
+        print(f"⚠️ Warning: recommended row gender ({row['gender']}) != expected gender "
+              f"({target_gender}) - double-check which matched_rows was passed in.")
+
+    base = pil_image.convert("RGB")
+    protect = get_face_protect_mask(base)  # (H, W) bool, at base's own resolution
+
+    canvas_w, canvas_h = canvas_size
+    head_w = int(canvas_w * head_width_frac)
+    scale = head_w / base.width
+    head_h = int(base.height * scale)
+
+    resized_face_crop = base.resize((head_w, head_h))
+    resized_protect_img = Image.fromarray(protect.astype(np.uint8) * 255).resize(
+        (head_w, head_h), resample=Image.NEAREST
+    )
+
+    canvas = Image.new("RGB", canvas_size, color=(128, 128, 128))
+    paste_x = (canvas_w - head_w) // 2
+    paste_y = int(canvas_h * 0.03)
+    canvas.paste(resized_face_crop, (paste_x, paste_y))
+
+    mask_arr = np.full((canvas_h, canvas_w), 255, dtype=np.uint8)  # 255 = let the model generate
+    protect_resized_arr = np.array(resized_protect_img) > 127
+    mask_arr[paste_y:paste_y + head_h, paste_x:paste_x + head_w][protect_resized_arr] = 0
+    mask = Image.fromarray(mask_arr).convert("L")
+
+    prompt = build_full_look_prompt(row)
+    generated = inpaint_pipe(
+        prompt=prompt,
+        image=canvas,
+        mask_image=mask,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        height=canvas_h,
+        width=canvas_w,
+    ).images[0]
+
+    # Hard-composite: guarantees zero change on the protected face/hair pixels
+    final_image = Image.composite(generated, canvas, mask)
+    return final_image, prompt
+
+
+# ---------------------------------------------------------------------------
+# SHOP LINKS (Part 3 logic, unchanged)
+# ---------------------------------------------------------------------------
+RETAILER_SEARCH_URLS = {
+    "zara":  "https://www.zara.com/us/en/search?searchTerm={query}&section={section}",
+    "hm":    "https://www2.hm.com/en_us/search-results.html?q={query}",
+    "asos":  "https://www.asos.com/us/{dept}/search/?q={query}",
+    "mango": "https://shop.mango.com/us/en/search?kw={query}",
+    "shein": "https://us.shein.com/pdsearch/{query}/",
+}
+
+
+def normalize_gender(gender):
+    g = str(gender).strip().lower() if gender is not None else ""
+    return "men" if g in ("male", "man", "men", "m") else "women"
+
+
+def to_shop_link(retailer, value, gender=None):
+    dept = normalize_gender(gender)
+    if retailer == "zara":
+        section = "MAN" if dept == "men" else "WOMAN"
+        return RETAILER_SEARCH_URLS["zara"].format(query=urllib.parse.quote(str(value)), section=section)
+    if retailer == "asos":
+        return RETAILER_SEARCH_URLS["asos"].format(query=urllib.parse.quote(str(value)), dept=dept)
+    gender_word = "men's" if dept == "men" else "women's"
+    return RETAILER_SEARCH_URLS[retailer].format(query=urllib.parse.quote(f"{gender_word} {value}"))
+
+
+def shop_links_markdown(row):
+    lines = []
+    for retailer, col in [("zara", "search_query_zara"), ("hm", "search_query_hm"),
+                           ("asos", "search_query_asos"), ("mango", "search_query_mango"),
+                           ("shein", "search_query_shein")]:
+        link = to_shop_link(retailer, row[col], gender=row["gender"])
+        lines.append(f"- **{retailer.upper()}**: [{row[col]}]({link})")
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# MAIN PIPELINE — shared by both input paths, runs all 3 GenAI patterns
+# ---------------------------------------------------------------------------
+def run_pipeline(matched_indices, matched_rows, matched_scores, base_image=None, question=None,
+                  expected_gender=None):
+    if len(matched_indices) == 0:
+        return [], None, "No matches found — try different filters.", "", ""
+
+    similarity = [round(1.0 - (d / 2.0), 3) for d in matched_scores]
+    gallery = [(images[idx], f"Match #{i + 1} ({similarity[i]})") for i, idx in enumerate(matched_indices)]
+
+    top_row = matched_rows.iloc[0]
+    # Pattern 3 needs a base photo: the user's own photo (Path A) or the top catalog match (Path B)
+    edit_base_image = base_image if base_image is not None else images[matched_indices[0]]
+    # If no explicit gender filter was given, fall back to the matched row's own gender, so
+    # the generated outfit's wording always matches the person actually shown in the photo.
+    gender_for_generation = expected_gender or top_row["gender"]
+
+    caption = generate_stylist_caption(top_row)                         # Pattern 1
+    new_image, _prompt = generate_new_outfit_image(
+        edit_base_image, top_row, target_gender=gender_for_generation
+    )                                                                     # Pattern 3
+
+    answer = ""
+    if question:
+        answer = answer_question_about_image(edit_base_image, question)  # Pattern 2
+
+    shop_md = "\n\n---\n\n".join(
+        f"**Match #{i + 1}** ({row['gender']}, {row['style_preference']}, colors: {row['recommended_colors']})\n"
+        + shop_links_markdown(row)
+        for i, (_, row) in enumerate(matched_rows.iterrows())
+    )
+    return gallery, new_image, caption, answer, shop_md
+
+
+def recommend_from_photo(photo, gender, age_group, question):
+    if photo is None:
+        return [], None, "Please upload a photo or pick a Quick Starter.", "", ""
+    query_emb = embed_query_image(photo)
+    idx, rows, scores = faiss_filtered_search(query_emb, gender=gender or None, age_group=age_group or None)
+    return run_pipeline(idx, rows, scores, base_image=photo, question=question, expected_gender=gender)
+
+
+def recommend_from_features(skin_tone, undertone, style, gender, age_group, question):
+    sentence = build_feature_sentence(skin_tone, undertone, style, gender, age_group)
+    query_emb = embed_query_text(sentence)
+    idx, rows, scores = faiss_filtered_search(query_emb, gender=gender or None, age_group=age_group or None)
+    return run_pipeline(idx, rows, scores, base_image=None, question=question, expected_gender=gender)
+
+
+# ---------------------------------------------------------------------------
+# GRADIO UI
+# ---------------------------------------------------------------------------
+with gr.Blocks(title="AI Fashion Stylist") as demo:
+    gr.Markdown(
+        "# 👗 AI Fashion Stylist\n"
+        "Upload a photo **or** describe your style — get 3 real catalog matches, shop links, "
+        "a brand-new AI-edited outfit photo, a stylist note, and answers to your styling questions."
+    )
+
+    with gr.Tab("📸 Upload a Photo"):
+        with gr.Row():
+            photo_in = gr.Image(type="pil", label="Your photo")
+            with gr.Column():
+                gender_a = gr.Dropdown(GENDERS, label="Gender (optional)")
+                age_a = gr.Dropdown(AGE_GROUPS, label="Age group (optional)")
+                question_a = gr.Textbox(label="Ask the stylist a question about your photo (optional)",
+                                         placeholder="e.g. What style would suit me best?")
+                btn_a = gr.Button("Find My Style", variant="primary")
+        gallery_a = gr.Gallery(label="Top 3 Matches", columns=3)
+        new_img_a = gr.Image(label="✨ New AI-Edited Outfit")
+        caption_a = gr.Textbox(label="Stylist Note")
+        answer_a = gr.Textbox(label="Answer to your question")
+        shop_a = gr.Markdown(label="Shop the Look")
+        btn_a.click(
+            recommend_from_photo,
+            [photo_in, gender_a, age_a, question_a],
+            [gallery_a, new_img_a, caption_a, answer_a, shop_a],
+        )
+
+        gr.Examples(
+            examples=[[p, None, None, "What style would suit me best?"] for p in SAMPLE_PHOTOS],
+            inputs=[photo_in, gender_a, age_a, question_a],
+            outputs=[gallery_a, new_img_a, caption_a, answer_a, shop_a],
+            fn=recommend_from_photo,
+            cache_examples=True,  # precomputed at Space startup -> truly "1-click and see a result"
+            label="Quick Starters",
+        )
+
+    with gr.Tab("✏️ Describe Your Style"):
+        with gr.Row():
+            with gr.Column():
+                skin_b = gr.Dropdown(SKIN_TONES, label="Skin tone", value=SKIN_TONES[0])
+                undertone_b = gr.Dropdown(UNDERTONES, label="Undertone", value=UNDERTONES[0])
+                style_b = gr.Dropdown(STYLES, label="Style preference", value=STYLES[0])
+                gender_b = gr.Dropdown(GENDERS, label="Gender", value=GENDERS[0])
+                age_b = gr.Dropdown(AGE_GROUPS, label="Age group", value=AGE_GROUPS[0])
+                question_b = gr.Textbox(label="Ask the stylist a question about the top match (optional)",
+                                         placeholder="e.g. Is this outfit formal or casual?")
+                btn_b = gr.Button("Find My Style", variant="primary")
+        gallery_b = gr.Gallery(label="Top 3 Matches", columns=3)
+        new_img_b = gr.Image(label="✨ New AI-Edited Outfit")
+        caption_b = gr.Textbox(label="Stylist Note")
+        answer_b = gr.Textbox(label="Answer to your question")
+        shop_b = gr.Markdown(label="Shop the Look")
+        btn_b.click(
+            recommend_from_features,
+            [skin_b, undertone_b, style_b, gender_b, age_b, question_b],
+            [gallery_b, new_img_b, caption_b, answer_b, shop_b],
+        )
+
+        gr.Examples(
+            examples=[
+                ["deep", "cool", "boho", "woman", "young adult", "Is this outfit formal or casual?"],
+                ["tan", "warm", "elegant", "man", "adult", "What occasion suits this outfit?"],
+                ["fair", "neutral", "minimalist", "woman", "teen", "What season is this outfit best for?"],
+            ],
+            inputs=[skin_b, undertone_b, style_b, gender_b, age_b, question_b],
+            outputs=[gallery_b, new_img_b, caption_b, answer_b, shop_b],
+            fn=recommend_from_features,
+            cache_examples=True,  # precomputed at Space startup -> truly "1-click and see a result"
+            label="Quick Starters",
+        )
+
+if __name__ == "__main__":
+    demo.launch()