| # /// script | |
| # requires-python = ">=3.10" | |
| # dependencies = [ | |
| # "torch", | |
| # "triton", | |
| # "kernels", | |
| # ] | |
| # /// | |
| """Minimal smoke test of the published kernel via get_kernel (run on a CUDA box).""" | |
| import torch | |
| from kernels import get_kernel | |
| kir = get_kernel("Molbap/kernel_image_resize", revision="main", trust_remote_code=True) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| images = [ | |
| torch.randint(0, 256, (3, h, w), dtype=torch.uint8, device=device) | |
| for h, w in [(640, 480), (800, 600), (384, 1024)] | |
| ] | |
| pixel_values = kir.resize_normalize( | |
| images, | |
| size=384, | |
| image_mean=[0.5, 0.5, 0.5], | |
| image_std=[0.5, 0.5, 0.5], | |
| rescale_factor=1 / 255, | |
| resample="bicubic", | |
| antialias=True, | |
| ) | |
| print(f"{len(images)} ragged images -> {tuple(pixel_values.shape)} {pixel_values.dtype}") | |