perf(nodes): avoid two full-size copies in Spandrel upscale - #9426
Open
dexhunter wants to merge 1 commit into
Open
perf(nodes): avoid two full-size copies in Spandrel upscale#9426dexhunter wants to merge 1 commit into
dexhunter wants to merge 1 commit into
Conversation
`upscale_image` held more host memory than the upscale needs: - `pil_to_tensor(image)` materialised the entire input as float32 up front, but the loop only ever uses one tile at a time. Crop each tile from the input image and convert it on demand instead. - The final `output_tensor.detach().numpy().astype(np.uint8)` copied the whole assembled image even though `output_tensor` is already uint8. `.numpy()` alone is a zero-copy view. For a 4x upscale of a 2048x2048 image (8192x8192 output, 192 MiB) the assembly's peak host memory drops from 639.9 MiB to 447.9 MiB; for a 1024x1024 input (4096x4096 output) from 159.9 MiB to 111.9 MiB. Tiling, the half-overlap trim, the clamp/uint8 conversion, cancellation and the step_callback sequence are unchanged, and the returned image is byte-for-byte identical, including for non-tile-aligned sizes and the untiled (tile_size=0) path. Co-Authored-By: Aiden <aiden@weco.ai>
dexhunter
requested review from
JPPhoto,
Pfannkuchensack,
blessedcoolant,
dunkeroni and
lstein
as code owners
July 31, 2026 21:38
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SpandrelImageToImageInvocation.upscale_imageholds more host memory than the upscaleactually needs. Two full-size buffers are avoidable:
pil_to_tensor(image)materialises theentire input as float32 before the loop, but the loop only ever slices one tile out
of it. This PR crops each tile from the input
PIL.Imageand converts it on demand,so only one tile-sized tensor exists at a time. The saving scales with the input.
output_tensor.detach().numpy().astype(np.uint8), butoutput_tensoris alreadytorch.uint8, andndarray.astypecopies unconditionally..numpy()on a CPUtensor is a zero-copy view, so the
astypeduplicated the whole image for nothing.Measured peak host RSS for the assembly, median of 5 runs each:
This is a host memory change, not a claim about inference speed. The model itself is
untouched and dominates the wall clock on a real upscale; what this removes is the
allocation and byte copying around it. I am not quoting a wall-clock figure because I
could not measure one under exclusive CPU conditions on my machine.
upscale_imageruns once perspandrel_image_to_imagenode, which is what the galleryUpscale button (
buildAdHocPostProcessingGraph.ts) and the canvasspandrel_filterboth enqueue. Neither sets
tile_size, so both use the default of 512. It also runsonce per iteration of
spandrel_image_to_image_autoscale, up to 5 of them.Related Issues / Discussions
None; found by profiling the upscale path.
QA Instructions
Upscale an image with any Spandrel model (e.g. a 4× ESRGAN) from the gallery Upscale
button, with and without tiling, and confirm the output is unchanged. The returned
image is byte-for-byte identical to before, including for non-tile-aligned sizes and
for the untiled
tile_size=0path. Tile geometry, the half-overlap trim, and theclamp/uint8 conversion are all unchanged, as are
is_canceled()handling and thestep_callbacksequence.Automated:
pytest tests/backend/tiles tests/app/invocationsgives 711 passed,1 skipped, 6 xfailed, identical to the pre-change baseline.
ruff checkandruff format --checkare clean.What I deliberately left out
clamp_()/mul_()on the model's output tensor. This reduces peak furtherand is the obvious next thing to try, but spandrel's
ImageModelDescriptor.__call__runs the model under
torch.inference_mode(), so everything it returns is aninference tensor, and mutating one from outside that mode raises
RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed. I hit thisrepeatedly while exploring, so I have left it alone.
Image.frombuffer(...)instead ofImage.fromarray(...). Avoids one more full-sizecopy, but makes the returned image alias the tensor's storage, which is a lifetime
hazard for a value that outlives this function.
numpybuffer wrapped bytorch.from_numpy. Scores slightlybetter again, but swaps the buffer's owning type in a way that reads less naturally
next to the rest of the method.
Merge Plan
Ordinary merge; no schema or migration impact.
Notes on how this was found
The change was found with an automated optimization search over this function
(trajectory: https://dashboard.weco.ai/share/NEJGgBZil8V-WOK5WgArl7XOZ1-jF7gE), scored
on peak resident memory with a byte-exactness gate against unmodified upstream output.
The link is a record of that search, not the author of the patch.
The diff here is hand-written and deliberately narrower than the search's best-scoring
candidate. That candidate also dropped the
tqdmprogress bar, replaced the zero-filledoutput buffer with an uninitialized
np.empty, hardcoded the channel count to 3, andadded a
gc.collect()call. It scored 1,352,282,112 bytes against the baseline's1,626,423,296 on the harness. I did not ship any of it: the uninitialized buffer is only
correct if every output pixel is written, and the hardcoded channel count would break any
input that is not 3-channel RGB.
Checklist
What's Newcopy (if doing a release after this PR) — n/a