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

invokeai/app/invocations/spandrel_image_to_image.py

@@ -1,7 +1,6 @@
 import functools
 from typing import Callable
 
-import numpy as np
 import torch
 from PIL import Image
 from tqdm import tqdm
@@ -90,20 +89,17 @@ def upscale_image(
         tiles = sorted(tiles, key=lambda x: x.coords.left)
         tiles = sorted(tiles, key=lambda x: x.coords.top)
 
-        # Prepare input image for inference.
-        image_tensor = SpandrelImageToImageModel.pil_to_tensor(image)
-
         # Scale the tiles for re-assembling the final image.
         scale = spandrel_model.scale
         scaled_tiles = [cls.scale_tile(tile, scale=scale) for tile in tiles]
 
         # Prepare the output tensor.
-        _, channels, height, width = image_tensor.shape
+        channels = len(image.getbands())
         output_tensor = torch.zeros(
-            (height * scale, width * scale, channels), dtype=torch.uint8, device=torch.device("cpu")
+            (image.height * scale, image.width * scale, channels), dtype=torch.uint8, device=torch.device("cpu")
         )
 
-        image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=spandrel_model.dtype)
+        device = TorchDevice.choose_torch_device()
 
         # Run the model on each tile.
         pbar = tqdm(list(zip(tiles, scaled_tiles, strict=True)), desc="Upscaling Tiles")
@@ -116,8 +112,21 @@ def upscale_image(
             if is_canceled():
                 raise CanceledException
 
-            # Extract the current tile from the input tensor.
-            input_tile = image_tensor[:, :, tile.coords.top : tile.coords.bottom, tile.coords.left : tile.coords.right]
+            # Crop the current tile from the input image and convert it on demand. Converting the
+            # whole image up front would keep a float32 copy of it in memory for the entire loop,
+            # even though only one tile is ever used at a time.
+            if (
+                tile.coords.top == 0
+                and tile.coords.bottom == image.height
+                and tile.coords.left == 0
+                and tile.coords.right == image.width
+            ):
+                input_image = image
+            else:
+                input_image = image.crop((tile.coords.left, tile.coords.top, tile.coords.right, tile.coords.bottom))
+            input_tile = SpandrelImageToImageModel.pil_to_tensor(input_image).to(
+                device=device, dtype=spandrel_model.dtype
+            )
 
             # Run the model on the tile.
             output_tile = spandrel_model.run(input_tile)
@@ -144,9 +153,9 @@ def upscale_image(
 
             step_callback(pbar.n + 1, pbar.total)
 
-        # Convert the output tensor to a PIL image.
-        np_image = output_tensor.detach().numpy().astype(np.uint8)
-        pil_image = Image.fromarray(np_image)
+        # Convert the output tensor to a PIL image. `output_tensor` is already uint8, so `.numpy()`
+        # is a zero-copy view; casting it to uint8 again here would copy the whole image for nothing.
+        pil_image = Image.fromarray(output_tensor.numpy())
 
         return pil_image
 

tests/app/invocations/test_spandrel_image_to_image.py

@@ -0,0 +1,35 @@
+import torch
+from PIL import Image
+
+from invokeai.app.invocations.spandrel_image_to_image import SpandrelImageToImageInvocation
+from invokeai.backend.util.devices import TorchDevice
+
+
+def test_upscale_image_does_not_crop_full_image_for_untiled_input(monkeypatch):
+    image = Image.new("RGB", (8, 8))
+
+    class IdentityModel:
+        scale = 1
+        dtype = torch.float32
+
+        @staticmethod
+        def run(image_tensor: torch.Tensor) -> torch.Tensor:
+            return image_tensor
+
+    def fail_crop(_image: Image.Image, _box: tuple[int, int, int, int]) -> Image.Image:
+        raise AssertionError("untiled input must not be copied with a full-image crop")
+
+    monkeypatch.setattr(Image.Image, "crop", fail_crop)
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", staticmethod(lambda: torch.device("cpu")))
+
+    result = SpandrelImageToImageInvocation.upscale_image(
+        image,
+        tile_size=0,
+        spandrel_model=IdentityModel(),
+        is_canceled=lambda: False,
+        step_callback=lambda *_: None,
+    )
+
+    assert result.size == image.size
+    assert result.mode == image.mode
+    assert result.tobytes() == image.tobytes()