"""Gradio demo: ViT image classification with pluggable LoRA adapters. Tabs: Classify - pick a task (ImageNet-1k base or any loaded LoRA adapter) and get top-k predictions with a confidence threshold. Compare: Base vs - side-by-side predictions from the same backbone with LoRA the adapter disabled vs enabled. Makes it obvious what LoRA added without touching the pretrained weights. """ from __future__ import annotations from pathlib import Path # Workaround for gradio_client<=1.3 schema bug that crashes on `additionalProperties: True` # and causes launch() to fail with "localhost is not accessible". # See: https://github.com/gradio-app/gradio/issues/10662 import gradio_client.utils as _gc_utils _orig_get_type = _gc_utils.get_type _orig_json_schema_to_python_type = _gc_utils._json_schema_to_python_type def _safe_get_type(schema): return _orig_get_type(schema) if isinstance(schema, dict) else "Any" def _safe_json_schema_to_python_type(schema, defs=None): if not isinstance(schema, dict): return "Any" return _orig_json_schema_to_python_type(schema, defs) _gc_utils.get_type = _safe_get_type _gc_utils._json_schema_to_python_type = _safe_json_schema_to_python_type import gradio as gr import torch from adapters import BASE_TASK_NAME, AdapterRegistry registry = AdapterRegistry() model = registry.model processor = registry.processor DEVICE = registry.device TASKS = registry.available_tasks() TASK_CHOICES = [(t.display_name, t.key) for t in TASKS] ADAPTER_TASKS = [t for t in TASKS if not t.is_base] ADAPTER_CHOICES = [(t.display_name, t.key) for t in ADAPTER_TASKS] EXAMPLES_DIR = Path(__file__).parent / "examples" FOOD_NAMES = {"baklava", "donut", "dumplings", "hotdog"} FOOD_TASK_KEY = "food101" if any(t.key == "food101" for t in TASKS) else BASE_TASK_NAME def _collect_examples() -> list[tuple[str, str]]: """Return (path, default_task_key) pairs for each example image.""" if not EXAMPLES_DIR.exists(): return [] pairs = [] for p in sorted(EXAMPLES_DIR.glob("*.jpg")): task_key = FOOD_TASK_KEY if p.stem.lower() in FOOD_NAMES else BASE_TASK_NAME pairs.append((str(p), task_key)) return pairs EXAMPLE_PAIRS = _collect_examples() EXAMPLE_IMAGES = [p for p, _ in EXAMPLE_PAIRS] def _to_inputs(image): if image.mode != "RGB": image = image.convert("RGB") return processor(images=image, return_tensors="pt").to(DEVICE) @torch.inference_mode() def classify(image, task_key, top_k, threshold): if image is None: return {} inputs = _to_inputs(image) with registry.use_task(task_key): logits = model(**inputs).logits[0] probs = torch.softmax(logits, dim=-1) id2label = registry.get_id2label(task_key) k = max(1, min(int(top_k), probs.shape[-1])) top = torch.topk(probs, k=k) return { id2label[idx.item()]: float(score) for score, idx in zip(top.values, top.indices) if float(score) >= threshold } def classify_compare(image, adapter_key, top_k, threshold): if image is None or not adapter_key: return {}, {} base = classify(image, BASE_TASK_NAME, top_k, threshold) adapter = classify(image, adapter_key, top_k, threshold) return base, adapter default_task = ADAPTER_CHOICES[0][1] if ADAPTER_CHOICES else BASE_TASK_NAME adapter_names_loaded = [a.display_name for a in ADAPTER_TASKS] with gr.Blocks(title="ViT + LoRA image classifier", theme=gr.themes.Soft()) as demo: gr.Markdown( rf""" # ViT image classifier with swappable LoRA adapters Backbone: [`{registry.base_model_id}`](https://huggingface.co/{registry.base_model_id}) (ViT-tiny, ~5.7M params). Device: **{DEVICE.upper()}**. Loaded adapters: **{', '.join(adapter_names_loaded) if adapter_names_loaded else 'none (base model only)'}**. Pick a task on the left to run top-k classification. The base weights are never mutated: the adapter adds a low-rank \(\Delta W\) and a new classification head, toggled on/off per request. """ ) with gr.Tabs(): with gr.Tab("Classify"): with gr.Row(): with gr.Column(scale=1): image = gr.Image( type="pil", label="Input image", sources=["upload", "clipboard", "webcam"], height=360, ) task = gr.Radio( choices=TASK_CHOICES, value=BASE_TASK_NAME, label="Task", info="Switch between the base model and any loaded LoRA adapter.", ) with gr.Row(): top_k = gr.Slider( minimum=1, maximum=10, value=5, step=1, label="Top-k", info="Number of predictions to show", ) threshold = gr.Slider( minimum=0.0, maximum=1.0, value=0.0, step=0.01, label="Confidence threshold", info="Hide predictions below this probability", ) submit = gr.Button("Classify", variant="primary") with gr.Column(scale=1): output = gr.Label(num_top_classes=10, label="Predictions") if EXAMPLE_PAIRS: gr.Examples( examples=[[p, tk, 5, 0.0] for p, tk in EXAMPLE_PAIRS], inputs=[image, task, top_k, threshold], outputs=output, fn=classify, cache_examples=False, label="Example images (animals default to ImageNet, foods to Food-101)", ) classify_args = dict(fn=classify, inputs=[image, task, top_k, threshold], outputs=output) submit.click(**classify_args) image.change(**classify_args) task.change(**classify_args) top_k.change(**classify_args) threshold.change(**classify_args) if ADAPTER_TASKS: with gr.Tab("Compare: Base vs LoRA"): gr.Markdown( "Runs the **same image** through the shared backbone twice: " "once with the LoRA adapter disabled (original ImageNet-1k head) " "and once with the selected adapter enabled." ) with gr.Row(): with gr.Column(scale=1): image2 = gr.Image( type="pil", label="Input image", sources=["upload", "clipboard", "webcam"], height=360, ) adapter_sel = gr.Dropdown( choices=ADAPTER_CHOICES, value=default_task, label="Adapter", interactive=True, ) with gr.Row(): top_k2 = gr.Slider( minimum=1, maximum=10, value=5, step=1, label="Top-k", ) threshold2 = gr.Slider( minimum=0.0, maximum=1.0, value=0.0, step=0.01, label="Confidence threshold", ) submit2 = gr.Button("Compare", variant="primary") with gr.Row(): base_out = gr.Label(num_top_classes=10, label="Base (ImageNet-1k)") adapter_out = gr.Label(num_top_classes=10, label="Adapter prediction") if EXAMPLE_IMAGES: gr.Examples( examples=[[p, default_task, 5, 0.0] for p in EXAMPLE_IMAGES], inputs=[image2, adapter_sel, top_k2, threshold2], outputs=[base_out, adapter_out], fn=classify_compare, cache_examples=False, label="Example images", ) compare_args = dict( fn=classify_compare, inputs=[image2, adapter_sel, top_k2, threshold2], outputs=[base_out, adapter_out], ) submit2.click(**compare_args) image2.change(**compare_args) adapter_sel.change(**compare_args) top_k2.change(**compare_args) threshold2.change(**compare_args) else: with gr.Tab("Compare: Base vs LoRA"): gr.Markdown( "_No LoRA adapters are currently loaded._\n\n" "Train one with `python train_lora.py --push-to-hub /` " "and register it in `adapters.py`." ) if __name__ == "__main__": demo.launch( server_name="0.0.0.0", server_port=7860, show_api=False, share=False, )