#!/usr/bin/env python3 """Explicit-setup OpenCLIP vision embedding worker. `--setup` is the only mode allowed to obtain model artifacts. `--embed` runs with outbound sockets denied and checks the setup manifest digest before it loads the model, so a REST embedding request can never silently provision. """ import argparse import base64 import hashlib import io import json import os import socket import sys MODEL = "ViT-B-32" PRETRAINED = "laion2b_s34b_b79k" MODEL_REPOSITORY = "laion/CLIP-ViT-B-32-laion2B-s34B-b79K" MODEL_REVISION = "1a25a446712ba5ee05982a381eed697ef9b435cf" MODEL_WEIGHTS_FILENAME = "open_clip_pytorch_model.bin" MODEL_WEIGHTS_BYTES = 605219813 MODEL_WEIGHTS_SHA256 = "1bd3c7172de5b207ceac554f5ab5266166f3b9baccc9af5989bc801016d080ad" MODEL_DIGEST = "sha256:" + MODEL_WEIGHTS_SHA256 EMBEDDING_DIMENSION = 512 def emit(value): print(json.dumps(value, separators=(",", ":")), flush=True) def deny_network(): def denied(*_args, **_kwargs): raise RuntimeError("OpenCLIP inference is offline; run explicit vision setup to obtain artifacts") socket.create_connection = denied socket.socket.connect = denied socket.socket.connect_ex = denied def sha256_file(path): digest = hashlib.sha256() with open(path, "rb") as source: while True: block = source.read(1024 * 1024) if not block: break digest.update(block) return digest.hexdigest() def model_weights_path(cache_dir): return os.path.join(cache_dir, MODEL_WEIGHTS_FILENAME) def verify_model_weights(cache_dir): path = model_weights_path(cache_dir) if not os.path.isfile(path): raise RuntimeError("Pinned OpenCLIP weight is missing; run explicit vision setup") actual_bytes = os.path.getsize(path) if actual_bytes != MODEL_WEIGHTS_BYTES: raise RuntimeError( "Pinned OpenCLIP weight has an unexpected size " "(expected %d bytes, got %d); run explicit vision setup" % (MODEL_WEIGHTS_BYTES, actual_bytes) ) actual_sha256 = sha256_file(path) if actual_sha256 != MODEL_WEIGHTS_SHA256: raise RuntimeError( "Pinned OpenCLIP weight checksum mismatch " "(expected %s, got %s); run explicit vision setup" % (MODEL_WEIGHTS_SHA256, actual_sha256) ) return path def artifact_entries(cache_dir): path = verify_model_weights(cache_dir) return [{ "path": MODEL_WEIGHTS_FILENAME, "bytes": os.path.getsize(path), "sha256": MODEL_WEIGHTS_SHA256, }] def download_pinned_model_weights(cache_dir): os.makedirs(cache_dir, exist_ok=True) try: verified = verify_model_weights(cache_dir) except RuntimeError: # An interrupted/corrupt setup may leave a file at the managed target. # Only explicit setup is permitted to replace it. force_download = True else: return verified from huggingface_hub import hf_hub_download downloaded = hf_hub_download( repo_id=MODEL_REPOSITORY, filename=MODEL_WEIGHTS_FILENAME, revision=MODEL_REVISION, local_dir=cache_dir, local_files_only=False, force_download=force_download, ) expected_path = model_weights_path(cache_dir) if os.path.abspath(downloaded) != os.path.abspath(expected_path): raise RuntimeError("Pinned OpenCLIP download did not materialize in the managed cache") return verify_model_weights(cache_dir) def load_model(weights_path, offline): if offline: deny_network() os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" import torch import open_clip if not torch.cuda.is_available(): raise RuntimeError("JetPack CUDA Torch is unavailable for OpenCLIP") if os.environ.get("CUDA_VISIBLE_DEVICES") != "0": raise RuntimeError("OpenCLIP requires CUDA_VISIBLE_DEVICES=0") if torch.cuda.current_device() != 0: raise RuntimeError("OpenCLIP did not bind to logical CUDA device 0") device = torch.device("cuda:0") model, _, preprocess = open_clip.create_model_and_transforms( MODEL, pretrained=weights_path, device=device, ) model.eval() return torch, model, preprocess, device def setup(cache_dir): _weights_path = download_pinned_model_weights(cache_dir) import torch if not torch.cuda.is_available(): raise RuntimeError("JetPack CUDA Torch is unavailable for OpenCLIP") if os.environ.get("CUDA_VISIBLE_DEVICES") != "0": raise RuntimeError("OpenCLIP setup requires CUDA_VISIBLE_DEVICES=0") if torch.cuda.current_device() != 0: raise RuntimeError("OpenCLIP setup did not bind to logical CUDA device 0") entries = artifact_entries(cache_dir) major, minor = torch.cuda.get_device_capability(0) emit({ "type": "setup", "success": True, "model": MODEL, "pretrained": PRETRAINED, "model_repository": MODEL_REPOSITORY, "model_revision": MODEL_REVISION, "model_weights_sha256": MODEL_WEIGHTS_SHA256, "backend": "open-clip", "embedding_dimension": EMBEDDING_DIMENSION, "normalization": "l2", "device": torch.cuda.get_device_name(0), "compute_capability": "%d.%d" % (major, minor), "torch_cuda_version": str(torch.version.cuda), "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "artifacts": entries, "model_digest": MODEL_DIGEST, }) def image_from_input(input_): from PIL import Image path = input_.get("path") encoded = input_.get("bytesBase64") if isinstance(path, str) and path: return Image.open(path).convert("RGB") if isinstance(encoded, str) and encoded: return Image.open(io.BytesIO(base64.b64decode(encoded, validate=True))).convert("RGB") raise ValueError("Provide path or bytesBase64") def embed(cache_dir, expected_digest, expected_revision): if expected_digest != MODEL_DIGEST or expected_revision != MODEL_REVISION: raise RuntimeError("OpenCLIP model identity does not match the pinned setup contract") weights_path = verify_model_weights(cache_dir) input_ = json.loads(sys.stdin.read() or "{}") torch, model, preprocess, device = load_model(weights_path, True) image = image_from_input(input_) with torch.inference_mode(): tensor = preprocess(image).unsqueeze(0).to(device) vector = model.encode_image(tensor).float() vector = vector / vector.norm(dim=-1, keepdim=True).clamp_min(1e-12) values = vector[0].detach().cpu().numpy().astype("float32").tolist() if len(values) != EMBEDDING_DIMENSION or not all(isinstance(value, float) for value in values): raise RuntimeError("OpenCLIP returned an invalid image embedding") emit({ "type": "embedding", "success": True, "model": MODEL, "pretrained": PRETRAINED, "model_repository": MODEL_REPOSITORY, "model_revision": MODEL_REVISION, "model_weights_sha256": MODEL_WEIGHTS_SHA256, "backend": "open-clip", "model_digest": MODEL_DIGEST, "embedding": values, "dimension": EMBEDDING_DIMENSION, "normalization": "l2", }) def main(): parser = argparse.ArgumentParser() parser.add_argument("--setup", action="store_true") parser.add_argument("--embed", action="store_true") parser.add_argument("--cache-dir", required=True) parser.add_argument("--model-digest") parser.add_argument("--model-revision") args = parser.parse_args() if args.setup == args.embed: raise ValueError("Choose exactly one of --setup or --embed") if args.setup: setup(args.cache_dir) else: if not args.model_digest or not args.model_revision: raise ValueError("--model-digest and --model-revision are required for offline embedding") embed(args.cache_dir, args.model_digest, args.model_revision) if __name__ == "__main__": try: main() except Exception as error: emit({"type": "error", "success": False, "error": str(error)}) raise SystemExit(1)