| """Run one image through a published MoonViT K3.1 Pro snapshot and its Projector. |
| |
| Run from this downloaded bundle, or pass --model-dir for another pinned snapshot. |
| """ |
|
|
| import argparse |
| import hashlib |
| import json |
| import time |
| from pathlib import Path |
|
|
| import torch |
| from PIL import Image, ImageOps |
| from safetensors.torch import load_file |
| from transformers import AutoImageProcessor, AutoModel |
|
|
|
|
| import sys |
| sys.path.insert(0, str(Path(__file__).absolute().parents[1])) |
| from projector import Projector |
|
|
|
|
| def square_image(image, resolution): |
| width, height = image.size |
| size = (resolution, int(resolution * height / width)) if width < height else ( |
| int(resolution * width / height), resolution |
| ) |
| image = image.resize(size, Image.Resampling.BICUBIC) |
| left, top = round((size[0] - resolution) / 2), round((size[1] - resolution) / 2) |
| return image.crop((left, top, left + resolution, top + resolution)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("image", type=Path) |
| parser.add_argument("--model-dir", type=Path, default=Path(__file__).absolute().parents[1]) |
| parser.add_argument("--device", default="cuda:0") |
| parser.add_argument("--resolution", type=int, default=448) |
| args = parser.parse_args() |
| root = args.model_dir.absolute() |
| if args.resolution % 28: |
| parser.error("resolution must be divisible by 28 for 2x2 patch merging") |
| for name in ("config.json", "model.safetensors", "projector_config.json", "projector.safetensors"): |
| if not (root / name).is_file(): |
| parser.error(f"missing snapshot file: {root / name}") |
| dtype = torch.bfloat16 |
| model = AutoModel.from_pretrained(root, trust_remote_code=True, local_files_only=True, dtype=dtype) |
| model = model.to(args.device).eval() |
| processor = AutoImageProcessor.from_pretrained(root, trust_remote_code=True, local_files_only=True) |
| projector_config = json.loads((root / "projector_config.json").read_text()) |
| weights = load_file(root / "projector.safetensors") |
| if set(weights) != {"proj.0.weight", "proj.2.weight", "post_norm.weight"}: |
| raise ValueError("expected the three published Projector tensors") |
| if any(value.dtype != dtype for value in weights.values()): |
| raise ValueError("published Projector tensors must remain BF16") |
| projector = Projector(projector_config).to(dtype=dtype) |
| projector.load_state_dict(weights, strict=True) |
| projector = projector.to(args.device).eval() |
| with Image.open(args.image) as original: |
| original_size = list(original.size) |
| image = square_image(ImageOps.exif_transpose(original).convert("RGB"), args.resolution) |
| inputs = processor(images=[image], return_tensors="pt") |
| grid = inputs["grid_thws"] |
| expected_grid = [1, args.resolution // model.config.patch_size, args.resolution // model.config.patch_size] |
| if grid.tolist() != [expected_grid]: |
| raise ValueError(f"unexpected single-image patch grid: {grid.tolist()}") |
| inputs["pixel_values"] = inputs["pixel_values"].to(device=args.device, dtype=dtype) |
| captured = {} |
|
|
| def capture_final_tower(module, arguments, output): |
| captured["tower"] = output.detach().unsqueeze(0) |
|
|
| handle = model.encoder.register_forward_hook(capture_final_tower) |
| started = time.perf_counter() |
| try: |
| with torch.inference_mode(): |
| packed = model(**inputs) |
| if len(packed) != 1: |
| raise ValueError("this example requires one independent image") |
| merged = packed[0].reshape(1, -1, projector_config["input_size"]) |
| stages = captured | {"merged": merged, "projected": projector(merged)} |
| tokens = expected_grid[1] * expected_grid[2] |
| expected = {"tower": (1, tokens, 1024), "merged": (1, tokens // 4, 4096), |
| "projected": (1, tokens // 4, 7168)} |
| for name, value in stages.items(): |
| if tuple(value.shape) != expected[name] or value.dtype != dtype or str(value.device) != args.device: |
| raise ValueError(f"unexpected {name} output: {value.shape}, {value.dtype}, {value.device}") |
| if not torch.isfinite(value).all(): |
| raise ValueError(f"nonfinite {name} output") |
| if args.device.startswith("cuda"): |
| torch.cuda.synchronize() |
| finally: |
| handle.remove() |
| print(json.dumps({ |
| "model_dir": str(root), "image": str(args.image.absolute()), |
| "image_sha256": hashlib.sha256(args.image.read_bytes()).hexdigest(), |
| "image_original_size": original_size, "device": args.device, "dtype": str(dtype), |
| "batch_size": 1, "torch_threads": torch.get_num_threads(), "torch_version": torch.__version__, |
| "preprocessing": {"resize_short_edge": args.resolution, "interpolation": "PIL bicubic", |
| "center_crop": list(image.size), "processor": type(processor).__name__, |
| "image_mean": processor.image_mean, "image_std": processor.image_std, |
| "grid_thws": grid.tolist()}, |
| "stages": {name: {"shape": list(value.shape), "device": str(value.device), "dtype": str(value.dtype)} |
| for name, value in stages.items()}, |
| "forward_seconds": time.perf_counter() - started, |
| "scope": "functional inference; does not establish official production TORCH_SDPA BF16 parity", |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
