diff --git a/invokeai/app/invocations/spandrel_image_to_image.py b/invokeai/app/invocations/spandrel_image_to_image.py index fb870c429c4..1b4b165140b 100644 --- a/invokeai/app/invocations/spandrel_image_to_image.py +++ b/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 diff --git a/tests/app/invocations/test_spandrel_image_to_image.py b/tests/app/invocations/test_spandrel_image_to_image.py new file mode 100644 index 00000000000..79f65a889d6 --- /dev/null +++ b/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()