add prototype editing feature

app/model_io.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, Tuple
 
 import torch
 
+# ZeroGPU on HF Spaces: CUDA must not be initialized in the main process
+_ON_SPACES = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE"))
+
 
 @dataclass(frozen=True)
 class LoadedModel:
@@ -21,6 +25,9 @@ class LoadedModel:
 def _pick_device(device: str) -> torch.device:
     if device.strip().lower() == "cpu":
         return torch.device("cpu")
+    if _ON_SPACES:
+        # ZeroGPU: can't init CUDA in main process
+        return torch.device("cpu")
     if torch.cuda.is_available():
         return torch.device("cuda")
     return torch.device("cpu")
@@ -41,9 +48,13 @@ def load_style_model(
     if not ckpt_path.exists():
         raise FileNotFoundError(str(ckpt_path))
 
-    dev = _pick_device("cpu" if device == "auto" else device)
-    if device == "auto":
+    # On Spaces, always use CPU (ZeroGPU forbids CUDA in main process)
+    if _ON_SPACES:
+        dev = torch.device("cpu")
+    elif device == "auto":
         dev = _pick_device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        dev = _pick_device(device)
 
     ck = torch.load(str(ckpt_path), map_location="cpu")
     meta = ck.get("meta", {}) if isinstance(ck, dict) else {}
@@ -103,8 +114,14 @@ def embed_triview(
     # Use lazy dtype detection to avoid CUDA init at import time (ZeroGPU compatibility)
     import train_style_ddp as _ts
     _dtype = _ts._get_amp_dtype() if hasattr(_ts, "_get_amp_dtype") else torch.float16
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=_dtype, enabled=(lm.device.type == "cuda")):
-        z, _, _ = lm.model(views, masks)
+    # On CPU or Spaces, skip autocast entirely to avoid touching CUDA
+    use_amp = (lm.device.type == "cuda") and not _ON_SPACES
+    if use_amp:
+        with torch.no_grad(), torch.amp.autocast("cuda", dtype=_dtype, enabled=True):
+            z, _, _ = lm.model(views, masks)
+    else:
+        with torch.no_grad():
+            z, _, _ = lm.model(views, masks)
     z = torch.nn.functional.normalize(z.float(), dim=1)
     return z.squeeze(0).detach().cpu()
 

app/proto_db.py

@@ -43,6 +43,35 @@ class PrototypeDB:
         self.labels = torch.cat([self.labels, torch.tensor([lid], dtype=torch.long)], dim=0)
         return lid
 
+    def remove_prototype(self, proto_idx: int) -> bool:
+        """Remove a prototype by its index. Returns True if removed."""
+        if proto_idx < 0 or proto_idx >= self.centers.shape[0]:
+            return False
+        mask = torch.ones(self.centers.shape[0], dtype=torch.bool)
+        mask[proto_idx] = False
+        self.centers = self.centers[mask]
+        self.labels = self.labels[mask]
+        return True
+
+    def list_prototypes(self) -> List[Tuple[int, str, int]]:
+        """Returns list of (proto_idx, label_name, label_id) for all prototypes."""
+        result = []
+        for i in range(self.centers.shape[0]):
+            lid = int(self.labels[i].item())
+            name = self.id_to_name(lid)
+            result.append((i, name, lid))
+        return result
+
+    def get_label_summary(self) -> Dict[str, int]:
+        """Returns {label_name: count} for prototype counts per label."""
+        from collections import Counter
+        counts: Dict[str, int] = {}
+        for i in range(self.labels.shape[0]):
+            lid = int(self.labels[i].item())
+            name = self.id_to_name(lid)
+            counts[name] = counts.get(name, 0) + 1
+        return counts
+
     def save(self, path: Optional[str | Path] = None) -> Path:
         out = Path(path) if path is not None else self.source_path
         if out is None:

webui_gradio.py

@@ -18,6 +18,20 @@ except Exception:  # noqa: BLE001
 # Detect if running on HF Spaces (ZeroGPU requires special handling)
 _ON_SPACES = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE"))
 
+# ─────────────────────────────────────────────────────────────────────────────
+# Authentication for prototype management
+# Set ADMIN_PASSWORD env var, or it defaults to "admin" (change in production!)
+# ─────────────────────────────────────────────────────────────────────────────
+ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "admin")
+APPROVED_USERS: dict[str, str] = {
+    "admin": ADMIN_PASSWORD,
+}
+# Add more users via APPROVED_USERS_CSV env var: "user1:pass1,user2:pass2"
+for pair in os.getenv("APPROVED_USERS_CSV", "").split(","):
+    if ":" in pair:
+        u, p = pair.split(":", 1)
+        APPROVED_USERS[u.strip()] = p.strip()
+
 def _patch_fastapi_starlette_middleware_unpack() -> None:
     """
     Work around FastAPI/Starlette version mismatches where Starlette's Middleware
@@ -383,6 +397,37 @@ def save_db_as(path_text: str) -> str:
     return f"✅ Saved prototype DB to `{out_path}`"
 
 
+def list_prototypes() -> Tuple[str, List[List]]:
+    """Returns (status, table_rows) with all prototypes."""
+    if APP_STATE.db is None:
+        return "❌ Click **Load** first.", []
+    db = APP_STATE.db
+    protos = db.list_prototypes()
+    # Group by label for summary
+    summary = db.get_label_summary()
+    summary_text = ", ".join(f"{k}: {v}" for k, v in sorted(summary.items()))
+    rows = [[i, name, lid] for (i, name, lid) in protos]
+    return f"✅ {len(protos)} prototypes. Summary: {summary_text}", rows
+
+
+def remove_prototype(proto_idx: int, save_back: bool) -> Tuple[str, List[List]]:
+    """Remove a prototype by index. Returns (status, updated_table)."""
+    if APP_STATE.db is None:
+        return "❌ Click **Load** first.", []
+    db = APP_STATE.db
+    idx = int(proto_idx)
+    if not db.remove_prototype(idx):
+        return f"❌ Invalid index: {idx}", [[i, name, lid] for (i, name, lid) in db.list_prototypes()]
+
+    msg = f"✅ Removed prototype at index {idx}. DB now N={db.centers.shape[0]}."
+    if save_back and APP_STATE.proto_path:
+        db.save(APP_STATE.proto_path)
+        msg += f" Saved to `{APP_STATE.proto_path}`."
+
+    rows = [[i, name, lid] for (i, name, lid) in db.list_prototypes()]
+    return msg, rows
+
+
 def build_ui() -> gr.Blocks:
     ckpts = _list_ckpt_files(CKPT_DIR)
     protos = _list_proto_files(CKPT_DIR)
@@ -412,16 +457,63 @@ def build_ui() -> gr.Blocks:
             table = gr.Dataframe(headers=["label", "cosine_sim"], datatype=["str", "number"], interactive=False)
             run_btn.click(classify, inputs=[whole, topk], outputs=[out_status, table, face_prev, eyes_prev])
 
-        with gr.Tab("Add prototype"):
+        with gr.Tab("Manage Prototypes"):
             gr.Markdown(
-                "Add a new prototype to the loaded prototype DB by averaging embeddings of uploaded whole images.\n"
+                "### Prototype Management\n"
+                "**Authentication required** to add or remove prototypes. "
+                "Enter your credentials below."
+            )
+
+            with gr.Row():
+                auth_user = gr.Textbox(label="Username", placeholder="admin")
+                auth_pass = gr.Textbox(label="Password", type="password", placeholder="password")
+
+            def _check_auth(user: str, pwd: str) -> bool:
+                u = (user or "").strip()
+                p = (pwd or "").strip()
+                return APPROVED_USERS.get(u) == p and p != ""
+
+            # ─────────────── List Prototypes ───────────────
+            gr.Markdown("---\n#### List Prototypes")
+            list_btn = gr.Button("Refresh List", variant="secondary")
+            list_status = gr.Markdown("")
+            proto_table = gr.Dataframe(
+                headers=["index", "label", "label_id"],
+                datatype=["number", "str", "number"],
+                interactive=False,
+            )
+            list_btn.click(list_prototypes, inputs=[], outputs=[list_status, proto_table])
+
+            # ─────────────── Remove Prototype ───────────────
+            gr.Markdown("---\n#### Remove Prototype")
+            with gr.Row():
+                remove_idx = gr.Number(label="Prototype index to remove", precision=0)
+                remove_save = gr.Checkbox(value=True, label="Save after removal")
+            remove_btn = gr.Button("Remove Prototype", variant="stop")
+            remove_status = gr.Markdown("")
+
+            def _remove_with_auth(user, pwd, idx, save_back):
+                if not _check_auth(user, pwd):
+                    return "❌ Authentication failed. Check username/password.", []
+                return remove_prototype(idx, save_back)
+
+            remove_btn.click(
+                _remove_with_auth,
+                inputs=[auth_user, auth_pass, remove_idx, remove_save],
+                outputs=[remove_status, proto_table],
+            )
+
+            # ─────────────── Add Prototype ───────────────
+            gr.Markdown("---\n#### Add New Prototype")
+            gr.Markdown(
+                "Add a new prototype by averaging embeddings of uploaded whole images. "
                 "Multiple prototypes per label are allowed."
             )
             label = gr.Textbox(label="Label name (artist)", placeholder="e.g. new_artist")
             imgs = gr.Gallery(label="Whole images (1+)", columns=4, rows=2, height=240, allow_preview=True)
             uploader = gr.Files(label="Upload image files (whole)", file_types=["image"], file_count="multiple")
-            save_back = gr.Checkbox(value=True, label="Save back to selected prototype DB file after adding")
-            add_btn = gr.Button("Add prototype", variant="primary")
+            add_save_back = gr.Checkbox(value=True, label="Save back to prototype DB after adding")
+            add_btn = gr.Button("Add Prototype", variant="primary")
             add_status = gr.Markdown("")
 
             def _files_to_gallery(files):
@@ -436,13 +528,33 @@ def build_ui() -> gr.Blocks:
                         continue
                 return out
 
+            def _add_with_auth(user, pwd, lbl, images, save_back):
+                if not _check_auth(user, pwd):
+                    return "❌ Authentication failed. Check username/password."
+                return add_prototype(lbl, images, save_back)
+
             uploader.change(_files_to_gallery, inputs=[uploader], outputs=[imgs])
-            add_btn.click(add_prototype, inputs=[label, imgs, save_back], outputs=[add_status])
+            add_btn.click(
+                _add_with_auth,
+                inputs=[auth_user, auth_pass, label, imgs, add_save_back],
+                outputs=[add_status],
+            )
 
-            gr.Markdown("Save DB as (optional):")
-            save_path = gr.Textbox(label="Output path (relative paths go under ./checkpoints_style/)", placeholder="prototypes_custom.pt")
-            save_btn = gr.Button("Save As")
-            save_btn.click(save_db_as, inputs=[save_path], outputs=[add_status])
+            # ─────────────── Save DB As ───────────────
+            gr.Markdown("---\n#### Save DB As (optional)")
+            with gr.Row():
+                save_path = gr.Textbox(
+                    label="Output path (relative paths go under ./checkpoints_style/)",
+                    placeholder="prototypes_custom.pt",
+                )
+                save_btn = gr.Button("Save As")
+
+            def _save_with_auth(user, pwd, path):
+                if not _check_auth(user, pwd):
+                    return "❌ Authentication failed. Check username/password."
+                return save_db_as(path)
+
+            save_btn.click(_save_with_auth, inputs=[auth_user, auth_pass, save_path], outputs=[add_status])
 
     return demo