[LLM runner] Support HWC inputs in make_image_input (pytorch#22888)

## Summary

- add proper HWC and 1xHWC support to `make_image_input`
- convert HWC tensors to CHW internally
- keep existing CHW behavior and improve validation/docs
- add coverage for all supported layouts, dtypes, and invalid inputs

Fixes pytorch#22685

## Test plan

- `python -m pytest -q
extension/llm/runner/test/test_runner_pybindings.py`
- 14 passed

cc @larryliu0820 @mergennachin @cccclai @helunwencser @jackzhxng
@digantdesai

Author: fatih-uzlmz

extension/llm/runner/README.md

@@ -274,10 +274,14 @@ token_ids = [1, 15043, 445, 2420]  # Example token IDs
 token_input = make_token_input(token_ids)
 
 # 3. Image input from torch tensor
-# Supports channels-first (C,H,W), or (1,C,H,W) with batch size 1
+# CHW is the default, including when width is 3 or 4
 image_chw = torch.randint(0, 255, (3, 224, 224), dtype=torch.uint8)
 image_input = make_image_input(image_chw)
 
+# HWC inputs require an explicit layout, including when height is 3 or 4
+image_hwc = torch.randint(0, 255, (224, 224, 3), dtype=torch.uint8)
+image_input_hwc = make_image_input(image_hwc, layout="HWC")
+
 # Float tensors also supported for normalized images
 image_float = torch.rand(3, 224, 224, dtype=torch.float32)  # CHW, normalized
 image_input_float = make_image_input(image_float)

extension/llm/runner/_llm_runner.pyi

@@ -553,18 +553,25 @@ def make_text_input(text: str) -> MultimodalInput:
     """
     ...
 
-def make_image_input(image_tensor: torch.Tensor) -> MultimodalInput:
+def make_image_input(
+    image_tensor: torch.Tensor, layout: str = "CHW"
+) -> MultimodalInput:
     """
     Create an image input from a torch tensor.
 
     Args:
-        image_tensor: Torch tensor with shape (H, W, C), (1, H, W, C), (C, H, W), or (1, C, H, W)
+        image_tensor: Contiguous uint8 or float32 tensor with an optional batch
+            dimension of size 1.
+        layout: Tensor layout, either "CHW" or "HWC". Defaults to "CHW" to
+            preserve existing behavior. HWC inputs are converted to CHW
+            internally, including when height is 3 or 4.
 
     Returns:
         A MultimodalInput containing the image
 
     Raises:
-        RuntimeError: If the tensor has invalid dimensions or number of channels
+        RuntimeError: If the tensor has invalid dimensions, layout, dtype, or
+            number of channels
     """
     ...
 

extension/llm/runner/pybindings.cpp

@@ -571,7 +571,8 @@ PYBIND11_MODULE(_llm_runner, m) {
 
   m.def(
       "make_image_input",
-      [](torch::Tensor image_tensor) -> MultimodalInput {
+      [](torch::Tensor image_tensor,
+         const std::string& layout) -> MultimodalInput {
         if (image_tensor.dim() == 4) {
           if (image_tensor.size(0) != 1) {
             throw std::runtime_error(
@@ -582,30 +583,31 @@ PYBIND11_MODULE(_llm_runner, m) {
 
         if (image_tensor.dim() != 3) {
           throw std::runtime_error(
-              "Image tensor must be 3-dimensional (H, W, C) or 4-dimensional (1, H, W, C)");
+              "Image tensor must have shape (H, W, C), (1, H, W, C), (C, H, W), or (1, C, H, W)");
+        }
+
+        if (!image_tensor.is_contiguous()) {
+          throw std::runtime_error("Image tensor must be contiguous");
         }
 
         int64_t height, width, channels;
-        // Check for memory format and permute to CHW if necessary
-        if (image_tensor.is_contiguous(at::MemoryFormat::ChannelsLast)) {
-          // Input is HWC, permute to CHW
+        if (layout == "CHW") {
+          channels = image_tensor.size(0);
+          height = image_tensor.size(1);
+          width = image_tensor.size(2);
+        } else if (layout == "HWC") {
           height = image_tensor.size(0);
           width = image_tensor.size(1);
           channels = image_tensor.size(2);
           image_tensor = image_tensor.permute({2, 0, 1});
-        } else if (image_tensor.is_contiguous(at::MemoryFormat::Contiguous)) {
-          // Input is CHW
-          channels = image_tensor.size(0);
-          height = image_tensor.size(1);
-          width = image_tensor.size(2);
         } else {
           throw std::runtime_error(
-              "Image tensor must be contiguous in either channels last (H, W, C) or contiguous (C, H, W) format.");
+              "Image tensor layout must be 'CHW' or 'HWC'");
         }
 
         if (channels != 3 && channels != 4) {
           throw std::runtime_error(
-              "Image must have 3 (RGB) or 4 (RGBA) channels");
+              "Image must have 3 (RGB) or 4 (RGBA) channels in the dimension specified by layout");
         }
 
         image_tensor = image_tensor.contiguous();
@@ -630,8 +632,9 @@ PYBIND11_MODULE(_llm_runner, m) {
               "Unsupported image tensor dtype. Only uint8 and float32 are supported.");
         }
       },
-      "Create an image input from a torch tensor (H, W, C), (1, H, W, C), (C, H, W), or (1, C, H, W)",
-      py::arg("image_tensor"));
+      "Create an image input from a contiguous uint8 or float32 torch tensor. layout must be 'CHW' or 'HWC' and defaults to 'CHW'.",
+      py::arg("image_tensor"),
+      py::arg("layout") = "CHW");
 
   m.def(
       "make_audio_input",

extension/llm/runner/test/test_runner_pybindings.py

@@ -183,7 +183,7 @@ def test_invalid_image_array(self):
         # Wrong dimensions (expects 3D or 4D tensor)
         with self.assertRaises(RuntimeError) as cm:
             make_image_input(torch.ones((100,), dtype=torch.uint8))
-        self.assertIn("3-dimensional", str(cm.exception))
+        self.assertIn("must have shape", str(cm.exception))
 
         # Wrong number of channels
         with self.assertRaises(RuntimeError) as cm:
@@ -252,15 +252,73 @@ def test_make_text_input(self):
         self.assertEqual(text_input.get_text(), "Hello")
 
     def test_make_image_input(self):
-        """Test make_image_input helper."""
-        # Create a test image tensor (RGB, CHW)
-        img_tensor = torch.zeros((3, 100, 150), dtype=torch.uint8)
-        img_tensor[0, :, :] = 255  # Red channel
+        """Test supported image layouts and dtypes."""
+        dimensions = ((3, 7), (4, 7), (5, 3), (5, 4), (5, 7))
+        for height, width in dimensions:
+            for channels in (3, 4):
+                for dtype in (torch.uint8, torch.float32):
+                    chw = torch.arange(channels * height * width, dtype=dtype).reshape(
+                        channels, height, width
+                    )
+                    hwc = chw.permute(1, 2, 0).contiguous()
+                    layouts = {
+                        "chw_default": (chw, None),
+                        "1xchw": (chw.unsqueeze(0), "CHW"),
+                        "hwc": (hwc, "HWC"),
+                        "1xhwc": (hwc.unsqueeze(0), "HWC"),
+                    }
+
+                    for name, (tensor, layout) in layouts.items():
+                        with self.subTest(
+                            channels=channels,
+                            dtype=dtype,
+                            height=height,
+                            width=width,
+                            layout=name,
+                        ):
+                            image_input = (
+                                make_image_input(tensor)
+                                if layout is None
+                                else make_image_input(tensor, layout=layout)
+                            )
+                            self.assertTrue(image_input.is_image())
+                            image = image_input.get_image()
+                            self.assertEqual(image.channels, channels)
+                            self.assertEqual(image.height, height)
+                            self.assertEqual(image.width, width)
+                            expected_data = chw.flatten().tolist()
+                            actual_data = (
+                                image.uint8_data
+                                if dtype == torch.uint8
+                                else image.float_data
+                            )
+                            self.assertEqual(actual_data, expected_data)
+
+    def test_make_image_input_rejects_invalid_inputs(self):
+        """Test invalid image dimensions, layouts, and dtypes."""
+        with self.assertRaisesRegex(RuntimeError, "must have shape"):
+            make_image_input(torch.ones((100,), dtype=torch.uint8))
 
-        image_input = make_image_input(img_tensor)
-        self.assertTrue(image_input.is_image())
+        with self.assertRaisesRegex(RuntimeError, "Batch size"):
+            make_image_input(torch.ones((2, 3, 5, 7), dtype=torch.uint8))
+
+        with self.assertRaisesRegex(RuntimeError, r"3 \(RGB\) or 4 \(RGBA\)"):
+            make_image_input(torch.ones((2, 5, 7), dtype=torch.uint8))
+
+        with self.assertRaisesRegex(RuntimeError, r"3 \(RGB\) or 4 \(RGBA\)"):
+            make_image_input(torch.ones((5, 7, 2), dtype=torch.uint8), layout="HWC")
 
-        # Test with RGBA (CHW)
-        img_tensor_rgba = torch.ones((4, 50, 50), dtype=torch.uint8) * 128
-        image_input_rgba = make_image_input(img_tensor_rgba)
-        self.assertTrue(image_input_rgba.is_image())
+        with self.assertRaisesRegex(RuntimeError, "must be 'CHW' or 'HWC'"):
+            make_image_input(torch.ones((3, 5, 7), dtype=torch.uint8), layout="NHWC")
+
+        with self.assertRaisesRegex(RuntimeError, "Only uint8 and float32"):
+            make_image_input(torch.ones((3, 5, 7), dtype=torch.int32))
+
+        noncontiguous_tensors = (
+            (torch.ones((3, 5, 7), dtype=torch.uint8).transpose(1, 2), "CHW"),
+            (torch.ones((5, 7, 3), dtype=torch.uint8).transpose(0, 1), "HWC"),
+        )
+        for tensor, layout in noncontiguous_tensors:
+            with self.subTest(shape=tuple(tensor.shape)):
+                with self.assertRaisesRegex(RuntimeError, "must be contiguous"):
+                    make_image_input(tensor, layout=layout)