Multithreading: num_threads and no-GIL

pillow_jxl/JpegXLImagePlugin.py

@@ -96,6 +96,7 @@ def _save(im, fp, filename, save_all=False):
     use_container = info.get("use_container", False)
     use_original_profile = info.get("use_original_profile", False)
     jpeg_encode = info.get("lossless_jpeg", True)
+    num_threads = info.get("num_threads", -1)
 
     enc = Encoder(
         mode=im.mode,
@@ -105,6 +106,7 @@ def _save(im, fp, filename, save_all=False):
         effort=effort,
         use_container=use_container,
         use_original_profile=use_original_profile,
+        num_threads=num_threads,
     )
     # FIXME (Isotr0py): im.filename maybe None if parse stream
     # TODO (Isotr0py): This part should be refactored in the near future

src/decode.rs

@@ -69,13 +69,16 @@ pub fn convert_pixels(pixels: Pixels) -> Vec<u8> {
 }
 
 #[pyclass(module = "pillow_jxl")]
-pub struct Decoder;
+pub struct Decoder {
+    num_threads: isize,
+}
 
 #[pymethods]
 impl Decoder {
     #[new]
-    fn new() -> Self {
-        Self
+    #[pyo3(signature = (num_threads = -1))]
+    fn new(num_threads: isize) -> Self {
+        Self { num_threads }
     }
 
     #[pyo3(signature = (data))]
@@ -84,7 +87,24 @@ impl Decoder {
         _py: Python,
         data: &[u8],
     ) -> (bool, ImageInfo, Cow<'_, [u8]>, Cow<'_, [u8]>) {
-        let parallel_runner = ThreadsRunner::default();
+        _py.allow_threads(|| self.call_inner(data))
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("Decoder"))
+    }
+}
+
+impl Decoder {
+    fn call_inner(&self, data: &[u8]) -> (bool, ImageInfo, Cow<'_, [u8]>, Cow<'_, [u8]>) {
+        let parallel_runner = ThreadsRunner::new(
+            None,
+            if self.num_threads < 0 {
+                None
+            } else {
+                Some(self.num_threads as usize)
+            },
+        ).unwrap();
         let decoder = decoder_builder()
             .icc_profile(true)
             .parallel_runner(&parallel_runner)
@@ -106,8 +126,4 @@ impl Decoder {
             Cow::Owned(icc_profile),
         )
     }
-
-    fn __repr__(&self) -> PyResult<String> {
-        Ok(format!("Decoder"))
-    }
 }

src/encode.rs

@@ -16,12 +16,13 @@ pub struct Encoder {
     effort: u32,
     use_container: bool,
     use_original_profile: bool,
+    num_threads: isize,
 }
 
 #[pymethods]
 impl Encoder {
     #[new]
-    #[pyo3(signature = (mode, lossless=false, quality=1.0, decoding_speed=0, effort=7, use_container=true, use_original_profile=false))]
+    #[pyo3(signature = (mode, lossless=false, quality=1.0, decoding_speed=0, effort=7, use_container=false, use_original_profile=false, num_threads=-1))]
     fn new(
         mode: &str,
         lossless: bool,
@@ -30,6 +31,7 @@ impl Encoder {
         effort: u32,
         use_container: bool,
         use_original_profile: bool,
+        num_threads: isize,
     ) -> Self {
         let (num_channels, has_alpha) = match mode {
             "RGBA" => (4, true),
@@ -58,13 +60,14 @@ impl Encoder {
             effort,
             use_container,
             use_original_profile,
+            num_threads,
         }
     }
 
     #[pyo3(signature = (data, width, height, jpeg_encode, exif=None, jumb=None, xmp=None))]
     fn __call__(
         &self,
-        _py: Python,
+        py: Python,
         data: &[u8],
         width: u32,
         height: u32,
@@ -73,7 +76,36 @@ impl Encoder {
         jumb: Option<&[u8]>,
         xmp: Option<&[u8]>,
     ) -> Cow<'_, [u8]> {
-        let parallel_runner = ThreadsRunner::default();
+        py.allow_threads(|| self.call_inner(data, width, height, jpeg_encode, exif, jumb, xmp))
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!(
+            "Encoder(has_alpha={}, lossless={}, quality={}, decoding_speed={}, effort={}, num_threads={})",
+            self.has_alpha, self.lossless, self.quality, self.decoding_speed, self.effort, self.num_threads
+        ))
+    }
+}
+
+impl Encoder {
+    fn call_inner(
+        &self,
+        data: &[u8],
+        width: u32,
+        height: u32,
+        jpeg_encode: bool,
+        exif: Option<&[u8]>,
+        jumb: Option<&[u8]>,
+        xmp: Option<&[u8]>,
+    ) -> Cow<'_, [u8]> {
+        let parallel_runner = ThreadsRunner::new(
+            None,
+            if self.num_threads < 0 {
+                None
+            } else {
+                Some(self.num_threads as usize)
+            },
+        ).unwrap();
         let mut encoder = encoder_builder()
             .parallel_runner(&parallel_runner)
             .jpeg_quality(self.quality)
@@ -125,11 +157,4 @@ impl Encoder {
         };
         Cow::Owned(buffer.data)
     }
-
-    fn __repr__(&self) -> PyResult<String> {
-        Ok(format!(
-            "Encoder(has_alpha={}, lossless={}, quality={}, decoding_speed={}, effort={})",
-            self.has_alpha, self.lossless, self.quality, self.decoding_speed, self.effort
-        ))
-    }
 }