Skip to content

perf(nodes): avoid two full-size copies in Spandrel upscale - #9426

Open
dexhunter wants to merge 1 commit into
invoke-ai:mainfrom
dexhunter:perf/spandrel-upscale-host-memory
Open

perf(nodes): avoid two full-size copies in Spandrel upscale#9426
dexhunter wants to merge 1 commit into
invoke-ai:mainfrom
dexhunter:perf/spandrel-upscale-host-memory

Conversation

@dexhunter

@dexhunter dexhunter commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Summary

SpandrelImageToImageInvocation.upscale_image holds more host memory than the upscale
actually needs. Two full-size buffers are avoidable:

  1. A float32 copy of the whole input image. pil_to_tensor(image) materialises the
    entire 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.Image and converts it on demand,
    so only one tile-sized tensor exists at a time. The saving scales with the input.
  2. A full-size copy of the assembled output. The final line was
    output_tensor.detach().numpy().astype(np.uint8), but output_tensor is already
    torch.uint8, and ndarray.astype copies unconditionally. .numpy() on a CPU
    tensor is a zero-copy view, so the astype duplicated the whole image for nothing.

Measured peak host RSS for the assembly, median of 5 runs each:

output size before after saved
8192×8192×3 (192 MiB) — a 4× upscale of a 2048×2048 image 639.9 MiB 447.9 MiB 192.0 MiB
4096×4096×3 (48 MiB) — a 4× upscale of a 1024×1024 image 159.9 MiB 111.9 MiB 48.0 MiB

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_image runs once per spandrel_image_to_image node, which is what the gallery
Upscale button (buildAdHocPostProcessingGraph.ts) and the canvas spandrel_filter
both enqueue. Neither sets tile_size, so both use the default of 512. It also runs
once 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=0 path. Tile geometry, the half-overlap trim, and the
clamp/uint8 conversion are all unchanged, as are is_canceled() handling and the
step_callback sequence.

Automated: pytest tests/backend/tiles tests/app/invocations gives 711 passed,
1 skipped, 6 xfailed, identical to the pre-change baseline. ruff check and
ruff format --check are clean.

What I deliberately left out

  • In-place clamp_()/mul_() on the model's output tensor. This reduces peak further
    and is the obvious next thing to try, but spandrel's ImageModelDescriptor.__call__
    runs the model under torch.inference_mode(), so everything it returns is an
    inference tensor, and mutating one from outside that mode raises RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed. I hit this
    repeatedly while exploring, so I have left it alone.
  • Image.frombuffer(...) instead of Image.fromarray(...). Avoids one more full-size
    copy, but makes the returned image alias the tensor's storage, which is a lifetime
    hazard for a value that outlives this function.
  • Assembling into a numpy buffer wrapped by torch.from_numpy. Scores slightly
    better 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 tqdm progress bar, replaced the zero-filled
output buffer with an uninitialized np.empty, hardcoded the channel count to 3, and
added a gc.collect() call. It scored 1,352,282,112 bytes against the baseline's
1,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

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable) — behaviour is unchanged; existing tests cover it
  • ❗Changes to a redux slice have a corresponding migration — n/a
  • Documentation added / updated (if applicable) — n/a
  • Updated What's New copy (if doing a release after this PR) — n/a

`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>
@github-actions github-actions Bot added python PRs that change python files invocations PRs that change invocations labels Jul 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

invocations PRs that change invocations python PRs that change python files

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant