app.py

"""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 <user>/<repo>` "
                    "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,
    )