"""ProgResViT — progressive-resolution / progressive-width adaptive ViT. Interactive ImageNet-1K classification demo that exposes the paper's input-adaptive routing: round 1 runs a narrow subnetwork on a low-resolution image, and only uncertain images continue to round 2 at higher resolution and wider width. Paper: https://huggingface.co/papers/2609.03216 Code: https://github.com/ds-kiel/ProgResViT """ import json import os import time import spaces # must precede torch import torch import gradio as gr from PIL import Image from huggingface_hub import hf_hub_download from safetensors.torch import load_file from timm.data.transforms_factory import create_transform from timm.models import create_model # --------------------------------------------------------------------------- # Model registry # --------------------------------------------------------------------------- # GMACs come from the authors' measured sweeps (results/RESULTS.md in the # upstream repo): the threshold=0 row is the full two-round cost, the # threshold=10 row (every image exits after round 1) is the round-1 cost. VARIANTS = { "160 → 384 · KD (84.9% top-1)": { "repo": "NCPS/progresvit-deit-s-160-384-kd-imagenet1k", "sizes": (160, 384), "gmacs": (0.615, 16.152), "top1": (73.940, 84.894), "amp": True, "threshold": 0.226, }, "160 → 384 (83.7% top-1)": { "repo": "NCPS/progresvit-deit-s-160-384-imagenet1k", "sizes": (160, 384), "gmacs": (0.615, 16.152), "top1": (70.616, 83.714), "amp": False, "threshold": 0.267, }, "192 → 240 · KD (83.8% top-1)": { "repo": "NCPS/progresvit-deit-s-192-240-kd-imagenet1k", "sizes": (192, 240), "gmacs": (0.912, 6.267), "top1": (76.018, 83.794), "amp": False, "threshold": 0.209, }, "192 → 240 (82.2% top-1)": { "repo": "NCPS/progresvit-deit-s-192-240-imagenet1k", "sizes": (192, 240), "gmacs": (0.912, 6.267), "top1": (73.238, 82.202), "amp": False, "threshold": 0.356, }, } DEFAULT_VARIANT = "160 → 384 · KD (84.9% top-1)" DEFAULT_THRESHOLD = VARIANTS[DEFAULT_VARIANT]["threshold"] PROGRESS_STAGES = (3, 6) # attention heads active in round 1 / round 2 CACHE_VERSION = 1 with open(os.path.join(os.path.dirname(__file__), "imagenet_classes.json")) as f: _ID2LABEL = json.load(f) IMAGENET_CLASSES = [_ID2LABEL[str(i)] for i in range(1000)] MODELS = {} TRANSFORMS = {} CROPS = {} for _name, _spec in VARIANTS.items(): _cfg = json.load(open(hf_hub_download(_spec["repo"], "config.json"))) _model = create_model( "progresvit", pretrained=False, num_classes=_cfg["num_classes"], **_cfg["model_args"], ) _state = load_file(hf_hub_download(_spec["repo"], "model.safetensors")) _model.load_state_dict(_state, strict=True) _pc = _cfg["pretrained_cfg"] TRANSFORMS[_name] = create_transform( input_size=tuple(_pc["input_size"]), is_training=False, interpolation=_pc["interpolation"], mean=tuple(_pc["mean"]), std=tuple(_pc["std"]), crop_pct=_pc["crop_pct"], crop_mode=_pc["crop_mode"], crop_border_pixels=0, use_prefetcher=False, ) CROPS[_name] = int(_pc["input_size"][-1]) MODELS[_name] = _model.eval().to("cuda") print(f"loaded {_name} from {_spec['repo']} (eval crop {CROPS[_name]})", flush=True) def _topk_dict(logits: torch.Tensor, k: int = 5) -> dict: probs = logits.float().softmax(dim=-1)[0] values, indices = probs.topk(k) return {IMAGENET_CLASSES[int(i)]: float(v) for v, i in zip(values, indices)} @spaces.GPU(duration=15) def classify( image: Image.Image, variant: str = DEFAULT_VARIANT, threshold: float = DEFAULT_THRESHOLD, ) -> tuple: """Classify an image with ProgResViT's progressive, input-adaptive rounds. Args: image: input photograph to classify against the 1000 ImageNet-1K classes. variant: which ProgResViT DeiT-S checkpoint to use (resolution schedule and whether it was trained with knowledge distillation). threshold: routing threshold on the round-1 top-10 prediction entropy. The image exits after the cheap first round when its entropy falls below this value; higher values exit more images and save more compute. Returns: A tuple of (final top-5 prediction, routing report in markdown, round-1 top-5 prediction, round-2 top-5 prediction). """ if image is None: raise gr.Error("Please provide an image.") spec = VARIANTS[variant] model = MODELS[variant] sizes = spec["sizes"] g1, g2 = spec["gmacs"] x = TRANSFORMS[variant](image.convert("RGB")).unsqueeze(0).to("cuda") started = time.perf_counter() with torch.inference_mode(): if spec["amp"]: ctx = torch.autocast("cuda", dtype=torch.bfloat16) else: ctx = torch.autocast("cuda", enabled=False) with ctx: tokens1, logits1 = model._forward_stage( x, 0, None, PROGRESS_STAGES, sizes ) _, logits2 = model._forward_stage( x, 1, tokens1, PROGRESS_STAGES, sizes ) entropy = float(model.entropy(logits1.float())[0, 0]) elapsed = time.perf_counter() - started exited_early = entropy < threshold final_logits = logits1 if exited_early else logits2 used_gmacs = g1 if exited_early else g2 saving = 100.0 * (1.0 - used_gmacs / g2) round1 = _topk_dict(logits1) round2 = _topk_dict(logits2) final = _topk_dict(final_logits) if exited_early: decision = ( f"**Exited after round 1.** Entropy `{entropy:.3f}` is below the " f"threshold `{threshold:.3f}`, so the {sizes[1]} px round was skipped." ) else: decision = ( f"**Continued to round 2.** Entropy `{entropy:.3f}` is at or above the " f"threshold `{threshold:.3f}`, so round 1's tokens were recycled and " f"refined at {sizes[1]} px." ) report = f"""### Routing {decision} | | Round 1 | Round 2 | This image | |---|---|---|---| | Input resolution | {sizes[0]} px | {sizes[1]} px | **{sizes[0] if exited_early else sizes[1]} px** | | Active attention heads | {PROGRESS_STAGES[0]} / 6 | {PROGRESS_STAGES[1]} / 6 | **{PROGRESS_STAGES[0] if exited_early else PROGRESS_STAGES[1]} / 6** | | Cumulative GMACs | {g1:.3f} | {g2:.3f} | **{used_gmacs:.3f}** | | ImageNet top-1 if always stopped here | {spec['top1'][0]:.2f}% | {spec['top1'][1]:.2f}% | — | Compute saved versus always running both rounds: **{saving:.1f}%** · inference {elapsed * 1000:.0f} ms """ return final, report, round1, round2 CSS = """ #col-container { max-width: 1180px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ # Ordered so the first rows tell the story: `red_fox` stays uncertain after round 1 # (which calls it a kit fox) and gets corrected in round 2, while `acoustic_guitar` # is confident enough to exit after the cheap first round. EXAMPLES = [ ["examples/red_fox.jpg"], ["examples/acoustic_guitar.jpg"], ["examples/husky_dog.jpg"], ["examples/pizza_board.jpg"], ["examples/bird_kingfisher.jpg"], ["examples/chameleon.jpg"], ["examples/hot_air_balloon.jpg"], ["examples/vintage_camera.jpg"], ["examples/library_interior.jpg"], ["examples/spiral_staircase.jpg"], ["examples/monstera_plant.jpg"], ] with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """# ProgResViT — adaptive-compute image classification An input-adaptive Vision Transformer that classifies progressively: round 1 runs a **narrow** subnetwork on a **low-resolution** image, and only images whose prediction is still uncertain continue to round 2 at **higher resolution** with a **wider** subnetwork, reusing the tokens produced in round 1. [Paper](https://huggingface.co/papers/2609.03216) · [Code](https://github.com/ds-kiel/ProgResViT) · [Checkpoints](https://huggingface.co/NCPS) """ ) with gr.Row(): with gr.Column(): image = gr.Image(label="Image", type="pil", height=340) run = gr.Button("Classify", variant="primary") variant = gr.Dropdown( label="Checkpoint", choices=list(VARIANTS), value=DEFAULT_VARIANT, ) threshold = gr.Slider( label="Routing threshold (round-1 entropy)", minimum=0.0, maximum=2.0, step=0.001, value=DEFAULT_THRESHOLD, info="0 = always run both rounds · higher = exit more images early", ) with gr.Column(): final_out = gr.Label(label="Prediction", num_top_classes=5) report_out = gr.Markdown(label="Routing report") with gr.Accordion("Round-by-round predictions", open=False): with gr.Row(): round1_out = gr.Label(label="Round 1 (low-res, narrow)", num_top_classes=5) round2_out = gr.Label(label="Round 2 (high-res, wide)", num_top_classes=5) gr.Examples( examples=EXAMPLES, inputs=[image], outputs=[final_out, report_out, round1_out, round2_out], fn=classify, cache_examples=True, cache_mode="lazy", examples_per_page=12, ) def _sync_threshold(name: str) -> float: """Reset the routing threshold to the checkpoint's reported operating point.""" return VARIANTS[name]["threshold"] variant.change(_sync_threshold, inputs=variant, outputs=threshold) run.click( classify, inputs=[image, variant, threshold], outputs=[final_out, report_out, round1_out, round2_out], api_name="classify", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)