[fbsync] Reapply "Fix memory leak in decode_webp (#8712)" (#8723) (#8724)

Summary: Co-authored-by: atalman <[email protected]>

Reviewed By: vmoens

Differential Revision: D68021982

fbshipit-source-id: 5db51ce48cba0c2e51e3ac85dd1a1a907bb8e306

Author: NicolasHug

packaging/pre_build_script.sh

@@ -32,7 +32,9 @@ else
     conda install -yq ffmpeg=4.2 libjpeg-turbo -c pytorch-nightly
   fi
 
-  yum install -y libjpeg-turbo-devel libwebp-devel freetype gnutls
+  conda install libwebp -yq
+  conda install libjpeg-turbo -c pytorch
+  yum install -y freetype gnutls
   pip install auditwheel
 fi
 

test/test_image.py

@@ -45,6 +45,8 @@
 IS_MACOS = sys.platform == "darwin"
 PILLOW_VERSION = tuple(int(x) for x in PILLOW_VERSION.split("."))
 WEBP_TEST_IMAGES_DIR = os.environ.get("WEBP_TEST_IMAGES_DIR", "")
+# See https://github.com/pytorch/vision/pull/8724#issuecomment-2503964558
+ROCM_WEBP_MESSAGE = "ROCM not built with webp support."
 
 # Hacky way of figuring out whether we compiled with libavif/libheif (those are
 # currenlty disabled by default)

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

@@ -3,6 +3,7 @@
 
 #if WEBP_FOUND
 #include "webp/decode.h"
+#include "webp/types.h"
 #endif // WEBP_FOUND
 
 namespace vision {
@@ -44,10 +45,12 @@ torch::Tensor decode_webp(
 
   auto decoded_data =
       decoding_func(encoded_data_p, encoded_data_size, &width, &height);
+
   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}, torch::kUInt8);
+      decoded_data, {height, width, num_channels}, deleter, torch::kUInt8);
 
   return out.permute({2, 0, 1});
 }