Update app.py

app.py

@@ -1,350 +1,294 @@
-import os, json, gc, random
-from pathlib import Path
-from typing import List, Tuple, Dict
-
-import gradio as gr
-import numpy as np
+import os, io, math, json, random, numpy as np
+from typing import List, Tuple, Dict, Any
 from PIL import Image
 
 import torch
-import open_clip
+import torch.nn.functional as F
+
+import gradio as gr
 from datasets import load_dataset
 from sklearn.neighbors import NearestNeighbors
-from diffusers import StableDiffusionImageVariationPipeline
-
-# -----------------------------
-# Config (CPU-friendly)
-# -----------------------------
-DATASET_ID = "tukey/human_face_emotions_roboflow"
-EMB_MODEL_NAME = "ViT-H-14"            # open_clip model name
-EMB_PRETRAINED = "laion2b_s32b_b79k"   # laion/CLIP-ViT-H-14-laion2B-s32B-b79K
-GEN_MODEL_ID = "lambdalabs/sd-image-variations-diffusers"
-
-CACHE_DIR = Path("./cache"); CACHE_DIR.mkdir(parents=True, exist_ok=True)
-EMB_MEMMAP_PATH = CACHE_DIR / "clip_vith14_laion2b.float32.memmap"
-LABELS_MEMMAP_PATH = CACHE_DIR / "labels.U32.memmap"
-KNN_META_PATH = CACHE_DIR / "knn_meta.json"
-
-# tiny index + light generation
-INDEX_MAX_DEFAULT = 80
-BATCH_SIZE_DEFAULT = 32
-N_SYN_DEFAULT = 3
-STEPS_DEFAULT = 12
-GUIDANCE_SCALES = [2.5, 3.0, 3.5, 4.0]
-NUM_SYN_TO_SHOW = 5
-
-DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else "cpu")
-
-# -----------------------------
-# Canonical labels + stress
-# -----------------------------
-CANON = {"anger","disgust","fear","happy","neutral","sad","surprise","contempt"}
-CANON_MAP = {
-    "angry": "anger", "happiness": "happy", "happy": "happy",
-    "sadness": "sad", "sad": "sad", "surprised": "surprise", "surprise": "surprise",
-    "content": "neutral", "calm": "neutral", "neutral": "neutral",
-    "contempt": "contempt", "fear": "fear", "disgust": "disgust", "anger": "anger",
-}
-STRESS_W = {"anger":0.95,"fear":0.90,"sad":0.80,"disgust":0.70,"contempt":0.65,"surprise":0.45,"neutral":0.20,"happy":0.05}
-def _bucket(pct: float) -> str: return "Low" if pct < 33 else ("Medium" if pct < 66 else "High")
-def stress_from_top3(res: List[Dict]) -> Tuple[float, str]:
-    probs = {}
-    for r in res:
-        lbl = CANON_MAP.get(str(r["emotion"]).lower(), str(r["emotion"]).lower())
-        if lbl not in CANON: continue
-        probs[lbl] = probs.get(lbl, 0.0) + float(r["confidence_pct"]) / 100.0
-    Z = sum(probs.values()) or 1.0
-    for k in list(probs): probs[k] /= Z
-    s01 = sum(probs.get(k, 0.0) * STRESS_W.get(k, 0.0) for k in probs)
-    s01 = max(0.0, min(1.0, s01))
-    pct = round(s01 * 100.0, 2)
-    return pct, _bucket(pct)
-
-# -----------------------------
-# Lazy globals
-# -----------------------------
-_openclip_model = None
-_preprocess = None
-_nn = None
-_X = None
-_labels_source = None
-_gen_pipe = None
-_dataset_for_labels = None
-
-# -----------------------------
-# Init / cache helpers
-# -----------------------------
-def _load_openclip():
-    global _openclip_model, _preprocess
-    if _openclip_model is not None: return _openclip_model, _preprocess
-    model, _, preprocess = open_clip.create_model_and_transforms(
-        model_name=EMB_MODEL_NAME, pretrained=EMB_PRETRAINED, device=DEVICE
+from transformers import pipeline
+
+# =============== CONFIG ===============
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Embeddings (your chosen model)
+OPENCLIP_BACKBONE = "ViT-H-14"
+OPENCLIP_PRETRAIN = "laion2B-s32B-b79K"  # laion/CLIP-ViT-H-14-laion2B-s32B-b79K
+INDEX_SIZE = int(os.getenv("INDEX_SIZE", 400))  # כמה תמונות מהדאטהסט לאינדוקס
+TOPK_NEAREST = 5
+
+# Emotion model (your chosen model)
+EMO_MODEL = "prithivMLmods/Facial-Emotion-Detection-SigLIP2"
+
+# Optional image-variation generator (your chosen model)
+USE_SD_VARIATIONS = True
+SD_MODEL = "lambdalabs/sd-image-variations-diffusers"
+# =====================================
+
+# ---------- Load OpenCLIP for image embeddings ----------
+try:
+    import open_clip
+    _openclip_model, _, _openclip_preprocess = open_clip.create_model_and_transforms(
+        OPENCLIP_BACKBONE, pretrained=OPENCLIP_PRETRAIN
     )
-    model.eval()
-    _openclip_model, _preprocess = model, preprocess
-    return _openclip_model, _preprocess
-
-def _fit_knn(X): return NearestNeighbors(metric="cosine", algorithm="brute").fit(X)
-
-def _ensure_knn_index(index_max: int, batch_size: int, progress: gr.Progress | None = None):
-    """Build (first run) or load a tiny memmap + KNN over a subset of the dataset."""
-    global _nn, _X, _labels_source, _dataset_for_labels
-
-    if _nn is not None and _X is not None:
-        return
-
-    dataset = load_dataset(DATASET_ID, split="train")
-    if index_max:
-        dataset = dataset.select(range(min(index_max, len(dataset))))
-    _dataset_for_labels = dataset
-    N = len(dataset)
-
-    if EMB_MEMMAP_PATH.exists() and KNN_META_PATH.exists():
-        meta = json.load(open(KNN_META_PATH))
-        if int(meta.get("N", -1)) == N:
-            D = int(meta["D"])
-            X = np.memmap(EMB_MEMMAP_PATH, mode="r", dtype="float32", shape=(N, D))
-            labels = np.memmap(LABELS_MEMMAP_PATH, mode="r", dtype="U32", shape=(N,)) if LABELS_MEMMAP_PATH.exists() else None
-            _X = X; _labels_source = labels; _nn = _fit_knn(X)
-            return
-
-    model, preprocess = _load_openclip()
-    labels_mm = np.memmap(LABELS_MEMMAP_PATH, mode="w+", dtype="U32", shape=(N,))
-    X_w = None; D = None
-
-    with torch.no_grad():
-        for start in range(0, N, batch_size):
-            end = min(start + batch_size, N)
-            imgs = [dataset[i]["image"].convert("RGB") for i in range(start, end)]
-            x = torch.stack([preprocess(im) for im in imgs])
-            if DEVICE in ("cuda", "mps"): x = x.to(DEVICE)
-            v = model.encode_image(x).float()
-            v = v / v.norm(dim=-1, keepdim=True)
-            if X_w is None:
-                D = v.shape[1]
-                X_w = np.memmap(EMB_MEMMAP_PATH, mode="w+", dtype="float32", shape=(N, D))
-            X_w[start:end] = v.detach().cpu().numpy()
-            for i in range(start, end):
-                try: labels_mm[i] = str(dataset[i]["qa"][0]["answer"] or "")
-                except Exception: labels_mm[i] = ""
-            if progress: progress(((end)/N), desc=f"Building index {end}/{N}")
-
-    del X_w; gc.collect()
-    json.dump({"N": int(N), "D": int(D)}, open(KNN_META_PATH, "w"))
-    X = np.memmap(EMB_MEMMAP_PATH, mode="r", dtype="float32", shape=(N, D))
-    labels = np.memmap(LABELS_MEMMAP_PATH, mode="r", dtype="U32", shape=(N,))
-    _X = X; _labels_source = labels; _nn = _fit_knn(X)
-
-def _label_by_idx(i: int):
-    global _labels_source, _dataset_for_labels
-    if _labels_source is not None:
-        lab = str(_labels_source[i]);  return lab if lab else None
-    try: return _dataset_for_labels[i]["qa"][0]["answer"]
-    except Exception: return None
-
-# -----------------------------
-# Embedding + inference utils
-# -----------------------------
+    _openclip_model = _openclip_model.to(DEVICE).eval()
+except Exception as e:
+    raise RuntimeError(
+        f"Failed to load OpenCLIP ({OPENCLIP_BACKBONE} / {OPENCLIP_PRETRAIN}). "
+        f"Install 'open_clip_torch' and verify CUDA if available. Error: {e}"
+    )
+
+@torch.inference_mode()
 def embed_image(img: Image.Image) -> np.ndarray:
-    model, preprocess = _load_openclip()
-    with torch.no_grad():
-        x = preprocess(img.convert("RGB")).unsqueeze(0)
-        if DEVICE in ("cuda", "mps"): x = x.to(DEVICE)
-        v = model.encode_image(x).float()
-        v = v / v.norm(dim=-1, keepdim=True)
-    return v.detach().cpu().numpy().squeeze()
-
-def _top3_emotions_weighted_from_embed(q: np.ndarray,
-                                       start_k: int = 30, step: int = 30,
-                                       method: str = "softmax", tau: float = 0.1):
-    max_k = _X.shape[0]; k = min(start_k, max_k)
-    while True:
-        dist, idx = _nn.kneighbors(q.reshape(1, -1), n_neighbors=k)
-        idx, dist = idx[0], dist[0]
-        sims = np.clip(1.0 - dist, 0.0, None)
-        w = np.exp(sims / tau) if method == "softmax" else sims
-        scores: Dict[str, float] = {}
-        for i, wi in zip(idx, w):
-            lab = _label_by_idx(int(i))
-            if lab is None: continue
-            lab = CANON_MAP.get(str(lab).lower(), str(lab).lower())
-            scores[lab] = scores.get(lab, 0.0) + float(wi)
-        if len([k for k in scores if k in CANON]) >= 3 or k == max_k:
-            break
-        k = min(k + step, max_k)
-    scores = {k: v for k, v in scores.items() if k in CANON and v > 0}
-    if not scores: return []
-    top_items = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:3]
-    vals = np.array([v for _, v in top_items], dtype=np.float32)
-    pct = (vals / vals.sum()) * 100.0 if vals.sum() > 0 else np.zeros_like(vals)
-    return [{"rank": i+1, "emotion": lab, "confidence_pct": int(round(p))}
-            for i, ((lab, _), p) in enumerate(zip(top_items, pct))]
-
-def analyze_face(image: Image.Image):
-    q = embed_image(image)
-    top3 = _top3_emotions_weighted_from_embed(q)
-    stress_pct, stress_lbl = stress_from_top3(top3)
-    return top3, f"{stress_pct}% ({stress_lbl})", q
-
-# ----- Nearest neighbors images from dataset -----
-def _get_dataset_image(i: int) -> Image.Image:
-    return _dataset_for_labels[int(i)]["image"].convert("RGB")
-
-def nearest_k_images_from_dataset(q_emb: np.ndarray, k: int = 5):
-    dist, idx = _nn.kneighbors(q_emb.reshape(1, -1), n_neighbors=k)
-    dist, idx = dist[0], idx[0]
-    sims = (1.0 - dist).tolist()
+    img = img.convert("RGB")
+    tens = _openclip_preprocess(img).unsqueeze(0).to(DEVICE)
+    feats = _openclip_model.encode_image(tens)
+    feats = F.normalize(feats, dim=-1).squeeze(0).detach().cpu().numpy().astype(np.float32)
+    return feats  # shape [D]
+
+# ---------- Dataset + index ----------
+DATASET_NAME = "tukey/human_face_emotions_roboflow"
+DATASET_SPLIT = "train"
+
+def _load_images_for_index(n: int) -> List[Image.Image]:
+    ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
+    n = min(n, len(ds))
+    imgs = []
+    for i in range(n):
+        # column is usually "image"
+        im = ds[i].get("image")
+        if isinstance(im, Image.Image):
+            imgs.append(im.copy())
+    return imgs
+
+def build_index(imgs: List[Image.Image]) -> Tuple[NearestNeighbors, np.ndarray]:
+    vecs = []
+    for im in imgs:
+        vecs.append(embed_image(im))
+    X = np.stack(vecs, axis=0)
+    nn = NearestNeighbors(metric="cosine", n_neighbors=min(TOPK_NEAREST, len(imgs)))
+    nn.fit(X)
+    return nn, X
+
+print("Loading dataset & building index (first time only)...")
+DATASET_IMAGES: List[Image.Image] = _load_images_for_index(INDEX_SIZE)
+NN_MODEL, EMB_MATRIX = build_index(DATASET_IMAGES)
+print(f"Index ready with {len(DATASET_IMAGES)} images.")
+
+def nearest5(pil_img: Image.Image) -> List[Tuple[Image.Image, str]]:
+    q = embed_image(pil_img).reshape(1, -1)
+    dists, idxs = NN_MODEL.kneighbors(q, n_neighbors=min(5, len(DATASET_IMAGES)))
+    # cosine distance -> similarity = 1 - dist
     out = []
-    for i, s in zip(idx, sims):
-        img = _get_dataset_image(int(i))
-        out.append((img, float(s), int(i)))
+    for rank, (dist, idx) in enumerate(zip(dists[0], idxs[0]), start=1):
+        sim = 1.0 - float(dist)
+        im = DATASET_IMAGES[int(idx)]
+        caption = f"#{rank}  sim={sim:.3f}  idx={int(idx)}"
+        out.append((im, caption))
+    return out  # list of (PIL, caption)
+
+# ---------- Emotion & Stress ----------
+EMO_MAP = {
+    "anger": "anger", "angry": "anger",
+    "disgust": "disgust",
+    "fear": "fear",
+    "happy": "happy", "happiness": "happy",
+    "neutral": "neutral", "calm": "neutral",
+    "sad": "sad", "sadness": "sad",
+    "surprise": "surprise",
+    "contempt": "contempt",
+}
+
+# higher == more stressed
+STRESS_WEIGHTS = {
+    "anger": 0.95,
+    "fear": 0.90,
+    "disgust": 0.70,
+    "sad": 0.80,
+    "surprise": 0.55,
+    "neutral": 0.30,
+    "contempt": 0.65,
+    "happy": 0.10,
+}
+
+def _bucket(p: float) -> str:
+    return "Low" if p < 33 else ("Medium" if p < 66 else "High")
+
+emo_pipe = pipeline("image-classification", model=EMO_MODEL, device=0 if DEVICE == "cuda" else -1)
+
+def _pipe_to_probs(res: List[Dict[str, Any]]) -> Dict[str, float]:
+    acc: Dict[str, float] = {}
+    for r in res:
+        label = (r.get("label") or r.get("emotion") or "").lower()
+        if not label:
+            continue
+        key = EMO_MAP.get(label, label)
+        score = float(r.get("score") or r.get("confidence") or r.get("confidence_pct", 0.0) / 100.0)
+        acc[key] = acc.get(key, 0.0) + score
+    Z = sum(acc.values()) or 1.0
+    for k in list(acc.keys()):
+        acc[k] = acc[k] / Z
+    return acc
+
+def emotions_top3(pil_img: Image.Image) -> List[List[Any]]:
+    res = emo_pipe(pil_img)
+    probs = _pipe_to_probs(res)
+    items = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)[:3]
+    table = []
+    for i, (emo, p) in enumerate(items, start=1):
+        table.append([i, emo, round(100.0 * p, 2)])
+    return table  # [[rank, emotion, pct]]
+
+def stress_index(pil_img: Image.Image) -> Tuple[str, float]:
+    res = emo_pipe(pil_img)
+    probs = _pipe_to_probs(res)
+    raw = 0.0
+    for k, v in probs.items():
+        w = STRESS_WEIGHTS.get(k, 0.5)
+        raw += v * w
+    pct = max(0.0, min(100.0, 100.0 * raw))
+    return f"{pct:.1f}%  ({_bucket(pct)})", pct
+
+# ---------- Optional: SD image variations (1 image only) ----------
+sd_pipe = None
+if USE_SD_VARIATIONS:
+    try:
+        from diffusers import StableDiffusionImageVariationPipeline
+        sd_pipe = StableDiffusionImageVariationPipeline.from_pretrained(
+            SD_MODEL, torch_dtype=torch.float32
+        )
+        sd_pipe = sd_pipe.to(DEVICE)
+    except Exception as e:
+        print(f"[WARN] Could not load {SD_MODEL}. Generation disabled. Error: {e}")
+        sd_pipe = None
+
+def generate_one_variation(pil_img: Image.Image, steps: int) -> Image.Image:
+    if sd_pipe is None:
+        raise gr.Error("Image-variation pipeline is not available on this Space.")
+    pil_img = pil_img.convert("RGB")
+    out = sd_pipe(pil_img, guidance_scale=3.0, num_inference_steps=int(steps)).images[0]
     return out
 
-# -----------------------------
-# Generator (optional)
-# -----------------------------
-def _get_gen_pipe():
-    global _gen_pipe
-    if _gen_pipe is not None: return _gen_pipe
-    gen_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-    pipe = StableDiffusionImageVariationPipeline.from_pretrained(
-        GEN_MODEL_ID, revision="v2.0", torch_dtype=gen_dtype
-    ).to(DEVICE)
-    _gen_pipe = pipe
-    return _gen_pipe
-
-def generate_synthetics(base_image: Image.Image, base_embed: np.ndarray, n_syn: int, steps: int, progress: gr.Progress):
-    pipe = _get_gen_pipe()
-    base_gen = torch.Generator(device="cpu").manual_seed(42)
-    records = []
-    for _ in progress.tqdm(range(n_syn), desc="Generating"):
-        seed = int(torch.randint(0, 2**31 - 1, (1,), generator=base_gen).item())
-        gs = random.choice(GUIDANCE_SCALES)
-        g = torch.Generator(device="cpu").manual_seed(seed)
-        out = pipe(image=base_image.convert("RGB"), guidance_scale=gs, num_inference_steps=steps, generator=g)
-        img = out.images[0]
-        emb = embed_image(img)
-        sim = float(np.dot(emb, base_embed))
-        top3_syn = _top3_emotions_weighted_from_embed(emb)
-        stress_pct, stress_lbl = stress_from_top3(top3_syn)
-        records.append({"image": img, "similarity": sim, "top3": top3_syn, "stress": f"{stress_pct}% ({stress_lbl})"})
-    records.sort(key=lambda r: r["similarity"], reverse=True)
-    return records[:NUM_SYN_TO_SHOW]
-
-# -----------------------------
-# UI
-# -----------------------------
-def _format_top3_for_table(top3: List[Dict]) -> List[List]:
-    return [[r["rank"], r["emotion"], r["confidence_pct"]] for r in top3]
-
-with gr.Blocks(title="Face Emotions + Stress (CPU Fast)") as demo:
+# =====================  GRADIO UI  =====================
+CSS = """
+.box { border: 1px solid #e5e7eb; border-radius: 12px; padding: 10px; }
+"""
+
+with gr.Blocks(title="Face Emotion & Stress Analyzer — CPU-friendly", css=CSS, fill_height=False) as demo:
     gr.Markdown(
-        "## Face Emotion & Stress Analyzer — CPU-friendly\n"
+        "### Face Emotion & Stress Analyzer — CPU-friendly\n"
         "- Embeddings: **laion/CLIP-ViT-H-14-laion2B-s32B-b79K** (open_clip)\n"
-        "- Optional SD variations: **lambdalabs/sd-image-variations-diffusers**\n"
-        "- Also shows **nearest 5 images from the dataset** for 1-click results.\n"
+        "- Emotion model: **prithivMLmods/Facial-Emotion-Detection-SigLIP2**\n"
+        "- Optional SD variations: **lambdalabs/sd-image-variations-diffusers** (1 synthetic only)\n"
+        "- Right column shows nearest 5 images from the dataset (clickable)."
     )
 
+    # ---- Row 1: upload + (top3_emotion_original | stress_original) ----
     with gr.Row():
-        inp = gr.Image(type="pil", label="Upload a face image", sources=["upload", "webcam"])
-        idx_cap = gr.Slider(20, 200, value=INDEX_MAX_DEFAULT, step=10, label="Index size (smaller = faster)")
-        bs = gr.Slider(8, 64, value=BATCH_SIZE_DEFAULT, step=8, label="Batch size (build)")
-
-    analyze_btn = gr.Button("Analyze (no synthetics)")
-
-    with gr.Row():
-        with gr.Column():
-            top3_tbl = gr.Dataframe(
+        with gr.Column(scale=2):
+            upload_image = gr.Image(label="Upload face image", type="pil")
+        with gr.Column(scale=1):
+            top3_emotion_original = gr.Dataframe(
                 headers=["Rank", "Emotion", "Confidence (%)"],
                 datatype=["number", "str", "number"],
-                interactive=False, row_count=(3, "fixed"), col_count=(3, "fixed"),
-                label="Top-3 emotions (original image)"
+                interactive=False, label="Top-3 emotions (original image)",
+                value=[]
             )
-            stress_txt = gr.Label(label="Stress index (original)")
-        with gr.Column():
-            # Nearest 5 from dataset (one-click examples)
-            nn_gal = gr.Gallery(
-                label="Nearest 5 from dataset (click one)",
-                columns=[5], height=220, preview=True
+        with gr.Column(scale=1):
+            stress_original = gr.Label(label="Stress index (original)")
+
+    gr.Markdown("#### Analyze (no synthetics)")
+
+    with gr.Row(equal_height=False):
+        # ---------- LEFT COLUMN ----------
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown("**gen_variations_control** — generate only **one** synthetic")
+                steps = gr.Slider(8, 40, value=12, step=1, label="Diffusion steps (higher=slower/better)")
+                gen_btn = gr.Button("Generate 1 synthetic", variant="primary")
+                picked_synth = gr.Image(label="Synthetic preview")
+            top3_emotion_synth = gr.Dataframe(
+                headers=["Rank", "Emotion", "Confidence (%)"],
+                datatype=["number", "str", "number"],
+                interactive=False, label="top3_emotion_synth",
+                value=[]
             )
-            nn_stress = gr.Label(label="Stress (nearest image)")
-            nn_top3 = gr.JSON(label="Top-3 emotions (nearest image)")
-
-            # Optional generator
-            n_syn = gr.Slider(0, 5, value=N_SYN_DEFAULT, step=1, label="How many SD variations to generate")
-            steps = gr.Slider(8, 30, value=STEPS_DEFAULT, step=2, label="Diffusion steps (higher = slower/better)")
-            gen_btn = gr.Button("Generate variations (optional)")
-            gal = gr.Gallery(label="Synthetic variations (click one)", columns=[5], height=220, preview=True)
-            syn_stress = gr.Label(label="Stress (selected synthetic)")
-            syn_top3 = gr.JSON(label="Top-3 emotions (selected synthetic)")
-
-    status = gr.Markdown(visible=False)
-
-    # State
-    syn_state = gr.State([])    # generated variations
-    q_state = gr.State(None)    # embedding of original image
-    img_state = gr.State(None)  # original image
-
-    # ---- Analyze ----
-    def do_analyze(image: Image.Image, cap: int, batch: int, progress=gr.Progress(track_tqdm=True)):
-        try:
-            _ensure_knn_index(index_max=int(cap), batch_size=int(batch), progress=progress)
-            top3, stress, q = analyze_face(image)
-
-            # nearest 5 images from dataset
-            neigh = nearest_k_images_from_dataset(np.array(q, dtype=np.float32), k=5)
-            nn_items = [(im, f"sim={sim:.3f} • idx={idx}") for im, sim, idx in neigh]
-
-            # return: top3, stress, nn gallery, (empty SD gallery), syn_state, q, img, status
-            return (_format_top3_for_table(top3), stress,
-                    nn_items, [], [], q, image, gr.update(visible=False))
-        except Exception as e:
-            return None, None, [], [], [], None, None, gr.update(visible=True, value=f"**Error:** {e}")
-
-    analyze_btn.click(
-        do_analyze,
-        inputs=[inp, idx_cap, bs],
-        outputs=[top3_tbl, stress_txt, nn_gal, gal, syn_state, q_state, img_state, status]
-    )
+            stress_synth = gr.Label(label="stress_synth")
 
-    # ---- One-click on a nearest image ----
-    def on_nn_select(evt: gr.SelectData, q):
-        if q is None:
-            return gr.update(value="Analyze first"), None
-        neigh = nearest_k_images_from_dataset(np.array(q, dtype=np.float32), k=5)
-        i = max(0, min(int(evt.index), len(neigh)-1))
-        img, _, _ = neigh[i]
-        emb = embed_image(img)
-        top3 = _top3_emotions_weighted_from_embed(emb)
-        stress_pct, stress_lbl = stress_from_top3(top3)
-        return f"{stress_pct}% ({stress_lbl})", top3
-
-    nn_gal.select(fn=on_nn_select, inputs=[q_state], outputs=[nn_stress, nn_top3])
-
-    # ---- Optional: generate SD variations ----
-    def do_generate(n: int, s: int, q, img, progress=gr.Progress()):
-        if q is None or img is None:
-            return [], [], gr.update(visible=True, value="**Error:** Analyze first."), None
-        try:
-            recs = generate_synthetics(img, np.array(q, dtype=np.float32), n_syn=int(n), steps=int(s), progress=progress)
-            items = [(r["image"], f"sim={r['similarity']:.3f}") for r in recs]
-            return items, recs, gr.update(visible=False), None
-        except Exception as e:
-            return [], [], gr.update(visible=True, value=f"**Error:** {e}"), None
+        # ---------- RIGHT COLUMN ----------
+        with gr.Column(scale=1):
+            nearest_images_5 = gr.Gallery(
+                label="nearest_images_5 (1-click on 5 examples)",
+                columns=5, rows=1, height=200, allow_preview=False, show_label=True
+            )
+            tops_emotion_nearest = gr.Dataframe(
+                headers=["Rank", "Emotion", "Confidence (%)"],
+                datatype=["number", "str", "number"],
+                interactive=False, label="tops_emotion_nearest_image",
+                value=[]
+            )
+            stress_nearest = gr.Label(label="stress_nearest_image")
+
+    # --------- Hidden states ---------
+    gallery_images_state = gr.State([])   # store PILs
+    gallery_index_state = gr.State([])    # store dataset indexes (ints)
+
+    # ================= Callbacks =================
+    def on_upload(img: Image.Image):
+        if img is None:
+            return gr.update(), gr.update(value=""), [], [], []
+        # original
+        t3 = emotions_top3(img)
+        s_label, _ = stress_index(img)
+        # nearest gallery
+        gal = nearest5(img)  # list[(PIL, caption)]
+        gal_imgs = [g[0] for g in gal]
+        gal_caps = [g[1] for g in gal]
+        # gr.Gallery accepts [(img, caption), ...]
+        gallery = [(im, cap) for im, cap in zip(gal_imgs, gal_caps)]
+        # return
+        return t3, s_label, gallery, gal_imgs, list(range(len(gal_imgs)))
+
+    upload_image.change(
+        fn=on_upload,
+        inputs=[upload_image],
+        outputs=[top3_emotion_original, stress_original, nearest_images_5, gallery_images_state, gallery_index_state]
+    )
 
-    gen_btn.click(
-        do_generate,
-        inputs=[n_syn, steps, q_state, img_state],
-        outputs=[gal, syn_state, status, syn_top3]
+    def on_gallery_select(evt: gr.SelectData, imgs: List[Image.Image], idxs: List[int]):
+        # evt.index is the clicked cell
+        if imgs is None or not imgs:
+            return [], ""
+        i = int(evt.index) if evt is not None else 0
+        i = max(0, min(i, len(imgs)-1))
+        im = imgs[i]
+        t3 = emotions_top3(im)
+        s_label, _ = stress_index(im)
+        return t3, s_label
+
+    nearest_images_5.select(
+        fn=on_gallery_select,
+        inputs=[gallery_images_state, gallery_index_state],
+        outputs=[tops_emotion_nearest, stress_nearest]
     )
 
-    # select from generated synthetics
-    def on_gallery_select(evt: gr.SelectData, syn_records: List[Dict]):
-        if not syn_records or evt is None: return gr.update(value=None), gr.update(value=None)
-        i = int(evt.index); rec = syn_records[i]
-        return gr.update(value=rec["stress"]), gr.update(value=rec["top3"])
+    def on_generate(img: Image.Image, steps_val: int):
+        if img is None:
+            raise gr.Error("Upload an image first.")
+        if sd_pipe is None:
+            raise gr.Error("Synthetic generation is disabled on this Space.")
+        synth = generate_one_variation(img, steps_val)
+        t3 = emotions_top3(synth)
+        s_label, _ = stress_index(synth)
+        return synth, t3, s_label
 
-    gal.select(fn=on_gallery_select, inputs=[syn_state], outputs=[syn_stress, syn_top3])
+    gen_btn.click(
+        fn=on_generate,
+        inputs=[upload_image, steps],
+        outputs=[picked_synth, top3_emotion_synth, stress_synth]
+    )
 
 if __name__ == "__main__":
     demo.launch()