"""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()