Update app.py

app.py

@@ -1,256 +1,315 @@
-import os, glob, math
-import numpy as np
+import os, json, tempfile
+import cv2, numpy as np, gradio as gr
 from PIL import Image
-import gradio as gr
 
-# Your hairstyles live here:
-HAIR_DIR = "hair"
-
-# Optional heavy deps (installed via requirements.txt)
+# ===================== Paths & Assets =====================
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+CANDIDATES = [
+    os.path.join(BASE_DIR, "hair"),  # your folder
+    os.path.join(BASE_DIR, "assets", "hairstyles"),
+    os.path.join(BASE_DIR, "assets", "Hairstyles"),
+    os.path.join(BASE_DIR, "hairstyles"),
+]
+HAIR_DIR = None
+for p in CANDIDATES:
+    if os.path.isdir(p):
+        HAIR_DIR = p
+        break
+if HAIR_DIR is None:
+    HAIR_DIR = os.path.join(BASE_DIR, "hair")
+    os.makedirs(HAIR_DIR, exist_ok=True)
+
+META_PATH = os.path.join(HAIR_DIR, "meta.json")  # optional per-style anchors
+
+# ===================== Dependencies =====================
 try:
-    import cv2
     import mediapipe as mp
-    MP_AVAILABLE = True
-except Exception:
-    MP_AVAILABLE = False
-
-APP_TITLE = "Salon Hairstyle Virtual Try-On — Pro Demo"
-
-# --------------------------
-# Asset loading / utilities
-# --------------------------
-def list_hairstyle_files():
-    files = sorted(glob.glob(os.path.join(HAIR_DIR, "*.png")))
-    labels = [os.path.basename(f).replace(".png","").replace("_"," ").title() for f in files]
-    thumbs = [f for f in files]  # Gradio Gallery accepts path strings
-    return labels, files, thumbs
-
-def ensure_rgba(im: Image.Image) -> Image.Image:
-    return im if im.mode == "RGBA" else im.convert("RGBA")
-
-def overlay_rgba(base: Image.Image, overlay: Image.Image, x: int, y: int) -> Image.Image:
-    base = ensure_rgba(base)
-    overlay = ensure_rgba(overlay)
-    bw, bh = base.size
-    ow, oh = overlay.size
-    if ow <= 0 or oh <= 0:
-        return base
-    x1, y1 = max(0, x), max(0, y)
-    x2, y2 = min(bw, x + ow), min(bh, y + oh)
-    if x1 >= x2 or y1 >= y2:
-        return base
-    crop = overlay.crop((x1 - x, y1 - y, x2 - x, y2 - y))
-    region = base.crop((x1, y1, x2, y2))
-    region = Image.alpha_composite(region, crop)
-    base.paste(region, (x1, y1))
-    return base
-
-# --------------------------
-# Face detection / anchoring
-# --------------------------
-def detect_face_keypoints(np_img):
-    """
-    Returns dict:
-      bbox: (x1,y1,x2,y2)
-      temples: (lx,ly, rx,ry)
-      forehead_y: int
-    """
-    h, w = np_img.shape[:2]
-
-    def fallback():
-        cx, cy = w // 2, int(h * 0.42)
-        bw, bh = int(w * 0.4), int(h * 0.45)
-        x1, y1 = max(0, cx - bw // 2), max(0, cy - bh // 2)
-        x2, y2 = min(w, x1 + bw), min(h, y1 + bh)
-        return {
-            "bbox": (x1, y1, x2, y2),
-            "temples": (x1, (y1 + y2) // 2, x2, (y1 + y2) // 2),
-            "forehead_y": max(0, int(y1 - 0.1 * (y2 - y1))),
-        }
-
-    if not MP_AVAILABLE:
-        return fallback()
-
-    mpfm = mp.solutions.face_mesh
-    with mpfm.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=False) as fm:
-        rgb = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB) if np_img.shape[2] == 3 else np_img[..., :3]
-        res = fm.process(rgb)
+except Exception as e:
+    raise RuntimeError(f"Mediapipe import failed. Check requirements pins. Details: {e}")
+
+mp_face_mesh = mp.solutions.face_mesh
+mp_selfie_seg = mp.solutions.selfie_segmentation
+LM = {"left_eye_outer": 33, "right_eye_outer": 263, "mid_forehead": 10}
+
+# ===================== Utilities =====================
+def load_hairstyles():
+    try:
+        files = [f for f in os.listdir(HAIR_DIR) if f.lower().endswith(".png")]
+    except FileNotFoundError:
+        files = []
+    files.sort()
+    return files
+
+def load_meta():
+    if os.path.exists(META_PATH):
+        try:
+            with open(META_PATH, "r") as f:
+                m = json.load(f)
+            return m if isinstance(m, dict) else {}
+        except Exception:
+            return {}
+    return {}
+
+def premultiply_alpha(bgra):
+    """Nicer blending (removes gray halos)."""
+    bgr = bgra[:, :, :3].astype(np.float32) / 255.0
+    a = (bgra[:, :, 3:4].astype(np.float32) / 255.0)
+    bgr_pm = (bgr * a * 255.0).astype(np.uint8)
+    return np.dstack([bgr_pm, bgra[:, :, 3]])
+
+def load_hair_png(name):
+    path = os.path.join(HAIR_DIR, name)
+    hair = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # BGRA
+    if hair is None or hair.shape[2] != 4:
+        raise ValueError(f"Invalid hair asset: {name} (must be RGBA PNG)")
+    return premultiply_alpha(hair)
+
+def detect_face_keypoints(img_bgr):
+    h, w = img_bgr.shape[:2]
+    with mp_face_mesh.FaceMesh(
+        static_image_mode=True, max_num_faces=1, refine_landmarks=True,
+        min_detection_confidence=0.6
+    ) as fm:
+        res = fm.process(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
         if not res.multi_face_landmarks:
-            return fallback()
-
+            return None
         lm = res.multi_face_landmarks[0].landmark
-        xs = np.array([p.x for p in lm]) * w
-        ys = np.array([p.y for p in lm]) * h
-        x1, y1, x2, y2 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
-
-        def safe_idx(i):
-            i = int(i)
-            i = max(0, min(len(lm) - 1, i))
-            return int(xs[i]), int(ys[i])
-
-        lx, ly = safe_idx(127)  # approx left temple
-        rx, ry = safe_idx(356)  # approx right temple
-        _, fy = safe_idx(10)    # approx forehead
-        y1 = max(0, int(y1 - 0.12 * (y2 - y1)))  # include hairline
-        return {
-            "bbox": (x1, y1, x2, y2),
-            "temples": (lx, ly, rx, ry),
-            "forehead_y": fy,
-        }
-
-def place_and_render(base_np, hair_path, scale, lr_shift_pct, ud_shift_pct, rotation_deg):
-    if base_np is None:
-        raise gr.Error("Please upload or capture a photo first.")
-    base = Image.fromarray(base_np.astype("uint8")).convert("RGBA")
-    key = detect_face_keypoints(np.array(base.convert("RGB")))
-    x1, y1, x2, y2 = key["bbox"]
-    lx, ly, rx, ry = key["temples"]
-    forehead_y = key["forehead_y"]
-
-    hair = Image.open(hair_path).convert("RGBA")
-
-    # Scale from temple distance
-    temple_dx = rx - lx
-    temple_dy = ry - ly
-    temple_dist = max(1, (temple_dx ** 2 + temple_dy ** 2) ** 0.5)
-    target_w = max(1, int(temple_dist * 2.0 * scale))
-    ratio = target_w / hair.width
-    target_h = max(1, int(hair.height * ratio))
-    hair_resized = hair.resize((target_w, target_h), Image.LANCZOS)
-
-    # Auto-rotation + manual extra rotation
-    auto_deg = math.degrees(math.atan2(temple_dy, temple_dx))
-    rot_total = auto_deg + rotation_deg
-    hair_resized = hair_resized.rotate(rot_total, expand=True)
-
-    # Anchor at temple midpoint, above forehead
-    midx = int((lx + rx) / 2)
-    anchor_x = int(midx - hair_resized.width / 2)
-    anchor_y = int(forehead_y - hair_resized.height * 0.45)
-
-    # Apply shifts (percent of image size)
-    img_w, img_h = base.size
-    anchor_x += int(lr_shift_pct * img_w / 100.0)  # +right, -left
-    anchor_y += int(ud_shift_pct * img_h / 100.0)  # +down,  -up
-
-    out = overlay_rgba(base, hair_resized, anchor_x, anchor_y).convert("RGB")
+        def xy(i): return np.array([lm[i].x*w, lm[i].y*h], dtype=np.float32)
+        return np.stack([xy(LM["left_eye_outer"]), xy(LM["right_eye_outer"]), xy(LM["mid_forehead"])])
+
+def person_mask(img_bgr):
+    with mp_selfie_seg.SelfieSegmentation(model_selection=1) as seg:
+        rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        m = seg.process(rgb).segmentation_mask
+        mask = (m > 0.5).astype(np.float32)
+        mask = cv2.GaussianBlur(mask, (35, 35), 0)
+        return mask
+
+def hair_reference_points(hair_bgra, filename, meta):
+    h, w = hair_bgra.shape[:2]
+    if filename in meta:
+        pts = np.array(meta[filename], dtype=np.float32)
+        if pts.shape == (3, 2):
+            return pts
+    # Defaults (OK for many styles). For perfect fit, add 3 points per file to meta.json.
+    pL = np.array([0.30*w, 0.60*h], dtype=np.float32)
+    pR = np.array([0.70*w, 0.60*h], dtype=np.float32)
+    pM = np.array([0.50*w, 0.40*h], dtype=np.float32)
+    return np.stack([pL, pR, pM], axis=0)
+
+def warp_and_alpha_blend(base_bgr, hair_bgra, M, opacity=1.0):
+    H, W = base_bgr.shape[:2]
+    hair_rgb = hair_bgra[:, :, :3]
+    hair_a = hair_bgra[:, :, 3] / 255.0
+    hair_warp = cv2.warpAffine(hair_rgb, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
+    a_warp = cv2.warpAffine(hair_a, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
+    a = np.clip(a_warp * opacity, 0, 1)[..., None]
+    out = (a * hair_warp + (1 - a) * base_bgr).astype(np.uint8)
     return out
 
-# --------------------------
-# Gradio callbacks
-# --------------------------
-def refresh_assets():
-    labels, files, thumbs = list_hairstyle_files()
-    if not files:
-        return (
-            gr.update(value=None, choices=[]),
-            gr.update(value=None),
-            "No PNGs found in the 'hair/' folder.",
-        )
-    dd = gr.update(choices=labels, value=labels[0])
-    gallery = gr.update(value=[[p, l] for p, l in zip(thumbs, labels)])
-    return dd, gallery, f"Found {len(files)} hairstyles."
-
-def refresh_dropdown_only():
-    labels, files, _ = list_hairstyle_files()
-    if not files:
-        return gr.update(value=None, choices=[])
-    return gr.update(choices=labels, value=labels[0])
-
-def pick_from_gallery(gallery_select, current_dropdown):
-    # When user clicks a gallery item, we set the dropdown to match
-    labels, _, _ = list_hairstyle_files()
-    if not labels:
-        raise gr.Error("No hairstyles available.")
-    if gallery_select is None:
-        return current_dropdown
-    idx = gallery_select[0]
-    idx = max(0, min(len(labels) - 1, idx))
-    return labels[idx]
-
-def apply_from_dropdown(image, hairstyle_label, scale, lr, ud, rot):
-    labels, files, _ = list_hairstyle_files()
-    if not labels:
-        raise gr.Error("No hairstyle assets found.")
-    if hairstyle_label not in labels:
-        raise gr.Error("Choose a hairstyle first.")
-    path = files[labels.index(hairstyle_label)]
-    return place_and_render(image, path, scale, lr, ud, rot)
-
-def save_image(img_np):
-    if img_np is None:
-        raise gr.Error("Nothing to save. Generate a preview first.")
-    out_path = "output_tryon.png"
-    Image.fromarray(img_np).save(out_path)
+def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity, meta):
+    if image is None:
+        return None, "Upload a photo or enable webcam."
+    if not hairstyle:
+        return np.array(image), "Pick a hairstyle first."
+
+    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+    kpts = detect_face_keypoints(img_bgr)
+    if kpts is None:
+        return image, "No face detected. Try a brighter, front-facing photo."
+
+    hair = load_hair_png(hairstyle)
+    hair_pts = hair_reference_points(hair, hairstyle, meta)
+
+    # Destination points (with user nudges)
+    dst = kpts.copy()
+    dst[:, 0] += dx
+    dst[:, 1] += dy
+
+    # Scale + rotate around hair anchor centroid
+    center = hair_pts.mean(axis=0)
+    theta = np.deg2rad(rot_deg)
+    s = max(0.5, scale_pct / 100.0)
+    R = np.array([[np.cos(theta), -np.sin(theta)],
+                  [np.sin(theta),  np.cos(theta)]], dtype=np.float32)
+    hair_pts_adj = (hair_pts - center) @ R.T * s + center
+
+    M, _ = cv2.estimateAffinePartial2D(hair_pts_adj, dst, method=cv2.LMEDS)
+    if M is None:
+        return image, "Could not compute alignment for this image/style."
+
+    out = warp_and_alpha_blend(img_bgr, hair, M, opacity=opacity)
+
+    # Restrict to head region for cleaner look
+    head = person_mask(img_bgr)
+    head3 = head[..., None]
+    out = (head3 * out + (1 - head3) * img_bgr).astype(np.uint8)
+
+    out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
+    return out_rgb, "OK"
+
+def save_png_to_tmp(img, filename="output_tryon.png"):
+    """Create a file in /tmp and return the path."""
+    if img is None:
+        raise gr.Error("No image to download. Click Apply first.")
+    out_path = os.path.join(tempfile.gettempdir(), filename)
+    if isinstance(img, np.ndarray):
+        Image.fromarray(img).save(out_path)
+    else:
+        img.save(out_path)
     return out_path
 
-# --------------------------
-# UI
-# --------------------------
-with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"# {APP_TITLE}")
-    gr.Markdown("Upload a photo or use webcam. Put **transparent PNGs** in the hair/ folder, then click **Refresh**.")
-
-    with gr.Tabs():
-        # -------- Upload tab --------
-        with gr.Tab("📷 Photo (Upload)"):
-            with gr.Row():
-                with gr.Column():
-                    img = gr.Image(label="Photo", sources=["upload"], type="numpy", height=420)
-
-                    hair_dd = gr.Dropdown(label="Hairstyle", choices=[], interactive=True)
-                    refresh = gr.Button("🔄 Refresh")
-                    status = gr.Markdown("Drop PNGs into hair/ and press Refresh.")
-                    gallery = gr.Gallery(
-                        label="Hairstyles (click to choose)",
-                        columns=4,
-                        height=220,
-                        allow_preview=False,
-                        interactive=True,
+# ----- WHITE background thumbnails (no checkerboard) -----
+def thumb_on_white(hair_bgra, max_h=220):
+    h, w = hair_bgra.shape[:2]
+    scale = min(1.0, max_h / h)
+    hair_bgra = cv2.resize(hair_bgra, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_LINEAR)
+    h, w = hair_bgra.shape[:2]
+    bg_rgb = np.full((h, w, 3), 255, dtype=np.uint8)  # white background
+    a = (hair_bgra[:, :, 3:4] / 255.0)
+    comp = (a * hair_bgra[:, :, :3] + (1 - a) * bg_rgb).astype(np.uint8)
+    return cv2.cvtColor(comp, cv2.COLOR_BGR2RGB)
+
+def build_gallery_items(files):
+    items = []
+    for fname in files:
+        try:
+            img = load_hair_png(fname)
+            items.append((thumb_on_white(img), fname))  # (image, caption)
+        except Exception:
+            continue
+    return items
+
+# ===================== UI =====================
+def build_ui():
+    META = load_meta()
+    HAIR_FILES = load_hairstyles()
+
+    with gr.Blocks(title="Salon Hairstyle Virtual Try-On", css="""
+    .gradio-container {max-width: 1200px; margin:auto;}  /* wider so more tiles fit */
+    @media (max-width: 768px){ .gradio-container {padding: 8px;} }
+    """) as demo:
+        gr.Markdown("Upload a photo or use webcam. Put transparent **PNGs** in the **`hair/`** folder, then click **Refresh**.")
+
+        files_state = gr.State(HAIR_FILES)   # keep filenames
+        meta_state  = gr.State(META)
+
+        with gr.Tabs():
+            # ---------------- Photo Tab ----------------
+            with gr.Tab("📷 Photo (Upload)"):
+                with gr.Row():
+                    in_img = gr.Image(label="Input photo (JPEG/PNG)", type="pil", height=360, sources=["upload"])
+                    out_img = gr.Image(label="Preview", height=360)
+                with gr.Row():
+                    hair_sel = gr.Dropdown(
+                        choices=HAIR_FILES,
+                        value=(HAIR_FILES[0] if HAIR_FILES else None),
+                        label="Selected hairstyle",
+                        interactive=True
                     )
+                    apply_btn = gr.Button("✨ Apply (Align & Overlay)")
+                    # One-click download + visible link fallback
+                    download_btn = gr.DownloadButton("⬇️ Download")
+                    download_file = gr.File(label="download (fallback)", visible=False)
+                status = gr.Markdown()
 
-                    with gr.Accordion("Fine-tune placement", open=False):
-                        scale = gr.Slider(0.6, 2.2, value=1.3, step=0.01, label="Scale (× temple distance)")
-                        lr = gr.Slider(-30, 30, value=0, step=0.5, label="Left ↔ Right shift (%)")
-                        ud = gr.Slider(-30, 30, value=0, step=0.5, label="Up ↕ Down shift (%)")
-                        rot = gr.Slider(-30, 30, value=0, step=0.5, label="Extra rotation (°)")
-
-                    run_dd = gr.Button("✨ Apply")
-
-                with gr.Column():
-                    out = gr.Image(label="Preview", height=480)
-                    save_btn = gr.Button("💾 Save result")
-                    file_out = gr.File(label="Download")
-
-            # Wiring
-            refresh.click(fn=refresh_assets, inputs=None, outputs=[hair_dd, gallery, status])
-            demo.load(fn=refresh_assets, inputs=None, outputs=[hair_dd, gallery, status])
-
-            # Clicking gallery picks the dropdown value; one Apply button only
-            gallery.select(fn=pick_from_gallery, inputs=[gallery, hair_dd], outputs=hair_dd)
-            run_dd.click(fn=apply_from_dropdown, inputs=[img, hair_dd, scale, lr, ud, rot], outputs=out)
-            save_btn.click(fn=save_image, inputs=out, outputs=file_out)
-
-        # -------- Webcam tab --------
-        with gr.Tab("🎥 Webcam (Live Beta)"):
-            gr.Markdown("Live mode processes frames continuously. Keep resolution small on CPU Spaces.")
-            cam = gr.Image(sources=["webcam"], streaming=True, type="numpy", label="Webcam")
-            hair_dd2 = gr.Dropdown(label="Hairstyle", choices=[], interactive=True)
-            scale2 = gr.Slider(0.6, 2.2, value=1.25, step=0.01, label="Scale")
-            lr2 = gr.Slider(-30, 30, value=0, step=0.5, label="Left ↔ Right shift (%)")
-            ud2 = gr.Slider(-30, 30, value=0, step=0.5, label="Up ↕ Down shift (%)")
-            rot2 = gr.Slider(-30, 30, value=0, step=0.5, label="Rotation (°)")
-            out_live = gr.Image(label="Live Preview", interactive=False, height=420)
-
-            def live_process(frame, label, s, lrp, udp, r):
-                labels, files, _ = list_hairstyle_files()
-                if frame is None or not labels or label not in labels:
-                    return frame
-                path = files[labels.index(label)]
-                return place_and_render(frame, path, s, lrp, udp, r)
-
-            cam.stream(fn=live_process, inputs=[cam, hair_dd2, scale2, lr2, ud2, rot2], outputs=out_live)
-            demo.load(fn=refresh_dropdown_only, inputs=None, outputs=hair_dd2)
-
+                with gr.Row():
+                    refresh = gr.Button("🔄 Refresh")
+                count_md = gr.Markdown(f"Found {len(HAIR_FILES)} hairstyles.")
+                gallery = gr.Gallery(
+                    label="Hairstyles (click to choose)",
+                    value=build_gallery_items(HAIR_FILES),
+                    columns=6, rows=2, height=340,     # up to 12 visible at once
+                    allow_preview=False, object_fit="contain", show_label=True
+                )
+
+                with gr.Accordion("Fine-tune placement", open=True):
+                    with gr.Row():
+                        scale = gr.Slider(50, 200, 100, 1, label="Scale (≈ temple distance %)")
+                        rot   = gr.Slider(-30, 30, 0, 1, label="Extra rotation (°)")
+                    with gr.Row():
+                        dx    = gr.Slider(-200, 200, 0, 1, label="Left ↔ Right shift (px)")
+                        dy    = gr.Slider(-200, 200, 0, 1, label="Up ↕ Down shift (px)")
+                    opacity = gr.Slider(0.2, 1.0, 1.0, 0.05, label="Hair opacity")
+
+                # --- Callbacks ---
+                def do_apply(im, hfile, s, r, dxv, dyv, op, meta):
+                    return apply_tryon(im, hfile, s, r, dxv, dyv, op, meta)
+
+                apply_btn.click(
+                    fn=do_apply,
+                    inputs=[in_img, hair_sel, scale, rot, dx, dy, opacity, meta_state],
+                    outputs=[out_img, status]
+                )
+
+                # Return path for one-click download AND show fallback file link
+                def prepare_download_dual(im):
+                    path = save_png_to_tmp(im, "output_tryon.png")
+                    return path, gr.File.update(value=path, visible=True)
+
+                download_btn.click(fn=prepare_download_dual, inputs=[out_img], outputs=[download_btn, download_file])
+
+                def do_refresh():
+                    files = load_hairstyles()
+                    items = build_gallery_items(files)
+                    msg = f"Found {len(files)} hairstyles."
+                    return items, gr.update(choices=files, value=(files[0] if files else None)), files, msg
+
+                refresh.click(fn=do_refresh, inputs=[], outputs=[gallery, hair_sel, files_state, count_md])
+
+                # Clicking a tile sets the dropdown to that filename
+                def on_gallery_select(evt, files):
+                    idx = getattr(evt, "index", None)
+                    if idx is None or not files:
+                        return gr.update()
+                    if idx >= len(files):
+                        idx = len(files) - 1
+                    return gr.update(value=files[idx])
+
+                gallery.select(on_gallery_select, inputs=[files_state], outputs=[hair_sel])
+
+            # ---------------- Webcam Tab ----------------
+            with gr.Tab("📹 Webcam (Live Beta)"):
+                cam = gr.Image(sources=["webcam"], streaming=True, type="pil", label="Enable camera")
+                hair2 = gr.Dropdown(choices=HAIR_FILES, value=(HAIR_FILES[0] if HAIR_FILES else None), label="Selected hairstyle")
+                with gr.Row():
+                    scale2 = gr.Slider(50, 200, 100, 1, label="Scale %")
+                    rot2   = gr.Slider(-25, 25, 0, 1, label="Rotate (°)")
+                with gr.Row():
+                    dx2    = gr.Slider(-150, 150, 0, 1, label="Left ↔ Right (px)")
+                    dy2    = gr.Slider(-150, 150, 0, 1, label="Up ↕ Down (px)")
+                opacity2 = gr.Slider(0.2, 1.0, 0.95, 0.05, label="Hair opacity")
+                out2 = gr.Image(label="Live result", height=360)
+                state_live = gr.State(None)
+                snap = gr.Button("📸 Snapshot")
+                save_live_btn = gr.DownloadButton("⬇️ Download Snapshot")
+                save_live_file = gr.File(label="snapshot (fallback)", visible=False)
+
+                def live(im, h, s, r, dxv, dyv, op, meta):
+                    res, _ = apply_tryon(im, h, s, r, dxv, dyv, op, meta)
+                    return res, res
+
+                cam.stream(
+                    fn=live,
+                    inputs=[cam, hair2, scale2, rot2, dx2, dy2, opacity2, meta_state],
+                    outputs=[out2, state_live]
+                )
+
+                snap.click(lambda x: x, inputs=[state_live], outputs=[out2])
+
+                def prepare_webcam_download_dual(im):
+                    path = save_png_to_tmp(im, "tryon_webcam.png")
+                    return path, gr.File.update(value=path, visible=True)
+
+                save_live_btn.click(fn=prepare_webcam_download_dual, inputs=[state_live], outputs=[save_live_btn, save_live_file])
+
+        return demo
+
+# Export for Spaces autostart
+app = build_ui()
+demo = app
+
+# Local dev
 if __name__ == "__main__":
-    demo.launch()
+    app.launch()