| import argparse |
| from pathlib import Path |
| import sys |
|
|
| from PIL import Image |
| import torch |
| from huggingface_hub import snapshot_download |
| from torchvision import transforms |
|
|
| sys.path.insert(0, str(Path(__file__).absolute().parents[1])) |
| from vision import load_vision |
|
|
|
|
| MODEL_ID = "immanuelpeter/C-RADIOv4-H-3.5-Super" |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Extract Nemotron Super visual features.") |
| parser.add_argument("image", type=Path) |
| parser.add_argument("--model", default=MODEL_ID) |
| args = parser.parse_args() |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 |
| bundle = Path(args.model) if Path(args.model).is_dir() else Path(snapshot_download(args.model)) |
| model = load_vision(bundle, dtype=dtype).to(device) |
|
|
| image = Image.open(args.image).convert("RGB") |
| pixels = transforms.Compose([ |
| transforms.Resize(448, interpolation=transforms.InterpolationMode.BICUBIC), |
| transforms.CenterCrop(448), |
| transforms.ToTensor(), |
| ])(image).unsqueeze(0).to(device) |
|
|
| for stage, tokens in model(pixels).items(): |
| print(stage, tuple(tokens.shape)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
