From ba76065fec7c6ce28e1cff6a24c0420bd4b22482 Mon Sep 17 00:00:00 2001 From: adabeyta Date: Mon, 6 Jul 2026 18:09:59 +0000 Subject: [PATCH] Port the CPU jpeg encoder (encode_jpeg) to the stable ABI. --- CMakeLists.txt | 1 + setup.py | 15 +++++++ torchvision/csrc/io/image/cpu/encode_jpeg.cpp | 43 +++++++++++++------ torchvision/csrc/io/image/cpu/encode_jpeg.h | 6 +-- torchvision/csrc/io/image/image.cpp | 1 - torchvision/csrc/io/image/image.h | 1 - 6 files changed, 49 insertions(+), 18 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 33ffad9ec59..e28254de879 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -105,6 +105,7 @@ set(TORCHVISION_STABLE_SOURCES ${TVCPP}/io/image/cuda/decode_jpegs_cuda.cpp ${TVCPP}/io/image/cuda/encode_jpegs_cuda.cpp ${TVCPP}/io/image/cpu/encode_png.cpp + ${TVCPP}/io/image/cpu/encode_jpeg.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/setup.py b/setup.py index fe8065ad7f2..fecf4934ea1 100644 --- a/setup.py +++ b/setup.py @@ -166,6 +166,7 @@ def get_macros_and_flags(): 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", } STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/nms_kernel.hip" if IS_ROCM else "ops/cuda/nms_kernel.cu")) STABLE_SOURCES.add( @@ -470,6 +471,8 @@ def make_image_stable_extension(): + _stable(image_dir.glob("cpu/*.cpp")) + _stable(image_dir.glob("hip/*.cpp" if IS_ROCM else "cuda/*.cpp")) ) + # Shared libjpeg glue. Legacy decode_jpeg still needs it, so it builds into both extensions. + sources.append(image_dir / "cpu/common_jpeg.cpp") Extension = CppExtension @@ -484,6 +487,18 @@ def make_image_stable_extension(): else: warnings.warn("Building torchvision without PNG support") + if USE_JPEG: + jpeg_found, jpeg_include_dir, jpeg_library_dir = find_library(header="jpeglib.h") + if jpeg_found: + print("Building torchvision with JPEG image support") + if jpeg_include_dir is not None and jpeg_library_dir is not None: + include_dirs.append(jpeg_include_dir) + library_dirs.append(jpeg_library_dir) + libraries.append("jpeg") + define_macros += [("JPEG_FOUND", 1)] + else: + warnings.warn("Building torchvision without JPEG 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: diff --git a/torchvision/csrc/io/image/cpu/encode_jpeg.cpp b/torchvision/csrc/io/image/cpu/encode_jpeg.cpp index d29f3f4f481..fa647565aab 100644 --- a/torchvision/csrc/io/image/cpu/encode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/encode_jpeg.cpp @@ -1,9 +1,12 @@ #include "encode_jpeg.h" +#include +#include #include #include +#include "../common_stable.h" #include "common_jpeg.h" namespace vision { @@ -11,7 +14,9 @@ namespace image { #if !JPEG_FOUND -torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { +torch::stable::Tensor encode_jpeg( + const torch::stable::Tensor& data, + int64_t quality) { STD_TORCH_CHECK( false, "encode_jpeg: torchvision not compiled with libjpeg support"); } @@ -28,9 +33,9 @@ using JpegSizeType = size_t; using namespace detail; -torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { - C10_LOG_API_USAGE_ONCE( - "torchvision.csrc.io.image.cpu.encode_jpeg.encode_jpeg"); +torch::stable::Tensor encode_jpeg( + const torch::stable::Tensor& data, + int64_t quality) { // Define compression structures and error handling struct jpeg_compress_struct cinfo {}; struct torch_jpeg_error_mgr jerr {}; @@ -43,7 +48,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { // unwind C++ stack frames, so destructors of objects created after setjmp // won't run. We use std::optional to declare tensors before setjmp while // deferring construction, and explicitly reset them on the error path. - std::optional input; + std::optional input; cinfo.err = jpeg_std_error(&jerr.pub); jerr.pub.error_exit = torch_jpeg_error_exit; @@ -63,12 +68,12 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { } // Check that the input tensor is on CPU - STD_TORCH_CHECK( - data.device() == torch::kCPU, "Input tensor should be on CPU"); + STD_TORCH_CHECK(data.is_cpu(), "Input tensor should be on CPU"); // Check that the input tensor dtype is uint8 STD_TORCH_CHECK( - data.dtype() == torch::kU8, "Input tensor dtype should be uint8"); + data.scalar_type() == torch::headeronly::ScalarType::Byte, + "Input tensor dtype should be uint8"); // Check that the input tensor is 3-dimensional STD_TORCH_CHECK( @@ -78,7 +83,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { int channels = data.size(0); int height = data.size(1); int width = data.size(2); - input = data.permute({1, 2, 0}).contiguous(); + input = torch::stable::contiguous(stable_permute(data, {1, 2, 0})); STD_TORCH_CHECK( channels == 1 || channels == 3, @@ -104,7 +109,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { jpeg_start_compress(&cinfo, TRUE); auto stride = width * channels; - auto ptr = input->data_ptr(); + auto ptr = input->mutable_data_ptr(); // Encode JPEG file while (cinfo.next_scanline < cinfo.image_height) { @@ -115,13 +120,25 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { jpeg_finish_compress(&cinfo); jpeg_destroy_compress(&cinfo); - torch::TensorOptions options = torch::TensorOptions{torch::kU8}; - auto out_tensor = - torch::from_blob(jpegBuf, {(long)jpegSize}, ::free, options); + auto out_tensor = torch::stable::from_blob( + jpegBuf, + {static_cast(jpegSize)}, + {1}, + torch::stable::Device(torch::headeronly::DeviceType::CPU), + torch::headeronly::ScalarType::Byte, + ::free); jpegBuf = nullptr; return out_tensor; } #endif +STABLE_TORCH_LIBRARY_FRAGMENT(image, m) { + m.def("encode_jpeg(Tensor data, int quality) -> Tensor"); +} + +STABLE_TORCH_LIBRARY_IMPL(image, CompositeExplicitAutograd, m) { + m.impl("encode_jpeg", TORCH_BOX(&encode_jpeg)); +} + } // namespace image } // namespace vision diff --git a/torchvision/csrc/io/image/cpu/encode_jpeg.h b/torchvision/csrc/io/image/cpu/encode_jpeg.h index 25084e154d6..ed2c02725c2 100644 --- a/torchvision/csrc/io/image/cpu/encode_jpeg.h +++ b/torchvision/csrc/io/image/cpu/encode_jpeg.h @@ -1,12 +1,12 @@ #pragma once -#include +#include namespace vision { namespace image { -C10_EXPORT torch::Tensor encode_jpeg( - const torch::Tensor& data, +torch::stable::Tensor encode_jpeg( + const torch::stable::Tensor& data, int64_t quality); } // namespace image diff --git a/torchvision/csrc/io/image/image.cpp b/torchvision/csrc/io/image/image.cpp index 4fb457d5a99..22f3dbe2183 100644 --- a/torchvision/csrc/io/image/image.cpp +++ b/torchvision/csrc/io/image/image.cpp @@ -14,7 +14,6 @@ static auto registry = &decode_jpeg) .op("image::decode_webp(Tensor encoded_data, int mode) -> Tensor", &decode_webp) - .op("image::encode_jpeg", &encode_jpeg) .op("image::read_file", &read_file) .op("image::write_file", &write_file) .op("image::decode_image(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor", diff --git a/torchvision/csrc/io/image/image.h b/torchvision/csrc/io/image/image.h index 1cf49876b89..23a897f73af 100644 --- a/torchvision/csrc/io/image/image.h +++ b/torchvision/csrc/io/image/image.h @@ -5,5 +5,4 @@ #include "cpu/decode_jpeg.h" #include "cpu/decode_png.h" #include "cpu/decode_webp.h" -#include "cpu/encode_jpeg.h" #include "cpu/read_write_file.h"