diff --git a/pillow_jxl/JpegXLImagePlugin.py b/pillow_jxl/JpegXLImagePlugin.py index 61415d9..a07aa41 100644 --- a/pillow_jxl/JpegXLImagePlugin.py +++ b/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 diff --git a/src/decode.rs b/src/decode.rs index 34864a0..27e70d4 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -69,13 +69,16 @@ pub fn convert_pixels(pixels: Pixels) -> Vec { } #[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 { + 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(¶llel_runner) @@ -106,8 +126,4 @@ impl Decoder { Cow::Owned(icc_profile), ) } - - fn __repr__(&self) -> PyResult { - Ok(format!("Decoder")) - } } diff --git a/src/encode.rs b/src/encode.rs index 1a2af2d..613feae 100644 --- a/src/encode.rs +++ b/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 { + 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(¶llel_runner) .jpeg_quality(self.quality) @@ -125,11 +157,4 @@ impl Encoder { }; Cow::Owned(buffer.data) } - - fn __repr__(&self) -> PyResult { - Ok(format!( - "Encoder(has_alpha={}, lossless={}, quality={}, decoding_speed={}, effort={})", - self.has_alpha, self.lossless, self.quality, self.decoding_speed, self.effort - )) - } }