From a6d6dd1b684a9942bb0c1b8df943c495fcaa8369 Mon Sep 17 00:00:00 2001
From: Frankie Dintino <fdintino@gmail.com>
Date: Wed, 3 Jul 2024 11:23:10 -0400
Subject: [PATCH] Allow users to pass max_threads to the avif encoder via
 Image.save (#54)

* Allow users to pass max_threads to the avif encoder via Image.save
* Distinguish global default max threads from encoder param
* Allow overriding DEFAULT_MAX_THREADS on AvifImagePlugin module

---------

Co-authored-by: Wyatt Borsos <yitb@meta.com>
---
 src/pillow_avif/AvifImagePlugin.py |  3 +++
 src/pillow_avif/_avif.c            | 17 +++++++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/pillow_avif/AvifImagePlugin.py b/src/pillow_avif/AvifImagePlugin.py
index 6aa6bd7..80dfb52 100644
--- a/src/pillow_avif/AvifImagePlugin.py
+++ b/src/pillow_avif/AvifImagePlugin.py
@@ -16,6 +16,7 @@
 # to Image.open (see https://github.com/python-pillow/Pillow/issues/569)
 DECODE_CODEC_CHOICE = "auto"
 CHROMA_UPSAMPLING = "auto"
+DEFAULT_MAX_THREADS = 0
 
 _VALID_AVIF_MODES = {"RGB", "RGBA"}
 
@@ -138,6 +139,7 @@ def _save(im, fp, filename, save_all=False):
     duration = info.get("duration", 0)
     subsampling = info.get("subsampling", "4:2:0")
     speed = info.get("speed", 6)
+    max_threads = info.get("max_threads", DEFAULT_MAX_THREADS)
     codec = info.get("codec", "auto")
     range_ = info.get("range", "full")
     tile_rows_log2 = info.get("tile_rows", 0)
@@ -196,6 +198,7 @@ def _save(im, fp, filename, save_all=False):
         qmax,
         quality,
         speed,
+        max_threads,
         codec,
         range_,
         tile_rows_log2,
diff --git a/src/pillow_avif/_avif.c b/src/pillow_avif/_avif.c
index 68c4d01..59516f8 100644
--- a/src/pillow_avif/_avif.c
+++ b/src/pillow_avif/_avif.c
@@ -46,7 +46,7 @@ typedef struct {
 
 static PyTypeObject AvifDecoder_Type;
 
-static int max_threads = 0;
+static int default_max_threads = 0;
 
 static void
 init_max_threads(void) {
@@ -85,7 +85,7 @@ init_max_threads(void) {
         goto error;
     }
 
-    max_threads = (int)num_cpus;
+    default_max_threads = (int)num_cpus;
 
 done:
     Py_XDECREF(os);
@@ -321,6 +321,7 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
     int quality = 75;
     int speed = 8;
     int exif_orientation = 0;
+    int max_threads = default_max_threads;
     PyObject *icc_bytes;
     PyObject *exif_bytes;
     PyObject *xmp_bytes;
@@ -336,7 +337,7 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "IIsiiiissiiOOSSiSO",
+            "IIsiiiiissiiOOSSiSO",
             &width,
             &height,
             &subsampling,
@@ -344,6 +345,7 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
             &qmax,
             &quality,
             &speed,
+            &max_threads,
             &codec,
             &range,
             &tile_rows_log2,
@@ -445,7 +447,10 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
         encoder = avifEncoderCreate();
 
         if (max_threads == 0) {
-            init_max_threads();
+            if (default_max_threads == 0) {
+                init_max_threads();
+            }
+            max_threads = default_max_threads;
         }
 
         int is_aom_encode = strcmp(codec, "aom") == 0 ||
@@ -777,10 +782,10 @@ AvifDecoderNew(PyObject *self_, PyObject *args) {
 
     self->decoder = avifDecoderCreate();
 #if AVIF_VERSION >= 80400
-    if (max_threads == 0) {
+    if (default_max_threads == 0) {
         init_max_threads();
     }
-    self->decoder->maxThreads = max_threads;
+    self->decoder->maxThreads = default_max_threads;
 #endif
     self->decoder->codecChoice = codec;