Deploy MobileCLIP classifier with Gradio interface

README.md

@@ -1,12 +1,53 @@
 ---
-title: Clipspace
-emoji: 📉
-colorFrom: indigo
-colorTo: red
+title: MobileCLIP Image Classifier
+emoji: 📸
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 5.44.1
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 📸 MobileCLIP-B Image Classifier
+
+Interactive web interface for Apple's MobileCLIP-B zero-shot image classification model.
+
+## Features
+
+- 🖼️ **Image Classification**: Upload any image for instant classification
+- 🏷️ **Dynamic Labels**: Add and manage classification labels on-the-fly
+- 📊 **Visual Results**: See confidence scores with interactive charts
+- ⚡ **Fast Inference**: Optimized for < 30ms latency on GPU
+- 🔒 **Admin Panel**: Secure label management interface
+
+## Environment Variables
+
+Configure these in your Space Settings (Settings → Variables and secrets):
+
+| Variable | Description | Required |
+|----------|-------------|----------|
+| `ADMIN_TOKEN` | Secret token for admin operations | Yes (for admin features) |
+| `HF_LABEL_REPO` | Hub dataset repo for label storage (e.g., `username/mobileclip-labels`) | No |
+| `HF_WRITE_TOKEN` | Hugging Face token with write permissions | No |
+| `HF_READ_TOKEN` | Hugging Face token with read permissions | No |
+
+## How It Works
+
+1. **Model**: Uses MobileCLIP-B with re-parameterized MobileOne blocks for efficient inference
+2. **Labels**: Loads from `items.json` or dynamically from Hub repository
+3. **Processing**: Pre-computes text embeddings for fast classification
+4. **Interface**: Gradio provides the web UI with image upload and admin controls
+
+## Admin Features
+
+With proper authentication, admins can:
+- Add new classification labels without redeploying
+- Reload specific label versions from the Hub
+- View current statistics and label information
+
+## License
+
+- Model weights: Apple Sample Code License (ASCL)
+- Interface code: MIT

app.py

@@ -0,0 +1,280 @@
+import gradio as gr
+import base64
+import json
+import os
+from PIL import Image
+import io
+from handler import EndpointHandler
+
+handler = EndpointHandler()
+
+def classify_image(image, top_k=10):
+    """
+    Main classification function for public interface.
+    """
+    if image is None:
+        return None, "Please upload an image"
+    
+    try:
+        # Convert PIL image to base64
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        img_b64 = base64.b64encode(buffered.getvalue()).decode()
+        
+        # Call handler
+        result = handler({
+            "inputs": {
+                "image": img_b64,
+                "top_k": int(top_k)
+            }
+        })
+        
+        # Format results for display
+        if isinstance(result, list):
+            # Create formatted output
+            output_text = "**Top {} Classifications:**\n\n".format(len(result))
+            
+            # Create a dictionary for the bar chart
+            chart_data = {}
+            
+            for i, item in enumerate(result, 1):
+                score_pct = item['score'] * 100
+                output_text += f"{i}. **{item['label']}** (ID: {item['id']}): {score_pct:.2f}%\n"
+                chart_data[item['label']] = item['score']
+            
+            return chart_data, output_text
+        else:
+            return None, f"Error: {result.get('error', 'Unknown error')}"
+            
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def upsert_labels_admin(admin_token, new_items_json):
+    """
+    Admin function to add new labels.
+    """
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        # Parse the JSON input
+        items = json.loads(new_items_json) if new_items_json else []
+        
+        result = handler({
+            "inputs": {
+                "op": "upsert_labels",
+                "token": admin_token,
+                "items": items
+            }
+        })
+        
+        if result.get("status") == "ok":
+            return f"✅ Success! Added {result.get('added', 0)} new labels. Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        else:
+            return f"❌ Error: {result.get('detail', result.get('error', 'Unknown error'))}"
+            
+    except json.JSONDecodeError:
+        return "❌ Error: Invalid JSON format"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def reload_labels_admin(admin_token, version):
+    """
+    Admin function to reload a specific label version.
+    """
+    if not admin_token:
+        return "Error: Admin token required"
+    
+    try:
+        result = handler({
+            "inputs": {
+                "op": "reload_labels",
+                "token": admin_token,
+                "version": int(version) if version else 1
+            }
+        })
+        
+        if result.get("status") == "ok":
+            return f"✅ Labels reloaded successfully! Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("status") == "nochange":
+            return f"ℹ️ No change needed. Current version: {result.get('labels_version', 'unknown')}"
+        elif result.get("error") == "unauthorized":
+            return "❌ Error: Invalid admin token"
+        elif result.get("error") == "invalid_version":
+            return "❌ Error: Invalid version number"
+        else:
+            return f"❌ Error: {result.get('error', 'Unknown error')}"
+            
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def get_current_stats():
+    """
+    Get current label statistics.
+    """
+    try:
+        num_labels = len(handler.class_ids) if hasattr(handler, 'class_ids') else 0
+        version = getattr(handler, 'labels_version', 1)
+        device = handler.device if hasattr(handler, 'device') else "unknown"
+        
+        stats = f"""
+        **Current Statistics:**
+        - Number of labels: {num_labels}
+        - Labels version: {version}
+        - Device: {device}
+        - Model: MobileCLIP-B
+        """
+        
+        if hasattr(handler, 'class_names') and len(handler.class_names) > 0:
+            stats += f"\n- Sample labels: {', '.join(handler.class_names[:5])}"
+            if len(handler.class_names) > 5:
+                stats += "..."
+        
+        return stats
+    except Exception as e:
+        return f"Error getting stats: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="MobileCLIP Image Classifier") as demo:
+    gr.Markdown("""
+    # 🖼️ MobileCLIP-B Zero-Shot Image Classifier
+    
+    Upload an image to classify it using MobileCLIP-B model with dynamic label management.
+    """)
+    
+    with gr.Tab("🔍 Image Classification"):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    type="pil",
+                    label="Upload Image"
+                )
+                top_k_slider = gr.Slider(
+                    minimum=1,
+                    maximum=50,
+                    value=10,
+                    step=1,
+                    label="Number of top results to show"
+                )
+                classify_btn = gr.Button("🚀 Classify Image", variant="primary")
+            
+            with gr.Column():
+                output_chart = gr.BarPlot(
+                    label="Classification Confidence",
+                    x_label="Label",
+                    y_label="Confidence",
+                    vertical=False,
+                    height=400
+                )
+                output_text = gr.Markdown(label="Classification Results")
+        
+        gr.Examples(
+            examples=[
+                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/image_classifier/examples/cheetah.jpg"],
+                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/image_classifier/examples/elephant.jpg"],
+                ["https://raw.githubusercontent.com/gradio-app/gradio/main/demo/image_classifier/examples/giraffe.jpg"]
+            ],
+            inputs=input_image,
+            label="Example Images"
+        )
+        
+        classify_btn.click(
+            classify_image,
+            inputs=[input_image, top_k_slider],
+            outputs=[output_chart, output_text]
+        )
+    
+    with gr.Tab("🔧 Admin Panel"):
+        gr.Markdown("""
+        ### Admin Functions
+        **Note:** Requires admin token (set via environment variable `ADMIN_TOKEN`)
+        """)
+        
+        with gr.Row():
+            admin_token_input = gr.Textbox(
+                label="Admin Token",
+                type="password",
+                placeholder="Enter admin token"
+            )
+        
+        with gr.Accordion("📊 Current Statistics", open=True):
+            stats_display = gr.Markdown(value=get_current_stats())
+            refresh_stats_btn = gr.Button("🔄 Refresh Stats")
+            refresh_stats_btn.click(
+                get_current_stats,
+                outputs=stats_display
+            )
+        
+        with gr.Accordion("➕ Add New Labels", open=False):
+            gr.Markdown("""
+            Add new labels by providing JSON array:
+            ```json
+            [
+                {"id": 100, "name": "new_object", "prompt": "a photo of a new_object"},
+                {"id": 101, "name": "another_object", "prompt": "a photo of another_object"}
+            ]
+            ```
+            """)
+            new_items_input = gr.Code(
+                label="New Items JSON",
+                language="json",
+                lines=5,
+                value='[\n  {"id": 100, "name": "example", "prompt": "a photo of example"}\n]'
+            )
+            upsert_btn = gr.Button("➕ Add Labels", variant="primary")
+            upsert_output = gr.Markdown()
+            
+            upsert_btn.click(
+                upsert_labels_admin,
+                inputs=[admin_token_input, new_items_input],
+                outputs=upsert_output
+            )
+        
+        with gr.Accordion("🔄 Reload Label Version", open=False):
+            gr.Markdown("Reload labels from a specific version stored in the Hub")
+            version_input = gr.Number(
+                label="Version Number",
+                value=1,
+                precision=0
+            )
+            reload_btn = gr.Button("🔄 Reload Version", variant="primary")
+            reload_output = gr.Markdown()
+            
+            reload_btn.click(
+                reload_labels_admin,
+                inputs=[admin_token_input, version_input],
+                outputs=reload_output
+            )
+    
+    with gr.Tab("ℹ️ About"):
+        gr.Markdown("""
+        ## About MobileCLIP-B Classifier
+        
+        This Space provides a web interface for Apple's MobileCLIP-B model, optimized for fast zero-shot image classification.
+        
+        ### Features:
+        - 🚀 **Fast inference**: < 30ms on GPU
+        - 🏷️ **Dynamic labels**: Add/update labels without redeployment
+        - 🔄 **Version control**: Track and reload label versions
+        - 📊 **Visual results**: Bar charts and confidence scores
+        
+        ### Environment Variables (set in Space Settings):
+        - `ADMIN_TOKEN`: Secret token for admin operations
+        - `HF_LABEL_REPO`: Hub repository for label storage (e.g., "username/labels")
+        - `HF_WRITE_TOKEN`: Token with write permissions to label repo
+        - `HF_READ_TOKEN`: Token with read permissions (optional, defaults to write token)
+        
+        ### Model Details:
+        - **Architecture**: MobileCLIP-B with MobileOne blocks
+        - **Text Encoder**: Transformer-based, 77 token context
+        - **Image Size**: 224x224
+        - **Embedding Dim**: 512
+        
+        ### License:
+        Model weights are licensed under Apple Sample Code License (ASCL).
+        """)
+
+if __name__ == "__main__":
+    demo.launch()

handler.py

@@ -0,0 +1,385 @@
+import contextlib, io, base64, torch, json, os, threading
+from PIL import Image
+import open_clip
+from huggingface_hub import hf_hub_download, create_commit, CommitOperationAdd
+from safetensors.torch import save_file, load_file
+from reparam import reparameterize_model
+
+ADMIN_TOKEN    = os.getenv("ADMIN_TOKEN", "")
+HF_LABEL_REPO  = os.getenv("HF_LABEL_REPO", "")  # e.g. "org/mobileclip-labels"
+HF_WRITE_TOKEN = os.getenv("HF_WRITE_TOKEN", "")
+HF_READ_TOKEN  = os.getenv("HF_READ_TOKEN", HF_WRITE_TOKEN)
+
+
+def _fingerprint(device: str, dtype: torch.dtype) -> dict:
+    return {
+        "model_id": "MobileCLIP-B",
+        "pretrained": "datacompdr",
+        "open_clip": getattr(open_clip, "__version__", "unknown"),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda if torch.cuda.is_available() else None,
+        "dtype_runtime": str(dtype),
+        "text_norm": "L2",
+        "logit_scale": 100.0,
+    }
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+
+        # 1) Load model + transforms
+        model, _, self.preprocess = open_clip.create_model_and_transforms(
+            "MobileCLIP-B", pretrained="datacompdr"
+        )
+        model.eval()
+        model = reparameterize_model(model)
+        model.to(self.device)
+        if self.device == "cuda":
+            model = model.to(torch.float16)
+        self.model = model
+        self.tokenizer = open_clip.get_tokenizer("MobileCLIP-B")
+        self.fingerprint = _fingerprint(self.device, self.dtype)
+        self._lock = threading.Lock()
+
+        # 2) Try to load snapshot from Hub; else seed from items.json
+        loaded = False
+        if HF_LABEL_REPO:
+            with contextlib.suppress(Exception):
+                loaded = self._load_snapshot_from_hub_latest()
+        if not loaded:
+            with open(f"{path}/items.json", "r", encoding="utf-8") as f:
+                items = json.load(f)
+            prompts = [it["prompt"] for it in items]
+            self.class_ids = [int(it["id"]) for it in items]
+            self.class_names = [it["name"] for it in items]
+            with torch.no_grad():
+                toks = self.tokenizer(prompts).to(self.device)
+                feats = self.model.encode_text(toks)
+                feats = feats / feats.norm(dim=-1, keepdim=True)
+            self.text_features_cpu = feats.detach().cpu().to(torch.float32).contiguous()
+            self._to_device()
+            self.labels_version = 1
+
+    def __call__(self, data):
+        payload = data.get("inputs", data)
+
+        # Admin op: upsert_labels
+        op = payload.get("op")
+        if op == "upsert_labels":
+            if payload.get("token") != ADMIN_TOKEN:
+                return {"error": "unauthorized"}
+            items = payload.get("items", []) or []
+            added = self._upsert_items(items)
+            if added > 0:
+                new_ver = int(getattr(self, "labels_version", 1)) + 1
+                try:
+                    self._persist_snapshot_to_hub(new_ver)
+                    self.labels_version = new_ver
+                except Exception as e:
+                    return {"status": "error", "added": added, "detail": str(e)}
+            return {"status": "ok", "added": added, "labels_version": getattr(self, "labels_version", 1)}
+
+        # Admin op: reload_labels
+        if op == "reload_labels":
+            if payload.get("token") != ADMIN_TOKEN:
+                return {"error": "unauthorized"}
+            try:
+                ver = int(payload.get("version"))
+            except Exception:
+                return {"error": "invalid_version"}
+            ok = self._load_snapshot_from_hub_version(ver)
+            return {"status": "ok" if ok else "nochange", "labels_version": getattr(self, "labels_version", 0)}
+
+        # Freshness guard (optional)
+        min_ver = payload.get("min_labels_version")
+        if isinstance(min_ver, int) and min_ver > getattr(self, "labels_version", 0):
+            with contextlib.suppress(Exception):
+                self._load_snapshot_from_hub_version(min_ver)
+
+        # Classification path (unchanged contract)
+        img_b64 = payload["image"]
+        image = Image.open(io.BytesIO(base64.b64decode(img_b64))).convert("RGB")
+        img_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
+        if self.device == "cuda":
+            img_tensor = img_tensor.to(torch.float16)
+        with torch.no_grad():
+            img_feat = self.model.encode_image(img_tensor)
+            img_feat /= img_feat.norm(dim=-1, keepdim=True)
+            probs = (100.0 * img_feat @ self.text_features.T).softmax(dim=-1)[0]
+        results = zip(self.class_ids, self.class_names, probs.detach().cpu().tolist())
+        top_k = int(payload.get("top_k", len(self.class_ids)))
+        return sorted(
+            [{"id": i, "label": name, "score": float(p)} for i, name, p in results],
+            key=lambda x: x["score"],
+            reverse=True,
+        )[:top_k]
+
+    # ------------- helpers -------------
+    def _encode_text(self, prompts):
+        with torch.no_grad():
+            toks = self.tokenizer(prompts).to(self.device)
+            feats = self.model.encode_text(toks)
+            feats = feats / feats.norm(dim=-1, keepdim=True)
+            return feats
+
+    def _to_device(self):
+        self.text_features = self.text_features_cpu.to(
+            self.device, dtype=(torch.float16 if self.device == "cuda" else torch.float32)
+        )
+
+    def _upsert_items(self, new_items):
+        if not new_items:
+            return 0
+        with self._lock:
+            known = set(getattr(self, "class_ids", []))
+            batch = [it for it in new_items if int(it.get("id")) not in known]
+            if not batch:
+                return 0
+            prompts = [it["prompt"] for it in batch]
+            feats = self._encode_text(prompts).detach().cpu().to(torch.float32)
+            if not hasattr(self, "text_features_cpu"):
+                self.text_features_cpu = feats.contiguous()
+                self.class_ids = [int(it["id"]) for it in batch]
+                self.class_names = [it["name"] for it in batch]
+            else:
+                self.text_features_cpu = torch.cat([self.text_features_cpu, feats], dim=0).contiguous()
+                self.class_ids.extend([int(it["id"]) for it in batch])
+                self.class_names.extend([it["name"] for it in batch])
+            self._to_device()
+            return len(batch)
+
+    def _persist_snapshot_to_hub(self, version: int):
+        if not HF_LABEL_REPO:
+            raise RuntimeError("HF_LABEL_REPO not set")
+        if not HF_WRITE_TOKEN:
+            raise RuntimeError("HF_WRITE_TOKEN not set for publishing")
+
+        emb_path = "/tmp/embeddings.safetensors"
+        meta_path = "/tmp/meta.json"
+        latest_bytes = io.BytesIO(json.dumps({"version": int(version)}).encode("utf-8"))
+
+        save_file({"embeddings": self.text_features_cpu.to(torch.float32)}, emb_path)
+        meta = {
+            "items": [{"id": int(i), "name": n} for i, n in zip(self.class_ids, self.class_names)],
+            "fingerprint": self.fingerprint,
+            "dims": int(self.text_features_cpu.shape[1]),
+            "count": int(self.text_features_cpu.shape[0]),
+            "version": int(version),
+        }
+        with open(meta_path, "w", encoding="utf-8") as f:
+            json.dump(meta, f)
+
+        ops = [
+            CommitOperationAdd(
+                path_in_repo=f"snapshots/v{version}/embeddings.safetensors",
+                path_or_fileobj=emb_path,
+                lfs=True,
+            ),
+            CommitOperationAdd(
+                path_in_repo=f"snapshots/v{version}/meta.json",
+                path_or_fileobj=meta_path,
+            ),
+            CommitOperationAdd(
+                path_in_repo="snapshots/latest.json",
+                path_or_fileobj=latest_bytes,
+            ),
+        ]
+        create_commit(
+            repo_id=HF_LABEL_REPO,
+            repo_type="dataset",
+            operations=ops,
+            token=HF_WRITE_TOKEN,
+            commit_message=f"labels v{version}",
+        )
+
+    def _load_snapshot_from_hub_version(self, version: int) -> bool:
+        if not HF_LABEL_REPO:
+            return False
+        with self._lock:
+            emb_p = hf_hub_download(
+                HF_LABEL_REPO,
+                f"snapshots/v{version}/embeddings.safetensors",
+                repo_type="dataset",
+                token=HF_READ_TOKEN,
+                force_download=True,
+            )
+            meta_p = hf_hub_download(
+                HF_LABEL_REPO,
+                f"snapshots/v{version}/meta.json",
+                repo_type="dataset",
+                token=HF_READ_TOKEN,
+                force_download=True,
+            )
+            meta = json.load(open(meta_p, "r", encoding="utf-8"))
+            if meta.get("fingerprint") != self.fingerprint:
+                raise RuntimeError("Embedding/model fingerprint mismatch")
+            feats = load_file(emb_p)["embeddings"]  # float32 CPU
+            self.text_features_cpu = feats.contiguous()
+            self.class_ids = [int(x["id"]) for x in meta.get("items", [])]
+            self.class_names = [x["name"] for x in meta.get("items", [])]
+            self.labels_version = int(meta.get("version", version))
+            self._to_device()
+            return True
+
+    def _load_snapshot_from_hub_latest(self) -> bool:
+        if not HF_LABEL_REPO:
+            return False
+        try:
+            latest_p = hf_hub_download(
+                HF_LABEL_REPO,
+                "snapshots/latest.json",
+                repo_type="dataset",
+                token=HF_READ_TOKEN,
+            )
+        except Exception:
+            return False
+        latest = json.load(open(latest_p, "r", encoding="utf-8"))
+        ver = int(latest.get("version", 0))
+        if ver <= 0:
+            return False
+        return self._load_snapshot_from_hub_version(ver)
+
+# """
+# MobileCLIP‑B Zero‑Shot Image Classifier  (Hugging Face Inference Endpoint)
+# ===========================================================================
+
+# * One container instance is created per replica; the `EndpointHandler`
+#   object below is instantiated exactly **once** at start‑up.
+
+# * At request time (`__call__`) we receive a base‑64‑encoded image, run a
+#   **single forward pass**, and return class probabilities.
+
+# Design choices
+# --------------
+
+# 1. **Model & transform come from OpenCLIP**  
+#    This guarantees we apply **identical preprocessing** to what the model
+#    was trained with (224 × 224 crop + mean/std normalisation).
+
+# 2. **Re‑parameterisation for inference**  
+#    MobileCLIP uses MobileOne blocks that have extra convolution branches
+#    for training; `reparameterize_model` fuses them so inference is fast
+#    and deterministic.
+
+# 3. **Text embeddings are cached**  
+#    The class “prompts” (e.g. `"a photo of a cat"`) are encoded **once at
+#    start‑up**.  Each request therefore encodes *only* the image and
+#    performs a single matrix multiplication.
+
+# 4. **Mixed precision on GPU**  
+#    If the container has CUDA, we cast the model **and** inputs to
+#    `float16`.  That halves memory and roughly doubles throughput on most
+#    modern GPUs.  On CPU we stay in `float32` for numerical stability.
+# """
+
+# import contextlib, io, base64, json
+# from pathlib import Path
+# from typing import Any, Dict, List
+
+# import torch
+# from PIL import Image
+# import open_clip
+
+# from reparam import reparameterize_model   # local copy (~60 LoC) of Apple’s helper
+
+
+# class EndpointHandler:
+#     """
+#     Hugging Face entry‑point.  The toolkit will instantiate this class
+#     once and call it for every HTTP request.
+
+#     Parameters
+#     ----------
+#     path : str, optional
+#         Root directory of the repository.  HF mounts the code under
+#         `/repository`; we use this path to locate `items.json`.
+#     """
+
+#     # ------------------------------------------------------------------ #
+#     #                 INITIALISATION  (runs **once**)                     #
+#     # ------------------------------------------------------------------ #
+#     def __init__(self, path: str = "") -> None:
+#         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+#         # 1️⃣  Load MobileCLIP‑B weights & transforms -------------------
+#         #    `pretrained="datacompdr"` makes OpenCLIP download the
+#         #    official checkpoint from the Hub (cached in the image layer).
+#         model, _, self.preprocess = open_clip.create_model_and_transforms(
+#             "MobileCLIP-B", pretrained="datacompdr"
+#         )
+#         model.eval()                       # disable dropout / BN updates
+#         model = reparameterize_model(model)  # fuse MobileOne branches
+#         model.to(self.device)
+#         if self.device == "cuda":
+#             model = model.to(torch.float16)  # FP16 for throughput
+#         self.model = model                  # hold a reference
+
+#         # 2️⃣  Build the tokenizer once --------------------------------
+#         tokenizer = open_clip.get_tokenizer("MobileCLIP-B")
+
+#         # 3️⃣  Load class metadata -------------------------------------
+#         #     Expect JSON file: [{"id": 3, "name": "cat", "prompt": "cat"}, …]
+#         items_path = Path(path) / "items.json"
+#         with items_path.open("r", encoding="utf-8") as f:
+#             class_defs: List[Dict[str, Any]] = json.load(f)
+
+#         #     Extract the bits we need later
+#         prompts                 = [item["prompt"] for item in class_defs]
+#         self.class_ids:   List[int]   = [item["id"]   for item in class_defs]
+#         self.class_names: List[str]   = [item["name"] for item in class_defs]
+
+#         # 4️⃣  Encode all prompts once ---------------------------------
+#         with torch.no_grad():
+#             text_tokens  = tokenizer(prompts).to(self.device)
+#             text_feats   = self.model.encode_text(text_tokens)
+#             text_feats   = text_feats / text_feats.norm(dim=-1, keepdim=True)
+#         self.text_features = text_feats           # [num_classes, 512]
+
+#     # ------------------------------------------------------------------ #
+#     #                          INFERENCE CALL                            #
+#     # ------------------------------------------------------------------ #
+#     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+#         """
+#         Parameters
+#         ----------
+#         data : dict
+#             Either the raw payload `{"image": "<base64>"}` **or** the
+#             Hugging Face convention `{"inputs": {...}}`.
+
+#         Returns
+#         -------
+#         list of dict
+#             Sorted list of `{"id": int, "label": str, "score": float}`.
+#             Scores are the softmax probabilities over the *provided*
+#             class list (they sum to 1.0).
+#         """
+#         # 1️⃣  Unpack the request payload ------------------------------
+#         payload: Dict[str, Any] = data.get("inputs", data)
+#         img_b64: str = payload["image"]
+
+#         # 2️⃣  Decode + preprocess -------------------------------------
+#         image      = Image.open(io.BytesIO(base64.b64decode(img_b64))).convert("RGB")
+#         img_tensor = self.preprocess(image).unsqueeze(0).to(self.device)  # [1, 3, 224, 224]
+#         if self.device == "cuda":
+#             img_tensor = img_tensor.to(torch.float16)
+
+#         # 3️⃣  Forward pass (image only) -------------------------------
+#         with torch.no_grad():                    # no autograd graph
+#             img_feat = self.model.encode_image(img_tensor)            # [1, 512]
+#             img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True) # L2‑normalise
+
+#             # cosine similarity → logits → softmax probabilities
+#             probs = (100 * img_feat @ self.text_features.T).softmax(dim=-1)[0]  # [num_classes]
+
+#         # 4️⃣  Assemble JSON‑serialisable response ---------------------
+#         results = zip(self.class_ids, self.class_names, probs.cpu().tolist())
+#         return sorted(
+#             [{"id": cid, "label": name, "score": float(p)} for cid, name, p in results],
+#             key=lambda x: x["score"],
+#             reverse=True,
+#         )
+
+
+        

items.json

@@ -0,0 +1,902 @@
+[
+  {
+    "name": "Laptop",
+    "short_description": "Laptops are portable computers that can help you play games and learn new things. They can fold up just like a book!",
+    "category": "Office",
+    "rarity": "common",
+    "fun_fact": "Laptops have tiny fans to keep them cool while working.",
+    "id": 0,
+    "prompt": "a photo of a laptop computer (also called a notebook computer)",
+    "slug": "laptop"
+  },
+  {
+    "name": "Pen",
+    "short_description": "Pens are used to write and draw. They come in many colors.",
+    "category": "Office",
+    "rarity": "common",
+    "fun_fact": "Pens can write underwater using special waterproof ink.",
+    "id": 1,
+    "prompt": "a photo of an ink pen for writing",
+    "slug": "pen"
+  },
+  {
+    "name": "Notebook",
+    "short_description": "A notebook is a book with paper inside. You can write or draw in it!",
+    "category": "Office",
+    "rarity": "common",
+    "fun_fact": "Leonardo da Vinci used notebooks to sketch out their brilliant ideas, like flying machines!",
+    "id": 2,
+    "prompt": "a photo of a paper notebook for writing (stationery)",
+    "slug": "notebook"
+  },
+  {
+    "name": "Stapler",
+    "short_description": "A stapler helps you keep papers together by using metal clips.",
+    "category": "Office",
+    "rarity": "common",
+    "fun_fact": "The first stapler was made for King Louis XV in France.",
+    "id": 3,
+    "prompt": "a photo of a stapler (office paper stapler)",
+    "slug": "stapler"
+  },
+  {
+    "name": "Highlighter",
+    "short_description": "A highlighter is a bright pen used to mark important words or pictures.",
+    "category": "Office",
+    "rarity": "rare",
+    "fun_fact": "Highlighter ink glows under black light, like magic!",
+    "id": 4,
+    "prompt": "a photo of a highlighter pen (also called a fluorescent marker)",
+    "slug": "highlighter"
+  },
+  {
+    "name": "Paper Clip",
+    "short_description": "Paper clips are tiny helpers that keep papers together. They're shiny and colorful!",
+    "category": "Office",
+    "rarity": "rare",
+    "fun_fact": "The first paper clip was made over 100 years ago.",
+    "id": 5,
+    "prompt": "a photo of a paper clip (wire office clip)",
+    "slug": "paper_clip"
+  },
+  {
+    "name": "Desk Chair",
+    "short_description": "A desk chair is a seat for working or studying. It can roll and spin!",
+    "category": "Office",
+    "rarity": "rare",
+    "fun_fact": "Early chairs were just wooden and had no cushions.",
+    "id": 6,
+    "prompt": "a photo of a desk chair (office chair)",
+    "slug": "desk_chair"
+  },
+  {
+    "name": "Printer",
+    "short_description": "A printer makes copies of what's on a computer. It uses paper and ink.",
+    "category": "Office",
+    "rarity": "super rare",
+    "fun_fact": "Printers make pictures from millions of tiny dots.",
+    "id": 7,
+    "prompt": "a photo of a document printer (inkjet or laser printer)",
+    "slug": "printer"
+  },
+  {
+    "name": "Coffee Mug",
+    "short_description": "This is a cup with a handle to drink from. It's great for hot drinks!",
+    "category": "Office",
+    "rarity": "super rare",
+    "fun_fact": "Some mugs can change color or show pictures when they get hot.",
+    "id": 8,
+    "prompt": "a photo of a coffee mug (cup with a handle)",
+    "slug": "coffee_mug"
+  },
+  {
+    "name": "Pencil",
+    "short_description": "A pencil lets you draw and write. It's long, thin, and has an eraser.",
+    "category": "School",
+    "rarity": "common",
+    "fun_fact": "A single pencil can draw a line about 35 miles long.",
+    "id": 9,
+    "prompt": "a photo of a wooden graphite pencil (for writing)",
+    "slug": "pencil"
+  },
+  {
+    "name": "Backpack",
+    "short_description": "A backpack is a bag worn on your back. It holds all your school supplies!",
+    "category": "School",
+    "rarity": "common",
+    "fun_fact": "Some backpacks have hidden pockets for secret treasures.",
+    "id": 10,
+    "prompt": "a photo of a backpack (school bag)",
+    "slug": "backpack"
+  },
+  {
+    "name": "Book",
+    "short_description": "A book has pages with words and pictures. You can learn or enjoy stories.",
+    "category": "School",
+    "rarity": "common",
+    "fun_fact": "Book pages are called leaves, just like trees!",
+    "id": 11,
+    "prompt": "a photo of a printed book (paper book)",
+    "slug": "book"
+  },
+  {
+    "name": "Crayon",
+    "short_description": "Crayons are colorful sticks used to draw and color. They're easy to hold!",
+    "category": "School",
+    "rarity": "common",
+    "fun_fact": "Crayons are wax sticks; your warm hands help them draw smoothly.",
+    "id": 12,
+    "prompt": "a photo of a wax crayon (drawing crayon)",
+    "slug": "crayon"
+  },
+  {
+    "name": "Glue Stick",
+    "short_description": "A glue stick is used to stick paper and crafts. Just roll and glue!",
+    "category": "School",
+    "rarity": "common",
+    "fun_fact": "Some glue sticks glow in the dark for nighttime crafting fun!",
+    "id": 13,
+    "prompt": "a photo of a glue stick (solid glue)",
+    "slug": "glue_stick"
+  },
+  {
+    "name": "Ruler",
+    "short_description": "A ruler is used to measure how long things are.",
+    "category": "School",
+    "rarity": "rare",
+    "fun_fact": "Most rulers are 12 inches long — exactly one foot, just like your shoe!",
+    "id": 14,
+    "prompt": "a photo of a measuring ruler (straightedge)",
+    "slug": "ruler"
+  },
+  {
+    "name": "Lunchbox",
+    "short_description": "A lunchbox holds your yummy food for school. It keeps everything safe!",
+    "category": "School",
+    "rarity": "rare",
+    "fun_fact": "Thick, squishy walls help keep your lunch cool and fresh.",
+    "id": 15,
+    "prompt": "a photo of a lunchbox (school lunch box)",
+    "slug": "lunchbox"
+  },
+  {
+    "name": "Eraser",
+    "short_description": "Erasers help remove pencil marks. They're like magic fixers for your writing!",
+    "category": "School",
+    "rarity": "rare",
+    "fun_fact": "Erasers were first made from bread crumbs in ancient Japan.",
+    "id": 16,
+    "prompt": "a photo of an eraser (rubber eraser)",
+    "slug": "eraser"
+  },
+  {
+    "name": "Classroom Globe",
+    "short_description": "A globe shows us where countries and oceans are. It spins!",
+    "category": "School",
+    "rarity": "super rare",
+    "fun_fact": "A globe shows what our Earth looks like from space!",
+    "id": 17,
+    "prompt": "a photo of a classroom globe (world globe)",
+    "slug": "classroom_globe"
+  },
+  {
+    "name": "Daisy",
+    "short_description": "A daisy is a cute flower with white petals and a yellow center.",
+    "category": "Garden",
+    "rarity": "common",
+    "fun_fact": "Daisy flowers close their petals at night and open again in the morning.",
+    "id": 18,
+    "prompt": "a photo of a daisy (flower)",
+    "slug": "daisy"
+  },
+  {
+    "name": "Rock",
+    "short_description": "Rocks are hard and come in different shapes and sizes. You can find them almost anywhere!",
+    "category": "Garden",
+    "rarity": "common",
+    "fun_fact": "Some rocks were once melted lava from volcanoes, then cooled and hardened.",
+    "id": 19,
+    "prompt": "a photo of a rock (stone)",
+    "slug": "rock"
+  },
+  {
+    "name": "Leaf",
+    "short_description": "Leaves grow on trees and plants. They can be green, red, or yellow!",
+    "category": "Garden",
+    "rarity": "common",
+    "fun_fact": "Leaves make food for plants using sunlight, water, and air.",
+    "id": 20,
+    "prompt": "a photo of a leaf (plant leaf)",
+    "slug": "leaf"
+  },
+  {
+    "name": "Stick",
+    "short_description": "Sticks are long and hard, great for building or playing pretend!",
+    "category": "Garden",
+    "rarity": "common",
+    "fun_fact": "Sticks make great homes for small bugs and spiders.",
+    "id": 21,
+    "prompt": "a photo of a stick (tree branch or twig)",
+    "slug": "stick"
+  },
+  {
+    "name": "Watering Can",
+    "short_description": "A watering can helps us water plants! Pour water through the spout.",
+    "category": "Garden",
+    "rarity": "common",
+    "fun_fact": "The sprinkler cap at the end of the spout is called a 'rose'!",
+    "id": 22,
+    "prompt": "a photo of a watering can (garden watering can)",
+    "slug": "watering_can"
+  },
+  {
+    "name": "Snail",
+    "short_description": "Snails are little animals with shells on their backs. They move very slowly.",
+    "category": "Garden",
+    "rarity": "rare",
+    "fun_fact": "Snails make slippery slime so they can slide without scratching.",
+    "id": 23,
+    "prompt": "a photo of a snail (animal with a spiral shell)",
+    "slug": "snail"
+  },
+  {
+    "name": "Ladybug",
+    "short_description": "Ladybugs are small, red insects with black spots. They love to eat pesky bugs.",
+    "category": "Garden",
+    "rarity": "rare",
+    "fun_fact": "Ladybugs can have up to seven spots — count them if you can!",
+    "id": 24,
+    "prompt": "a photo of a ladybug (also called a ladybird beetle)",
+    "slug": "ladybug"
+  },
+  {
+    "name": "Birdhouse",
+    "short_description": "A birdhouse is a little house for birds. They stay safe and cozy inside.",
+    "category": "Garden",
+    "rarity": "super rare",
+    "fun_fact": "Some birds like to decorate their houses with shiny objects!",
+    "id": 25,
+    "prompt": "a photo of a birdhouse (nesting box)",
+    "slug": "birdhouse"
+  },
+  {
+    "name": "Garden Gnome",
+    "short_description": "A garden gnome is a small statue that lives in gardens and brings fun.",
+    "category": "Garden",
+    "rarity": "super rare",
+    "fun_fact": "Garden gnomes were once believed to guard treasures buried underground.",
+    "id": 26,
+    "prompt": "a photo of a garden gnome (garden statue)",
+    "slug": "garden_gnome"
+  },
+  {
+    "name": "Swing",
+    "short_description": "A swing is a seat that hangs from ropes or chains. It's super fun to swing high!",
+    "category": "Playground",
+    "rarity": "common",
+    "fun_fact": "Pumping your legs makes you go higher — each pump gives the swing a push.",
+    "id": 27,
+    "prompt": "a photo of a swing (playground swing)",
+    "slug": "swing"
+  },
+  {
+    "name": "Slide",
+    "short_description": "Slides are fun surfaces to quickly slide down. They can be found in playgrounds.",
+    "category": "Playground",
+    "rarity": "common",
+    "fun_fact": "Slides can be found in shapes like spirals and waves.",
+    "id": 28,
+    "prompt": "a photo of a slide (playground slide)",
+    "slug": "slide"
+  },
+  {
+    "name": "Ball",
+    "short_description": "Balls are round toys you can throw, catch, or kick. They're lots of fun!",
+    "category": "Playground",
+    "rarity": "common",
+    "fun_fact": "Balls are spheres, so they roll easily in any direction.",
+    "id": 29,
+    "prompt": "a photo of a ball (toy ball)",
+    "slug": "ball"
+  },
+  {
+    "name": "Sandbox",
+    "short_description": "A sandbox is a fun place where kids can play with sand. You can build castles and dig holes!",
+    "category": "Playground",
+    "rarity": "common",
+    "fun_fact": "Sand can be different colors like yellow, white, and even pink!",
+    "id": 30,
+    "prompt": "a photo of a sandbox (box of play sand)",
+    "slug": "sandbox"
+  },
+  {
+    "name": "Seesaw",
+    "short_description": "A seesaw is a long board that goes up and down. Friends can play on it together!",
+    "category": "Playground",
+    "rarity": "common",
+    "fun_fact": "Playing on a seesaw helps improve your balance and teamwork.",
+    "id": 31,
+    "prompt": "a photo of a seesaw (teeter-totter)",
+    "slug": "seesaw"
+  },
+  {
+    "name": "Hula Hoop",
+    "short_description": "A hula hoop is a big round toy you can spin around your waist. It's super fun!",
+    "category": "Playground",
+    "rarity": "rare",
+    "fun_fact": "The world record is over 100 hours of hula hooping!",
+    "id": 32,
+    "prompt": "a photo of a hula hoop (toy hoop)",
+    "slug": "hula_hoop"
+  },
+  {
+    "name": "Frisbee",
+    "short_description": "A frisbee is a round, flying disc that you can throw and catch with friends.",
+    "category": "Playground",
+    "rarity": "rare",
+    "fun_fact": "There's a sport called Ultimate Frisbee, like soccer with Frisbees.",
+    "id": 33,
+    "prompt": "a photo of a frisbee (flying disc)",
+    "slug": "frisbee"
+  },
+  {
+    "name": "Climbing Frame",
+    "short_description": "A climbing frame is a structure to climb and play on. It's super fun!",
+    "category": "Playground",
+    "rarity": "super rare",
+    "fun_fact": "Climbing frames can be shaped like rockets or castles!",
+    "id": 34,
+    "prompt": "a photo of a climbing frame (jungle gym)",
+    "slug": "climbing_frame"
+  },
+  {
+    "name": "Bouncer",
+    "short_description": "A bouncer is a fun thing to jump on. It helps you bounce up and down!",
+    "category": "Playground",
+    "rarity": "super rare",
+    "fun_fact": "Springs or bouncy air push you back up after each jump—boing, boing!",
+    "id": 35,
+    "prompt": "a photo of a playground bouncer (spring rider toy)",
+    "slug": "bouncer"
+  },
+  {
+    "name": "Dog",
+    "short_description": "A dog is a fun, furry friend. They love to play and wag their tails!",
+    "category": "Animals",
+    "rarity": "common",
+    "fun_fact": "Dogs have a super sense of smell, much better than ours!",
+    "id": 36,
+    "prompt": "a photo of a dog (domestic dog)",
+    "slug": "dog"
+  },
+  {
+    "name": "Cat",
+    "short_description": "Cats are furry animals with whiskers and tails. They love to purr and nap.",
+    "category": "Animals",
+    "rarity": "common",
+    "fun_fact": "Whiskers help cats feel around in the dark or tight spaces.",
+    "id": 37,
+    "prompt": "a photo of a cat (house cat)",
+    "slug": "cat"
+  },
+  {
+    "name": "Bird",
+    "short_description": "Birds have wings and feathers. They love to chirp and fly!",
+    "category": "Animals",
+    "rarity": "common",
+    "fun_fact": "Penguins are birds but can't fly; they are great swimmers!",
+    "id": 38,
+    "prompt": "a photo of a bird (animal with feathers)",
+    "slug": "bird"
+  },
+  {
+    "name": "Fly",
+    "short_description": "Flies are tiny insects that love to buzz around. They have wings and big eyes!",
+    "category": "Animals",
+    "rarity": "common",
+    "fun_fact": "A fly's wings beat over 200 times every second!",
+    "id": 39,
+    "prompt": "a photo of a fly (housefly insect)",
+    "slug": "fly"
+  },
+  {
+    "name": "Fish",
+    "short_description": "Fish live in water and come in many colors and sizes. They have fins to swim.",
+    "category": "Animals",
+    "rarity": "common",
+    "fun_fact": "Fish can sleep with their eyes open—because they have no eyelids!",
+    "id": 40,
+    "prompt": "a photo of a fish (aquarium or pond fish)",
+    "slug": "fish"
+  },
+  {
+    "name": "Rabbit",
+    "short_description": "Rabbits are fluffy animals with long ears and love to hop around.",
+    "category": "Animals",
+    "rarity": "rare",
+    "fun_fact": "A rabbit's teeth never stop growing, so they chew a lot!",
+    "id": 41,
+    "prompt": "a photo of a rabbit (bunny)",
+    "slug": "rabbit"
+  },
+  {
+    "name": "Frog",
+    "short_description": "Frogs are jumpy animals that live in ponds. They make funny ribbit sounds!",
+    "category": "Animals",
+    "rarity": "rare",
+    "fun_fact": "Frogs can jump over 20 times their own body length!",
+    "id": 42,
+    "prompt": "a photo of a frog (amphibian)",
+    "slug": "frog"
+  },
+  {
+    "name": "Duck",
+    "short_description": "Ducks are birds that love to swim and quack. They have webbed feet!",
+    "category": "Animals",
+    "rarity": "super rare",
+    "fun_fact": "Ducks have waterproof feathers thanks to special oils they produce.",
+    "id": 43,
+    "prompt": "a photo of a duck (water bird)",
+    "slug": "duck"
+  },
+  {
+    "name": "Squirrel",
+    "short_description": "Squirrels are small, bushy-tailed animals that love to climb trees. They gather nuts!",
+    "category": "Animals",
+    "rarity": "super rare",
+    "fun_fact": "Squirrels plant thousands of trees by forgetting where they buried nuts.",
+    "id": 44,
+    "prompt": "a photo of a squirrel (tree squirrel)",
+    "slug": "squirrel"
+  },
+  {
+    "name": "Car",
+    "short_description": "A car has wheels and a steering wheel. It can take you places!",
+    "category": "Transportation",
+    "rarity": "common",
+    "fun_fact": "The first cars had no steering wheels—drivers used levers instead.",
+    "id": 45,
+    "prompt": "a photo of a car (automobile)",
+    "slug": "car"
+  },
+  {
+    "name": "Bus",
+    "short_description": "A bus is a big vehicle that takes people to places. It's like a car, but much bigger!",
+    "category": "Transportation",
+    "rarity": "common",
+    "fun_fact": "The first buses were pulled by horses before engines were invented!",
+    "id": 46,
+    "prompt": "a photo of a bus (passenger bus)",
+    "slug": "bus"
+  },
+  {
+    "name": "Bicycle",
+    "short_description": "Bicycles have pedals and handlebars you use to ride around. They're great for fun and exercise!",
+    "category": "Transportation",
+    "rarity": "common",
+    "fun_fact": "The first bicycles were called \"velocipedes\" and had wooden wheels!",
+    "id": 47,
+    "prompt": "a photo of a bicycle (bike)",
+    "slug": "bicycle"
+  },
+  {
+    "name": "Scooter",
+    "short_description": "A scooter is a ride-on toy with two wheels. You can push with your feet!",
+    "category": "Transportation",
+    "rarity": "common",
+    "fun_fact": "Early scooters were made of wood with roller skate wheels.",
+    "id": 48,
+    "prompt": "a photo of a scooter (kick scooter)",
+    "slug": "scooter"
+  },
+  {
+    "name": "Train",
+    "short_description": "Trains are big, fast vehicles that travel on tracks. They carry people and goods.",
+    "category": "Transportation",
+    "rarity": "common",
+    "fun_fact": "The longest train ever was over 4.5 miles long!",
+    "id": 49,
+    "prompt": "a photo of a train (railway train)",
+    "slug": "train"
+  },
+  {
+    "name": "Fire Truck",
+    "short_description": "A fire truck helps put out fires. It carries firefighters and water.",
+    "category": "Transportation",
+    "rarity": "rare",
+    "fun_fact": "Some fire trucks have ladders that reach over 100 feet high.",
+    "id": 50,
+    "prompt": "a photo of a fire truck (fire engine)",
+    "slug": "fire_truck"
+  },
+  {
+    "name": "Tractor",
+    "short_description": "Tractors are big machines that help farmers. They pull heavy things and plow fields.",
+    "category": "Transportation",
+    "rarity": "rare",
+    "fun_fact": "Some tractors have wheels taller than a person!",
+    "id": 51,
+    "prompt": "a photo of a tractor (farm tractor)",
+    "slug": "tractor"
+  },
+  {
+    "name": "Plane",
+    "short_description": "Planes are big flying machines. They take us on long trips in the sky!",
+    "category": "Transportation",
+    "rarity": "super rare",
+    "fun_fact": "The fastest passenger plane flew faster than the speed of sound.",
+    "id": 52,
+    "prompt": "a photo of an airplane (passenger plane)",
+    "slug": "plane"
+  },
+  {
+    "name": "Boat",
+    "short_description": "A boat is a fun way to travel on water. It can be big or small!",
+    "category": "Transportation",
+    "rarity": "super rare",
+    "fun_fact": "Viking boats were called longships and had dragon heads!",
+    "id": 53,
+    "prompt": "a photo of a boat (watercraft)",
+    "slug": "boat"
+  },
+  {
+    "name": "Toothbrush",
+    "short_description": "A toothbrush helps clean your teeth. Use it with toothpaste for a shiny smile!",
+    "category": "House",
+    "rarity": "common",
+    "fun_fact": "The first toothbrushes were made from sticks and twigs!",
+    "id": 54,
+    "prompt": "a photo of a toothbrush (manual toothbrush)",
+    "slug": "toothbrush"
+  },
+  {
+    "name": "Towel",
+    "short_description": "A towel is soft and helps you dry off after a bath or swim.",
+    "category": "House",
+    "rarity": "common",
+    "fun_fact": "Towels soak up water with thousands of tiny fabric loops.",
+    "id": 55,
+    "prompt": "a photo of a towel (bath towel)",
+    "slug": "towel"
+  },
+  {
+    "name": "Lamp",
+    "short_description": "A lamp lights up a room. You turn it on with a switch!",
+    "category": "House",
+    "rarity": "common",
+    "fun_fact": "The first electric lamp was made over 140 years ago by Thomas Edison.",
+    "id": 56,
+    "prompt": "a photo of a lamp (table or floor lamp)",
+    "slug": "lamp"
+  },
+  {
+    "name": "Chair",
+    "short_description": "A chair is a seat with four legs and a back. It's great for sitting!",
+    "category": "House",
+    "rarity": "common",
+    "fun_fact": "In ancient Egypt, only kings and queens had chairs with backs.",
+    "id": 57,
+    "prompt": "a photo of a chair (wooden or cushioned chair)",
+    "slug": "chair"
+  },
+  {
+    "name": "Cushion",
+    "short_description": "A cushion is a soft pad you can sit on. It makes couches more comfy!",
+    "category": "House",
+    "rarity": "common",
+    "fun_fact": "In ancient times, cushions were signs of royalty and power.",
+    "id": 58,
+    "prompt": "a photo of a cushion (sofa cushion)",
+    "slug": "cushion"
+  },
+  {
+    "name": "Fan",
+    "short_description": "A fan makes the air move and keeps us cool. It's like a mini wind!",
+    "category": "House",
+    "rarity": "rare",
+    "fun_fact": "Electric fans were invented over 130 years ago!",
+    "id": 59,
+    "prompt": "a photo of a fan (electric fan)",
+    "slug": "fan"
+  },
+  {
+    "name": "Clock",
+    "short_description": "A clock tells the time. It has numbers and hands to point.",
+    "category": "House",
+    "rarity": "rare",
+    "fun_fact": "Big Ben is a famous clock tower in London, England.",
+    "id": 60,
+    "prompt": "a photo of a clock (analog or wall clock)",
+    "slug": "clock"
+  },
+  {
+    "name": "Painting",
+    "short_description": "Paintings are cool pictures made with colors on paper or canvas.",
+    "category": "House",
+    "rarity": "super rare",
+    "fun_fact": "Famous artist Leonardo da Vinci loved painting with eggs.",
+    "id": 61,
+    "prompt": "a photo of a painting (art on canvas or paper)",
+    "slug": "painting"
+  },
+  {
+    "name": "Plant Pot",
+    "short_description": "A plant pot holds soil for growing plants. It's like a tiny garden!",
+    "category": "House",
+    "rarity": "super rare",
+    "fun_fact": "Using big pots lets plants grow big and tall!",
+    "id": 62,
+    "prompt": "a photo of a plant pot (flower pot)",
+    "slug": "plant_pot"
+  },
+  {
+    "name": "T-Shirt",
+    "short_description": "A T-shirt is a piece of clothing with short sleeves. It's soft and comfy!",
+    "category": "Clothing",
+    "rarity": "common",
+    "fun_fact": "T-shirts got their name because of their T-shape.",
+    "id": 63,
+    "prompt": "a photo of a t-shirt (short-sleeved shirt)",
+    "slug": "t_shirt"
+  },
+  {
+    "name": "Socks",
+    "short_description": "Socks are soft cloth tubes for your feet. They come in lots of colors!",
+    "category": "Clothing",
+    "rarity": "common",
+    "fun_fact": "There are special socks for running with extra cushion.",
+    "id": 64,
+    "prompt": "a photo of socks (pair of socks)",
+    "slug": "socks"
+  },
+  {
+    "name": "Shoes",
+    "short_description": "Shoes protect your feet and look cool! They come in all colors.",
+    "category": "Clothing",
+    "rarity": "common",
+    "fun_fact": "Shoes were first invented over 5,000 years ago!",
+    "id": 65,
+    "prompt": "a photo of shoes (pair of shoes)",
+    "slug": "shoes"
+  },
+  {
+    "name": "Hat",
+    "short_description": "A hat is something you wear on your head. It can keep you warm or cool!",
+    "category": "Clothing",
+    "rarity": "common",
+    "fun_fact": "Some hats are made just for magic tricks!",
+    "id": 66,
+    "prompt": "a photo of a hat (headwear)",
+    "slug": "hat"
+  },
+  {
+    "name": "Jacket",
+    "short_description": "A jacket keeps you warm when it's chilly. It's like a cozy hug.",
+    "category": "Clothing",
+    "rarity": "common",
+    "fun_fact": "Winter jackets are often stuffed with fluffy feathers!",
+    "id": 67,
+    "prompt": "a photo of a jacket (outerwear)",
+    "slug": "jacket"
+  },
+  {
+    "name": "Scarf",
+    "short_description": "A scarf is a soft piece of clothing you wear around your neck. It keeps you warm.",
+    "category": "Clothing",
+    "rarity": "rare",
+    "fun_fact": "In ancient Rome, scarves were used to wipe sweat, not for warmth.",
+    "id": 68,
+    "prompt": "a photo of a scarf (neckwear)",
+    "slug": "scarf"
+  },
+  {
+    "name": "Gloves",
+    "short_description": "Gloves are like tiny jackets for your hands. They keep you warm or safe!",
+    "category": "Clothing",
+    "rarity": "rare",
+    "fun_fact": "Astronauts wear special gloves to hold tools in space!",
+    "id": 69,
+    "prompt": "a photo of gloves (pair of gloves)",
+    "slug": "gloves"
+  },
+  {
+    "name": "Sunglasses",
+    "short_description": "Sunglasses protect your eyes from the bright sun. They come in cool shapes!",
+    "category": "Clothing",
+    "rarity": "rare",
+    "fun_fact": "In ancient China, judges wore dark glasses to hide their eyes in court.",
+    "id": 70,
+    "prompt": "a photo of sunglasses (pair of sunglasses)",
+    "slug": "sunglasses"
+  },
+  {
+    "name": "Umbrella",
+    "short_description": "Umbrellas keep you dry in the rain. They open wide like a big flower!",
+    "category": "Clothing",
+    "rarity": "super rare",
+    "fun_fact": "Umbrellas were once used as sun protection, not for rain.",
+    "id": 71,
+    "prompt": "a photo of an umbrella (rain umbrella)",
+    "slug": "umbrella"
+  },
+  {
+    "name": "Sofa",
+    "short_description": "A sofa is a big, comfy seat. You can sit or lie on it.",
+    "category": "Living Room",
+    "rarity": "common",
+    "fun_fact": "The word 'sofa' comes from an Arabic word meaning 'bench with cushions'.",
+    "id": 72,
+    "prompt": "a photo of a sofa (couch)",
+    "slug": "sofa"
+  },
+  {
+    "name": "TV",
+    "short_description": "A TV shows your favorite cartoons and movies with sound and pictures.",
+    "category": "Living Room",
+    "rarity": "common",
+    "fun_fact": "The first TV shows were black and white.",
+    "id": 73,
+    "prompt": "a photo of a television (TV set)",
+    "slug": "tv"
+  },
+  {
+    "name": "Remote Control",
+    "short_description": "A remote helps you change channels on the TV. It has buttons to press!",
+    "category": "Living Room",
+    "rarity": "rare",
+    "fun_fact": "Remote controls use light beams called infrared to work.",
+    "id": 74,
+    "prompt": "a photo of a remote control (TV remote)",
+    "slug": "remote_control"
+  },
+  {
+    "name": "Rug",
+    "short_description": "Rugs are soft mats that cover the floor. They come in many colors and shapes.",
+    "category": "Living Room",
+    "rarity": "common",
+    "fun_fact": "Some handmade rugs can take years to finish.",
+    "id": 75,
+    "prompt": "a photo of a rug (floor carpet)",
+    "slug": "rug"
+  },
+  {
+    "name": "Bookshelf",
+    "short_description": "A bookshelf holds your books, keeping them tidy and easy to find.",
+    "category": "Living Room",
+    "rarity": "rare",
+    "fun_fact": "Some bookshelves have secret compartments for hiding small treasures.",
+    "id": 76,
+    "prompt": "a photo of a bookshelf (bookcase)",
+    "slug": "bookshelf"
+  },
+  {
+    "name": "Blanket",
+    "short_description": "A blanket is soft and keeps you warm. It's great for snuggling!",
+    "category": "Living Room",
+    "rarity": "rare",
+    "fun_fact": "Blankets keep you warm by trapping your body heat inside.",
+    "id": 77,
+    "prompt": "a photo of a blanket (throw blanket)",
+    "slug": "blanket"
+  },
+  {
+    "name": "Candle",
+    "short_description": "Candles are sticks of wax with a wick that can be lit to make light. They also smell nice!",
+    "category": "Living Room",
+    "rarity": "super rare",
+    "fun_fact": "The earliest candles were made from animal fat, not wax.",
+    "id": 78,
+    "prompt": "a photo of a candle (wax candle)",
+    "slug": "candle"
+  },
+  {
+    "name": "Record Player",
+    "short_description": "A record player plays music from vinyl records. It uses a needle to read grooves.",
+    "category": "Living Room",
+    "rarity": "super rare",
+    "fun_fact": "A needle travels in grooves to create the music we hear.",
+    "id": 79,
+    "prompt": "a photo of a record player (turntable)",
+    "slug": "record_player"
+  },
+  {
+    "name": "Coffee Table",
+    "short_description": "A coffee table sits in front of the sofa and holds drinks, snacks, or board games.",
+    "category": "Living Room",
+    "rarity": "common",
+    "fun_fact": "Coffee tables can be square, round, or even shaped like animals!",
+    "id": 80,
+    "prompt": "a photo of a coffee table (low table)",
+    "slug": "coffee_table"
+  },
+  {
+    "name": "Teddy Bear",
+    "short_description": "Teddy bears are soft, cuddly toys. They love hugs and snuggles!",
+    "category": "Toys & Crafts",
+    "rarity": "common",
+    "fun_fact": "Teddy bears are named after U.S. President Theodore Roosevelt.",
+    "id": 81,
+    "prompt": "a photo of a teddy bear (stuffed toy)",
+    "slug": "teddy_bear"
+  },
+  {
+    "name": "Toy Car",
+    "short_description": "A toy car is fun to zoom around. You can pretend it's on a racetrack!",
+    "category": "Toys & Crafts",
+    "rarity": "common",
+    "fun_fact": "Some tiny toy cars can roll over 20 feet on one push!",
+    "id": 82,
+    "prompt": "a photo of a toy car (miniature car)",
+    "slug": "toy_car"
+  },
+  {
+    "name": "LEGO",
+    "short_description": "LEGO bricks are colorful blocks that you can use to build anything you imagine!",
+    "category": "Toys & Crafts",
+    "rarity": "common",
+    "fun_fact": "The word LEGO means 'play well' in Danish.",
+    "id": 83,
+    "prompt": "a photo of lego bricks (building blocks)",
+    "slug": "lego"
+  },
+  {
+    "name": "Paintbrush",
+    "short_description": "A paintbrush is used to make art by spreading paint on surfaces. It has soft bristles!",
+    "category": "Toys & Crafts",
+    "rarity": "common",
+    "fun_fact": "The oldest known paintbrushes were made from animal hair tied to sticks.",
+    "id": 84,
+    "prompt": "a photo of a paintbrush (art brush)",
+    "slug": "paintbrush"
+  },
+  {
+    "name": "Doll",
+    "short_description": "A doll is a fun toy that looks like a person. You can dress it up and play with it!",
+    "category": "Toys & Crafts",
+    "rarity": "rare",
+    "fun_fact": "Ancient dolls were made from clay, wood, and even ivory.",
+    "id": 85,
+    "prompt": "a photo of a doll (toy doll)",
+    "slug": "doll"
+  },
+  {
+    "name": "Videogame Console",
+    "short_description": "A console lets you play video games on your TV. How fun!",
+    "category": "Toys & Crafts",
+    "rarity": "common",
+    "fun_fact": "The first home console was called the Magnavox Odyssey.",
+    "id": 86,
+    "prompt": "a photo of a video game console (gaming console)",
+    "slug": "videogame_console"
+  },
+  {
+    "name": "Puzzle Piece",
+    "short_description": "A puzzle piece fits with others to make a big picture. It's like a fun mystery!",
+    "category": "Toys & Crafts",
+    "rarity": "rare",
+    "fun_fact": "The first jigsaw puzzles were made from maps to teach kids geography.",
+    "id": 87,
+    "prompt": "a photo of a puzzle piece (jigsaw puzzle piece)",
+    "slug": "puzzle_piece"
+  },
+  {
+    "name": "Stickers",
+    "short_description": "Stickers are like colorful, sticky friends you can put on things! They make everything fun.",
+    "category": "Toys & Crafts",
+    "rarity": "rare",
+    "fun_fact": "Some artists use thousands of tiny stickers to create huge, amazing pictures.",
+    "id": 88,
+    "prompt": "a photo of stickers (adhesive stickers)",
+    "slug": "stickers"
+  },
+  {
+    "name": "Magic Wand",
+    "short_description": "A magic wand is a stick that wizards use to do magic! It sparkles and shines.",
+    "category": "Toys & Crafts",
+    "rarity": "super rare",
+    "fun_fact": "Wizards wave their wands in special ways for magic to work!",
+    "id": 89,
+    "prompt": "a photo of a magic wand (toy wand)",
+    "slug": "magic_wand"
+  }
+]

reparam.py

@@ -0,0 +1,341 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2024 Apple Inc. All Rights Reserved.
+#
+from typing import Union, Tuple
+
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ["MobileOneBlock", "reparameterize_model"]
+
+
+class SEBlock(nn.Module):
+    """Squeeze and Excite module.
+
+    Pytorch implementation of `Squeeze-and-Excitation Networks` -
+    https://arxiv.org/pdf/1709.01507.pdf
+    """
+
+    def __init__(self, in_channels: int, rd_ratio: float = 0.0625) -> None:
+        """Construct a Squeeze and Excite Module.
+
+        Args:
+            in_channels: Number of input channels.
+            rd_ratio: Input channel reduction ratio.
+        """
+        super(SEBlock, self).__init__()
+        self.reduce = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=int(in_channels * rd_ratio),
+            kernel_size=1,
+            stride=1,
+            bias=True,
+        )
+        self.expand = nn.Conv2d(
+            in_channels=int(in_channels * rd_ratio),
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            bias=True,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """Apply forward pass."""
+        b, c, h, w = inputs.size()
+        x = F.avg_pool2d(inputs, kernel_size=[h, w])
+        x = self.reduce(x)
+        x = F.relu(x)
+        x = self.expand(x)
+        x = torch.sigmoid(x)
+        x = x.view(-1, c, 1, 1)
+        return inputs * x
+
+
+class MobileOneBlock(nn.Module):
+    """MobileOne building block.
+
+    This block has a multi-branched architecture at train-time
+    and plain-CNN style architecture at inference time
+    For more details, please refer to our paper:
+    `An Improved One millisecond Mobile Backbone` -
+    https://arxiv.org/pdf/2206.04040.pdf
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        inference_mode: bool = False,
+        use_se: bool = False,
+        use_act: bool = True,
+        use_scale_branch: bool = True,
+        num_conv_branches: int = 1,
+        activation: nn.Module = nn.GELU(),
+    ) -> None:
+        """Construct a MobileOneBlock module.
+
+        Args:
+            in_channels: Number of channels in the input.
+            out_channels: Number of channels produced by the block.
+            kernel_size: Size of the convolution kernel.
+            stride: Stride size.
+            padding: Zero-padding size.
+            dilation: Kernel dilation factor.
+            groups: Group number.
+            inference_mode: If True, instantiates model in inference mode.
+            use_se: Whether to use SE-ReLU activations.
+            use_act: Whether to use activation. Default: ``True``
+            use_scale_branch: Whether to use scale branch. Default: ``True``
+            num_conv_branches: Number of linear conv branches.
+        """
+        super(MobileOneBlock, self).__init__()
+        self.inference_mode = inference_mode
+        self.groups = groups
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_conv_branches = num_conv_branches
+
+        # Check if SE-ReLU is requested
+        if use_se:
+            self.se = SEBlock(out_channels)
+        else:
+            self.se = nn.Identity()
+
+        if use_act:
+            self.activation = activation
+        else:
+            self.activation = nn.Identity()
+
+        if inference_mode:
+            self.reparam_conv = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias=True,
+            )
+        else:
+            # Re-parameterizable skip connection
+            self.rbr_skip = (
+                nn.BatchNorm2d(num_features=in_channels)
+                if out_channels == in_channels and stride == 1
+                else None
+            )
+
+            # Re-parameterizable conv branches
+            if num_conv_branches > 0:
+                rbr_conv = list()
+                for _ in range(self.num_conv_branches):
+                    rbr_conv.append(
+                        self._conv_bn(kernel_size=kernel_size, padding=padding)
+                    )
+                self.rbr_conv = nn.ModuleList(rbr_conv)
+            else:
+                self.rbr_conv = None
+
+            # Re-parameterizable scale branch
+            self.rbr_scale = None
+            if not isinstance(kernel_size, int):
+                kernel_size = kernel_size[0]
+            if (kernel_size > 1) and use_scale_branch:
+                self.rbr_scale = self._conv_bn(kernel_size=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply forward pass."""
+        # Inference mode forward pass.
+        if self.inference_mode:
+            return self.activation(self.se(self.reparam_conv(x)))
+
+        # Multi-branched train-time forward pass.
+        # Skip branch output
+        identity_out = 0
+        if self.rbr_skip is not None:
+            identity_out = self.rbr_skip(x)
+
+        # Scale branch output
+        scale_out = 0
+        if self.rbr_scale is not None:
+            scale_out = self.rbr_scale(x)
+
+        # Other branches
+        out = scale_out + identity_out
+        if self.rbr_conv is not None:
+            for ix in range(self.num_conv_branches):
+                out += self.rbr_conv[ix](x)
+
+        return self.activation(self.se(out))
+
+    def reparameterize(self):
+        """Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
+        https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
+        architecture used at training time to obtain a plain CNN-like structure
+        for inference.
+        """
+        if self.inference_mode:
+            return
+        kernel, bias = self._get_kernel_bias()
+        self.reparam_conv = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            bias=True,
+        )
+        self.reparam_conv.weight.data = kernel
+        self.reparam_conv.bias.data = bias
+
+        # Delete un-used branches
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__("rbr_conv")
+        self.__delattr__("rbr_scale")
+        if hasattr(self, "rbr_skip"):
+            self.__delattr__("rbr_skip")
+
+        self.inference_mode = True
+
+    def _get_kernel_bias(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Method to obtain re-parameterized kernel and bias.
+        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L83
+
+        Returns:
+            Tuple of (kernel, bias) after fusing branches.
+        """
+        # get weights and bias of scale branch
+        kernel_scale = 0
+        bias_scale = 0
+        if self.rbr_scale is not None:
+            kernel_scale, bias_scale = self._fuse_bn_tensor(self.rbr_scale)
+            # Pad scale branch kernel to match conv branch kernel size.
+            pad = self.kernel_size // 2
+            kernel_scale = torch.nn.functional.pad(kernel_scale, [pad, pad, pad, pad])
+
+        # get weights and bias of skip branch
+        kernel_identity = 0
+        bias_identity = 0
+        if self.rbr_skip is not None:
+            kernel_identity, bias_identity = self._fuse_bn_tensor(self.rbr_skip)
+
+        # get weights and bias of conv branches
+        kernel_conv = 0
+        bias_conv = 0
+        if self.rbr_conv is not None:
+            for ix in range(self.num_conv_branches):
+                _kernel, _bias = self._fuse_bn_tensor(self.rbr_conv[ix])
+                kernel_conv += _kernel
+                bias_conv += _bias
+
+        kernel_final = kernel_conv + kernel_scale + kernel_identity
+        bias_final = bias_conv + bias_scale + bias_identity
+        return kernel_final, bias_final
+
+    def _fuse_bn_tensor(
+        self, branch: Union[nn.Sequential, nn.BatchNorm2d]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Method to fuse batchnorm layer with preceeding conv layer.
+        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L95
+
+        Args:
+            branch: Sequence of ops to be fused.
+
+        Returns:
+            Tuple of (kernel, bias) after fusing batchnorm.
+        """
+        if isinstance(branch, nn.Sequential):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, "id_tensor"):
+                input_dim = self.in_channels // self.groups
+
+                kernel_size = self.kernel_size
+                if isinstance(self.kernel_size, int):
+                    kernel_size = (self.kernel_size, self.kernel_size)
+
+                kernel_value = torch.zeros(
+                    (self.in_channels, input_dim, kernel_size[0], kernel_size[1]),
+                    dtype=branch.weight.dtype,
+                    device=branch.weight.device,
+                )
+                for i in range(self.in_channels):
+                    kernel_value[
+                        i, i % input_dim, kernel_size[0] // 2, kernel_size[1] // 2
+                    ] = 1
+                self.id_tensor = kernel_value
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def _conv_bn(self, kernel_size: int, padding: int) -> nn.Sequential:
+        """Helper method to construct conv-batchnorm layers.
+
+        Args:
+            kernel_size: Size of the convolution kernel.
+            padding: Zero-padding size.
+
+        Returns:
+            Conv-BN module.
+        """
+        mod_list = nn.Sequential()
+        mod_list.add_module(
+            "conv",
+            nn.Conv2d(
+                in_channels=self.in_channels,
+                out_channels=self.out_channels,
+                kernel_size=kernel_size,
+                stride=self.stride,
+                padding=padding,
+                groups=self.groups,
+                bias=False,
+            ),
+        )
+        mod_list.add_module("bn", nn.BatchNorm2d(num_features=self.out_channels))
+        return mod_list
+
+
+def reparameterize_model(model: torch.nn.Module) -> nn.Module:
+    """Method returns a model where a multi-branched structure
+        used in training is re-parameterized into a single branch
+        for inference.
+
+    Args:
+        model: MobileOne model in train mode.
+
+    Returns:
+        MobileOne model in inference mode.
+    """
+    # Avoid editing original graph
+    model = copy.deepcopy(model)
+    for module in model.modules():
+        if hasattr(module, "reparameterize"):
+            module.reparameterize()
+    return model

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=4.0.0
+Pillow
+torch>=2.0.0
+torchvision
+open_clip_torch
+huggingface_hub>=0.23.0
+safetensors>=0.4.3