fix: improve python bindings and sanity check

flake.lock

@@ -73,11 +73,11 @@
         "nixpkgs": "nixpkgs"
       },
       "locked": {
-        "lastModified": 1750234878,
-        "narHash": "sha256-q9DRC9zdpzUf88qqg1qbhP1qgJbE2cMtn8oUmosuyT8=",
+        "lastModified": 1754038838,
+        "narHash": "sha256-oHigCT4z0ayyLyEuxdZooSXRAZP8lfOkZHzY1lx1U50=",
         "owner": "huggingface",
         "repo": "hf-nix",
-        "rev": "c7132f90763d756da3e77da62e01be0a4546dc57",
+        "rev": "336f781fa284e193baa3d4c3ce3f95fb34e9ffad",
         "type": "github"
       },
       "original": {
@@ -98,11 +98,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1750790603,
-        "narHash": "sha256-m7FoTYWDV811Y7FiuJPa/uCOV63rf6LHxWportuI9h0=",
+        "lastModified": 1755181472,
+        "narHash": "sha256-xOXjhehC5xi/XB4fXZ5c0L2sSyDjJQdlH7/BcdHLBaM=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "37cad313efea84e213b2fc13b2ec808d273a126d",
+        "rev": "85da46f660c1c43b40771c3df3b223bb3fa39bec",
         "type": "github"
       },
       "original": {
@@ -113,17 +113,17 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1747820358,
-        "narHash": "sha256-fTqsZsUX6M3yeEvgyQvXcbGmT2CaRVyVwsi8eK29Oj4=",
-        "owner": "danieldk",
+        "lastModified": 1752785354,
+        "narHash": "sha256-Y33ryUz7MPqKrZwlbQcsYCUz2jAJCacRf8jbs0tYUlA=",
+        "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "d3c1681180717528068082103bf323147de6ab0b",
+        "rev": "d38025438a6ee456758dc03188ca6873a415463b",
         "type": "github"
       },
       "original": {
-        "owner": "danieldk",
-        "ref": "cudatoolkit-12.9-kernel-builder",
+        "owner": "nixos",
         "repo": "nixpkgs",
+        "rev": "d38025438a6ee456758dc03188ca6873a415463b",
         "type": "github"
       }
     },

scripts/sanity.py

@@ -6,18 +6,14 @@ import numpy as np
 
 print(dir(img2gray))
 
-img = Image.open("/home/ubuntu/Projects/img2gray/kernel-builder-logo-color.png").convert("RGB")
+img = Image.open("kernel-builder-logo-color.png").convert("RGB")
 img = np.array(img)
-img_tensor = torch.from_numpy(img)
+img_tensor = torch.from_numpy(img).cuda()
 print(img_tensor.shape)  # HWC
-img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).contiguous().cuda()  # BCHW
-print(img_tensor.shape)  # BCHW
 
 gray_tensor = img2gray.img2gray(img_tensor).squeeze()
-print(gray_tensor.shape)  # B1HW
+print(gray_tensor.shape)  # HW
 
 # save the output image
-gray_img = gray_tensor.cpu().numpy()  # 1HW -> HW
-gray_img = Image.fromarray(gray_img.astype(np.uint8), mode="L")
-
-gray_img.save("/home/ubuntu/Projects/img2gray/kernel-builder-logo-gray.png")
+gray_img = Image.fromarray(gray_tensor.cpu().numpy().astype(np.uint8), mode="L")
+gray_img.save("kernel-builder-logo-gray.png")

torch-ext/img2gray/__init__.py

@@ -2,17 +2,13 @@ import torch
 
 from ._ops import ops
 
-def img2gray(input: torch.Tensor) -> torch.Tensor:
-    # we expect input to be in BCHW format
-    batch, channels, height, width = input.shape
 
+def img2gray(input: torch.Tensor) -> torch.Tensor:
+    # we expect input to be in CHW format
+    height, width, channels = input.shape
     assert channels == 3, "Input image must have 3 channels (RGB)"
 
-    output = torch.empty((batch, 1, height, width), device=input.device, dtype=input.dtype)
-
-    for b in range(batch):
-        single_image = input[b].permute(1, 2, 0).contiguous()  # HWC
-        single_output = output[b].reshape(height, width)  # HW
-        ops.img2gray(single_image, single_output)
+    output = torch.empty((height, width), device=input.device, dtype=input.dtype)
+    ops.img2gray(input, output)
 
-    return output
+    return output

torch-ext/torch_binding.cpp

@@ -1,12 +1,11 @@
+// torch-ext/torch_binding.cpp
 #include <torch/library.h>
-
-#include "registration.h"
-#include "torch_binding.h"
-
+#include "registration.h" // included in the build
+#include "torch_binding.h" // Declares our img2gray_cuda function
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-    ops.def("img2gray(Tensor input, Tensor output) -> ()");
-    ops.impl("img2gray", torch::kCUDA, &img2gray_cuda);
+  ops.def("img2gray(Tensor input, Tensor! output) -> ()");
+  ops.impl("img2gray", torch::kCUDA, &img2gray_cuda);
 }
-  
-REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)