diff --git a/CMakeLists.txt b/CMakeLists.txt index 39a9fe58fca..1d9ab7b24d2 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/packaging/wheel/relocate.py b/packaging/wheel/relocate.py index 59f8688d931..b83811047a7 100644 --- a/packaging/wheel/relocate.py +++ b/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) diff --git a/setup.py b/setup.py index 6f95bbbfd36..fb57c645fdd 100644 --- a/setup.py +++ b/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(), ] diff --git a/torchvision/csrc/io/image/common.cpp b/torchvision/csrc/io/image/common.cpp deleted file mode 100644 index 0be5f67532f..00000000000 --- a/torchvision/csrc/io/image/common.cpp +++ /dev/null @@ -1,51 +0,0 @@ - -#include "common.h" - -// If we are in a Windows environment, we need to define -// initialization functions for the _custom_ops extension -#ifdef _WIN32 -void* PyInit_image(void) { - return nullptr; -} -#endif - -namespace vision { -namespace image { - -void validate_encoded_data(const torch::Tensor& encoded_data) { - STD_TORCH_CHECK( - encoded_data.is_contiguous(), "Input tensor must be contiguous."); - STD_TORCH_CHECK( - encoded_data.dtype() == torch::kU8, - "Input tensor must have uint8 data type, got ", - encoded_data.dtype()); - STD_TORCH_CHECK( - encoded_data.dim() == 1 && encoded_data.numel() > 0, - "Input tensor must be 1-dimensional and non-empty, got ", - encoded_data.dim(), - " dims and ", - encoded_data.numel(), - " numels."); -} - -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 image -} // namespace vision diff --git a/torchvision/csrc/io/image/common.h b/torchvision/csrc/io/image/common.h deleted file mode 100644 index 1b459c5d7ea..00000000000 --- a/torchvision/csrc/io/image/common.h +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace vision { -namespace image { - -/* Should be kept in-sync with Python ImageReadMode enum */ -using ImageReadMode = int64_t; -const ImageReadMode IMAGE_READ_MODE_UNCHANGED = 0; -const ImageReadMode IMAGE_READ_MODE_GRAY = 1; -const ImageReadMode IMAGE_READ_MODE_GRAY_ALPHA = 2; -const ImageReadMode IMAGE_READ_MODE_RGB = 3; -const ImageReadMode IMAGE_READ_MODE_RGB_ALPHA = 4; - -void validate_encoded_data(const torch::Tensor& encoded_data); - -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); - -} // namespace image -} // namespace vision diff --git a/torchvision/csrc/io/image/cpu/decode_image.cpp b/torchvision/csrc/io/image/cpu/decode_image.cpp index 42d64bb0216..3a0a50a7afd 100644 --- a/torchvision/csrc/io/image/cpu/decode_image.cpp +++ b/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 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(stack[0]); -} - -} // namespace - torch::stable::Tensor decode_image( const torch::stable::Tensor& data, ImageReadMode mode, diff --git a/torchvision/csrc/io/image/cpu/decode_webp.cpp b/torchvision/csrc/io/image/cpu/decode_webp.cpp index d5c949e77eb..b0366dd8371 100644 --- a/torchvision/csrc/io/image/cpu/decode_webp.cpp +++ b/torchvision/csrc/io/image/cpu/decode_webp.cpp @@ -1,5 +1,9 @@ #include "decode_webp.h" -#include "../common.h" + +#include +#include + +#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(); + auto encoded_data_p = encoded_data.const_data_ptr(); 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(__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 diff --git a/torchvision/csrc/io/image/cpu/decode_webp.h b/torchvision/csrc/io/image/cpu/decode_webp.h index d5c81547c42..7511f70d678 100644 --- a/torchvision/csrc/io/image/cpu/decode_webp.h +++ b/torchvision/csrc/io/image/cpu/decode_webp.h @@ -1,13 +1,13 @@ #pragma once -#include -#include "../common.h" +#include +#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 diff --git a/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp b/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp index c87029c42b2..775b095bf14 100644 --- a/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp +++ b/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp @@ -3,6 +3,10 @@ #include #include +// In ROCm builds, the rocJPEG implementation will register this op. +// So we keep this translation unit for nvJPEG and the no-GPU builds only. +#if !ROCJPEG_FOUND + #if !NVJPEG_FOUND namespace vision { namespace image { @@ -618,13 +622,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(image, m) { "decode_jpegs_cuda(Tensor[] encoded_images, int mode, Device device) -> Tensor[]"); } -// In ROCm builds, the rocJPEG implementation will register this op. -// So we keep this registration for nvJPEG and the no-GPU builds only. -#if !ROCJPEG_FOUND STABLE_TORCH_LIBRARY_IMPL(image, CompositeExplicitAutograd, m) { m.impl("decode_jpegs_cuda", TORCH_BOX(&decode_jpegs_cuda)); } -#endif } // namespace image } // namespace vision + +#endif // !ROCJPEG_FOUND diff --git a/torchvision/csrc/io/image/hip/decode_jpegs_cuda.cpp b/torchvision/csrc/io/image/hip/decode_jpegs_cuda.cpp index 043560a36d8..34ed16d749a 100644 --- a/torchvision/csrc/io/image/hip/decode_jpegs_cuda.cpp +++ b/torchvision/csrc/io/image/hip/decode_jpegs_cuda.cpp @@ -200,6 +200,11 @@ std::vector RocJpegDecoder::decode_images( return output_tensors; } +STABLE_TORCH_LIBRARY_FRAGMENT(image, m) { + m.def( + "decode_jpegs_cuda(Tensor[] encoded_images, int mode, Device device) -> Tensor[]"); +} + STABLE_TORCH_LIBRARY_IMPL(image, CompositeExplicitAutograd, m) { m.impl("decode_jpegs_cuda", TORCH_BOX(&decode_jpegs_cuda)); } diff --git a/torchvision/csrc/io/image/image.cpp b/torchvision/csrc/io/image/image.cpp deleted file mode 100644 index 2ef160b1835..00000000000 --- a/torchvision/csrc/io/image/image.cpp +++ /dev/null @@ -1,13 +0,0 @@ -#include "image.h" - -#include - -namespace vision { -namespace image { - -static auto registry = torch::RegisterOperators().op( - "image::decode_webp(Tensor encoded_data, int mode) -> Tensor", - &decode_webp); - -} // namespace image -} // namespace vision diff --git a/torchvision/csrc/io/image/image.h b/torchvision/csrc/io/image/image.h deleted file mode 100644 index 7968e1d0cab..00000000000 --- a/torchvision/csrc/io/image/image.h +++ /dev/null @@ -1,3 +0,0 @@ -#pragma once - -#include "cpu/decode_webp.h" diff --git a/torchvision/io/image.py b/torchvision/io/image.py index ffc01470593..9f242b65e7d 100644 --- a/torchvision/io/image.py +++ b/torchvision/io/image.py @@ -11,16 +11,12 @@ def _has_image_ops(): return False -if _load_library("image"): +if _load_library("image_stable"): def _has_image_ops(): # noqa: F811 return True -# Stable-ABI image extension (migrated ops) -_load_library("image_stable") - - def _assert_has_image_ops(): if not _has_image_ops(): raise RuntimeError(