Add Gradio app for inference and model selection

- Introduced app.py for a Gradio interface allowing users to upload images and select between baseline and GNN models for inference.
- Implemented ModelRegistry for caching loaded models and managing checkpoints.
- Added README_SPACE.md to guide users on setting up the Hugging Face Space with the necessary files and configurations.
- Created requirements-space.txt for dependencies specific to the Gradio app.
- Developed scripts for single-image inference and comparison of baseline vs GNN model metrics.
- Enhanced existing README.md with updated training and evaluation instructions for both models.

README.md

@@ -105,14 +105,25 @@ Important: the processing pipeline is strict chair-only (`primary_object_class:
 
 ### 2) Train/Evaluate Models
 
-This repo provides dataset + experiment configs. Connect model training to your entrypoints (for example `scripts/train.py` and `scripts/eval.py`).
+#### Baseline (CLIP/DINO + MLP)
 
-Typical flow:
+- Smoke run (quick sanity check):
 
-1. Run baseline experiments from `configs/experiment.yaml`.
-2. Run the proposed GNN experiment.
-3. Run ablation experiments.
-4. Compare all runs in one metrics table.
+  `python scripts/train_baseline.py --run-name smoke_baseline --mode smoke --device auto`
+
+- Full run:
+
+  `python scripts/train_baseline.py --run-name baseline_full --mode full --device auto`
+
+#### GNN (CLIP/DINO + bipartite native GNN)
+
+- GNN smoke run:
+
+  `python scripts/train_gnn.py --run-name smoke_gnn --mode smoke --device auto`
+
+- Compare baseline vs GNN (after both smoke runs complete):
+
+  `python scripts/compare_baseline_vs_gnn.py --baseline-run smoke_baseline --gnn-run smoke_gnn`
 
 ## Expected Outputs
 

README_SPACE.md

@@ -0,0 +1,45 @@
+## Hugging Face Space Setup
+
+### 1. Prepare Repository
+
+Push this project (including `app.py`, `requirements-space.txt`, and checkpoints) to a Git repository that you will connect to your Space.
+
+Recommended checkpoint locations:
+
+- Baseline: `outputs/smoke_baseline/best.pt`
+- GNN: `outputs/smoke_gnn/best.pt`
+
+### 2. Create a New Space
+
+1. On Hugging Face, create a new **Space**.
+2. Select **Gradio** as the SDK.
+3. Point the Space to this repository (or upload the files manually).
+
+### 3. Configure Runtime
+
+- Hardware: start with **CPU**; upgrade to GPU if needed.
+- Python dependencies:
+  - Set the Space to use `requirements-space.txt` as the dependency file.
+
+### 4. Local Test Before Pushing
+
+Create and activate a clean virtualenv, then run:
+
+```bash
+pip install -r requirements-space.txt
+python app.py
+```
+
+Open the printed URL in your browser, upload a chair image, and verify that:
+
+- Baseline and GNN models both run.
+- Switching model type updates which checkpoint is used.
+
+### 5. Expected Behavior in the Space
+
+- The app starts from `app.py`.
+- Default checkpoint paths:
+  - Baseline: `outputs/smoke_baseline/best.pt`
+  - GNN: `outputs/smoke_gnn/best.pt`
+- You can override checkpoint paths via the text box in the UI.
+

app.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict
+
+import gradio as gr
+import torch
+from PIL import Image
+
+from infer.pipeline import LoadedModel, load_checkpoint, predict_image
+
+
+DEFAULT_BASELINE_CKPT = Path("outputs/smoke_baseline/best.pt")
+DEFAULT_GNN_CKPT = Path("outputs/smoke_gnn/best.pt")
+
+
+class ModelRegistry:
+    def __init__(self) -> None:
+        self._cache: Dict[str, LoadedModel] = {}
+
+    def get(self, key: str, path: Path, device: torch.device) -> LoadedModel:
+        cached = self._cache.get(key)
+        if cached is not None:
+            return cached
+        loaded = load_checkpoint(path, device=device)
+        self._cache[key] = loaded
+        return loaded
+
+
+registry = ModelRegistry()
+
+
+def infer_gradio(
+    image: Image.Image,
+    model_type: str,
+    checkpoint_path: str,
+    top_k: int,
+    threshold: float,
+) -> Any:
+    if image is None:
+        return [], "No image provided."
+    path = Path(checkpoint_path).expanduser()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    try:
+        key = f"{model_type}:{path}"
+        loaded = registry.get(key, path, device)
+        # Save temp image to disk-agnostic path is unnecessary; use pipeline logic directly.
+        from infer.pipeline import predict_image as predict_from_loaded
+        from infer.pipeline import LoadedModel as _LM  # type: ignore
+
+        assert isinstance(loaded, _LM)
+        # Wrap PIL image call using temporary in-memory path emulation.
+        # For simplicity, reuse adapter directly.
+        result = loaded.adapter.predict(
+            image=image,
+            label_vocab=loaded.label_vocab,
+            top_k=top_k,
+            threshold=threshold,
+        )
+        rows = [
+            {"label": label, "score": round(score, 4), "positive@thr": label in result.positives}
+            for label, score in zip(result.labels, result.scores)
+        ]
+        message = f"Positives (>= {threshold:.2f}): {', '.join(result.positives) if result.positives else 'none'}"
+        return rows, message
+    except Exception as exc:  # noqa: BLE001
+        return [], f"Error: {exc}"
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Adjective-Aware Chair Attributes\nSelect baseline or GNN, upload an image, and view attribute scores.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Chair image")
+            model_type = gr.Radio(
+                choices=["baseline", "gnn"],
+                value="baseline",
+                label="Model type",
+            )
+            checkpoint = gr.Textbox(
+                value=str(DEFAULT_BASELINE_CKPT),
+                label="Checkpoint path",
+            )
+
+            def _sync_ckpt(choice: str) -> str:
+                return str(DEFAULT_GNN_CKPT if choice == "gnn" else DEFAULT_BASELINE_CKPT)
+
+            model_type.change(_sync_ckpt, inputs=model_type, outputs=checkpoint)
+
+            top_k = gr.Slider(1, 20, value=5, step=1, label="Top-K")
+            threshold = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Threshold")
+            run_btn = gr.Button("Run inference")
+
+        with gr.Column(scale=2):
+            table = gr.Dataframe(
+                headers=["label", "score", "positive@thr"],
+                datatype=["str", "number", "bool"],
+                label="Attribute scores",
+            )
+            message = gr.Markdown()
+
+    run_btn.click(
+        infer_gradio,
+        inputs=[image_input, model_type, checkpoint, top_k, threshold],
+        outputs=[table, message],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()
+

requirements-space.txt

@@ -0,0 +1,9 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+
+torch
+torchvision
+torchaudio
+timm
+Pillow
+gradio
+

requirements.txt

@@ -10,3 +10,4 @@ timm
 scikit-learn
 Pillow
 tqdm
+numpy

scripts/compare_baseline_vs_gnn.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_metrics(run_name: str, output_dir: Path) -> Dict[str, Any]:
+    metrics_path = output_dir / run_name / "metrics.json"
+    if not metrics_path.exists():
+        raise RuntimeError(f"Metrics not found: {metrics_path}")
+    with metrics_path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Compare baseline vs GNN runs on test metrics.")
+    parser.add_argument("--baseline-run", type=str, required=True, help="Baseline run_name (e.g. smoke_baseline).")
+    parser.add_argument("--gnn-run", type=str, required=True, help="GNN run_name (e.g. smoke_gnn).")
+    parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "outputs")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    baseline = load_metrics(args.baseline_run, args.output_dir)
+    gnn = load_metrics(args.gnn_run, args.output_dir)
+
+    def extract(row: Dict[str, Any]) -> Dict[str, float]:
+        test = row.get("test_metrics", {})
+        return {
+            "map": float(test.get("map", 0.0)),
+            "macro_f1": float(test.get("macro_f1", 0.0)),
+            "micro_f1": float(test.get("micro_f1", 0.0)),
+        }
+
+    b = extract(baseline)
+    g = extract(gnn)
+    delta = {k: g[k] - b[k] for k in b.keys()}
+
+    report = {
+        "baseline_run": args.baseline_run,
+        "gnn_run": args.gnn_run,
+        "baseline_test": b,
+        "gnn_test": g,
+        "delta": delta,
+    }
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()
+

scripts/infer.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import torch
+
+from infer.pipeline import load_checkpoint, predict_image
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run single-image inference for baseline or GNN checkpoints.")
+    parser.add_argument("--model-type", type=str, choices=["auto", "baseline", "gnn"], default="auto")
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=Path("outputs/smoke_baseline/best.pt"),
+        help="Path to checkpoint (.pt).",
+    )
+    parser.add_argument(
+        "--image",
+        type=Path,
+        required=True,
+        help="Path to input image.",
+    )
+    parser.add_argument("--top-k", type=int, default=5)
+    parser.add_argument("--threshold", type=float, default=0.5)
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"])
+    parser.add_argument("--benchmark", action="store_true", help="Print simple latency timing.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    loaded = load_checkpoint(args.checkpoint, device=device)
+
+    if args.benchmark:
+        import time
+
+        start = time.time()
+        result = predict_image(
+            image_path=args.image,
+            loaded=loaded,
+            top_k=args.top_k,
+            threshold=args.threshold,
+        )
+        elapsed = time.time() - start
+    else:
+        result = predict_image(
+            image_path=args.image,
+            loaded=loaded,
+            top_k=args.top_k,
+            threshold=args.threshold,
+        )
+        elapsed = None
+
+    payload: Dict[str, Any] = {
+        "image": str(args.image),
+        "top_k": args.top_k,
+        "threshold": args.threshold,
+        "labels": result.labels,
+        "scores": result.scores,
+        "positives": result.positives,
+    }
+    if elapsed is not None:
+        payload["latency_sec"] = elapsed
+
+    print(json.dumps(payload, indent=2))
+
+
+if __name__ == "__main__":
+    main()
+

scripts/train_baseline.py

@@ -96,7 +96,7 @@ def build_loaders(cfg: Dict[str, Any], batch_size: int, eval_batch_size: int, sm
 
     num_workers = int(cfg["dataset"].get("num_workers", 0))
     if smoke:
-        num_workers = min(num_workers, 2)
+        num_workers = min(num_workers, 8)
 
     train_loader = build_dataloader(
         dataset=train_ds,
@@ -189,8 +189,8 @@ def train() -> None:
 
     smoke = args.mode == "smoke"
     train_epochs = 2 if smoke else int(cfg["training"]["stage1"].get("epochs", 20))
-    batch_size = 8 if smoke else int(cfg["batching"].get("batch_size", 32))
-    eval_batch_size = 8 if smoke else int(cfg["batching"].get("eval_batch_size", 64))
+    batch_size = 4 if smoke else int(cfg["batching"].get("batch_size", 32))
+    eval_batch_size = 4 if smoke else int(cfg["batching"].get("eval_batch_size", 64))
 
     cuda_available = torch.cuda.is_available()
     selected_device = "cuda" if (args.device == "auto" and cuda_available) else args.device

scripts/train_gnn.py

@@ -0,0 +1,361 @@
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+from data.io_utils import write_json  # noqa: E402
+from train.config_utils import (  # noqa: E402
+    apply_dot_overrides,
+    deep_merge,
+    load_yaml,
+    parse_overrides_json,
+)
+from train.dataset import (  # noqa: E402
+    ChairAttributeDataset,
+    build_dataloader,
+    class_pos_weights,
+    load_json,
+    split_paths,
+)
+from train.encoders import build_backbones  # noqa: E402
+from train.graph_builder import build_bipartite_batch  # noqa: E402
+from train.graph_models import NativeGNNClassifier  # noqa: E402
+from train.losses import bce_logits_loss  # noqa: E402
+from train.metrics import compute_metrics  # noqa: E402
+
+
+def log(message: str) -> None:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    print(f"[{ts}] [train_gnn] {message}", flush=True)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Train GNN model (bipartite graph).")
+    parser.add_argument("--dataset-config", type=Path, default=REPO_ROOT / "configs" / "dataset.yaml")
+    parser.add_argument("--model-config", type=Path, default=REPO_ROOT / "configs" / "model.yaml")
+    parser.add_argument("--train-config", type=Path, default=REPO_ROOT / "configs" / "train.yaml")
+    parser.add_argument("--eval-config", type=Path, default=REPO_ROOT / "configs" / "eval.yaml")
+    parser.add_argument("--run-name", type=str, required=True)
+    parser.add_argument("--overrides-json", type=str, default="")
+    parser.add_argument("--mode", type=str, default="full", choices=["smoke", "full"])
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"])
+    return parser.parse_args()
+
+
+def set_seed(seed: int, deterministic: bool) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = deterministic
+    torch.backends.cudnn.benchmark = not deterministic
+
+
+def load_combined_config(args: argparse.Namespace) -> Dict[str, Any]:
+    dataset_cfg = load_yaml(args.dataset_config)
+    model_cfg = load_yaml(args.model_config)
+    train_cfg = load_yaml(args.train_config)
+    eval_cfg = load_yaml(args.eval_config)
+    combined = deep_merge(dataset_cfg, model_cfg)
+    combined = deep_merge(combined, train_cfg)
+    combined = deep_merge(combined, eval_cfg)
+    overrides = parse_overrides_json(args.overrides_json)
+    if overrides:
+        combined = apply_dot_overrides(combined, overrides)
+    return combined
+
+
+def build_loaders(cfg: Dict[str, Any], batch_size: int, eval_batch_size: int, smoke: bool) -> Tuple[Any, Any, Any, Dict[str, int]]:
+    processed_root = REPO_ROOT / cfg["dataset"]["processed_dir"]
+    label_vocab_path = processed_root / cfg["processing"]["label_vocab_file"]
+    label_freq_path = processed_root / cfg["processing"]["label_frequency_file"]
+    cache_dir = REPO_ROOT / cfg["dataset"]["root_dir"] / "images"
+    splits = split_paths(processed_root=processed_root, split_dir_name=cfg["processing"]["split_dir"])
+
+    label_vocab = load_json(label_vocab_path)
+    label_freq = load_json(label_freq_path)
+
+    train_ds = ChairAttributeDataset(splits["train"], label_vocab_path, REPO_ROOT, cache_dir)
+    val_ds = ChairAttributeDataset(splits["val"], label_vocab_path, REPO_ROOT, cache_dir)
+    test_ds = ChairAttributeDataset(splits["test"], label_vocab_path, REPO_ROOT, cache_dir)
+
+    if smoke:
+        # Cap samples for faster smoke training.
+        train_ds.samples = train_ds.samples[:512]
+        val_ds.samples = val_ds.samples[:128]
+        test_ds.samples = test_ds.samples[:128]
+
+    num_workers = int(cfg["dataset"].get("num_workers", 0))
+    if smoke:
+        num_workers = min(num_workers, 8)
+
+    train_loader = build_dataloader(
+        dataset=train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
+    )
+    val_loader = build_dataloader(
+        dataset=val_ds,
+        batch_size=eval_batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
+    )
+    test_loader = build_dataloader(
+        dataset=test_ds,
+        batch_size=eval_batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
+    )
+    return train_loader, val_loader, test_loader, {"vocab": label_vocab, "freq": label_freq}
+
+
+def build_backbone_and_gnn(
+    cfg: Dict[str, Any],
+    num_labels: int,
+    device: torch.device,
+) -> Tuple[torch.nn.Module, NativeGNNClassifier, int]:
+    use_clip = bool(cfg["encoders"]["clip"]["enabled"])
+    use_dino = bool(cfg["encoders"]["dino"]["enabled"])
+    clip_name = str(cfg["encoders"]["clip"].get("model_name", "vit_base_patch16_clip_224.openai"))
+    dino_name = str(cfg["encoders"]["dino"].get("model_name", "vit_base_patch14_dinov2.lvd142m"))
+    if clip_name == "ViT-B-32":
+        clip_name = "vit_base_patch16_clip_224.openai"
+    if dino_name == "vit_base_patch14_dinov2":
+        dino_name = "vit_base_patch14_dinov2.lvd142m"
+
+    clip_backbone, dino_backbone, feature_dim = build_backbones(
+        use_clip=use_clip,
+        use_dino=use_dino,
+        clip_model_name=clip_name,
+        dino_model_name=dino_name,
+        clip_pretrained=bool(cfg["encoders"]["clip"].get("pretrained", True)),
+        dino_pretrained=bool(cfg["encoders"]["dino"].get("pretrained", True)),
+    )
+
+    class BackboneWrapper(torch.nn.Module):
+        def __init__(self, clip_module: torch.nn.Module | None, dino_module: torch.nn.Module | None) -> None:
+            super().__init__()
+            self.clip_module = clip_module
+            self.dino_module = dino_module
+
+        def encode(self, images: List[Any], device: torch.device) -> torch.Tensor:
+            feats: List[torch.Tensor] = []
+            if self.clip_module is not None:
+                feats.append(self.clip_module.forward_pil(images, device=device))
+            if self.dino_module is not None:
+                feats.append(self.dino_module.forward_pil(images, device=device))
+            if len(feats) == 1:
+                return feats[0]
+            return torch.cat(feats, dim=1)
+
+    backbone = BackboneWrapper(clip_backbone, dino_backbone).to(device)
+
+    hidden_dims = [layer_cfg["out_dim"] for layer_cfg in cfg["gnn"]["layers"]]
+    dropout = float(cfg["gnn"].get("dropout", 0.2))
+    gnn_model = NativeGNNClassifier(
+        in_dim=feature_dim,
+        hidden_dims=hidden_dims,
+        num_attributes=num_labels,
+        dropout=dropout,
+    ).to(device)
+
+    return backbone, gnn_model, feature_dim
+
+
+def run_eval(
+    backbone: torch.nn.Module,
+    gnn_model: NativeGNNClassifier,
+    loader: Any,
+    device: torch.device,
+    pos_weight: torch.Tensor,
+) -> Dict[str, Any]:
+    backbone.eval()
+    gnn_model.eval()
+    all_logits: List[np.ndarray] = []
+    all_targets: List[np.ndarray] = []
+    losses: List[float] = []
+    with torch.no_grad():
+        for batch in loader:
+            targets = batch["targets"].to(device)
+            images = batch["images"]
+            feats = backbone.encode(images=images, device=device).unsqueeze(1)
+            graph = build_bipartite_batch(feats=feats, targets=targets)
+            logits = gnn_model(graph)
+            loss = bce_logits_loss(logits, targets, pos_weight=pos_weight)
+            losses.append(float(loss.item()))
+            all_logits.append(logits.detach().cpu().numpy())
+            all_targets.append(targets.detach().cpu().numpy())
+    logits_np = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, 1))
+    targets_np = np.concatenate(all_targets, axis=0) if all_targets else np.zeros((0, 1))
+    metrics = compute_metrics(targets=targets_np, logits=logits_np)
+    metrics["loss"] = float(np.mean(losses)) if losses else 0.0
+    return metrics
+
+
+def train() -> None:
+    args = parse_args()
+    log(f"Starting run_name={args.run_name} mode={args.mode} device_request={args.device}")
+    log(f"dataset_config={args.dataset_config}")
+    log(f"model_config={args.model_config}")
+    log(f"train_config={args.train_config}")
+    cfg = load_combined_config(args)
+    run_cfg = cfg.get("run", {})
+    seed = int(run_cfg.get("seed", 42))
+    deterministic = bool(run_cfg.get("deterministic", True))
+    set_seed(seed=seed, deterministic=deterministic)
+    log(f"Seed set to {seed}, deterministic={deterministic}")
+
+    smoke = args.mode == "smoke"
+    train_epochs = 1 if smoke else int(cfg["training"]["stage1"].get("epochs", 20))
+    batch_size = 4 if smoke else int(cfg["batching"].get("batch_size", 32))
+    eval_batch_size = 4 if smoke else int(cfg["batching"].get("eval_batch_size", 64))
+
+    cuda_available = torch.cuda.is_available()
+    selected_device = "cuda" if (args.device == "auto" and cuda_available) else args.device
+    if args.device == "auto" and not cuda_available:
+        selected_device = "cpu"
+    if selected_device == "cuda" and not cuda_available:
+        raise RuntimeError(
+            "CUDA requested but not available in torch build/environment. "
+            "Install CUDA-enabled PyTorch in the active virtualenv."
+        )
+    device = torch.device(selected_device)
+    log(
+        f"Resolved device={device.type}, torch_cuda_available={cuda_available}, "
+        f"torch_version={torch.__version__}, torch_cuda={torch.version.cuda}"
+    )
+    log("Building dataloaders from processed manifests...")
+    train_loader, val_loader, test_loader, label_info = build_loaders(
+        cfg=cfg,
+        batch_size=batch_size,
+        eval_batch_size=eval_batch_size,
+        smoke=smoke,
+    )
+    log(
+        f"Dataloader sizes -> train_batches={len(train_loader)} val_batches={len(val_loader)} "
+        f"test_batches={len(test_loader)}"
+    )
+    label_vocab = label_info["vocab"]
+    label_freq = label_info["freq"]
+    num_labels = len(label_vocab)
+    log(f"Label space size: {num_labels}")
+    log("Building encoder backbones and GNN head (first run may download pretrained weights)...")
+    backbone, gnn_model, _ = build_backbone_and_gnn(cfg=cfg, num_labels=num_labels, device=device)
+    log("Model construction complete")
+    pos_weight = class_pos_weights(label_frequencies=label_freq, label_vocab=label_vocab).to(device)
+
+    lr = float(cfg["training"]["stage1"].get("lr_head", 1e-3))
+    weight_decay = float(cfg["optimization"].get("weight_decay", 1e-4))
+    optimizer = torch.optim.AdamW(
+        list(gnn_model.parameters()),
+        lr=lr,
+        weight_decay=weight_decay,
+    )
+
+    output_root = REPO_ROOT / run_cfg.get("output_dir", "outputs")
+    run_dir = output_root / args.run_name
+    run_dir.mkdir(parents=True, exist_ok=True)
+
+    history: List[Dict[str, Any]] = []
+    best_map = -1.0
+    best_path = run_dir / "best.pt"
+    log(f"Training for epochs={train_epochs}, batch_size={batch_size}, eval_batch_size={eval_batch_size}")
+    for epoch in range(1, train_epochs + 1):
+        backbone.eval()
+        gnn_model.train()
+        train_losses: List[float] = []
+        for batch in tqdm(train_loader, desc=f"{args.run_name} epoch {epoch}/{train_epochs}"):
+            targets = batch["targets"].to(device)
+            images = batch["images"]
+            feats = backbone.encode(images=images, device=device).unsqueeze(1)
+            graph = build_bipartite_batch(feats=feats, targets=targets)
+            logits = gnn_model(graph)
+            loss = bce_logits_loss(logits, targets, pos_weight=pos_weight)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            train_losses.append(float(loss.item()))
+
+        val_metrics = run_eval(
+            backbone=backbone,
+            gnn_model=gnn_model,
+            loader=val_loader,
+            device=device,
+            pos_weight=pos_weight,
+        )
+        epoch_row = {
+            "epoch": epoch,
+            "train_loss": float(np.mean(train_losses)) if train_losses else 0.0,
+            "val_loss": val_metrics["loss"],
+            "val_map": val_metrics["map"],
+            "val_macro_f1": val_metrics["macro_f1"],
+            "val_micro_f1": val_metrics["micro_f1"],
+        }
+        history.append(epoch_row)
+        log(
+            f"Epoch {epoch}/{train_epochs} -> train_loss={epoch_row['train_loss']:.4f} "
+            f"val_map={epoch_row['val_map']:.4f} val_macro_f1={epoch_row['val_macro_f1']:.4f}"
+        )
+        if val_metrics["map"] > best_map:
+            best_map = val_metrics["map"]
+            torch.save(
+                {
+                    "backbone_state": backbone.state_dict(),
+                    "gnn_state": gnn_model.state_dict(),
+                    "label_vocab": label_vocab,
+                    "config": cfg,
+                },
+                best_path,
+            )
+            log(f"Saved new best checkpoint at epoch {epoch}: {best_path}")
+
+    log("Running final test evaluation using best checkpoint...")
+    checkpoint = torch.load(best_path, map_location=device)
+    backbone.load_state_dict(checkpoint["backbone_state"])
+    gnn_model.load_state_dict(checkpoint["gnn_state"])
+    test_metrics = run_eval(
+        backbone=backbone,
+        gnn_model=gnn_model,
+        loader=test_loader,
+        device=device,
+        pos_weight=pos_weight,
+    )
+    payload = {
+        "run_name": args.run_name,
+        "mode": args.mode,
+        "num_labels": num_labels,
+        "best_val_map": best_map,
+        "history": history,
+        "test_metrics": test_metrics,
+    }
+    write_json(run_dir / "metrics.json", payload)
+    write_json(run_dir / "label_vocab.json", label_vocab)
+    with (run_dir / "overrides.json").open("w", encoding="utf-8") as handle:
+        json.dump(parse_overrides_json(args.overrides_json), handle, indent=2, ensure_ascii=True)
+    log(f"Training complete for {args.run_name}")
+    print(json.dumps({"run_name": args.run_name, "test_map": test_metrics["map"]}, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    train()
+

src/infer/model_adapters.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+
+import torch
+from PIL import Image
+
+from train.encoders import build_backbones
+from train.graph_builder import build_bipartite_batch
+from train.graph_models import NativeGNNClassifier
+from train.models import BaselineClassifier
+
+
+@dataclass
+class InferenceResult:
+    labels: List[str]
+    scores: List[float]
+    positives: List[str]
+
+
+def _sorted_labels(
+    scores: torch.Tensor,
+    label_vocab: Dict[str, int],
+    top_k: int,
+    threshold: float,
+) -> InferenceResult:
+    inv_vocab = {idx: label for label, idx in label_vocab.items()}
+    probs = scores.sigmoid().detach().cpu().view(-1)
+    values, indices = torch.topk(probs, k=min(top_k, probs.numel()))
+    labels = [inv_vocab[int(i)] for i in indices]
+    positives = [label for label, score in zip(labels, values) if float(score) >= threshold]
+    return InferenceResult(
+        labels=labels,
+        scores=[float(v) for v in values],
+        positives=positives,
+    )
+
+
+class BaseAdapter:
+    def __init__(self, device: torch.device) -> None:
+        self.device = device
+
+    def predict(
+        self,
+        image: Image.Image,
+        label_vocab: Dict[str, int],
+        top_k: int = 5,
+        threshold: float = 0.5,
+    ) -> InferenceResult:
+        raise NotImplementedError
+
+
+class BaselineAdapter(BaseAdapter):
+    def __init__(self, checkpoint: Dict[str, Any], device: torch.device) -> None:
+        super().__init__(device)
+        cfg = checkpoint["config"]
+        num_labels = len(checkpoint["label_vocab"])
+        use_clip = bool(cfg["encoders"]["clip"]["enabled"])
+        use_dino = bool(cfg["encoders"]["dino"]["enabled"])
+        clip_name = str(cfg["encoders"]["clip"].get("model_name", "vit_base_patch16_clip_224.openai"))
+        dino_name = str(cfg["encoders"]["dino"].get("model_name", "vit_base_patch14_dinov2.lvd142m"))
+        if clip_name == "ViT-B-32":
+            clip_name = "vit_base_patch16_clip_224.openai"
+        if dino_name == "vit_base_patch14_dinov2":
+            dino_name = "vit_base_patch14_dinov2.lvd142m"
+
+        clip_backbone, dino_backbone, feature_dim = build_backbones(
+            use_clip=use_clip,
+            use_dino=use_dino,
+            clip_model_name=clip_name,
+            dino_model_name=dino_name,
+            clip_pretrained=bool(cfg["encoders"]["clip"].get("pretrained", True)),
+            dino_pretrained=bool(cfg["encoders"]["dino"].get("pretrained", True)),
+        )
+        hidden_dims = list(cfg["baseline_heads"]["mlp"].get("hidden_dims", [512, 256]))
+        dropout = float(cfg["baseline_heads"]["mlp"].get("dropout", 0.2))
+        model = BaselineClassifier(
+            clip_backbone=clip_backbone,
+            dino_backbone=dino_backbone,
+            feature_dim=feature_dim,
+            hidden_dims=hidden_dims,
+            dropout=dropout,
+            num_labels=num_labels,
+        )
+        model.load_state_dict(checkpoint["model_state"])
+        model.eval().to(device)
+        self.model = model
+
+    def predict(
+        self,
+        image: Image.Image,
+        label_vocab: Dict[str, int],
+        top_k: int = 5,
+        threshold: float = 0.5,
+    ) -> InferenceResult:
+        with torch.no_grad():
+            logits = self.model([image], device=self.device)[0]
+        return _sorted_labels(logits, label_vocab, top_k, threshold)
+
+
+class GNNAdapter(BaseAdapter):
+    def __init__(self, checkpoint: Dict[str, Any], device: torch.device) -> None:
+        super().__init__(device)
+        cfg = checkpoint["config"]
+        num_labels = len(checkpoint["label_vocab"])
+
+        use_clip = bool(cfg["encoders"]["clip"]["enabled"])
+        use_dino = bool(cfg["encoders"]["dino"]["enabled"])
+        clip_name = str(cfg["encoders"]["clip"].get("model_name", "vit_base_patch16_clip_224.openai"))
+        dino_name = str(cfg["encoders"]["dino"].get("model_name", "vit_base_patch14_dinov2.lvd142m"))
+        if clip_name == "ViT-B-32":
+            clip_name = "vit_base_patch16_clip_224.openai"
+        if dino_name == "vit_base_patch14_dinov2":
+            dino_name = "vit_base_patch14_dinov2.lvd142m"
+
+        clip_backbone, dino_backbone, feature_dim = build_backbones(
+            use_clip=use_clip,
+            use_dino=use_dino,
+            clip_model_name=clip_name,
+            dino_model_name=dino_name,
+            clip_pretrained=bool(cfg["encoders"]["clip"].get("pretrained", True)),
+            dino_pretrained=bool(cfg["encoders"]["dino"].get("pretrained", True)),
+        )
+
+        class BackboneWrapper(torch.nn.Module):
+            def __init__(self, clip_module: torch.nn.Module | None, dino_module: torch.nn.Module | None) -> None:
+                super().__init__()
+                self.clip_module = clip_module
+                self.dino_module = dino_module
+
+            def encode(self, images: List[Image.Image], device: torch.device) -> torch.Tensor:
+                feats: List[torch.Tensor] = []
+                if self.clip_module is not None:
+                    feats.append(self.clip_module.forward_pil(images, device=device))
+                if self.dino_module is not None:
+                    feats.append(self.dino_module.forward_pil(images, device=device))
+                if len(feats) == 1:
+                    return feats[0]
+                return torch.cat(feats, dim=1)
+
+        backbone = BackboneWrapper(clip_backbone, dino_backbone).to(device)
+        self.backbone = backbone
+
+        hidden_dims = [layer_cfg["out_dim"] for layer_cfg in cfg["gnn"]["layers"]]
+        dropout = float(cfg["gnn"].get("dropout", 0.2))
+        gnn_model = NativeGNNClassifier(
+            in_dim=feature_dim,
+            hidden_dims=hidden_dims,
+            num_attributes=num_labels,
+            dropout=dropout,
+        )
+        gnn_model.load_state_dict(checkpoint["gnn_state"])
+        gnn_model.eval().to(device)
+        self.gnn_model = gnn_model
+
+    def predict(
+        self,
+        image: Image.Image,
+        label_vocab: Dict[str, int],
+        top_k: int = 5,
+        threshold: float = 0.5,
+    ) -> InferenceResult:
+        with torch.no_grad():
+            feats = self.backbone.encode([image], device=self.device).unsqueeze(1)
+            # Single-sample multi-hot target is unknown at inference; use zeros to build graph shape.
+            dummy_targets = torch.zeros(1, len(label_vocab), device=self.device)
+            graph = build_bipartite_batch(feats=feats, targets=dummy_targets)
+            logits = self.gnn_model(graph)[0]
+        return _sorted_labels(logits, label_vocab, top_k, threshold)
+

src/infer/pipeline.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict
+
+import torch
+from PIL import Image
+
+from train.dataset import load_json
+from .model_adapters import BaselineAdapter, GNNAdapter, InferenceResult
+
+
+@dataclass
+class LoadedModel:
+    adapter: Any
+    label_vocab: Dict[str, int]
+
+
+def _detect_model_type(state: Dict[str, Any]) -> str:
+    if "model_state" in state:
+        return "baseline"
+    if "gnn_state" in state and "backbone_state" in state:
+        return "gnn"
+    raise RuntimeError("Unrecognized checkpoint format (expected baseline or GNN keys).")
+
+
+def load_checkpoint(path: Path, device: torch.device | None = None) -> LoadedModel:
+    if device is None:
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        else:
+            device = torch.device("cpu")
+    state: Dict[str, Any] = torch.load(path, map_location=device)
+    model_type = _detect_model_type(state)
+    label_vocab = state["label_vocab"]
+    if model_type == "baseline":
+        adapter = BaselineAdapter(state, device=device)
+    else:
+        adapter = GNNAdapter(state, device=device)
+    return LoadedModel(adapter=adapter, label_vocab=label_vocab)
+
+
+def predict_image(
+    image_path: Path,
+    loaded: LoadedModel,
+    top_k: int = 5,
+    threshold: float = 0.5,
+) -> InferenceResult:
+    image = Image.open(image_path).convert("RGB")
+    return loaded.adapter.predict(
+        image=image,
+        label_vocab=loaded.label_vocab,
+        top_k=top_k,
+        threshold=threshold,
+    )
+

src/train/graph_builder.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+
+import torch
+
+from .graph_models import BipartiteGraphBatch
+
+
+def build_bipartite_batch(
+    feats: torch.Tensor,
+    targets: torch.Tensor,
+) -> BipartiteGraphBatch:
+    """Construct a simple bipartite batch graph from features and multi-hot targets.
+
+    Args:
+        feats:   [B, O, D] fused object features (one object per image for now).
+        targets: [B, A]   multi-hot attribute labels.
+
+    Returns:
+        BipartiteGraphBatch with:
+          - object_feats: feats
+          - attr_feats:   trainable attribute embeddings will be attached later
+          - edge_index:   edges from object indices to attribute indices
+          - edge_weight:  all ones
+    """
+    device = feats.device
+    bsz, num_objects, _ = feats.shape
+    num_attrs = targets.shape[1]
+
+    # For smoke training we assume one object per image (num_objects == 1).
+    # We still keep the general formulation for clarity.
+    object_indices: List[int] = []
+    attr_indices: List[int] = []
+
+    # Flatten batch/object into a single dimension of size B * O.
+    for b in range(bsz):
+        for o in range(num_objects):
+            flat_obj_idx = b * num_objects + o
+            active_attrs = (targets[b] > 0.5).nonzero(as_tuple=False).view(-1)
+            for a in active_attrs.tolist():
+                object_indices.append(flat_obj_idx)
+                attr_indices.append(b * num_attrs + a)
+
+    if not object_indices:
+        # No positive labels in batch; create a dummy self-loop to keep shapes valid.
+        object_indices = [0]
+        attr_indices = [0]
+
+    edge_index = torch.tensor(
+        [object_indices, attr_indices],
+        dtype=torch.long,
+        device=device,
+    )
+    edge_weight = torch.ones(edge_index.shape[1], device=device)
+
+    # Attribute features are left as zeros here; they will be replaced by the model
+    # with either CLIP text or trainable embeddings. For smoke we just keep zeros.
+    attr_feats = torch.zeros(bsz, num_attrs, feats.shape[-1], device=device)
+
+    return BipartiteGraphBatch(
+        object_feats=feats,
+        attr_feats=attr_feats,
+        edge_index=edge_index,
+        edge_weight=edge_weight,
+    )
+

src/train/graph_models.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import torch
+from torch import nn
+
+
+@dataclass
+class BipartiteGraphBatch:
+    """Lightweight container for a batch of bipartite graphs.
+
+    All tensors live on the same device. Shapes:
+      - object_feats:  [B, num_objects, in_dim]
+      - attr_feats:    [B, num_attrs, attr_dim]
+      - edge_index:    [2, num_edges] with edges from object -> attribute
+      - edge_weight:   [num_edges]
+    """
+
+    object_feats: torch.Tensor
+    attr_feats: torch.Tensor
+    edge_index: torch.Tensor
+    edge_weight: torch.Tensor
+
+
+class BipartiteMessagePassingLayer(nn.Module):
+    """Single message passing layer for object <-> attribute bipartite graphs.
+
+    For smoke training we keep a simple formulation:
+      1) Aggregate attribute messages into each object using weighted mean.
+      2) Project and combine with previous object features via residual MLP.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int, attr_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.attr_to_obj = nn.Linear(attr_dim, in_dim)
+        self.proj = nn.Linear(in_dim, out_dim)
+        self.update = nn.Sequential(
+            nn.Linear(in_dim + out_dim, out_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+        )
+
+    def forward(
+        self,
+        object_feats: torch.Tensor,
+        attr_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_weight: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # object_feats: [B, O, Din], attr_feats: [B, A, Din_attr]
+        bsz, num_objects, _ = object_feats.shape
+        device = object_feats.device
+
+        if edge_weight is None:
+            edge_weight = torch.ones(edge_index.shape[1], device=device)
+
+        src_obj = edge_index[0]  # indices in [0, B*O)
+        src_attr = edge_index[1]  # indices in [0, B*A)
+
+        # Flatten batch/object and batch/attr dimensions for gathering.
+        flat_objects = object_feats.reshape(bsz * num_objects, -1)
+        flat_attrs = attr_feats.reshape(bsz * attr_feats.shape[1], -1)
+
+        # Messages from attributes to objects.
+        attr_msgs = flat_attrs.index_select(0, src_attr)  # [E, Din_attr]
+        attr_msgs = self.attr_to_obj(attr_msgs)  # [E, Din]
+        w = edge_weight.view(-1, 1)
+        weighted_msgs = attr_msgs * w
+
+        # Aggregate messages per object index.
+        agg = torch.zeros_like(flat_objects, device=device)
+        agg.index_add_(0, src_obj, weighted_msgs)
+
+        # Normalize by total incoming weight per object to compute mean.
+        weight_sums = torch.zeros(flat_objects.shape[0], device=device)
+        weight_sums.index_add_(0, src_obj, edge_weight)
+        weight_sums = weight_sums.clamp_min(1e-6).view(-1, 1)
+        agg = agg / weight_sums
+
+        # Project aggregated messages and combine with original object features.
+        proj_msgs = self.proj(agg)
+        combined = torch.cat([flat_objects, proj_msgs], dim=-1)
+        updated = self.update(combined)
+
+        return updated.view(bsz, num_objects, self.out_dim)
+
+
+class NativeGNNClassifier(nn.Module):
+    """Simple bipartite GNN classifier for multi-label attribute prediction."""
+
+    def __init__(
+        self,
+        in_dim: int,
+        hidden_dims: List[int],
+        num_attributes: int,
+        dropout: float = 0.2,
+    ) -> None:
+        super().__init__()
+        layers: List[nn.Module] = []
+        dims = [in_dim] + hidden_dims
+        attr_dim = in_dim
+        for dim_in, dim_out in zip(dims[:-1], dims[1:]):
+            layers.append(BipartiteMessagePassingLayer(dim_in, dim_out, attr_dim=attr_dim, dropout=dropout))
+        self.layers = nn.ModuleList(layers)
+        self.classifier = nn.Linear(dims[-1], num_attributes)
+
+    def forward(
+        self,
+        graph: BipartiteGraphBatch,
+    ) -> torch.Tensor:
+        x = graph.object_feats
+        for layer in self.layers:
+            x = layer(
+                object_feats=x,
+                attr_feats=graph.attr_feats,
+                edge_index=graph.edge_index,
+                edge_weight=graph.edge_weight,
+            )
+        # Predict attributes for each object, then average over objects in batch.
+        logits_per_object = self.classifier(x)  # [B, O, num_attributes]
+        return logits_per_object.mean(dim=1)
+