Port CPU decode_webp to stable ABI and retire the unstable image extension

CMakeLists.txt

@@ -112,6 +112,7 @@ set(TORCHVISION_STABLE_SOURCES
   ${TVCPP}/io/image/cpu/decode_jpeg.cpp
   ${TVCPP}/io/image/cpu/common_jpeg.cpp
   ${TVCPP}/io/image/cpu/decode_gif.cpp
+  ${TVCPP}/io/image/cpu/decode_webp.cpp
   ${TVCPP}/io/image/common_stable.cpp
   ${TVCPP}/vision_stable.cpp)
 # Pin them to torch 2.11. Per source, not library-wide: the pin makes ATen headers

packaging/wheel/relocate.py

@@ -315,9 +315,8 @@ def patch_linux():
     wheels = glob.glob(osp.join(PACKAGE_ROOT, "dist", "*.whl"))
     output_dir = osp.join(PACKAGE_ROOT, "dist", ".wheel-process")
 
-    image_binary = "image.so"
     image_stable_binary = "image_stable.so"
-    torchvision_binaries = [image_binary, image_stable_binary]
+    torchvision_binaries = [image_stable_binary]
     for wheel in wheels:
         if osp.exists(output_dir):
             shutil.rmtree(output_dir)
@@ -351,9 +350,8 @@ def patch_win():
     wheels = glob.glob(osp.join(PACKAGE_ROOT, "dist", "*.whl"))
     output_dir = osp.join(PACKAGE_ROOT, "dist", ".wheel-process")
 
-    image_binary = "image.pyd"
     image_stable_binary = "image_stable.pyd"
-    torchvision_binaries = [image_binary, image_stable_binary]
+    torchvision_binaries = [image_stable_binary]
     for wheel in wheels:
         if osp.exists(output_dir):
             shutil.rmtree(output_dir)

setup.py

@@ -164,25 +164,8 @@ def get_macros_and_flags():
     CSRS_DIR / "ops/cpu/nms_kernel.cpp",
     CSRS_DIR / "ops/mps/nms_kernel.mm",
     CSRS_DIR / "ops/quantized/cpu/qnms_kernel.cpp",
-    CSRS_DIR / "io/image/common_stable.cpp",
-    CSRS_DIR / "io/image/cpu/encode_png.cpp",
-    CSRS_DIR / "io/image/cpu/encode_jpeg.cpp",
-    CSRS_DIR / "io/image/cpu/read_write_file.cpp",
-    CSRS_DIR / "io/image/cpu/decode_image.cpp",
-    CSRS_DIR / "io/image/cpu/decode_png.cpp",
-    CSRS_DIR / "io/image/cpu/decode_jpeg.cpp",
-    CSRS_DIR / "io/image/cpu/common_jpeg.cpp",
-    CSRS_DIR / "io/image/cpu/decode_gif.cpp",
-    CSRS_DIR / "io/image/cpu/giflib/dgif_lib.c",
-    CSRS_DIR / "io/image/cpu/giflib/gif_hash.c",
-    CSRS_DIR / "io/image/cpu/giflib/gifalloc.c",
-    CSRS_DIR / "io/image/cpu/giflib/openbsd-reallocarray.c",
 }
 STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/nms_kernel.hip" if IS_ROCM else "ops/cuda/nms_kernel.cu"))
-STABLE_SOURCES.add(
-    CSRS_DIR / ("io/image/hip/decode_jpegs_cuda.cpp" if IS_ROCM else "io/image/cuda/decode_jpegs_cuda.cpp")
-)
-STABLE_SOURCES.add(CSRS_DIR / "io/image/cuda/encode_jpegs_cuda.cpp")
 
 
 def _not_stable(paths):
@@ -388,58 +371,7 @@ def find_library(header):
     return False, None, None
 
 
-def make_image_extension():
-    print("Building image extension")
-
-    include_dirs = TORCHVISION_INCLUDE.copy()
-    library_dirs = TORCHVISION_LIBRARY.copy()
-
-    libraries = []
-    define_macros, extra_compile_args = get_macros_and_flags()
-
-    image_dir = CSRS_DIR / "io/image"
-    sources = (
-        _not_stable(image_dir.glob("*.cpp"))
-        + _not_stable(image_dir.glob("cpu/*.cpp"))
-        + _not_stable(image_dir.glob("cpu/giflib/*.c"))
-    )
-
-    if IS_ROCM:
-        sources += _not_stable(image_dir.glob("hip/*.cpp"))
-    else:
-        sources += _not_stable(image_dir.glob("cuda/*.cpp"))
-
-    Extension = CppExtension
-
-    if USE_WEBP:
-        webp_found, webp_include_dir, webp_library_dir = find_library(header="webp/decode.h")
-        if webp_found:
-            print("Building torchvision with WEBP support")
-            print(f"{webp_include_dir = }")
-            print(f"{webp_library_dir = }")
-            if webp_include_dir is not None and webp_library_dir is not None:
-                # if those are None it means they come from standard paths that are already in the search paths, which we don't need to re-add.
-                include_dirs.append(webp_include_dir)
-                library_dirs.append(webp_library_dir)
-            webp_library = "libwebp" if sys.platform == "win32" else "webp"
-            libraries.append(webp_library)
-            define_macros += [("WEBP_FOUND", 1)]
-        else:
-            warnings.warn("Building torchvision without WEBP support")
-
-    return Extension(
-        name="torchvision.image",
-        sources=sorted(str(s) for s in sources),
-        include_dirs=include_dirs,
-        library_dirs=library_dirs,
-        define_macros=define_macros,
-        libraries=libraries,
-        extra_compile_args=extra_compile_args,
-    )
-
-
 def make_image_stable_extension():
-    # Stable-ABI sibling of make_image_extension(): only the migrated (_stable) image TUs.
     print("Building image_stable extension")
 
     include_dirs = TORCHVISION_INCLUDE.copy()
@@ -449,10 +381,10 @@ def make_image_stable_extension():
 
     image_dir = CSRS_DIR / "io/image"
     sources = (
-        _stable(image_dir.glob("*.cpp"))
-        + _stable(image_dir.glob("cpu/*.cpp"))
-        + _stable(image_dir.glob("cpu/giflib/*.c"))
-        + _stable(image_dir.glob("hip/*.cpp" if IS_ROCM else "cuda/*.cpp"))
+        list(image_dir.glob("*.cpp"))
+        + list(image_dir.glob("cpu/*.cpp"))
+        + list(image_dir.glob("cpu/giflib/*.c"))
+        + list(image_dir.glob("hip/*.cpp" if IS_ROCM else "cuda/*.cpp"))
     )
 
     Extension = CppExtension
@@ -480,6 +412,22 @@ def make_image_stable_extension():
         else:
             warnings.warn("Building torchvision without JPEG support")
 
+    if USE_WEBP:
+        webp_found, webp_include_dir, webp_library_dir = find_library(header="webp/decode.h")
+        if webp_found:
+            print("Building torchvision with WEBP support")
+            print(f"{webp_include_dir = }")
+            print(f"{webp_library_dir = }")
+            if webp_include_dir is not None and webp_library_dir is not None:
+                # if those are None it means they come from standard paths that are already in the search paths, which we don't need to re-add.
+                include_dirs.append(webp_include_dir)
+                library_dirs.append(webp_library_dir)
+            webp_library = "libwebp" if sys.platform == "win32" else "webp"
+            libraries.append(webp_library)
+            define_macros += [("WEBP_FOUND", 1)]
+        else:
+            warnings.warn("Building torchvision without WEBP support")
+
     if USE_NVJPEG and (torch.cuda.is_available() or FORCE_CUDA):
         nvjpeg_found = CUDA_HOME is not None and (Path(CUDA_HOME) / "include/nvjpeg.h").exists()
         if nvjpeg_found:
@@ -536,7 +484,6 @@ def run(self):
     extensions = [
         make_C_extension(),
         make_C_stable_extension(),
-        make_image_extension(),
         make_image_stable_extension(),
     ]
 

torchvision/csrc/io/image/cpu/decode_image.cpp

@@ -8,28 +8,11 @@
 #include "decode_gif.h"
 #include "decode_jpeg.h"
 #include "decode_png.h"
+#include "decode_webp.h"
 
 namespace vision {
 namespace image {
 
-namespace {
-
-// Shim over the legacy image::decode_webp op
-// not yet on the stable ABI.
-// TODO(stable-abi): remove the shim once the decoder is ported.
-torch::stable::Tensor decode_webp(
-    const torch::stable::Tensor& data,
-    ImageReadMode mode) {
-  const auto num_args = 2;
-  std::array<StableIValue, num_args> stack{
-      torch::stable::detail::from(data), torch::stable::detail::from(mode)};
-  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
-      "image::decode_webp", "", stack.data(), TORCH_ABI_VERSION));
-  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
-}
-
-} // namespace
-
 torch::stable::Tensor decode_image(
     const torch::stable::Tensor& data,
     ImageReadMode mode,

torchvision/csrc/io/image/cpu/decode_webp.cpp

@@ -1,5 +1,9 @@
 #include "decode_webp.h"
-#include "../common.h"
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+
+#include "../common_stable.h"
 
 #if WEBP_FOUND
 #include "webp/decode.h"
@@ -10,20 +14,41 @@ namespace vision {
 namespace image {
 
 #if !WEBP_FOUND
-torch::Tensor decode_webp(
-    const torch::Tensor& encoded_data,
+torch::stable::Tensor decode_webp(
+    const torch::stable::Tensor& encoded_data,
     ImageReadMode mode) {
   STD_TORCH_CHECK(
       false, "decode_webp: torchvision not compiled with libwebp support");
 }
 #else
 
-torch::Tensor decode_webp(
-    const torch::Tensor& encoded_data,
+namespace {
+bool should_this_return_rgb_or_rgba_let_me_know_in_the_comments_down_below_guys_see_you_in_the_next_video(
+    ImageReadMode mode,
+    bool has_alpha) {
+  // Return true if the calling decoding function should return a 3D RGB tensor,
+  // and false if it should return a 4D RGBA tensor.
+  // This function ignores the requested "grayscale" modes and treats it as
+  // "unchanged", so it should only used on decoders who don't support grayscale
+  // outputs.
+
+  if (mode == IMAGE_READ_MODE_RGB) {
+    return true;
+  }
+  if (mode == IMAGE_READ_MODE_RGB_ALPHA) {
+    return false;
+  }
+  // From here we assume mode is "unchanged", even for grayscale ones.
+  return !has_alpha;
+}
+} // namespace
+
+torch::stable::Tensor decode_webp(
+    const torch::stable::Tensor& encoded_data,
     ImageReadMode mode) {
   validate_encoded_data(encoded_data);
 
-  auto encoded_data_p = encoded_data.data_ptr<uint8_t>();
+  auto encoded_data_p = encoded_data.const_data_ptr<uint8_t>();
   auto encoded_data_size = encoded_data.numel();
 
   WebPBitstreamFeatures features;
@@ -34,7 +59,12 @@ torch::Tensor decode_webp(
       !features.has_animation, "Animated webp files are not supported.");
 
   if (mode == IMAGE_READ_MODE_GRAY || mode == IMAGE_READ_MODE_GRAY_ALPHA) {
-    TORCH_WARN_ONCE(
+    // TODO(stable-abi): warn once per process currently not available, replace
+    // when a stable warn-once lands.
+    aoti_torch_warn(
+        __func__,
+        __FILE__,
+        static_cast<uint32_t>(__LINE__),
         "Webp does not support grayscale conversions. "
         "The returned tensor will be in the colorspace of the original image.");
   }
@@ -55,12 +85,25 @@ torch::Tensor decode_webp(
   STD_TORCH_CHECK(decoded_data != nullptr, "WebPDecodeRGB[A] failed.");
 
   auto deleter = [decoded_data](void*) { WebPFree(decoded_data); };
-  auto out = torch::from_blob(
-      decoded_data, {height, width, num_channels}, deleter, torch::kUInt8);
+  auto out = torch::stable::from_blob(
+      decoded_data,
+      {height, width, num_channels},
+      {width * num_channels, num_channels, 1},
+      torch::stable::Device(torch::headeronly::DeviceType::CPU),
+      torch::headeronly::ScalarType::Byte,
+      deleter);
 
-  return out.permute({2, 0, 1});
+  return stable_permute(out, {2, 0, 1});
 }
 #endif // WEBP_FOUND
 
+STABLE_TORCH_LIBRARY_FRAGMENT(image, m) {
+  m.def("decode_webp(Tensor encoded_data, int mode) -> Tensor");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(image, CompositeExplicitAutograd, m) {
+  m.impl("decode_webp", TORCH_BOX(&decode_webp));
+}
+
 } // namespace image
 } // namespace vision

torchvision/csrc/io/image/cpu/decode_webp.h

@@ -1,13 +1,13 @@
 #pragma once
 
-#include <torch/types.h>
-#include "../common.h"
+#include <torch/csrc/stable/tensor.h>
+#include "../common_stable.h"
 
 namespace vision {
 namespace image {
 
-C10_EXPORT torch::Tensor decode_webp(
-    const torch::Tensor& encoded_data,
+torch::stable::Tensor decode_webp(
+    const torch::stable::Tensor& encoded_data,
     ImageReadMode mode = IMAGE_READ_MODE_UNCHANGED);
 
 } // namespace image