invoke-ai--invokeai/tests/app/invocations/test_image.py

import importlib.util
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock

import numpy
import torch
from PIL import Image, ImageFilter

from invokeai.app.invocations.image import ImageField, OklabUnsharpMaskInvocation, OklchImageHueAdjustmentInvocation
from invokeai.app.invocations.primitives import ImageCollectionInvocation
from invokeai.backend.image_util.color_conversion import (
    linear_srgb_from_oklab,
    linear_srgb_from_oklch,
    linear_srgb_from_srgb,
    okhsl_from_srgb,
    oklab_from_linear_srgb,
    oklch_from_oklab,
    srgb_from_hsl,
    srgb_from_linear_srgb,
    srgb_from_okhsl,
)

_COMPOSITION_NODES_SPEC = importlib.util.spec_from_file_location(
    "invokeai.app.invocations.composition_nodes",
    Path(__file__).resolve().parents[3] / "invokeai/app/invocations/composition-nodes.py",
)
assert _COMPOSITION_NODES_SPEC is not None
assert _COMPOSITION_NODES_SPEC.loader is not None
composition_nodes = importlib.util.module_from_spec(_COMPOSITION_NODES_SPEC)
_COMPOSITION_NODES_SPEC.loader.exec_module(composition_nodes)
InvokeAdjustImageHuePlusInvocation = composition_nodes.InvokeAdjustImageHuePlusInvocation
InvokeImageBlendInvocation = composition_nodes.InvokeImageBlendInvocation


def _build_context(input_image: Image.Image) -> MagicMock:
    context = MagicMock()
    context.images.get_pil.return_value = input_image
    context.images.save.side_effect = lambda image: SimpleNamespace(
        image_name="out", width=image.width, height=image.height
    )
    return context


def _max_abs_diff_uint8(left: Image.Image, right: Image.Image) -> int:
    left_arr = numpy.asarray(left, dtype=numpy.int16)
    right_arr = numpy.asarray(right, dtype=numpy.int16)
    return int(numpy.abs(left_arr - right_arr).max())


def test_image_collection_invocation_preserves_existing_collection_values() -> None:
    images = [ImageField(image_name="first"), ImageField(image_name="second")]

    output = ImageCollectionInvocation(collection=images).invoke(MagicMock())

    assert output.collection == images


def test_image_collection_invocation_appends_direct_images_after_chained_collection() -> None:
    chained_images = [ImageField(image_name="chained")]
    direct_images = [ImageField(image_name="direct_1"), ImageField(image_name="direct_2")]

    output = ImageCollectionInvocation(collection=chained_images, images=direct_images).invoke(MagicMock())

    assert output.collection == [*chained_images, *direct_images]


def test_image_collection_invocation_supports_empty_direct_images() -> None:
    chained_images = [ImageField(image_name="chained")]

    output = ImageCollectionInvocation(collection=chained_images, images=None).invoke(MagicMock())

    assert output.collection == chained_images


def test_image_collection_invocation_outputs_empty_collection_when_inputs_are_empty() -> None:
    output = ImageCollectionInvocation(collection=None, images=None).invoke(MagicMock())

    assert output.collection == []


def test_oklab_unsharp_mask_invocation_preserves_alpha_and_sharpens_lightness_only() -> None:
    input_image = Image.new("RGBA", (3, 1))
    input_image.putdata(
        [
            (255, 0, 0, 32),
            (0, 255, 0, 128),
            (0, 0, 255, 224),
        ]
    )

    context = _build_context(input_image)

    invocation = OklabUnsharpMaskInvocation(image=ImageField(image_name="in"), radius=1.0, strength=50.0)
    output = invocation.invoke(context)
    saved_image = context.images.save.call_args.kwargs["image"]

    assert output.image.image_name == "out"
    assert output.width == 3
    assert output.height == 1
    assert numpy.asarray(saved_image.getchannel("A")).reshape(-1).tolist() == [32, 128, 224]

    rgb = torch.from_numpy(numpy.asarray(input_image.convert("RGB"), dtype=numpy.float32) / 255.0).permute(2, 0, 1)
    blurred_rgb = torch.from_numpy(
        numpy.asarray(input_image.convert("RGB").filter(ImageFilter.GaussianBlur(radius=1.0)), dtype=numpy.float32)
        / 255.0
    ).permute(2, 0, 1)

    rgb_unsharp = torch.clamp(rgb + (rgb - blurred_rgb) * 0.5, 0.0, 1.0)
    rgb_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(rgb))
    blurred_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(blurred_rgb))
    expected_oklab = rgb_oklab.clone()
    expected_oklab[0, ...] = torch.clamp(
        rgb_oklab[0, ...] + (rgb_oklab[0, ...] - blurred_oklab[0, ...]) * 0.5,
        -1.0,
        1.0,
    )
    oklab_unsharp = srgb_from_linear_srgb(linear_srgb_from_oklab(expected_oklab))

    assert not torch.allclose(oklab_unsharp, rgb_unsharp, atol=1e-3)
    assert numpy.allclose(
        numpy.asarray(saved_image.convert("RGB"), dtype=numpy.float32) / 255.0,
        oklab_unsharp.permute(1, 2, 0).numpy(),
        atol=1 / 255.0,
    )


def test_oklch_hue_adjustment_invocation_preserves_alpha_and_rotates_hue_in_oklch() -> None:
    input_image = Image.new("RGBA", (2, 1))
    input_image.putdata(
        [
            (210, 80, 30, 64),
            (40, 160, 220, 192),
        ]
    )

    context = _build_context(input_image)

    invocation = OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=180)
    output = invocation.invoke(context)
    saved_image = context.images.save.call_args.kwargs["image"]

    rgb = torch.from_numpy(numpy.asarray(input_image.convert("RGB"), dtype=numpy.float32) / 255.0).permute(2, 0, 1)
    oklch = oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(rgb)))
    rotated_oklch = oklch.clone()
    rotated_oklch[2, ...] = (rotated_oklch[2, ...] + 180.0) % 360.0
    expected_rgb = srgb_from_linear_srgb(linear_srgb_from_oklch(rotated_oklch))

    assert output.image.image_name == "out"
    assert output.width == 2
    assert output.height == 1
    assert numpy.asarray(saved_image.getchannel("A")).reshape(-1).tolist() == [64, 192]
    assert numpy.allclose(
        numpy.asarray(saved_image.convert("RGB"), dtype=numpy.float32) / 255.0,
        expected_rgb.permute(1, 2, 0).numpy(),
        atol=1 / 255.0,
    )


def test_oklab_unsharp_mask_invocation_zero_strength_returns_original_image() -> None:
    input_image = Image.new("RGBA", (2, 2))
    input_image.putdata(
        [
            (12, 34, 56, 78),
            (90, 123, 45, 67),
            (210, 40, 80, 90),
            (255, 200, 10, 255),
        ]
    )
    context = _build_context(input_image)

    invocation = OklabUnsharpMaskInvocation(image=ImageField(image_name="in"), radius=1.5, strength=0.0)
    invocation.invoke(context)
    saved_image = context.images.save.call_args.kwargs["image"]

    assert _max_abs_diff_uint8(saved_image, input_image) <= 1


def test_oklab_unsharp_mask_invocation_does_not_introduce_color_on_grayscale_image() -> None:
    input_image = Image.new("RGB", (3, 1))
    input_image.putdata([(32, 32, 32), (128, 128, 128), (224, 224, 224)])
    context = _build_context(input_image)

    invocation = OklabUnsharpMaskInvocation(image=ImageField(image_name="in"), radius=1.0, strength=80.0)
    invocation.invoke(context)
    saved_image = context.images.save.call_args.kwargs["image"]
    saved_rgb = numpy.asarray(saved_image.convert("RGB"), dtype=numpy.uint8)

    assert numpy.abs(saved_rgb[..., 0].astype(numpy.int16) - saved_rgb[..., 1].astype(numpy.int16)).max() <= 1
    assert numpy.abs(saved_rgb[..., 1].astype(numpy.int16) - saved_rgb[..., 2].astype(numpy.int16)).max() <= 1


def test_oklab_unsharp_mask_invocation_clips_extreme_values_to_valid_rgb_range() -> None:
    input_image = Image.new("RGB", (3, 1))
    input_image.putdata([(255, 255, 255), (0, 0, 0), (255, 255, 255)])
    context = _build_context(input_image)

    invocation = OklabUnsharpMaskInvocation(image=ImageField(image_name="in"), radius=2.0, strength=500.0)
    invocation.invoke(context)
    saved_rgb = numpy.asarray(context.images.save.call_args.kwargs["image"].convert("RGB"), dtype=numpy.uint8)

    assert saved_rgb.min() >= 0
    assert saved_rgb.max() <= 255


def test_oklch_hue_adjustment_invocation_wraps_hue_values_and_supports_rgb_input() -> None:
    input_image = Image.new("RGB", (2, 1))
    input_image.putdata([(210, 80, 30), (40, 160, 220)])

    base_context = _build_context(input_image)
    zero_output = OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=0).invoke(base_context)
    zero_saved = base_context.images.save.call_args.kwargs["image"]

    full_turn_context = _build_context(input_image)
    full_turn_output = OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=360).invoke(
        full_turn_context
    )
    full_turn_saved = full_turn_context.images.save.call_args.kwargs["image"]

    negative_context = _build_context(input_image)
    OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=-180).invoke(negative_context)
    negative_saved = negative_context.images.save.call_args.kwargs["image"]

    positive_context = _build_context(input_image)
    OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=180).invoke(positive_context)
    positive_saved = positive_context.images.save.call_args.kwargs["image"]

    assert zero_output.width == 2
    assert zero_output.height == 1
    assert full_turn_output.width == 2
    assert full_turn_output.height == 1
    assert _max_abs_diff_uint8(zero_saved, input_image) <= 1
    assert _max_abs_diff_uint8(full_turn_saved, input_image) <= 1
    assert _max_abs_diff_uint8(negative_saved, positive_saved) <= 1


def test_new_oklab_nodes_preserve_alpha_for_non_rgba_alpha_modes() -> None:
    la_image = Image.new("LA", (2, 1))
    la_image.putdata([(32, 64), (192, 224)])

    unsharp_context = _build_context(la_image)
    OklabUnsharpMaskInvocation(image=ImageField(image_name="in"), radius=1.0, strength=25.0).invoke(unsharp_context)
    unsharp_saved = unsharp_context.images.save.call_args.kwargs["image"]

    hue_context = _build_context(la_image)
    OklchImageHueAdjustmentInvocation(image=ImageField(image_name="in"), hue=45).invoke(hue_context)
    hue_saved = hue_context.images.save.call_args.kwargs["image"]

    assert unsharp_saved.mode == "LA"
    assert hue_saved.mode == "LA"
    assert numpy.asarray(unsharp_saved.getchannel("A")).reshape(-1).tolist() == [64, 224]
    assert numpy.asarray(hue_saved.getchannel("A")).reshape(-1).tolist() == [64, 224]


def test_hue_adjust_plus_oklch_uses_degree_based_oklch_contract() -> None:
    input_image = Image.new("RGB", (2, 1))
    input_image.putdata([(210, 80, 30), (40, 160, 220)])

    context = _build_context(input_image)
    invocation = InvokeAdjustImageHuePlusInvocation(
        image=ImageField(image_name="in"),
        space="*Oklch / Oklab",
        degrees=180.0,
        ok_adaptive_gamut=0.0,
    )

    output = invocation.invoke(context)
    saved_image = context.images.save.call_args.args[0]

    rgb = torch.from_numpy(numpy.asarray(input_image, dtype=numpy.float32) / 255.0).permute(2, 0, 1)
    oklch = oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(rgb)))
    rotated_oklch = oklch.clone()
    rotated_oklch[2, ...] = (rotated_oklch[2, ...] + 180.0) % 360.0
    expected_rgb = srgb_from_linear_srgb(linear_srgb_from_oklch(rotated_oklch))

    assert output.width == 2
    assert output.height == 1
    assert numpy.allclose(
        numpy.asarray(saved_image.convert("RGB"), dtype=numpy.float32) / 255.0,
        expected_rgb.permute(1, 2, 0).numpy(),
        atol=1 / 255.0,
    )


def test_hue_adjust_plus_hsv_uses_degree_hue_contract() -> None:
    input_image = Image.new("RGB", (2, 1))
    input_image.putdata([(210, 80, 30), (40, 160, 220)])

    context = _build_context(input_image)
    invocation = InvokeAdjustImageHuePlusInvocation(
        image=ImageField(image_name="in"),
        space="HSV / HSL / RGB",
        degrees=90.0,
    )

    output = invocation.invoke(context)
    saved_image = context.images.save.call_args.args[0]

    hsv = numpy.asarray(input_image.convert("HSV"), dtype=numpy.float32) / 255.0
    hsv[..., 0] = ((hsv[..., 0] * 360.0) + 90.0) % 360.0 / 360.0
    expected_rgb = Image.fromarray((hsv * 255.0).astype(numpy.uint8), mode="HSV").convert("RGB")

    assert output.width == 2
    assert output.height == 1
    assert _max_abs_diff_uint8(saved_image.convert("RGB"), expected_rgb) <= 1


def test_hue_adjust_plus_okhsl_uses_degree_hue_contract() -> None:
    input_image = Image.new("RGB", (2, 1))
    input_image.putdata([(210, 80, 30), (40, 160, 220)])

    context = _build_context(input_image)
    invocation = InvokeAdjustImageHuePlusInvocation(
        image=ImageField(image_name="in"),
        space="Okhsl",
        degrees=90.0,
        ok_adaptive_gamut=0.0,
    )

    output = invocation.invoke(context)
    saved_image = context.images.save.call_args.args[0]

    rgb = torch.from_numpy(numpy.asarray(input_image, dtype=numpy.float32) / 255.0).permute(2, 0, 1)
    okhsl = okhsl_from_srgb(rgb)
    rotated_okhsl = okhsl.clone()
    rotated_okhsl[0, ...] = (rotated_okhsl[0, ...] + 90.0) % 360.0
    expected_rgb = srgb_from_okhsl(rotated_okhsl)

    assert output.width == 2
    assert output.height == 1
    assert numpy.allclose(
        numpy.asarray(saved_image.convert("RGB"), dtype=numpy.float32) / 255.0,
        expected_rgb.permute(1, 2, 0).numpy(),
        atol=1 / 255.0,
    )


def test_image_blend_oklch_subtract_wraps_hue_in_degrees() -> None:
    invocation = InvokeImageBlendInvocation(
        layer_upper=ImageField(image_name="upper"),
        layer_base=ImageField(image_name="base"),
        blend_mode="Subtract",
        color_space="Oklch (Oklab)",
        opacity=1.0,
        adaptive_gamut=0.0,
    )

    upper_oklch = torch.tensor([[[0.0]], [[0.0]], [[20.0]]], dtype=torch.float32)
    lower_oklch = torch.tensor([[[0.6]], [[0.18]], [[350.0]]], dtype=torch.float32)
    expected_linear_srgb = linear_srgb_from_oklch(torch.tensor([[[0.6]], [[0.18]], [[330.0]]], dtype=torch.float32))

    blank_rgb = torch.zeros((3, 1, 1), dtype=torch.float32)
    blank_alpha = torch.ones((1, 1), dtype=torch.float32)
    image_tensors = (
        blank_rgb,
        blank_rgb,
        blank_rgb,
        blank_rgb,
        blank_alpha,
        blank_alpha,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        upper_oklch,
        lower_oklch,
        None,
        None,
        None,
        None,
    )

    blended = invocation.apply_blend(image_tensors)

    assert torch.allclose(blended, expected_linear_srgb, atol=1e-5)


def test_image_blend_hsl_subtract_wraps_hue_in_degrees() -> None:
    invocation = InvokeImageBlendInvocation(
        layer_upper=ImageField(image_name="upper"),
        layer_base=ImageField(image_name="base"),
        blend_mode="Subtract",
        color_space="HSL (RGB)",
        opacity=1.0,
        adaptive_gamut=0.0,
    )

    upper_hsl = torch.tensor([[[20.0]], [[0.0]], [[0.0]]], dtype=torch.float32)
    lower_hsl = torch.tensor([[[350.0]], [[1.0]], [[0.5]]], dtype=torch.float32)
    expected_linear_srgb = linear_srgb_from_srgb(
        srgb_from_hsl(torch.tensor([[[330.0]], [[1.0]], [[0.5]]], dtype=torch.float32))
    )

    blank_rgb = torch.zeros((3, 1, 1), dtype=torch.float32)
    blank_alpha = torch.ones((1, 1), dtype=torch.float32)
    image_tensors = (
        blank_rgb,
        blank_rgb,
        blank_rgb,
        blank_rgb,
        blank_alpha,
        blank_alpha,
        None,
        None,
        None,
        upper_hsl,
        lower_hsl,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    )

    blended = invocation.apply_blend(image_tensors)

    assert torch.allclose(blended, expected_linear_srgb, atol=1e-5)