add app.py

Dockerfile

@@ -0,0 +1,35 @@
+FROM python:3.11-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PORT=7860 \
+    HF_HOME=/home/user/.cache/huggingface \
+    TORCH_HOME=/home/user/.cache/torch \
+    OMP_NUM_THREADS=1 \
+    MKL_NUM_THREADS=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    ca-certificates \
+    libgomp1 \
+    git \
+ && rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -u 1000 user
+USER user
+
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user:user requirements.txt .
+RUN pip install --upgrade pip && pip install -r requirements.txt
+
+COPY --chown=user:user . .
+
+RUN git clone --depth 1 https://github.com/facebookresearch/dinov2.git
+
+EXPOSE 7860
+
+CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port ${PORT:-7860} --workers 1"]

app.py

@@ -0,0 +1,648 @@
+# app.py
+import io
+import os
+import uuid
+import threading
+import hashlib
+from contextvars import ContextVar
+from typing import Optional, Dict, Any, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from fastapi import FastAPI, UploadFile, File, Query, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+# ============================================================
+# Config
+# ============================================================
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+MODEL_IMG_SIZE = 518
+ALLOW_ORIGINS = os.environ.get("ALLOW_ORIGINS", "*").split(",")
+
+# RAD-DINO checkpoint en HF
+RAD_BACKBONE_REPO_ID = "microsoft/rad-dino"
+RAD_BACKBONE_FILENAME = "backbone_compatible.safetensors"
+
+# Heads
+RAD_HEAD_CKPT_PATH = os.environ.get("RAD_HEAD_CKPT_PATH", "rad_dino_chestmnist_head.pt")
+DINO_HEAD_CKPT_PATH = os.environ.get("DINO_HEAD_CKPT_PATH", "dino_chestmnist_head.pt")
+
+# Normalización
+RAD_MEAN = torch.tensor([0.5307, 0.5307, 0.5307], dtype=torch.float32).view(3, 1, 1)
+RAD_STD  = torch.tensor([0.2583, 0.2583, 0.2583], dtype=torch.float32).view(3, 1, 1)
+
+# DINOv2 usual / ImageNet normalization
+DINO_MEAN = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(3, 1, 1)
+DINO_STD  = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(3, 1, 1)
+
+DEFAULT_LABEL_NAMES = [
+    "atelectasis", "cardiomegaly", "effusion", "infiltration",
+    "mass", "nodule", "pneumonia", "pneumothorax",
+    "consolidation", "edema", "emphysema", "fibrosis",
+    "pleural", "hernia"
+]
+
+MODEL_CONFIGS = {
+    "rad-dino": {
+        "backbone_type": "rad-dino",
+        "head_ckpt_path": RAD_HEAD_CKPT_PATH,
+        "model_name": "rad-dino-chestmnist",
+        "mean": RAD_MEAN,
+        "std": RAD_STD,
+    },
+    "dino": {
+        "backbone_type": "dino",
+        "head_ckpt_path": DINO_HEAD_CKPT_PATH,
+        "model_name": "dino-chestmnist",
+        "mean": DINO_MEAN,
+        "std": DINO_STD,
+    },
+}
+
+
+# ============================================================
+# Model definitions
+# ============================================================
+class MedicalHead(nn.Module):
+    def __init__(self, in_dim: int = 768, num_classes: int = 14, dropout: float = 0.1):
+        super().__init__()
+        self.drop = nn.Dropout(dropout)
+        self.fc = nn.Linear(in_dim, num_classes)
+
+    def forward(self, cls_token: torch.Tensor) -> torch.Tensor:
+        return self.fc(self.drop(cls_token))
+
+
+def round_tensor(t: torch.Tensor, decimals: int = 4) -> torch.Tensor:
+    s = 10 ** decimals
+    return torch.round(t * s) / s
+
+
+def preprocess_pil(pil_img: Image.Image, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
+    img = pil_img.convert("RGB").resize((MODEL_IMG_SIZE, MODEL_IMG_SIZE), Image.BICUBIC)
+    arr = np.array(img).astype("float32") / 255.0
+    x = torch.from_numpy(arr).permute(2, 0, 1)  # [3,H,W]
+    x = (x - mean) / std
+    return x.unsqueeze(0)  # [1,3,H,W]
+
+
+# ============================================================
+# Build backbones
+# ============================================================
+def ensure_local_dinov2_repo():
+    if not os.path.exists("./dinov2"):
+        raise FileNotFoundError(
+            "No encontré ./dinov2. Clona el repo primero con:\n"
+            "git clone https://github.com/facebookresearch/dinov2.git"
+        )
+
+
+def disable_fused_attn(model: nn.Module):
+    for blk in model.blocks:
+        if hasattr(blk.attn, "fused_attn"):
+            blk.attn.fused_attn = False
+
+
+def build_dinov2_backbone() -> nn.Module:
+    ensure_local_dinov2_repo()
+    model = torch.hub.load("./dinov2", "dinov2_vitb14", source="local")
+    model.eval().to(DEVICE)
+    disable_fused_attn(model)
+    return model
+
+
+def build_rad_dino_backbone() -> nn.Module:
+    model = build_dinov2_backbone()
+
+    backbone_path = hf_hub_download(
+        repo_id=RAD_BACKBONE_REPO_ID,
+        filename=RAD_BACKBONE_FILENAME
+    )
+    state = load_file(backbone_path)
+    model.load_state_dict(state, strict=True)
+    model.eval().to(DEVICE)
+    disable_fused_attn(model)
+    return model
+
+
+def build_head(head_ckpt_path: str) -> tuple[nn.Module, Dict[str, Any], List[str]]:
+    ckpt = torch.load(head_ckpt_path, map_location=DEVICE)
+    label_names = ckpt.get("label_names", DEFAULT_LABEL_NAMES)
+    num_classes = len(label_names)
+
+    head = MedicalHead(in_dim=768, num_classes=num_classes, dropout=0.1).to(DEVICE)
+    head.load_state_dict(ckpt["head_state_dict"])
+    head.eval()
+
+    return head, ckpt, label_names
+
+
+def build_model_bundle(model_key: str, cfg: Dict[str, Any]) -> Dict[str, Any]:
+    if cfg["backbone_type"] == "rad-dino":
+        backbone = build_rad_dino_backbone()
+    elif cfg["backbone_type"] == "dino":
+        backbone = build_dinov2_backbone()
+    else:
+        raise ValueError(f"backbone_type desconocido: {cfg['backbone_type']}")
+
+    head, ckpt, label_names = build_head(cfg["head_ckpt_path"])
+
+    bundle = {
+        "key": model_key,
+        "model_name": cfg["model_name"],
+        "backbone_type": cfg["backbone_type"],
+        "backbone": backbone,
+        "head": head,
+        "head_ckpt": ckpt,
+        "label_names": label_names,
+        "mean": cfg["mean"],
+        "std": cfg["std"],
+        "num_layers": len(backbone.blocks),
+        "num_heads": getattr(backbone.blocks[0].attn, "num_heads", None),
+        "current": {
+            "hash": None,
+            "attention_cls_full": None,
+            "logit_lens_full": None,
+        },
+        "results": {},
+        "lock": threading.Lock(),
+    }
+    return bundle
+
+
+# ============================================================
+# Hook registration per model
+# ============================================================
+def register_hooks(bundle: Dict[str, Any]):
+    _attn_in_var: ContextVar[Optional[list]] = ContextVar(
+        f"_attn_in_var_{bundle['key']}", default=None
+    )
+    _tok_var: ContextVar[Optional[list]] = ContextVar(
+        f"_tok_var_{bundle['key']}", default=None
+    )
+
+    def _save_attn_input(module, inp):
+        lst = _attn_in_var.get()
+        if lst is None:
+            return
+        if len(inp) == 0 or not torch.is_tensor(inp[0]):
+            return
+        # input to attn: [B, N, D]
+        lst.append(inp[0].detach())
+
+    def _save_block_out(module, inp, out):
+        lst = _tok_var.get()
+        if lst is None:
+            return
+        if torch.is_tensor(out):
+            # block output: [B, N, D]
+            lst.append(out.detach())
+
+    attn_hooks = []
+    tok_hooks = []
+
+    for blk in bundle["backbone"].blocks:
+        if not hasattr(blk, "attn"):
+            raise RuntimeError(f"No encontré blk.attn en backbone {bundle['key']}")
+        attn_hooks.append(blk.attn.register_forward_pre_hook(_save_attn_input))
+        tok_hooks.append(blk.register_forward_hook(_save_block_out))
+
+    bundle["_attn_in_var"] = _attn_in_var
+    bundle["_tok_var"] = _tok_var
+    bundle["_attn_hooks"] = attn_hooks
+    bundle["_tok_hooks"] = tok_hooks
+
+
+# ============================================================
+# Build all models
+# ============================================================
+MODELS: Dict[str, Dict[str, Any]] = {
+    key: build_model_bundle(key, cfg)
+    for key, cfg in MODEL_CONFIGS.items()
+}
+
+for _bundle in MODELS.values():
+    register_hooks(_bundle)
+
+for key, bundle in MODELS.items():
+    print(f"[server] model_key={key}")
+    print(f"[server] model_name={bundle['model_name']}")
+    print(f"[server] backbone_type={bundle['backbone_type']} device={DEVICE}")
+    print(f"[server] head_ckpt={MODEL_CONFIGS[key]['head_ckpt_path']}")
+    print(f"[server] num_layers={bundle['num_layers']} num_heads={bundle['num_heads']}")
+    print(f"[server] num_classes={len(bundle['label_names'])}")
+    if "best_val_auc" in bundle["head_ckpt"]:
+        print(f"[server] checkpoint best_val_auc={bundle['head_ckpt']['best_val_auc']:.4f}")
+
+
+# ============================================================
+# Inference helpers
+# ============================================================
+@torch.no_grad()
+def extract_cls(backbone: nn.Module, images: torch.Tensor) -> torch.Tensor:
+    feats = backbone.forward_features(images)
+    if "x_norm_clstoken" not in feats:
+        raise RuntimeError("forward_features no devolvió 'x_norm_clstoken'.")
+    return feats["x_norm_clstoken"]
+
+
+@torch.no_grad()
+def compute_logit_lens_from_tokens(tokens_per_layer: List[torch.Tensor], head: nn.Module):
+    logits_list = []
+    probs_list = []
+
+    for x_l in tokens_per_layer:
+        # x_l: [B, N, D]
+        cls_l = x_l[:, 0]            # [B, D]
+        logits_l = head(cls_l)       # [B, C]
+        probs_l = torch.sigmoid(logits_l)
+
+        logits_list.append(logits_l.detach().cpu())
+        probs_list.append(probs_l.detach().cpu())
+
+    logits_per_layer = torch.stack(logits_list, dim=0)  # [L, B, C]
+    probs_per_layer  = torch.stack(probs_list, dim=0)   # [L, B, C]
+    return logits_per_layer, probs_per_layer
+
+
+@torch.no_grad()
+def compute_cls_attention_from_inputs(backbone: nn.Module, attn_inputs: List[torch.Tensor]):
+    """
+    Reconstruct CLS->tokens attention per layer from the input to attention.
+    Returns list of [B, H, N], one per layer.
+    """
+    cls_attn_per_layer = []
+
+    for blk, x in zip(backbone.blocks, attn_inputs):
+        x = x.to(DEVICE)  # [B, N, D]
+
+        B, N, C = x.shape
+        num_heads = blk.attn.num_heads
+        head_dim = C // num_heads
+
+        qkv = blk.attn.qkv(x)  # [B, N, 3*C]
+        qkv = qkv.reshape(B, N, 3, num_heads, head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # [B, H, N, Hd]
+
+        attn = (q @ k.transpose(-2, -1)) * blk.attn.scale
+        attn = attn.softmax(dim=-1)
+
+        cls_attn = attn[:, :, 0, :].detach().cpu()  # [B, H, N]
+        cls_attn_per_layer.append(cls_attn)
+
+    return cls_attn_per_layer
+
+
+def analyze_image(bundle: Dict[str, Any], pil_img: Image.Image) -> Dict[str, Any]:
+    x = preprocess_pil(pil_img, bundle["mean"], bundle["std"]).to(DEVICE)
+
+    attn_inputs = []
+    layer_tokens = []
+
+    tok_token = bundle["_tok_var"].set(layer_tokens)
+    attn_token = bundle["_attn_in_var"].set(attn_inputs)
+
+    try:
+        with torch.no_grad():
+            with bundle["lock"]:
+                cls_final = extract_cls(bundle["backbone"], x)   # [1, 768]
+                logits_final = bundle["head"](cls_final)         # [1, C]
+
+        probs_final = torch.sigmoid(logits_final)[0].detach().cpu()
+        probs_final = round_tensor(probs_final, 6)
+
+        if len(layer_tokens) == 0:
+            raise RuntimeError("No se capturaron tokens por capa.")
+        if len(attn_inputs) == 0:
+            raise RuntimeError("No se capturaron entradas a atención por capa.")
+
+        logits_by_layer, probs_by_layer = compute_logit_lens_from_tokens(
+            layer_tokens, bundle["head"]
+        )
+        attn_maps = compute_cls_attention_from_inputs(bundle["backbone"], attn_inputs)
+
+        # ----------------------------------------------------
+        # attention_cls_full
+        # ----------------------------------------------------
+        attn_maps2 = [a.squeeze(0) for a in attn_maps]  # list of [H, N]
+
+        attn_serializable_all = []
+        attn_serializable_patches = []
+
+        for layer in attn_maps2:
+            layer_all = []
+            layer_patches = []
+            for head in layer:
+                head = round_tensor(head, 4)     # [N]
+                layer_all.append(head.tolist())
+                layer_patches.append(head[1:].tolist())  # remove CLS->CLS
+            attn_serializable_all.append(layer_all)
+            attn_serializable_patches.append(layer_patches)
+
+        num_tokens_all = len(attn_serializable_all[0][0])
+        num_patch_tokens = len(attn_serializable_patches[0][0])
+
+        export_attn = {
+            "model": bundle["model_name"],
+            "attention_type": "cls_only",
+            "num_layers": len(attn_serializable_all),
+            "num_heads": len(attn_serializable_all[0]),
+            "num_tokens_all": num_tokens_all,
+            "num_patch_tokens": num_patch_tokens,
+            "cls_index": 0,
+            "attention_cls_to_all_tokens": attn_serializable_all,
+            "attention_cls_to_patches": attn_serializable_patches,
+        }
+
+        # ----------------------------------------------------
+        # logit_lens_full
+        # ----------------------------------------------------
+        export_logit = {
+            "model": bundle["model_name"],
+            "num_layers": int(logits_by_layer.shape[0]),
+            "num_classes": int(logits_by_layer.shape[-1]),
+            "class_names": bundle["label_names"],
+            "checkpoint_best_val_auc": bundle["head_ckpt"].get("best_val_auc", None),
+            "final_probs": probs_final.tolist(),
+            "logits": [],
+            "probs_by_layer": [],
+        }
+
+        for l in range(logits_by_layer.shape[0]):
+            v_logits = round_tensor(logits_by_layer[l, 0], 4)
+            v_probs = round_tensor(probs_by_layer[l, 0], 6)
+            export_logit["logits"].append(v_logits.tolist())
+            export_logit["probs_by_layer"].append(v_probs.tolist())
+
+        return {
+            "attention_cls_full": export_attn,
+            "logit_lens_full": export_logit,
+        }
+
+    finally:
+        bundle["_tok_var"].reset(tok_token)
+        bundle["_attn_in_var"].reset(attn_token)
+        layer_tokens.clear()
+        attn_inputs.clear()
+
+
+# ============================================================
+# FastAPI app
+# ============================================================
+app = FastAPI(title="ChestMNIST Explainer API (RAD-DINO + DINO)", version="2.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=ALLOW_ORIGINS if ALLOW_ORIGINS != ["*"] else ["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def _no_store(resp: JSONResponse) -> JSONResponse:
+    resp.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, max-age=0"
+    resp.headers["Pragma"] = "no-cache"
+    return resp
+
+
+def get_model_bundle(model_key: str) -> Dict[str, Any]:
+    if model_key not in MODELS:
+        raise HTTPException(
+            status_code=404,
+            detail=f"Unknown model_key '{model_key}'. Available: {list(MODELS.keys())}"
+        )
+    return MODELS[model_key]
+
+
+# ============================================================
+# Root / health
+# ============================================================
+@app.get("/")
+def root():
+    return {
+        "status": "ok",
+        "device": DEVICE,
+        "available_models": list(MODELS.keys()),
+        "image_size": MODEL_IMG_SIZE,
+    }
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "device": DEVICE,
+        "available_models": list(MODELS.keys()),
+        "models": {
+            key: {
+                "model": bundle["model_name"],
+                "num_layers": bundle["num_layers"],
+                "num_heads": bundle["num_heads"],
+                "num_classes": len(bundle["label_names"]),
+                "class_names": bundle["label_names"],
+                "checkpoint_best_val_auc": bundle["head_ckpt"].get("best_val_auc", None),
+                "has_current": bundle["current"]["attention_cls_full"] is not None,
+            }
+            for key, bundle in MODELS.items()
+        }
+    }
+
+
+@app.get("/health/{model_key}")
+def health_model(model_key: str):
+    bundle = get_model_bundle(model_key)
+    return {
+        "status": "ok",
+        "device": DEVICE,
+        "model_key": model_key,
+        "model": bundle["model_name"],
+        "image_size": MODEL_IMG_SIZE,
+        "num_layers": bundle["num_layers"],
+        "num_heads": bundle["num_heads"],
+        "num_classes": len(bundle["label_names"]),
+        "class_names": bundle["label_names"],
+        "checkpoint_best_val_auc": bundle["head_ckpt"].get("best_val_auc", None),
+        "has_current": bundle["current"]["attention_cls_full"] is not None,
+    }
+
+
+# ============================================================
+# Legacy analyze with stored jobs
+# ============================================================
+@app.post("/analyze/{model_key}")
+async def analyze(
+    model_key: str,
+    file: UploadFile = File(...),
+    store: int = Query(0, description="1 => guarda resultados y entrega endpoints /results/{model_key}/{id}/..."),
+):
+    bundle = get_model_bundle(model_key)
+
+    if not file.content_type or not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Please upload an image file.")
+
+    raw = await file.read()
+    try:
+        img = Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Could not decode image.")
+
+    try:
+        out = analyze_image(bundle, img)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Model inference failed: {e}")
+
+    if store == 1:
+        job_id = str(uuid.uuid4())
+        bundle["results"][job_id] = out
+        return {
+            "model_key": model_key,
+            "job_id": job_id,
+            "endpoints": {
+                "attention_cls_full": f"/results/{model_key}/{job_id}/attention_cls_full.json",
+                "logit_lens_full": f"/results/{model_key}/{job_id}/logit_lens_full.json",
+            }
+        }
+
+    return out
+
+
+@app.get("/results/{model_key}/{job_id}/attention_cls_full.json")
+def get_attention(model_key: str, job_id: str):
+    bundle = get_model_bundle(model_key)
+    if job_id not in bundle["results"]:
+        raise HTTPException(status_code=404, detail="job_id not found")
+    return _no_store(JSONResponse(bundle["results"][job_id]["attention_cls_full"]))
+
+
+@app.get("/results/{model_key}/{job_id}/logit_lens_full.json")
+def get_logit(model_key: str, job_id: str):
+    bundle = get_model_bundle(model_key)
+    if job_id not in bundle["results"]:
+        raise HTTPException(status_code=404, detail="job_id not found")
+    return _no_store(JSONResponse(bundle["results"][job_id]["logit_lens_full"]))
+
+
+# ============================================================
+# Preferred: current endpoints per model
+# ============================================================
+@app.post("/analyze_current/{model_key}")
+async def analyze_current(model_key: str, file: UploadFile = File(...)):
+    bundle = get_model_bundle(model_key)
+
+    if not file.content_type or not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Please upload an image file.")
+
+    raw = await file.read()
+    img_hash = hashlib.sha256(raw).hexdigest()
+
+    if bundle["current"]["hash"] == img_hash and bundle["current"]["attention_cls_full"] is not None:
+        return {"status": "unchanged", "hash": img_hash, "model_key": model_key}
+
+    try:
+        img = Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Could not decode image.")
+
+    try:
+        out = analyze_image(bundle, img)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Model inference failed: {e}")
+
+    bundle["current"]["hash"] = img_hash
+    bundle["current"]["attention_cls_full"] = out["attention_cls_full"]
+    bundle["current"]["logit_lens_full"] = out["logit_lens_full"]
+
+    return {"status": "ok", "hash": img_hash, "model_key": model_key}
+
+
+@app.get("/{model_key}/attention_cls_full.json")
+def current_attention(model_key: str):
+    bundle = get_model_bundle(model_key)
+
+    if bundle["current"]["attention_cls_full"] is None:
+        raise HTTPException(
+            status_code=404,
+            detail=f"No current attention file for '{model_key}'. POST /analyze_current/{model_key} first."
+        )
+    return _no_store(JSONResponse(bundle["current"]["attention_cls_full"]))
+
+
+@app.get("/{model_key}/logit_lens_full.json")
+def current_logit(model_key: str):
+    bundle = get_model_bundle(model_key)
+
+    if bundle["current"]["logit_lens_full"] is None:
+        raise HTTPException(
+            status_code=404,
+            detail=f"No current logit file for '{model_key}'. POST /analyze_current/{model_key} first."
+        )
+    return _no_store(JSONResponse(bundle["current"]["logit_lens_full"]))
+
+
+# ============================================================
+# Optional backward-compatible aliases for RAD-DINO
+# ============================================================
+@app.post("/analyze_current")
+async def analyze_current_rad_default(file: UploadFile = File(...)):
+    return await analyze_current("rad-dino", file)
+
+
+@app.get("/attention_cls_full.json")
+def current_attention_rad_default():
+    return current_attention("rad-dino")
+
+
+@app.get("/logit_lens_full.json")
+def current_logit_rad_default():
+    return current_logit("rad-dino")
+
+
+@app.post("/analyze")
+async def analyze_rad_default(
+    file: UploadFile = File(...),
+    store: int = Query(0, description="1 => guarda resultados"),
+):
+    return await analyze("rad-dino", file, store)
+
+
+# ============================================================
+# Smoke test
+# ============================================================
+def smoke_test_local_image(image_path: str, model_key: str = "rad-dino"):
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"No existe la imagen: {image_path}")
+
+    bundle = get_model_bundle(model_key)
+    img = Image.open(image_path).convert("RGB")
+    out = analyze_image(bundle, img)
+
+    print(f"\n[smoke test] model_key={model_key} OK")
+    print("[smoke test] capas:", out["attention_cls_full"]["num_layers"])
+    print("[smoke test] heads:", out["attention_cls_full"]["num_heads"])
+    print("[smoke test] patch tokens:", out["attention_cls_full"]["num_patch_tokens"])
+
+    final_probs = out["logit_lens_full"]["final_probs"]
+    pairs = sorted(zip(bundle["label_names"], final_probs), key=lambda t: t[1], reverse=True)
+
+    print("\nTop-5 predicciones:")
+    for name, p in pairs[:5]:
+        print(f"  {name:<15} {p:.4f}")
+
+
+if __name__ == "__main__":
+    test_path = os.environ.get("TEST_IMAGE_PATH", "").strip()
+    test_model = os.environ.get("TEST_MODEL_KEY", "rad-dino").strip()
+
+    if test_path:
+        smoke_test_local_image(test_path, test_model)
+
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)

dino_chestmnist_head.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12e5409624f3333ac9cc00402d0d84d2a58bfd62aac0462dc3fe33bbda14c1fa
+size 133769

rad_dino_chestmnist_head.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b570912198f3cdf00fe2b01b7cae88cdc13539245173e7a13f665f4a241f5873
+size 45613

requirements.txt

@@ -0,0 +1,14 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+
+fastapi==0.115.6
+uvicorn[standard]==0.30.6
+python-multipart==0.0.9
+
+torch==2.5.1
+torchvision==0.20.1
+timm==0.9.16
+
+pillow==10.4.0
+numpy==2.1.3
+huggingface-hub==0.24.7
+safetensors==0.4.5