# /// script
# requires-python = "==3.10"
# dependencies = [
#     "kernels",
#     "numpy",
#     "pillow",
#     "torch",
# ]
# ///
import torch
from PIL import Image
import numpy as np
from kernels import get_kernel

# This downloads, caches, and loads the kernel library
# and makes the custom op available in torch.ops
img2gray_lib = get_kernel("drbh/img2gray")

img = Image.open("kernel-builder-logo-color.png").convert("RGB")
img = np.array(img)
img_tensor = torch.from_numpy(img).cuda()
print(img_tensor.shape)  # HWC

gray_tensor = img2gray_lib.img2gray(img_tensor).squeeze()
print(gray_tensor.shape)  # HW

# save the output image
gray_img = Image.fromarray(gray_tensor.cpu().numpy().astype(np.uint8))
gray_img.save("kernel-builder-logo-gray2.png")