tipsv2/run_reference.py

"""Run reference inference for all TIPSv2 HuggingFace models."""

import datetime
import gc
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import save_file
from torchvision import transforms
from transformers import AutoModel

try:
    from sklearn.decomposition import PCA
    HAS_SKLEARN = True
except ImportError:
    HAS_SKLEARN = False

try:
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

# Match TF32/cuDNN settings from transformers conftest.py so results are comparable
if hasattr(torch.backends, "cuda"):
    torch.backends.cuda.matmul.allow_tf32 = False
if hasattr(torch.backends.cudnn, "allow_tf32"):
    torch.backends.cudnn.allow_tf32 = False
if hasattr(torch.backends.cudnn, "conv") and hasattr(torch.backends.cudnn.conv, "fp32_precision"):
    torch.backends.cudnn.conv.fp32_precision = "ieee"

torch.set_printoptions(precision=5, sci_mode=False)

SCRIPT_DIR = Path(__file__).parent

MODELS = [
    "google/tipsv2-b14",
    "google/tipsv2-l14",
    "google/tipsv2-so400m14",
    "google/tipsv2-g14",
    "google/tipsv2-b14-dpt",
    "google/tipsv2-l14-dpt",
    "google/tipsv2-so400m14-dpt",
    "google/tipsv2-g14-dpt",
]

IMAGES = [
    SCRIPT_DIR / "reference_images/coco_000000039769.jpg",           # two cats on a sofa
    SCRIPT_DIR / "reference_images/depth_ade20k_00014.png",          # outdoor scene (ADE20K)
    SCRIPT_DIR / "reference_images/nyuv2_living_room_01260.jpg",     # indoor living room
    SCRIPT_DIR / "reference_images/pca_cph.jpeg",                    # city street (Copenhagen)
    SCRIPT_DIR / "reference_images/zeroseg_pascal_context_00049_image.png",  # bus on street
]

TEXT_QUERIES = ["two cats on a sofa", "a cat lying down", "a dog on a couch", "an empty room"]
IMAGE_SIZE = 448
# Identity normalization matching run_image_encoder_inference.py — no ImageNet normalization for TIPSv2.
IMAGE_MEAN = (0, 0, 0)
IMAGE_STD = (1.0, 1.0, 1.0)


def preprocess_image(path: Path, size: int = IMAGE_SIZE) -> torch.Tensor:
    transform = transforms.Compose([
        transforms.Resize((size, size)),
        transforms.ToTensor(),
        transforms.Normalize(IMAGE_MEAN, IMAGE_STD),
    ])
    return transform(Image.open(path).convert("RGB")).unsqueeze(0)


def tprint(msg: str, log_file) -> None:
    # print(msg)
    log_file.write(msg + "\n")
    log_file.flush()


def log_tensor(name: str, tensor: torch.Tensor, log_file) -> None:
    """Print tensor shape and a representative sample slice to stdout and log file."""
    ndim = tensor.ndim
    tprint(f"  {name}: shape={tuple(tensor.shape)}", log_file)
    if ndim >= 4:
        # 2D spatial tensor [batch, channel, H, W] — print 3x3 patch from first channel
        tprint(f"  {name}[0, 0, :3, :3] =\n{tensor[0, 0, :3, :3]}", log_file)
    elif ndim == 3:
        if tensor.is_floating_point():
            # Sequence-like [batch, seq, dim] — print first 5 features of first token
            tprint(f"  {name}[0, 0, :5] = {tensor[0, 0, :5]}", log_file)
        else:
            # Post-argmax spatial [batch, H, W] — print 3x3 patch
            tprint(f"  {name}[0, :3, :3] =\n{tensor[0, :3, :3]}", log_file)
    elif ndim == 2:
        if tensor.shape[1] == 1:
            # Column vector [queries, 1] — print all rows
            tprint(f"  {name}[:5, 0] = {tensor[:5, 0]}", log_file)
        else:
            # [batch, dim] or [queries, dim] — print first 5 values of first row
            tprint(f"  {name}[0, :5] = {tensor[0, :5]}", log_file)
    elif ndim == 1:
        tprint(f"  {name}[:5] = {tensor[:5]}", log_file)
    else:
        tprint(f"  {name} = {tensor.item()}", log_file)
    tprint("", log_file)


def log_base_model_outputs(
    forward: dict,
    postprocess: dict,
    text_emb: torch.Tensor,
    text_norm: torch.Tensor,
    log_file,
) -> None:
    tprint("=== backbone outputs ===", log_file)
    log_tensor("cls_token", forward["cls_token"], log_file)
    log_tensor("register_tokens", forward["register_tokens"], log_file)
    log_tensor("patch_tokens", forward["patch_tokens"], log_file)
    log_tensor("vision_pooler_output", forward["vision_pooler_output"], log_file)

    tprint("=== postprocess ===", log_file)
    log_tensor("cls_token_norm", postprocess["cls_token_norm"], log_file)

    tprint("  [classification scores: pre-argmax]", log_file)
    log_tensor("classification_scores", postprocess["classification_scores"], log_file)
    tprint("  [classification: post-argmax]", log_file)
    log_tensor("predicted_class", postprocess["predicted_class"], log_file)

    tprint("=== text embeddings (shared across images) ===", log_file)
    log_tensor("text_pooler_output", text_emb, log_file)
    log_tensor("text_embeddings_norm", text_norm, log_file)

    tprint("=== base model logits ===", log_file)
    if "loss" in postprocess:
        log_tensor("loss", postprocess["loss"], log_file)
    log_tensor("logits_per_image", postprocess["logits_per_image"], log_file)
    log_tensor("logits_per_text", postprocess["logits_per_text"], log_file)


def log_dpt_model_outputs(forward: dict, forward_postprocessed: dict, postprocess: dict, log_file) -> None:
    tprint("=== model outputs (native DPT resolution, no image_size) ===", log_file)
    log_tensor("depth", forward["depth"], log_file)
    log_tensor("normals", forward["normals"], log_file)
    tprint("  [segmentation logits: pre-argmax]", log_file)
    log_tensor("segmentation", forward["segmentation"], log_file)

    tprint("=== post-processed model outputs (upsampled to input size with image_size) ===", log_file)
    log_tensor("depth", forward_postprocessed["depth"], log_file)
    log_tensor("normals", forward_postprocessed["normals"], log_file)
    tprint("  [post-processed segmentation logits: pre-argmax]", log_file)
    log_tensor("segmentation", forward_postprocessed["segmentation"], log_file)

    tprint("=== postprocess ===", log_file)
    tprint("  [segmentation: post-argmax]", log_file)
    log_tensor("segmentation_labels", postprocess["segmentation_labels"], log_file)


def visualize_base_model(forward: dict, postprocess: dict, out_dir: Path) -> None:
    if not HAS_SKLEARN or not HAS_MATPLOTLIB:
        return

    patch_tokens = forward["patch_tokens"][0].numpy()  # sklearn requires numpy; (1024, D)
    grid = int(patch_tokens.shape[0] ** 0.5)
    patch_rgb = PCA(n_components=3, whiten=True).fit_transform(patch_tokens).reshape(grid, grid, 3)
    patch_rgb = 1 / (1 + np.exp(-2.0 * patch_rgb))
    Image.fromarray((patch_rgb * 255).astype(np.uint8)).save(out_dir / "patch_tokens_pca.png")

    similarity = postprocess["classification_scores"]  # (1, num_queries)

    fig, ax = plt.subplots(figsize=(10, 2))
    im = ax.imshow(similarity.numpy(), cmap="RdYlGn", aspect="auto", vmin=-1, vmax=1)
    ax.set_yticks([])
    ax.set_xticks(range(len(TEXT_QUERIES)))
    ax.set_xticklabels(TEXT_QUERIES, rotation=45, ha="right", fontsize=8)
    ax.set_title("Cls-Text Similarity")
    plt.colorbar(im, ax=ax, label="Cosine Similarity")
    plt.tight_layout()
    plt.savefig(out_dir / "text_similarity.png", dpi=100)
    plt.close()


def visualize_dpt_model(forward: dict, postprocess: dict, out_dir: Path) -> None:
    if not HAS_MATPLOTLIB:
        return

    depth = forward["depth"][0, 0]  # (H, W)
    depth_norm = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
    Image.fromarray((depth_norm.numpy() * 255).astype(np.uint8)).save(out_dir / "depth.png")

    normals_vis = ((forward["normals"][0].permute(1, 2, 0) + 1) / 2).clamp(0, 1)  # (H, W, 3)
    Image.fromarray((normals_vis.numpy() * 255).astype(np.uint8)).save(out_dir / "normals.png")

    seg_labels = postprocess["segmentation_labels"][0].numpy().astype(int)  # matplotlib cmap requires numpy
    cmap = cm.get_cmap("tab20", 150)
    seg_colored = cmap((seg_labels % 150) / 150.0)[:, :, :3]
    Image.fromarray((seg_colored * 255).astype(np.uint8)).save(out_dir / "segmentation.png")


def main() -> None:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    out_root = SCRIPT_DIR / "out"
    out_root.mkdir(exist_ok=True)

    now = datetime.datetime.now()
    date_str = now.strftime("%Y%m%d")
    time_str = now.strftime("%H%M%S")

    for model_id in MODELS:
        model_name = model_id.split("/")[-1]
        is_dpt = model_name.endswith("-dpt")
        model_dir = out_root / f"{date_str}_{time_str}_{model_name}_ref"

        print(f"\n=== {model_id} ===")
        model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
        model.eval().to(device)

        if is_dpt:
            _orig_forwards = {}
            _orig_F_normalize = torch.nn.functional.normalize

            def _identity_normalize(input, p=2.0, dim=1, eps=1e-12, out=None):
                return input

            for head_name in ("depth_head", "segmentation_head"):
                head = getattr(model, head_name)
                orig_fwd = head.forward
                _orig_forwards[head_name] = orig_fwd
                head.forward = lambda inputs, image_size=None, _fwd=orig_fwd: _fwd(inputs)

            # Normals head: also skip F.normalize so the no-image_size output is pre-normalization logits
            _normals_orig_fwd = model.normals_head.forward
            _orig_forwards["normals_head"] = _normals_orig_fwd

            def _normals_forward_no_normalize(inputs, image_size=None, _fwd=_normals_orig_fwd):
                torch.nn.functional.normalize = _identity_normalize
                try:
                    result = _fwd(inputs)
                finally:
                    torch.nn.functional.normalize = _orig_F_normalize
                return result

            model.normals_head.forward = _normals_forward_no_normalize

        try:
            text_emb = None
            text_norm = None
            if not is_dpt:
                with torch.no_grad():
                    text_emb = model.encode_text(TEXT_QUERIES).detach().cpu()
                text_norm = F.normalize(text_emb, dim=-1)
                model_dir.mkdir(parents=True, exist_ok=True)
                save_file(
                    {"text_pooler_output": text_emb, "text_embeddings_norm": text_norm},
                    model_dir / "text_embeddings.safetensors",
                )

            for image_path in IMAGES:
                image_stem = image_path.stem
                out_dir = model_dir / image_stem
                out_dir.mkdir(parents=True, exist_ok=True)

                pixel_values = preprocess_image(image_path).to(device)

                with torch.no_grad():
                    if is_dpt:
                        # Run without image_size (native DPT resolution, heads not upsampled)
                        outputs = model(pixel_values)
                        forward = {
                            "depth": outputs.depth.detach().cpu().contiguous(),
                            "normals": outputs.normals.detach().cpu().contiguous(),
                            "segmentation": outputs.segmentation.detach().cpu().contiguous(),
                        }
                        # Restore original forwards and run with image_size (upsampled to input size)
                        for head_name, orig_fwd in _orig_forwards.items():
                            getattr(model, head_name).forward = orig_fwd
                        outputs_pp = model(pixel_values)
                        forward_postprocessed = {
                            "depth": outputs_pp.depth.detach().cpu().contiguous(),
                            "normals": outputs_pp.normals.detach().cpu().contiguous(),
                            "segmentation": outputs_pp.segmentation.detach().cpu().contiguous(),
                        }
                        postprocess = {
                            "segmentation_labels": outputs_pp.segmentation.argmax(dim=1).detach().cpu().contiguous(),
                        }
                    else:
                        image_out = model.encode_image(pixel_values)
                        forward = {
                            "cls_token": image_out.cls_token.detach().cpu(),
                            "register_tokens": image_out.register_tokens.detach().cpu(),
                            "patch_tokens": image_out.patch_tokens.detach().cpu(),
                            "vision_pooler_output": image_out.cls_token[:, 0, :].detach().cpu(),
                        }
                        cls_norm = F.normalize(image_out.cls_token[:, 0, :], dim=-1).detach().cpu()
                        classification_scores = cls_norm @ text_norm.T  # (1, num_queries)
                        temperature = model.config.temperature
                        # Match transformers: logits_per_text = matmul(text_embeds, image_embeds.T) / temperature
                        logits_per_image = cls_norm @ text_norm.T / temperature  # (1, num_queries)
                        logits_per_text = text_norm @ cls_norm.T / temperature   # (num_queries, 1)
                        postprocess = {
                            "cls_token_norm": cls_norm,
                            "classification_scores": classification_scores,
                            "predicted_class": classification_scores.argmax(dim=-1),
                            "logits_per_image": logits_per_image,
                            "logits_per_text": logits_per_text,
                        }

                save_file(forward, out_dir / "forward.safetensors")
                save_file(postprocess, out_dir / "postprocess.safetensors")
                if is_dpt:
                    save_file(forward_postprocessed, out_dir / "forward_postprocessed.safetensors")

                log_path = out_dir / "inference_log.txt"
                with open(log_path, "w") as log_file:
                    tprint(f"model: {model_id}", log_file)
                    tprint(f"image: {image_path.name}", log_file)
                    tprint(f"device: {device}", log_file)
                    tprint("", log_file)
                    if is_dpt:
                        log_dpt_model_outputs(forward, forward_postprocessed, postprocess, log_file)
                    else:
                        log_base_model_outputs(forward, postprocess, text_emb, text_norm, log_file)

                if is_dpt:
                    visualize_dpt_model(forward_postprocessed, postprocess, out_dir)
                else:
                    visualize_base_model(forward, postprocess, out_dir)

                print(f"  {image_stem} -> {out_dir}")
        except Exception as exc:
            print(f"  ERROR: {exc}")
        finally:
            del model
            gc.collect()
            if device.type == "cuda":
                torch.cuda.empty_cache()


if __name__ == "__main__":
    main()