Just to mention for anyone else reading this - individual users can already vary their individual experience by setting MAXBLOCK themselves.

from PIL import ImageFile
ImageFile.MAXBLOCK = 65536 * 4

Have you tried using the Arrow interface, which is zero-copy?

Have you tried using the Arrow interface, which is zero-copy?

My use case still requires a numpy array as output (or at least access to the raw bytes). How would I do this from a user perspective with the Arrow interface? Currently I'm just using np.asarray(img) which calls .tobytes() under the hood.

I would guess

import numpy as np
import pyarrow as pa
from PIL import Image
img = Image.new("RGB", (12, 12))
np.array(pa.array(img))

but it doesn't seem faster to me.

np.array always does a copy so we wouldn't really gain much. pa.array(img).to_numpy(zero_copy_only=True) doesn't seem to work in my case so zero_copy_only=False or np.array() would be needed which seems to be very slow.

Here's a quick benchmark with this PR:

import numpy as np
import pyarrow as pa
from PIL import Image

rng = np.random.default_rng(42)

for size in (128, 256, 512, 1024, 2048, 4096, 8192, 16384):
    img = rng.integers(0, 256, size=(size, size, 3), dtype=np.uint8)
    img = Image.fromarray(img)

    print(f"{size}x{size}")
    %timeit img.tobytes()
    %timeit np.asarray(img)
    %timeit pa.array(img)
    %timeit pa.array(img).to_numpy(zero_copy_only=False)

128x128
4.61 μs ± 39 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
5.53 μs ± 31.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
1.28 μs ± 4.5 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
1.4 ms ± 3.29 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
256x256
13.1 μs ± 86.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
13.9 μs ± 19.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
1.3 μs ± 22.2 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
5.86 ms ± 14.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
512x512
45.8 μs ± 77.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
46.8 μs ± 95.1 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
1.27 μs ± 9.67 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
24 ms ± 107 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
1024x1024
246 μs ± 1.82 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
263 μs ± 6.77 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1.28 μs ± 5.84 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
97.8 ms ± 215 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
2048x2048
1.15 ms ± 13.2 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1.16 ms ± 15 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1.28 μs ± 6.15 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
395 ms ± 1.18 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

For image sizes of 4096x4096 it also seems like pyarrow fails with ValueError: Image is in multiple array blocks, use imaging_new_block for zero copy.

So in summary using pyarrow arrays are much faster, but going from pyarrow to numpy is very slow.

pa.array(img).to_numpy(zero_copy_only=True) doesn't seem to work in my case

Just in case it is something interesting that we should consider in the future, could you explain this slightly more?

pa.array(img).to_numpy(zero_copy_only=True) doesn't seem to work in my case

Just in case it is something interesting that we should consider in the future, could you explain this slightly more?

The following code would raise ArrowInvalid: Needed to copy 1 chunks with 0 nulls, but zero_copy_only was True

img = Image.fromarray(np.random.randint(0, 255, size=(128, 128, 3), dtype=np.uint8))
arr = pa.array(img)
np_arr = pa.array(img).to_numpy()

I haven't used pyarrow before, but judging from the docs it raises because "the conversion to a numpy array would require copying the underlying data (e.g. in presence of nulls, or for non-primitive types) " and since zero_copy_only was True.

Just to mention for anyone else reading this - individual users can already vary their individual experience by setting MAXBLOCK themselves.

The main benefit of this PR is that since the buffer size is dependent on the image size users will still have low memory usage for small images but benefit from a larger buffer size for large images.

Let me know if you have any concerns that would prevent merging of the PR.

@radarhere Any updates on when/whether this PR could be merged? Or are there any additional benchmarks that you would like me to run?

@lgeiger We'll probably need @wiredfool to look closer too.

The original point of this particular bit of code was to have predictable memory usage when running tobytes, in this case, either Image.MAX_BLOCK or at the very least, the size of one row. (as the shuffler is row based, and if the buffer is smaller than the row, no progress can be made.)

So, where previously we needed 2xImageMemory + 64k, now we need 3xImageMemory.

For smaller images, it's not a problem, but for larger images this may cause memory pressure where we didn't have it before. I'd consider that a regression.

One alternate here is to change the calculation so that it's the min of max(MAX_BLOCK, row_size) and the image size, at which point the MAX_BLOCK can be boosted without allocating excessive memory in the small image case.

There may other places where the MAX_BLOCK shouldn't be large, and they'd need some similar checks.

@wiredfool Thanks for taking a look. I agree increasing MAXBLOCK globally would be a problem since it increases memory usage, but for tobytes() I'm not sure this is actually the case.

The way I understand the code is the following: All chunks are appended to a list which is joined afterwards which causes the 2x ImageMemory usage on the Python side you mentioned above.

I don't think the max memory usage would include the additional 64k buffer, but I haven't looked at what the C code is actually doing so I might be wrong.

In any case, for this PR the output list only consists of a single item (assuming the buffer size was correct/large enough) preventing the need for allocating a new bytes object during the join. So the memory usage of the Python code would actually half. I thought about changing the code to directly return data but I wasn't sure whether I'm missing any edge cases where the the actual size would be larger than the buffer size estimate added here.

I double checked this with a memory profile and viewed the memory usage with memray summary <filename>:

import memray
import numpy as np
from PIL import Image

rng = np.random.default_rng(42)

def get_image(size):
    return Image.fromarray(rng.integers(0, 256, size=(size, size, 3), dtype=np.uint8))

for size in (512, 1024, 2048, 4096, 8192, 16384):
    img = get_image(size)
    with memray.Tracker(f"pr_{size}.bin"):
        img.tobytes()

And the results show that this PR halves the memory usage which matches my theory from above:

Ok, I was wrong here, the return b"".join(output) looks like it's just stringing together all the buffers with pointer magic. So the advantage you've got here is that you're basically turning that into a noop and making bigger allocations earlier. OTOH, with all the reallocs, I bet there's a buffer trimming going on in the C layer to make it joinable without doing a slice on the buffer in python space.

What might make sense is for there to be an encode_tobytes, similar to the other encode_to functions in encode.c, but essentially do all the looping in C space returning a single PyBytes. Not sure it would help completely, but it would take all that logic in python and move it down.

Aside -- I don't think the memray is measuring what you think it is -- I get different results when I replace

rng = np.random.default_rng(42)

def get_image(size):
    return Image.fromarray(rng.integers(0, 256, size=(size, size, 3), dtype=np.uint8))

with

def get_image(size):
    return Image.new('RGBA', (size,size), (0,1,2,3))

But I don't think it changes the conclusion there.

What might make sense is for there to be an encode_tobytes, similar to the other encode_to functions in encode.c, but essentially do all the looping in C space returning a single PyBytes. Not sure it would help completely, but it would take all that logic in python and move it down.

What would the benefit of this be compare to using a bufsize equal to the image size in encode? The best case memory usage of .tobytes() will always be 1x ImageMemory which is already very close to the memory measurements measured above. Or are there any benefits of doing the encoding in small chunks that I'm missing here? I actually tried something like this initially but I didn't really saw any benefits performance wise over this solution which requires much less code.

Aside -- I don't think the memray is measuring what you think it is -- I get different results when I replace

I'm not an expert with memray, could you elaborate? In your code example your using a 4-channel RGBA image where I tested a 3-channel RGB image from numpy. So if the memory usage in the Image.new('RGBA', (size,size), (0,1,2,3)) case is 4/3 higher that is expected. Or are you talking about other differences?

@wiredfool Friendly ping, do you mind having a look at my comments from above?

@aclark4life @wiredfool any updates on this?

If I understand correctly, this is not going to get merged without some comprehensive testing to ensure that there are no memory usage regressions. What may move this forward is some plan to do that testing.

If I understand correctly, this is not going to get merged without some comprehensive testing to ensure that there are no memory usage regressions. What may move this forward is some plan to do that testing.

@aclark4life As mentioned above, I did some memory profiling at #9220 (comment). In my measurements this PR actually halves the memory usage by preventing the need to concatenate byte chunks.

I've noticed the same thing - with a suitably large MAXBLOCK, encoding becomes faster (as more time is spent in C code) and doesn't require temporary memory, as CPython optimizes "".join([x]) to just increase the refcount of the immutable x, so it's essentially a free operation.

I shaved a bunch of time off a project (the same project that instigated #9649 and the rest of the perf series) with just setting MAXBLOCK to 16 megabytes (which is enough to fit a 1080p RGBA frame or two, and a pittance of memory on modern systems).

I think this PR has a good point, but it should probably only be limited to the "raw" encoder? With that encoder, the programmer's intent, after all, is "I need the raw bytes, and I know I will need the memory for the raw bytes". An alternate approach (that would IMO make a lot of sense for certain cases) would be to add a Image.iterbytes(encoder_name, *args, bufsize: int | None = None) -> Iterable[bytes] generator?