tipsv2/run_compare.py

"""Run inference for all TIPSv2 models and compare against latest reference results.

For each model × image, produces:
  - forward.safetensors / postprocess.safetensors  (new results)
  - inference_log.txt  (same as run_reference.py, with tensor comparison table appended)
  - {output_type}_composite.png  (ref | new | diff×10 side-by-side per visualised output)

At model level:
  - comparison.txt  (single table covering text embeddings + all image tensors)
  - comparison.png  (grid: rows=input images, column groups=output type×(ref, new, diff))
"""

import datetime
import gc
import io
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent / "transformers/src"))

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import load_file, save_file
from torchvision import transforms
from transformers import AutoModel, Tipsv2Tokenizer, Tipsv2DptImageProcessor, Tipsv2DptForDensePrediction

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

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 = [
    "guarin/tipsv2-b14",
    "guarin/tipsv2-l14",
    "guarin/tipsv2-so400m14",
    "guarin/tipsv2-g14",
    "guarin/tipsv2-b14-dpt",
    "guarin/tipsv2-l14-dpt",
    "guarin/tipsv2-so400m14-dpt",
    "guarin/tipsv2-g14-dpt",
]

IMAGES = [
    SCRIPT_DIR / "reference_images/coco_000000039769.jpg",
    SCRIPT_DIR / "reference_images/depth_ade20k_00014.png",
    SCRIPT_DIR / "reference_images/nyuv2_living_room_01260.jpg",
    SCRIPT_DIR / "reference_images/pca_cph.jpeg",
    SCRIPT_DIR / "reference_images/zeroseg_pascal_context_00049_image.png",
]

TEXT_QUERIES = ["two cats on a sofa", "a cat lying down", "a dog on a couch", "an empty room"]
IMAGE_SIZE = 448
IMAGE_MEAN = (0, 0, 0)
IMAGE_STD = (1.0, 1.0, 1.0)

DIFF_AMPLIFY = 10.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:
    log_file.write(msg + "\n")
    log_file.flush()


def log_tensor(name: str, tensor: torch.Tensor, log_file) -> None:
    ndim = tensor.ndim
    tprint(f"  {name}: shape={tuple(tensor.shape)}", log_file)
    if ndim >= 4:
        tprint(f"  {name}[0, 0, :3, :3] =\n{tensor[0, 0, :3, :3]}", log_file)
    elif ndim == 3:
        if tensor.is_floating_point():
            tprint(f"  {name}[0, 0, :5] = {tensor[0, 0, :5]}", log_file)
        else:
            tprint(f"  {name}[0, :3, :3] =\n{tensor[0, :3, :3]}", log_file)
    elif ndim == 2:
        if tensor.shape[1] == 1:
            tprint(f"  {name}[:5, 0] = {tensor[:5, 0]}", log_file)
        else:
            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, postprocess, text_emb, text_norm, 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, forward_postprocessed, postprocess, 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, postprocess, out_dir: Path) -> None:
    if not HAS_SKLEARN or not HAS_MATPLOTLIB:
        return
    patch_tokens = forward["patch_tokens"][0].numpy()
    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"]
    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, postprocess, out_dir: Path) -> None:
    if not HAS_MATPLOTLIB:
        return
    depth = forward["depth"].squeeze()  # handles (1,H,W) or (1,1,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)
    Image.fromarray((normals_vis.numpy() * 255).astype(np.uint8)).save(out_dir / "normals.png")

    seg_labels = postprocess["segmentation_labels"][0].numpy().astype(int)
    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")


# ---------------------------------------------------------------------------
# Reference directory lookup
# ---------------------------------------------------------------------------

def find_latest_ref_dir(out_root: Path, model_name: str) -> Path | None:
    """Return the most recent *_ref directory for model_name, or None."""
    candidates = sorted(out_root.glob(f"*_{model_name}_ref"))
    return candidates[-1] if candidates else None


# ---------------------------------------------------------------------------
# Tensor comparison helpers
# ---------------------------------------------------------------------------

_TABLE_HDR = f"  {'tensor':<40} {'equal':>6}  {'max_abs_diff':>14}  {'mean_abs_diff':>14}"
_TABLE_SEP = f"  {'-'*40} {'-'*6}  {'-'*14}  {'-'*14}"


def _compare_row(name: str, ref: torch.Tensor, new: torch.Tensor) -> tuple[str, bool, float, float]:
    """Return (formatted_line, equal, max_diff, mean_diff)."""
    if ref.shape != new.shape:
        line = f"  {name:<40} SHAPE MISMATCH: {tuple(ref.shape)} vs {tuple(new.shape)}"
        return line, False, float("nan"), float("nan")
    equal = torch.equal(ref, new)
    diff = (ref.float() - new.float()).abs()
    max_diff = diff.max().item()
    mean_diff = diff.mean().item()
    line = f"  {name:<40} {'True' if equal else 'False':>6}  {max_diff:>14.6f}  {mean_diff:>14.6f}"
    return line, equal, max_diff, mean_diff


def compare_tensor_dicts(
    section: str,
    ref_tensors: dict,
    new_tensors: dict,
    log_file,
) -> list[tuple[str, bool, float, float]]:
    """Log a comparison table for a pair of tensor dicts; return per-tensor stats."""
    tprint(f"=== comparison: {section} ===", log_file)
    tprint(_TABLE_HDR, log_file)
    tprint(_TABLE_SEP, log_file)
    rows = []
    all_keys = sorted(set(ref_tensors) | set(new_tensors))
    for key in all_keys:
        if key not in ref_tensors:
            tprint(f"  {key:<40} MISSING IN REFERENCE", log_file)
            rows.append((key, False, float("nan"), float("nan")))
            continue
        if key not in new_tensors:
            tprint(f"  {key:<40} MISSING IN NEW RUN", log_file)
            rows.append((key, False, float("nan"), float("nan")))
            continue
        line, equal, max_diff, mean_diff = _compare_row(key, ref_tensors[key], new_tensors[key])
        tprint(line, log_file)
        rows.append((key, equal, max_diff, mean_diff))
    tprint("", log_file)
    return rows


# ---------------------------------------------------------------------------
# Image composite helpers
# ---------------------------------------------------------------------------

def _load_rgb(path: Path, target_shape: tuple | None = None) -> np.ndarray | None:
    if not path.exists():
        return None
    img = np.array(Image.open(path).convert("RGB"))
    if target_shape is not None and img.shape != target_shape:
        pil = Image.fromarray(img).resize(
            (target_shape[1], target_shape[0]), Image.LANCZOS
        )
        img = np.array(pil)
    return img


def create_image_composite(
    ref_path: Path,
    new_path: Path,
    out_path: Path,
    amplify: float = DIFF_AMPLIFY,
) -> bool:
    """Save a (ref | new | diff×amplify) composite PNG. Returns True on success."""
    ref_img = _load_rgb(ref_path)
    new_img = _load_rgb(new_path, target_shape=ref_img.shape if ref_img is not None else None)
    if ref_img is None or new_img is None:
        return False
    diff = np.abs(ref_img.astype(np.float32) - new_img.astype(np.float32))
    diff_vis = np.clip(diff * amplify, 0, 255).astype(np.uint8)
    h, w = ref_img.shape[:2]
    composite = np.zeros((h, w * 3 + 4, 3), dtype=np.uint8)  # 2px separators
    composite[:, :w] = ref_img
    composite[:, w + 2: 2 * w + 2] = new_img
    composite[:, 2 * w + 4:] = diff_vis
    Image.fromarray(composite).save(out_path)
    return True


def create_model_comparison_png(
    model_dir: Path,
    ref_dir: Path,
    image_stems: list[str],
    is_dpt: bool,
    out_path: Path,
    amplify: float = DIFF_AMPLIFY,
) -> None:
    """Grid comparison PNG: rows = input images, column groups = output type × (ref, new, diff)."""
    if not HAS_MATPLOTLIB:
        return

    output_types = ["depth", "normals", "segmentation"] if is_dpt else ["patch_tokens_pca", "text_similarity"]
    nrows = len(image_stems)
    ncols = len(output_types) * 3

    fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 2.5, nrows * 2.5), squeeze=False)

    for row_idx, image_stem in enumerate(image_stems):
        for type_idx, output_type in enumerate(output_types):
            col_base = type_idx * 3
            ref_img = _load_rgb(ref_dir / image_stem / f"{output_type}.png")
            new_img = _load_rgb(
                model_dir / image_stem / f"{output_type}.png",
                target_shape=ref_img.shape if ref_img is not None else None,
            )

            if ref_img is not None and new_img is not None:
                diff = np.abs(ref_img.astype(np.float32) - new_img.astype(np.float32))
                diff_vis = np.clip(diff * amplify, 0, 255).astype(np.uint8)
                axes[row_idx, col_base].imshow(ref_img)
                axes[row_idx, col_base + 1].imshow(new_img)
                axes[row_idx, col_base + 2].imshow(diff_vis)
            else:
                for col_offset in range(3):
                    ax = axes[row_idx, col_base + col_offset]
                    ax.text(0.5, 0.5, "MISSING", ha="center", va="center",
                            transform=ax.transAxes, fontsize=8, color="red")

            if row_idx == 0:
                axes[row_idx, col_base].set_title(f"{output_type}\nref", fontsize=7)
                axes[row_idx, col_base + 1].set_title(f"{output_type}\nnew", fontsize=7)
                axes[row_idx, col_base + 2].set_title(f"{output_type}\ndiff×{int(amplify)}", fontsize=7)

        axes[row_idx, 0].set_ylabel(image_stem, fontsize=7, rotation=30, ha="right", va="center")

    for ax in axes.flat:
        ax.set_xticks([])
        ax.set_yticks([])

    ref_name = ref_dir.name
    plt.suptitle(
        f"{model_dir.name}\nreference: {ref_name}",
        fontsize=9,
        y=1.01,
    )
    plt.tight_layout()
    plt.savefig(out_path, dpi=100, bbox_inches="tight")
    plt.close()
    print(f"  comparison.png -> {out_path}")


# ---------------------------------------------------------------------------
# comparison.txt writer
# ---------------------------------------------------------------------------

LARGE_DIFF_THRESHOLD = 1e-4


def write_comparison_txt(
    out_path: Path,
    model_id: str,
    ref_dir: Path,
    is_dpt: bool,
    text_rows: list[tuple] | None,
    image_data: list[tuple],  # (image_stem, fwd_rows, fwd_pp_rows, post_rows)
) -> None:
    """Write a single comparison.txt covering text embeddings + all images."""
    header = _TABLE_HDR
    sep = _TABLE_SEP
    large_diffs = []  # (section_label, tensor_name, max_diff)

    def fmt_section(title, rows, section_label):
        lines = [f"--- {title} ---", header, sep]
        for name, equal, max_diff, mean_diff in rows:
            if np.isnan(max_diff):
                lines.append(f"  {name:<40} {'False':>6}  {'n/a':>14}  {'n/a':>14}")
            else:
                flag = " **" if max_diff >= LARGE_DIFF_THRESHOLD else ""
                lines.append(
                    f"  {name:<40} {'True' if equal else 'False':>6}  {max_diff:>14.6f}  {mean_diff:>14.6f}{flag}"
                )
                if max_diff >= LARGE_DIFF_THRESHOLD:
                    large_diffs.append((section_label, name, max_diff))
        return "\n".join(lines) + "\n"

    with open(out_path, "w") as fp:
        fp.write(f"model:     {model_id}\n")
        fp.write(f"reference: {ref_dir}\n\n")

        if text_rows:
            fp.write(fmt_section("text embeddings", text_rows, "text embeddings") + "\n")

        for image_stem, fwd_rows, fwd_pp_rows, post_rows in image_data:
            fp.write(f"{'='*70}\n")
            fp.write(f"image: {image_stem}\n\n")
            fp.write(fmt_section("forward", fwd_rows, f"{image_stem}/forward") + "\n")
            if fwd_pp_rows is not None:
                fp.write(fmt_section("forward_postprocessed", fwd_pp_rows, f"{image_stem}/forward_postprocessed") + "\n")
            fp.write(fmt_section("postprocess", post_rows, f"{image_stem}/postprocess") + "\n")

        if large_diffs:
            fp.write(f"{'='*70}\n")
            fp.write(f"NOTE: {len(large_diffs)} tensor(s) exceed the diff threshold ({LARGE_DIFF_THRESHOLD:.0e}):\n")
            for section_label, name, max_diff in large_diffs:
                fp.write(f"  ** [{section_label}] {name}: max_abs_diff={max_diff:.6f}\n")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

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}_cmp"

        ref_dir = find_latest_ref_dir(out_root, model_name)
        if ref_dir is None:
            print(f"\n=== {model_id} — no reference found, skipping ===")
            continue

        print(f"\n=== {model_id} ===")
        print(f"  reference: {ref_dir.name}")

        if "dpt" in model_id:
            model = Tipsv2DptForDensePrediction.from_pretrained(model_id)
        else:
            model = AutoModel.from_pretrained(model_id)
        model.eval().to(device)
        tokenizer = Tipsv2Tokenizer.from_pretrained(model_id) if not is_dpt else None
        image_processor = Tipsv2DptImageProcessor() if is_dpt else None

        text_rows = None
        image_data = []

        try:
            text_emb = None
            text_norm = None
            if not is_dpt:
                with torch.no_grad():
                    text_inputs = tokenizer(TEXT_QUERIES, return_tensors="pt", padding=True).to(device)
                    text_out = model.get_text_features(**text_inputs)
                    text_emb = text_out.pooler_output.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",
                )

                ref_text_path = ref_dir / "text_embeddings.safetensors"
                ref_text = None
                if ref_text_path.exists():
                    ref_text = load_file(ref_text_path)
                    new_text = {"text_pooler_output": text_emb, "text_embeddings_norm": text_norm}
                    _sink = io.StringIO()

                    class _SinkLF:
                        def write(self, s):
                            _sink.write(s)
                        def flush(self):
                            pass

                    text_rows = compare_tensor_dicts("text embeddings", ref_text, new_text, _SinkLF())

            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:
                        outputs = model(pixel_values)
                        forward = {
                            "depth": outputs.predicted_depth.unsqueeze(1).detach().cpu().contiguous(),
                            "normals": outputs.normals.detach().cpu().contiguous(),
                            "segmentation": outputs.segmentation_logits.detach().cpu().contiguous(),
                        }
                        target_sizes = [(IMAGE_SIZE, IMAGE_SIZE)]
                        pp_depth = image_processor.post_process_depth_estimation(outputs, target_sizes=target_sizes)[0]["predicted_depth"]
                        pp_normals = image_processor.post_process_normal_estimation(outputs, target_sizes=target_sizes)[0]["normals"]
                        pp_seg = F.interpolate(outputs.segmentation_logits.float(), size=target_sizes[0], mode="bilinear", align_corners=False)
                        pp_seg_labels = image_processor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)[0]
                        forward_postprocessed = {
                            "depth": pp_depth.unsqueeze(0).unsqueeze(0).detach().cpu().contiguous(),
                            "normals": pp_normals.unsqueeze(0).detach().cpu().contiguous(),
                            "segmentation": pp_seg.detach().cpu().contiguous(),
                        }
                        postprocess = {
                            "segmentation_labels": pp_seg_labels.unsqueeze(0).detach().cpu().contiguous(),
                        }
                    else:
                        vision_out = model.get_image_features(pixel_values=pixel_values)
                        num_reg = model.config.vision_config.num_register_tokens
                        seq = vision_out.last_hidden_state  # [batch, 1+num_reg+num_patches, dim]
                        forward = {
                            "cls_token": seq[:, 0:1, :].detach().cpu(),
                            "register_tokens": seq[:, 1:1 + num_reg, :].detach().cpu(),
                            "patch_tokens": seq[:, 1 + num_reg:, :].detach().cpu(),
                            "vision_pooler_output": vision_out.pooler_output.detach().cpu(),
                        }
                        cls_norm = F.normalize(vision_out.pooler_output, dim=-1).detach().cpu()
                        classification_scores = cls_norm @ text_norm.T
                        temperature = model.config.temperature_init_value
                        logits_per_image = cls_norm @ text_norm.T / temperature
                        logits_per_text = text_norm @ cls_norm.T / temperature
                        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")

                # ------ load reference tensors for this image ------
                ref_img_dir = ref_dir / image_stem
                ref_forward = load_file(ref_img_dir / "forward.safetensors") if (ref_img_dir / "forward.safetensors").exists() else {}
                ref_postprocess = load_file(ref_img_dir / "postprocess.safetensors") if (ref_img_dir / "postprocess.safetensors").exists() else {}
                ref_fwd_pp = load_file(ref_img_dir / "forward_postprocessed.safetensors") if is_dpt and (ref_img_dir / "forward_postprocessed.safetensors").exists() else {}

                # ------ write inference log with comparison table ------
                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(f"reference: {ref_dir.name}", 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)

                    fwd_rows = compare_tensor_dicts("forward", ref_forward, forward, log_file)
                    fwd_pp_rows = compare_tensor_dicts("forward_postprocessed", ref_fwd_pp, forward_postprocessed, log_file) if is_dpt else None
                    post_rows = compare_tensor_dicts("postprocess", ref_postprocess, postprocess, log_file)

                    if text_rows and ref_text is not None:
                        compare_tensor_dicts(
                            "text embeddings",
                            ref_text,
                            {"text_pooler_output": text_emb, "text_embeddings_norm": text_norm},
                            log_file,
                        )

                image_data.append((image_stem, fwd_rows, fwd_pp_rows, post_rows))

                # ------ visualise new outputs ------
                if is_dpt:
                    visualize_dpt_model(forward_postprocessed, postprocess, out_dir)
                else:
                    visualize_base_model(forward, postprocess, out_dir)

                # ------ per-image composites ------
                output_types = (
                    ["depth", "normals", "segmentation"]
                    if is_dpt
                    else ["patch_tokens_pca", "text_similarity"]
                )
                for output_type in output_types:
                    ref_png = ref_img_dir / f"{output_type}.png"
                    new_png = out_dir / f"{output_type}.png"
                    composite_path = out_dir / f"{output_type}_composite.png"
                    ok = create_image_composite(ref_png, new_png, composite_path)
                    if not ok:
                        print(f"  WARNING: could not create composite for {output_type} ({image_stem})")

                print(f"  {image_stem} -> {out_dir}")

            # ------ model-level comparison.txt ------
            model_dir.mkdir(parents=True, exist_ok=True)
            write_comparison_txt(
                model_dir / "comparison.txt",
                model_id,
                ref_dir,
                is_dpt,
                text_rows,
                image_data,
            )
            print(f"  comparison.txt -> {model_dir / 'comparison.txt'}")

            # ------ model-level comparison.png ------
            image_stems = [image_path.stem for image_path in IMAGES]
            create_model_comparison_png(
                model_dir,
                ref_dir,
                image_stems,
                is_dpt,
                model_dir / "comparison.png",
            )

        except Exception as exc:
            print(f"  ERROR: {exc}")
            import traceback
            traceback.print_exc()
        finally:
            del model
            gc.collect()
            if device.type == "cuda":
                torch.cuda.empty_cache()


if __name__ == "__main__":
    main()