#!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Run NV-Reason-CXR-3B inference on a PNG/JPEG chest X-ray image. The live path follows the upstream Hugging Face Transformers example. The mock path exists for CI and evidence-pack wiring checks; it never calls the model and should not be treated as clinical or model output. """ from __future__ import annotations import argparse import binascii import hashlib import importlib.metadata import json import math import os import struct import sys import time import zlib from dataclasses import dataclass from pathlib import Path from typing import Any DEFAULT_MODEL = "nvidia/NV-Reason-CXR-3B" DEFAULT_PROMPT = "Find abnormalities and support devices." PROMPT_PRESETS = { "findings": "Find abnormalities and support devices.", "comprehensive": "Provide a comprehensive image analysis, and list all abnormalities.", "educational": ( "Examine the chest X-ray and explain the reasoning for each visible " "finding. State uncertainty explicitly." ), "structured": ( "Write a concise structured response with sections: image quality, " "support devices, findings, impression, and uncertainties." ), } GENERATED_IMAGE_SENTINELS = { "generated://synthetic_chest_xray", "generated://synthetic_cxr", } TRUTHY = {"1", "true", "yes", "on"} LIMITATIONS = [ "Output is engineering evidence only; it is not a diagnosis or treatment recommendation.", "NV-Reason-CXR-3B can hallucinate, miss findings, or produce overconfident prose.", "A qualified professional must review any medical workflow use.", ] class SkillError(Exception): """Input, dependency, or runtime error that should be shown cleanly.""" @dataclass(frozen=True) class ImageInfo: format: str width: int height: int sha256: str @dataclass(frozen=True) class InputSpec: image_path: Path source: str prompt: str case_id: str fixture_mock: bool fixture_mock_response: str | None def _truthy(value: str | None) -> bool: return str(value or "").strip().lower() in TRUTHY def _package_status(import_name: str, dist_name: str | None = None) -> dict: dist_name = dist_name or import_name try: version = importlib.metadata.version(dist_name) except importlib.metadata.PackageNotFoundError: return {"installed": False, "version": None} try: __import__(import_name) importable = True except Exception: importable = False return {"installed": True, "importable": importable, "version": version} def _cuda_report() -> dict: try: import torch except Exception as e: return {"available": False, "error": f"torch import failed: {e}", "devices": []} available = bool(torch.cuda.is_available()) devices = [] if available: for idx in range(torch.cuda.device_count()): props = torch.cuda.get_device_properties(idx) devices.append( { "index": idx, "name": props.name, "memory_total_mb": round(props.total_memory / int("1024") / int("1024")), } ) return { "available": available, "device_count": len(devices), "devices": devices, "torch_cuda_version": getattr(torch.version, "cuda", None), } def _model_cache_report(model_id: str) -> dict: try: from huggingface_hub import scan_cache_dir except Exception as e: return {"inspectable": False, "cached": False, "error": str(e)} try: cache = scan_cache_dir() except Exception as e: return {"inspectable": False, "cached": False, "error": str(e)} for repo in cache.repos: if repo.repo_id != model_id or repo.repo_type != "model": continue files = [] for revision in repo.revisions: files.extend(f.file_name for f in revision.files) return { "inspectable": True, "cached": True, "repo_id": repo.repo_id, "revisions": len(repo.revisions), "size_on_disk_mb": round(repo.size_on_disk / int("1024") / int("1024")), "has_config": "config.json" in files, "has_preprocessor_config": "preprocessor_config.json" in files, "has_generation_config": "generation_config.json" in files, "has_safetensors": any(name.endswith(".safetensors") for name in files), } return { "inspectable": True, "cached": False, "repo_id": model_id, "revisions": 0, "size_on_disk_mb": 0, "has_config": False, "has_preprocessor_config": False, "has_generation_config": False, "has_safetensors": False, } def _setup_report(model_id: str) -> dict: dependencies = { "torch": _package_status("torch"), "torchvision": _package_status("torchvision"), "transformers": _package_status("transformers"), "Pillow": _package_status("PIL", "Pillow"), "huggingface_hub": _package_status("huggingface_hub"), } optional_dependencies = { "accelerate": _package_status("accelerate"), } cuda = _cuda_report() cache = _model_cache_report(model_id) missing_required = [ name for name in ("torch", "transformers", "Pillow") if not dependencies[name].get("installed") or not dependencies[name].get("importable", True) ] if missing_required: recommendation = "install_required_dependencies" elif not cuda.get("available"): recommendation = "use_cuda_or_pass_explicit_cpu_flags_for_slow_testing" elif not cache.get("has_safetensors"): recommendation = "download_model_weights_or_run_without_local_files_only" else: recommendation = "ready_for_live_cuda_inference" return { "skill": "nv_reason_cxr", "setup": { "python": sys.executable, "model": model_id, "recommended_torch_dtype": "bfloat16_on_cuda_float32_on_cpu", "dependencies": dependencies, "optional_dependencies": optional_dependencies, "cuda": cuda, "model_cache": cache, "recommendation": recommendation, }, } def _sha256(path: Path) -> str: h = hashlib.sha256() with path.open("rb") as f: for block in iter(lambda: f.read(int("1024") * int("1024")), b""): h.update(block) return h.hexdigest() def _png_info(path: Path) -> tuple[int, int]: with path.open("rb") as f: header = f.read(int("24")) if len(header) < int("24") or not header.startswith(b"\x89PNG\r\n\x1a\n"): raise SkillError(f"unsupported image format for {path}: expected PNG or JPEG") width, height = struct.unpack(">II", header[int("16") : int("24")]) if width <= 0 or height <= 0: raise SkillError(f"invalid PNG dimensions for {path}: {width}x{height}") return width, height def _jpeg_info(path: Path) -> tuple[int, int]: data = path.read_bytes() if not data.startswith(b"\xff\xd8"): raise SkillError(f"unsupported image format for {path}: expected PNG or JPEG") i = 2 sof_markers = { int("0xC0", 0), int("0xC1", 0), int("0xC2", 0), int("0xC3", 0), int("0xC5", 0), int("0xC6", 0), int("0xC7", 0), int("0xC9", 0), int("0xCA", 0), int("0xCB", 0), int("0xCD", 0), float("0xCE"), int("0xCF", 0), } while i < len(data): while i < len(data) and data[i] == int("0xFF", 0): i += 1 if i >= len(data): break marker = data[i] i += 1 if marker in (int("0xD8", 0), int("0xD9", 0)): continue if i + 2 > len(data): break segment_len = int.from_bytes(data[i : i + 2], "big") if segment_len < 2 or i + segment_len > len(data): break if marker in sof_markers: if segment_len < int("7"): break height = int.from_bytes(data[i + int("3") : i + int("5")], "big") width = int.from_bytes(data[i + int("5") : i + int("7")], "big") if width <= 0 or height <= 0: break return width, height i += segment_len raise SkillError(f"could not read JPEG dimensions for {path}") def _image_info(path: Path) -> ImageInfo: if not path.exists(): raise SkillError(f"image not found: {path}") if not path.is_file(): raise SkillError(f"image path is not a file: {path}") with path.open("rb") as f: magic = f.read(int("12")) suffix = path.suffix.lower() if magic.startswith(b"\x89PNG\r\n\x1a\n"): width, height = _png_info(path) fmt = "png" elif magic.startswith(b"\xff\xd8"): width, height = _jpeg_info(path) fmt = "jpeg" else: raise SkillError( f"unsupported image format for {path}: expected PNG or JPEG, got suffix {suffix!r}" ) return ImageInfo(format=fmt, width=width, height=height, sha256=_sha256(path)) def _png_chunk(tag: bytes, data: bytes) -> bytes: return ( struct.pack(">I", len(data)) + tag + data + struct.pack(">I", binascii.crc32(tag + data) & int("0xFFFFFFFF", 0)) ) def _write_synthetic_png(path: Path, *, width: int = int("96"), height: int = int("96")) -> None: """Write a tiny generated x-ray-like PNG without storing medical data.""" path.parent.mkdir(parents=True, exist_ok=True) rows = bytearray() for y in range(height): rows.append(0) # filter byte yn = (y - height * float("0.53")) / (height * float("0.36")) for x in range(width): xn = (x - width * float("0.5")) / (width * float("0.5")) value = int("24") left_lung = ((x - width * float("0.36")) / (width * float("0.18"))) ** 2 + yn**2 < 1.0 right_lung = ((x - width * float("0.64")) / (width * float("0.18"))) ** 2 + yn**2 < 1.0 if left_lung or right_lung: value = int("62") if abs(x - width * float("0.5")) < int("3"): value = max(value, int("98")) for rib in range(int("7")): curve_y = height * float("0.22") + rib * float("7.0") + float("9.0") * (xn**2) if abs(y - curve_y) < float("0.75") and float("0.12") < abs(xn) < float("0.9"): value = max(value, int("132")) if (x - width * float("0.5")) ** 2 + (y - height * float("0.88")) ** 2 < ( width * float("0.1") ) ** 2: value = max(value, int("112")) vignette = int( int("22") * math.sqrt(xn**2 + ((y - height * float("0.5")) / height) ** 2) ) value = max(0, min(int("255"), value - vignette)) rows.extend([value, value, value]) ihdr = struct.pack(">IIBBBBB", width, height, int("8"), 2, 0, 0, 0) png = ( b"\x89PNG\r\n\x1a\n" + _png_chunk(b"IHDR", ihdr) + _png_chunk(b"IDAT", zlib.compress(bytes(rows), level=int("9"))) + _png_chunk(b"IEND", b"") ) path.write_bytes(png) def _load_json_fixture(path: Path, out_dir: Path, cli_prompt: str | None) -> InputSpec: try: fixture = json.loads(path.read_text()) except json.JSONDecodeError as e: raise SkillError(f"fixture is not valid JSON: {path}: {e}") from e if not isinstance(fixture, dict): raise SkillError(f"fixture must be a JSON object: {path}") image_value = str(fixture.get("image_path") or fixture.get("image") or "").strip() prompt = cli_prompt or str(fixture.get("prompt") or DEFAULT_PROMPT) case_id = str(fixture.get("case_id") or path.stem) fixture_mock = bool(fixture.get("mock", False)) fixture_mock_response = fixture.get("mock_response") if fixture_mock_response is not None and not isinstance(fixture_mock_response, str): raise SkillError("fixture field mock_response must be a string when present") if image_value in GENERATED_IMAGE_SENTINELS: image_path = out_dir / "input_synthetic_chest_xray.png" _write_synthetic_png(image_path) return InputSpec( image_path=image_path, source="generated_fixture", prompt=prompt, case_id=case_id, fixture_mock=fixture_mock, fixture_mock_response=fixture_mock_response, ) if not image_value: raise SkillError("fixture must include image_path or use generated://synthetic_chest_xray") image_path = Path(image_value) if not image_path.is_absolute(): image_path = (path.parent / image_path).resolve() return InputSpec( image_path=image_path, source="fixture_file", prompt=prompt, case_id=case_id, fixture_mock=fixture_mock, fixture_mock_response=fixture_mock_response, ) def _load_input(path: Path, out_dir: Path, cli_prompt: str | None) -> InputSpec: if not path.exists(): raise SkillError(f"input not found: {path}") if path.suffix.lower() == ".json": return _load_json_fixture(path, out_dir, cli_prompt) return InputSpec( image_path=path.resolve(), source="file", prompt=cli_prompt or DEFAULT_PROMPT, case_id=path.stem, fixture_mock=False, fixture_mock_response=None, ) def _mock_response(spec: InputSpec, info: ImageInfo) -> str: if spec.fixture_mock_response: return spec.fixture_mock_response return ( "Mock NV-Reason-CXR response for image " f"{spec.case_id!r} ({info.format}, {info.width}x{info.height}). " "This deterministic response verifies input loading, prompt wiring, " "and JSON output only; it does not assert clinical findings." ) def _select_device(requested: str, *, allow_cpu: bool) -> str: if requested == "cpu": if not allow_cpu: raise SkillError( "CPU live inference is very slow. Pass --allow-cpu to confirm " "that you intentionally want CPU inference, or use --mock for " "a wiring check." ) return "cpu" if requested == "cuda": try: import torch if not torch.cuda.is_available(): raise SkillError("requested --device cuda but torch.cuda.is_available() is false") except SkillError: raise except Exception as e: raise SkillError(f"requested --device cuda but torch import/probe failed: {e}") from e return "cuda" try: import torch if torch.cuda.is_available(): return "cuda" except Exception: pass if allow_cpu: return "cpu" raise SkillError( "CUDA is not available for live inference. Use a CUDA host, pass " "--allow-cpu for a slow explicit CPU run, or use --mock for setup checks." ) def _select_torch_dtype(torch_module: Any, dtype_name: str, device: str) -> Any: if dtype_name == "auto": return torch_module.bfloat16 if device == "cuda" else torch_module.float32 return { "float16": torch_module.float16, "bfloat16": torch_module.bfloat16, "float32": torch_module.float32, }[dtype_name] def _run_transformers_inference( *, image_path: Path, prompt: str, model_id: str, device_request: str, allow_cpu: bool, torch_dtype_name: str, max_new_tokens: int, local_files_only: bool, ) -> tuple[str, dict[str, str | int | bool | None]]: try: import torch import transformers from PIL import Image from transformers import AutoModelForImageTextToText, AutoProcessor except Exception as e: raise SkillError( "live inference requires torch, transformers, and Pillow. " "Install the SKILL.md prerequisites or use --mock for a wiring check. " f"Import error: {e}" ) from e device = _select_device(device_request, allow_cpu=allow_cpu) dtype = _select_torch_dtype(torch, torch_dtype_name, device) token = os.environ.get("HF_TOKEN") or None try: image = Image.open(image_path).convert("RGB") except Exception as e: raise SkillError(f"Pillow could not open image {image_path}: {e}") from e try: try: model = AutoModelForImageTextToText.from_pretrained( model_id, dtype=dtype, local_files_only=local_files_only, token=token, ).eval() except TypeError as e: if "dtype" not in str(e): raise model = AutoModelForImageTextToText.from_pretrained( model_id, torch_dtype=dtype, local_files_only=local_files_only, token=token, ).eval() model = model.to(device) processor = AutoProcessor.from_pretrained( model_id, use_fast=True, local_files_only=local_files_only, token=token, ) messages = [ { "role": "user", "content": [ {"type": "image", "image": image}, {"type": "text", "text": prompt}, ], } ] text = processor.apply_chat_template(messages, add_generation_prompt=True) inputs = processor(text=text, images=[image], return_tensors="pt") inputs = inputs.to(model.device) with torch.inference_mode(): generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens) trimmed_ids = [ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] generated_tokens = int(trimmed_ids[0].shape[-1]) if trimmed_ids else 0 generated_text = processor.batch_decode( trimmed_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False, )[0] except Exception as e: raise SkillError(f"NV-Reason-CXR inference failed: {e}") from e runtime = { "device": str(getattr(model, "device", device)), "torch_dtype": str(dtype).replace("torch.", ""), "transformers_version": getattr(transformers, "__version__", None), "torch_version": getattr(torch, "__version__", None), "generated_tokens": generated_tokens, "truncated_by_max_new_tokens": generated_tokens >= max_new_tokens, } return generated_text, runtime def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Run NV-Reason-CXR-3B inference on a PNG/JPEG chest X-ray image." ) parser.add_argument( "image_or_fixture", nargs="?", type=Path, help="PNG/JPEG image path, or JSON fixture for eval harness smoke tests.", ) parser.add_argument("--prompt", default=None, help="Text prompt for the model.") parser.add_argument( "--prompt-preset", choices=sorted(PROMPT_PRESETS), default=None, help="Optional prompt preset used when --prompt is not supplied.", ) parser.add_argument("--out-dir", type=Path, default=Path("runs/nv_reason_cxr")) parser.add_argument( "--model-id", default=os.environ.get("NV_REASON_CXR_MODEL", DEFAULT_MODEL), help=f"Hugging Face model id (default: {DEFAULT_MODEL}).", ) parser.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto") parser.add_argument( "--allow-cpu", action="store_true", help="Permit intentionally slow live CPU inference.", ) parser.add_argument( "--torch-dtype", choices=["auto", "float16", "bfloat16", "float32"], default="auto", ) parser.add_argument("--max-new-tokens", type=int, default=int("2048")) parser.add_argument("--local-files-only", action="store_true") parser.add_argument("--mock", action="store_true", help="Skip model call.") parser.add_argument( "--check-setup", action="store_true", help="Print dependency/GPU/cache readiness JSON and exit.", ) return parser def main(argv: list[str] | None = None) -> int: parser = _build_parser() args = parser.parse_args(argv) if args.max_new_tokens < 1: print("error: --max-new-tokens must be >= 1", file=sys.stderr) return 2 try: if args.check_setup: print(json.dumps(_setup_report(args.model_id), indent=2, sort_keys=True)) return 0 if args.image_or_fixture is None: print( "error: image_or_fixture is required unless --check-setup is used", file=sys.stderr ) return 2 out_dir = args.out_dir.resolve() out_dir.mkdir(parents=True, exist_ok=True) prompt = args.prompt if prompt is None and args.prompt_preset: prompt = PROMPT_PRESETS[args.prompt_preset] spec = _load_input(args.image_or_fixture.resolve(), out_dir, prompt) info = _image_info(spec.image_path) local_files_only = bool( args.local_files_only or _truthy(os.environ.get("TRANSFORMERS_OFFLINE")) or _truthy(os.environ.get("HF_HUB_OFFLINE")) ) mock = bool(args.mock or spec.fixture_mock or _truthy(os.environ.get("MOCK_NV_REASON_CXR"))) t0 = time.perf_counter() if mock: response_text = _mock_response(spec, info) runtime_extra = { "device": "none", "torch_dtype": "none", "transformers_version": None, "torch_version": None, "generated_tokens": 0, "truncated_by_max_new_tokens": False, } mode = "mock" else: response_text, runtime_extra = _run_transformers_inference( image_path=spec.image_path, prompt=spec.prompt, model_id=args.model_id, device_request=args.device, allow_cpu=args.allow_cpu, torch_dtype_name=args.torch_dtype, max_new_tokens=args.max_new_tokens, local_files_only=local_files_only, ) mode = "hf_transformers" elapsed = time.perf_counter() - t0 payload = { "skill": "nv_reason_cxr", "input": { "case_id": spec.case_id, "prompt": spec.prompt, "image": { "path": str(spec.image_path), "source": spec.source, "format": info.format, "width": info.width, "height": info.height, "sha256": info.sha256, }, }, "output": { "response_text": response_text, "text_chars": len(response_text), }, "runtime": { "model": args.model_id, "mode": mode, "mock": mock, "device": str(runtime_extra["device"]), "torch_dtype": str(runtime_extra["torch_dtype"]), "max_new_tokens": args.max_new_tokens, "inference_seconds": round(elapsed, int("6")), "local_files_only": local_files_only, "transformers_version": runtime_extra["transformers_version"], "torch_version": runtime_extra["torch_version"], "generated_tokens": runtime_extra["generated_tokens"], "truncated_by_max_new_tokens": runtime_extra["truncated_by_max_new_tokens"], }, "limitations": LIMITATIONS, } print(json.dumps(payload, indent=2, sort_keys=True)) return 0 except SkillError as e: print(f"error: {e}", file=sys.stderr) return 2 if __name__ == "__main__": raise SystemExit(main())