chore: import upstream snapshot with attribution

tests/app/util/test_controlnet_utils.py

@@ -0,0 +1,50 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from invokeai.app.util.controlnet_utils import prepare_control_image
+from invokeai.backend.image_util.util import nms
+
+
+@pytest.mark.parametrize("num_channels", [1, 2, 3])
+def test_prepare_control_image_num_channels(num_channels):
+    """Test that the `num_channels` parameter is applied correctly in prepare_control_image(...)."""
+    np_image = np.zeros((256, 256, 3), dtype=np.uint8)
+    pil_image = Image.fromarray(np_image)
+
+    torch_image = prepare_control_image(
+        image=pil_image,
+        width=256,
+        height=256,
+        num_channels=num_channels,
+        device="cpu",
+        do_classifier_free_guidance=False,
+    )
+
+    assert torch_image.shape == (1, num_channels, 256, 256)
+
+
+@pytest.mark.parametrize("num_channels", [0, 4])
+def test_prepare_control_image_num_channels_too_large(num_channels):
+    """Test that an exception is raised in prepare_control_image(...) if the `num_channels` parameter is out of the
+    supported range.
+    """
+    np_image = np.zeros((256, 256, 3), dtype=np.uint8)
+    pil_image = Image.fromarray(np_image)
+
+    with pytest.raises(ValueError):
+        _ = prepare_control_image(
+            image=pil_image,
+            width=256,
+            height=256,
+            num_channels=num_channels,
+            device="cpu",
+            do_classifier_free_guidance=False,
+        )
+
+
+@pytest.mark.parametrize("threshold,sigma", [(None, 1.0), (1, None)])
+def test_nms_invalid_options(threshold: None | int, sigma: None | float):
+    """Test that an exception is raised in nms(...) if only one of the `threshold` or `sigma` parameters are provided."""
+    with pytest.raises(ValueError):
+        nms(np.zeros((256, 256, 3), dtype=np.uint8), threshold, sigma)

tests/app/util/test_custom_openapi.py

@@ -0,0 +1,61 @@
+from fastapi import FastAPI
+from pydantic import create_model
+
+from invokeai.app.invocations.baseinvocation import InvocationRegistry
+from invokeai.app.util.custom_openapi import get_openapi_func
+
+
+class _FakeOutput:
+    pass
+
+
+class _InvocationB:
+    __name__ = "InvocationB"
+
+    @classmethod
+    def model_json_schema(cls, mode: str, ref_template: str) -> dict:
+        return {"type": "object", "properties": {}}
+
+    @classmethod
+    def get_output_annotation(cls) -> type:
+        return _FakeOutput
+
+    @classmethod
+    def get_type(cls) -> str:
+        return "b_type"
+
+
+class _InvocationA:
+    __name__ = "InvocationA"
+
+    @classmethod
+    def model_json_schema(cls, mode: str, ref_template: str) -> dict:
+        return {"type": "object", "properties": {}}
+
+    @classmethod
+    def get_output_annotation(cls) -> type:
+        return _FakeOutput
+
+    @classmethod
+    def get_type(cls) -> str:
+        return "a_type"
+
+
+def test_invocation_output_map_required_is_sorted(monkeypatch: object) -> None:
+    """The 'required' list in InvocationOutputMap must be sorted so that the
+    generated openapi.json is deterministic regardless of set-iteration order."""
+
+    # A FastAPI app needs at least one route to produce a schema with 'components'.
+    DummyResponse = create_model("DummyResponse", ok=(bool, ...))
+    app = FastAPI(title="test")
+    app.get("/healthz", response_model=DummyResponse)(lambda: DummyResponse(ok=True))
+
+    monkeypatch.setattr(InvocationRegistry, "get_output_classes", classmethod(lambda cls: []))  # type: ignore[arg-type]
+    monkeypatch.setattr(  # type: ignore[arg-type]
+        InvocationRegistry, "get_invocation_classes", classmethod(lambda cls: [_InvocationB, _InvocationA])
+    )
+
+    schema = get_openapi_func(app)()
+    required = schema["components"]["schemas"]["InvocationOutputMap"]["required"]
+
+    assert required == ["a_type", "b_type"], f"Expected sorted required list, got: {required}"

tests/app/util/test_dynamicprompts.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import pytest
+
+from invokeai.app.util.dynamicprompts import find_missing_wildcards
+
+
+def test_find_missing_wildcards_detects_unknown_wildcard_in_variant() -> None:
+    # Regression: `__random__` inside a variant is parsed as a wildcard reference. Left unchecked it
+    # sends the combinatorial generator into an infinite loop, so it must be reported up front.
+    assert find_missing_wildcards("{__random__8chan|fenster|stuff}") == ["random"]
+
+
+def test_find_missing_wildcards_detects_unknown_wildcard_nested_in_sequence_in_variant() -> None:
+    # The wildcard hangs the generator even when wrapped in other text inside the variant value.
+    assert find_missing_wildcards("{a __nope__|b}") == ["nope"]
+
+
+@pytest.mark.parametrize("prompt", ["a __nope__ b", "__nope__", "a photo, __my_style__"])
+def test_find_missing_wildcards_ignores_wildcards_outside_variants(prompt: str) -> None:
+    # A wildcard used as plain literal text generates fine (no hang), so it must not be reported.
+    assert find_missing_wildcards(prompt) == []
+
+
+@pytest.mark.parametrize("prompt", ["plain text", "{a|b|c}", "a {2$$x|y|z}"])
+def test_find_missing_wildcards_ignores_prompts_without_wildcards(prompt: str) -> None:
+    assert find_missing_wildcards(prompt) == []
+
+
+def test_find_missing_wildcards_dedupes_repeated_unknown_wildcards() -> None:
+    assert find_missing_wildcards("{__nope__|a} {__nope__|b} {__other__|c}") == ["nope", "other"]

tests/app/util/test_step_callback.py

@@ -0,0 +1,119 @@
+"""Tests for diffusion step callback preview image generation."""
+
+import torch
+from PIL import Image
+
+from invokeai.app.util.step_callback import (
+    QWEN_IMAGE_LATENT_RGB_BIAS,
+    QWEN_IMAGE_LATENT_RGB_FACTORS,
+    sample_to_lowres_estimated_image,
+)
+
+
+class TestSampleToLowresEstimatedImage:
+    """Test the latent-to-preview-image conversion used during denoising."""
+
+    def test_qwen_image_preview_produces_valid_image(self):
+        """A synthetic Qwen latent tensor produces a valid RGB preview image."""
+        # Create a small 1x16x4x4 latent tensor (batch=1, channels=16, 4x4 spatial)
+        torch.manual_seed(42)
+        sample = torch.randn(1, 16, 4, 4)
+
+        factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype)
+        bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype)
+
+        image = sample_to_lowres_estimated_image(
+            samples=sample,
+            latent_rgb_factors=factors,
+            latent_rgb_bias=bias,
+        )
+
+        assert isinstance(image, Image.Image)
+        assert image.size == (4, 4)
+        assert image.mode == "RGB"
+
+    def test_qwen_image_preview_deterministic(self):
+        """The same input tensor always produces the same preview image."""
+        sample = torch.ones(1, 16, 2, 2)
+
+        factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype)
+        bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype)
+
+        image1 = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias)
+        image2 = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias)
+
+        assert list(image1.getdata()) == list(image2.getdata())
+
+    def test_qwen_image_preview_known_value(self):
+        """Verify the preview computation against a hand-calculated expected value.
+
+        With a 1x16x1x1 tensor of all ones:
+        - latent_image = [1,1,...,1] @ factors = sum of each column of factors
+        - R = sum(col 0) = 0.3677, G = sum(col 1) = 0.4577, B = sum(col 2) = 0.9101
+        - After bias: R = 0.1842, G = 0.3709, B = 0.5741
+        - After scale ((x+1)/2): R = 0.5921, G = 0.6855, B = 0.7871
+        - After quantize (*255): R = 151, G = 175, B = 201
+        """
+        sample = torch.ones(1, 16, 1, 1)
+
+        factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype)
+        bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype)
+
+        image = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias)
+
+        assert image.size == (1, 1)
+        pixel = image.getpixel((0, 0))
+
+        # Compute expected values
+        col_sums = [sum(row[c] for row in QWEN_IMAGE_LATENT_RGB_FACTORS) for c in range(3)]
+        expected = []
+        for c in range(3):
+            val = col_sums[c] + QWEN_IMAGE_LATENT_RGB_BIAS[c]
+            val = (val + 1) / 2  # scale from [-1,1] to [0,1]
+            val = max(0.0, min(1.0, val))  # clamp
+            expected.append(int(val * 255))
+
+        assert pixel == tuple(expected), f"Expected {tuple(expected)}, got {pixel}"
+
+    def test_qwen_image_preview_zeros_tensor(self):
+        """A zero tensor with bias produces a valid image reflecting just the bias."""
+        sample = torch.zeros(1, 16, 2, 2)
+
+        factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype)
+        bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype)
+
+        image = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias)
+
+        assert isinstance(image, Image.Image)
+        assert image.size == (2, 2)
+
+        # All pixels should be identical (uniform zero input)
+        pixels = [image.getpixel((x, y)) for y in range(image.height) for x in range(image.width)]
+        assert all(p == pixels[0] for p in pixels)
+
+        # With zero input, result = bias, scaled: ((bias + 1) / 2) * 255
+        expected = []
+        for c in range(3):
+            val = (QWEN_IMAGE_LATENT_RGB_BIAS[c] + 1) / 2
+            val = max(0.0, min(1.0, val))
+            expected.append(int(val * 255))
+        assert pixels[0] == tuple(expected)
+
+    def test_qwen_image_factors_have_correct_shape(self):
+        """Qwen Image uses 16 latent channels, so factors should be 16x3."""
+        assert len(QWEN_IMAGE_LATENT_RGB_FACTORS) == 16
+        for row in QWEN_IMAGE_LATENT_RGB_FACTORS:
+            assert len(row) == 3
+        assert len(QWEN_IMAGE_LATENT_RGB_BIAS) == 3
+
+    def test_3d_input_accepted(self):
+        """sample_to_lowres_estimated_image accepts 3D input (no batch dim)."""
+        sample = torch.randn(16, 4, 4)  # no batch dimension
+
+        factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype)
+        bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype)
+
+        image = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias)
+
+        assert isinstance(image, Image.Image)
+        assert image.size == (4, 4)

tests/app/util/test_torch_cuda_allocator.py

@@ -0,0 +1,127 @@
+import pytest
+import torch
+
+from tests.dangerously_run_function_in_subprocess import dangerously_run_function_in_subprocess
+
+# These tests are a bit fiddly, because the depend on the import behaviour of torch. They use subprocesses to isolate
+# the import behaviour of torch, and then check that the function behaves as expected. We have to hack in some logging
+# to check that the tested function is behaving as expected.
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
+def test_configure_torch_cuda_allocator_configures_backend():
+    """Test that configure_torch_cuda_allocator() raises a RuntimeError if the configured backend does not match the
+    expected backend."""
+
+    def test_func():
+        import os
+
+        # Unset the environment variable if it is set so that we can test setting it
+        try:
+            del os.environ["PYTORCH_CUDA_ALLOC_CONF"]
+        except KeyError:
+            pass
+
+        from unittest.mock import MagicMock
+
+        from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
+
+        mock_logger = MagicMock()
+
+        # Set the PyTorch CUDA memory allocator to cudaMallocAsync
+        configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=mock_logger)
+
+        # Verify that the PyTorch CUDA memory allocator was configured correctly
+        import torch
+
+        assert torch.cuda.get_allocator_backend() == "cudaMallocAsync"
+
+        # Verify that the logger was called with the correct message
+        mock_logger.info.assert_called_once()
+        args, _kwargs = mock_logger.info.call_args
+        logged_message = args[0]
+        print(logged_message)
+
+    stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
+    assert returncode == 0
+    assert "PyTorch CUDA memory allocator: cudaMallocAsync" in stdout
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
+def test_configure_torch_cuda_allocator_raises_if_torch_already_imported():
+    """Test that configure_torch_cuda_allocator() raises a RuntimeError if torch was already imported."""
+
+    def test_func():
+        from unittest.mock import MagicMock
+
+        # Import torch before calling configure_torch_cuda_allocator()
+        import torch  # noqa: F401
+
+        from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
+
+        try:
+            configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=MagicMock())
+        except RuntimeError as e:
+            print(e)
+
+    stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
+    assert returncode == 0
+    assert "configure_torch_cuda_allocator() must be called before importing torch." in stdout
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
+def test_configure_torch_cuda_allocator_warns_if_env_var_is_set_differently():
+    """Test that configure_torch_cuda_allocator() logs at WARNING level if PYTORCH_CUDA_ALLOC_CONF is set and doesn't
+    match the requested configuration."""
+
+    def test_func():
+        import os
+
+        # Explicitly set the environment variable
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:native"
+
+        from unittest.mock import MagicMock
+
+        from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
+
+        mock_logger = MagicMock()
+
+        # Set the PyTorch CUDA memory allocator a different configuration
+        configure_torch_cuda_allocator("backend:cudaMallocAsync", logger=mock_logger)
+
+        # Verify that the logger was called with the correct message
+        mock_logger.warning.assert_called_once()
+        args, _kwargs = mock_logger.warning.call_args
+        logged_message = args[0]
+        print(logged_message)
+
+    stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
+    assert returncode == 0
+    assert "Attempted to configure the PyTorch CUDA memory allocator with 'backend:cudaMallocAsync'" in stdout
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA device.")
+def test_configure_torch_cuda_allocator_logs_if_env_var_is_already_set_correctly():
+    """Test that configure_torch_cuda_allocator() logs at INFO level if PYTORCH_CUDA_ALLOC_CONF is set and matches the
+    requested configuration."""
+
+    def test_func():
+        import os
+
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:native"
+        from unittest.mock import MagicMock
+
+        from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
+
+        mock_logger = MagicMock()
+
+        configure_torch_cuda_allocator("backend:native", logger=mock_logger)
+
+        mock_logger.info.assert_called_once()
+        args, _kwargs = mock_logger.info.call_args
+        logged_message = args[0]
+        print(logged_message)
+
+    stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func)
+    assert returncode == 0
+    assert "PYTORCH_CUDA_ALLOC_CONF is already set to 'backend:native'" in stdout