wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
665 lines
26 KiB
Python
665 lines
26 KiB
Python
import inspect
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from invokeai.app.invocations.anima_denoise import AnimaDenoiseInvocation
|
|
from invokeai.app.invocations.cogview4_denoise import CogView4DenoiseInvocation
|
|
from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation
|
|
from invokeai.app.invocations.flux_denoise import FluxDenoiseInvocation
|
|
from invokeai.app.invocations.metadata_linked import FluxDenoiseLatentsMetaInvocation, ZImageDenoiseMetaInvocation
|
|
from invokeai.app.invocations.primitives import LatentsOutput
|
|
from invokeai.app.invocations.sd3_denoise import SD3DenoiseInvocation
|
|
from invokeai.app.invocations.z_image_denoise import ZImageDenoiseInvocation
|
|
from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional, get_schedule
|
|
from invokeai.backend.flux.schedulers import ANIMA_SCHEDULER_MAP, FLUX_SCHEDULER_MAP, ZIMAGE_SCHEDULER_MAP
|
|
from invokeai.backend.flux2.sampling_utils import get_schedule_flux2
|
|
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
|
|
|
|
|
|
def test_flux_prepare_noise_uses_external_noise():
|
|
invocation = FluxDenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
expected = torch.zeros(1, 16, 8, 8)
|
|
mock_context.tensors.load.return_value = expected
|
|
|
|
with patch("invokeai.app.invocations.flux_denoise.get_noise") as mock_get_noise:
|
|
noise = invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
assert torch.equal(noise, expected.to(dtype=torch.bfloat16))
|
|
mock_get_noise.assert_not_called()
|
|
|
|
|
|
def test_flux_prepare_noise_rejects_invalid_shape():
|
|
invocation = FluxDenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = torch.zeros(1, 15, 8, 8)
|
|
|
|
with pytest.raises(ValueError, match="Expected noise with shape"):
|
|
invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
|
|
def test_flux_add_noise_false_ignores_connected_noise():
|
|
invocation = FluxDenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
noise=MagicMock(latents_name="noise"),
|
|
add_noise=False,
|
|
width=64,
|
|
height=64,
|
|
num_steps=4,
|
|
denoising_start=0.25,
|
|
denoising_end=0.25,
|
|
positive_text_conditioning=MagicMock(conditioning_name="positive"),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
seed=123,
|
|
)
|
|
init_latents = torch.full((1, 16, 8, 8), 2.0)
|
|
dummy_conditioning = SimpleNamespace(
|
|
t5_embeds=torch.zeros(1, 4, 16),
|
|
clip_embeds=torch.zeros(1, 768),
|
|
to=lambda **_: dummy_conditioning,
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
mock_context.conditioning.load.return_value = SimpleNamespace(conditionings=[dummy_conditioning])
|
|
mock_context.models.get_config.return_value = SimpleNamespace(
|
|
base=BaseModelType.Flux, type=ModelType.Main, variant=None
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"invokeai.app.invocations.flux_denoise.TorchDevice.choose_torch_device", return_value=torch.device("cpu")
|
|
),
|
|
patch("invokeai.app.invocations.flux_denoise.FLUXConditioningInfo", object),
|
|
patch(
|
|
"invokeai.app.invocations.flux_denoise.RegionalPromptingExtension.from_text_conditioning",
|
|
return_value=MagicMock(),
|
|
),
|
|
patch.object(invocation, "_prepare_noise_tensor", side_effect=AssertionError("noise should be ignored")),
|
|
patch.object(invocation, "_load_redux_conditioning", return_value=[]),
|
|
patch("invokeai.app.invocations.flux_denoise.get_schedule", return_value=[0.75]),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
assert torch.equal(result, init_latents)
|
|
|
|
|
|
def test_flux2_prepare_noise_uses_external_noise():
|
|
invocation = Flux2DenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
expected = torch.zeros(1, 32, 8, 8)
|
|
mock_context.tensors.load.return_value = expected
|
|
|
|
with patch("invokeai.app.invocations.flux2_denoise.get_noise_flux2") as mock_get_noise:
|
|
noise = invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
assert torch.equal(noise, expected.to(dtype=torch.bfloat16))
|
|
mock_get_noise.assert_not_called()
|
|
|
|
|
|
def test_flux2_prepare_noise_rejects_invalid_shape():
|
|
invocation = Flux2DenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = torch.zeros(1, 16, 8, 8)
|
|
|
|
with pytest.raises(ValueError, match="Expected noise with shape"):
|
|
invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
|
|
def test_sd3_prepare_noise_uses_external_noise():
|
|
invocation = SD3DenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
expected = torch.zeros(1, 16, 8, 8)
|
|
mock_context.tensors.load.return_value = expected
|
|
|
|
with patch.object(invocation, "_get_noise") as mock_get_noise:
|
|
noise = invocation._prepare_noise_tensor(mock_context, 16, torch.bfloat16, torch.device("cpu"))
|
|
|
|
assert torch.equal(noise, expected.to(dtype=torch.bfloat16))
|
|
mock_get_noise.assert_not_called()
|
|
|
|
|
|
def test_sd3_prepare_noise_rejects_invalid_shape():
|
|
invocation = SD3DenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = torch.zeros(1, 8, 8, 8)
|
|
|
|
with pytest.raises(ValueError, match="Expected noise with shape"):
|
|
invocation._prepare_noise_tensor(mock_context, 16, torch.bfloat16, torch.device("cpu"))
|
|
|
|
|
|
def test_cogview4_prepare_noise_uses_external_noise():
|
|
invocation = CogView4DenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
expected = torch.zeros(1, 16, 8, 8)
|
|
mock_context.tensors.load.return_value = expected
|
|
|
|
with patch.object(invocation, "_get_noise") as mock_get_noise:
|
|
noise = invocation._prepare_noise_tensor(mock_context, 16, torch.bfloat16, torch.device("cpu"))
|
|
|
|
assert torch.equal(noise, expected.to(dtype=torch.bfloat16))
|
|
mock_get_noise.assert_not_called()
|
|
|
|
|
|
def test_cogview4_prepare_noise_rejects_invalid_shape():
|
|
invocation = CogView4DenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = torch.zeros(1, 4, 8, 8)
|
|
|
|
with pytest.raises(ValueError, match="Expected noise with shape"):
|
|
invocation._prepare_noise_tensor(mock_context, 16, torch.bfloat16, torch.device("cpu"))
|
|
|
|
|
|
def test_z_image_prepare_noise_uses_external_noise():
|
|
invocation = ZImageDenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
expected = torch.zeros(1, 16, 8, 8)
|
|
mock_context.tensors.load.return_value = expected
|
|
|
|
with patch.object(invocation, "_get_noise") as mock_get_noise:
|
|
noise = invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
assert torch.equal(noise, expected.to(dtype=torch.bfloat16))
|
|
mock_get_noise.assert_not_called()
|
|
|
|
|
|
def test_z_image_prepare_noise_rejects_invalid_shape():
|
|
invocation = ZImageDenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = torch.zeros(1, 8, 8, 8)
|
|
|
|
with pytest.raises(ValueError, match="Expected noise with shape"):
|
|
invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
|
|
def test_z_image_add_noise_false_ignores_connected_noise():
|
|
invocation = ZImageDenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
noise=MagicMock(latents_name="noise"),
|
|
add_noise=False,
|
|
width=64,
|
|
height=64,
|
|
steps=4,
|
|
denoising_start=0.0,
|
|
denoising_end=1.0,
|
|
positive_conditioning=SimpleNamespace(conditioning_name="positive", mask=None),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
seed=123,
|
|
scheduler="euler",
|
|
)
|
|
init_latents = torch.full((1, 16, 8, 8), 2.0)
|
|
dummy_conditioning = SimpleNamespace(prompt_embeds=torch.zeros(1, 4, 16))
|
|
dummy_conditioning.to = lambda **_: dummy_conditioning
|
|
regional_extension = SimpleNamespace(
|
|
regional_text_conditioning=SimpleNamespace(prompt_embeds=torch.zeros(1, 4, 16))
|
|
)
|
|
loaded_text_conditioning = [SimpleNamespace(prompt_embeds=torch.zeros(1, 4, 16), mask=None)]
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
mock_context.conditioning.load.return_value = SimpleNamespace(conditionings=[dummy_conditioning])
|
|
|
|
with (
|
|
patch(
|
|
"invokeai.app.invocations.z_image_denoise.TorchDevice.choose_torch_device", return_value=torch.device("cpu")
|
|
),
|
|
patch(
|
|
"invokeai.app.invocations.z_image_denoise.TorchDevice.choose_bfloat16_safe_dtype",
|
|
return_value=torch.bfloat16,
|
|
),
|
|
patch("invokeai.app.invocations.z_image_denoise.ZImageConditioningInfo", object),
|
|
patch(
|
|
"invokeai.app.invocations.z_image_denoise.ZImageRegionalPromptingExtension.from_text_conditionings",
|
|
return_value=regional_extension,
|
|
),
|
|
patch.object(invocation, "_load_text_conditioning", return_value=loaded_text_conditioning),
|
|
patch.object(invocation, "_prepare_noise_tensor", side_effect=AssertionError("noise should be ignored")),
|
|
patch.object(invocation, "_get_sigmas", return_value=[0.75]),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
assert torch.equal(result, init_latents)
|
|
|
|
|
|
def test_anima_prepare_noise_uses_external_noise():
|
|
invocation = AnimaDenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
expected = torch.zeros(1, 16, 1, 8, 8)
|
|
mock_context.tensors.load.return_value = expected
|
|
|
|
with patch.object(invocation, "_get_noise") as mock_get_noise:
|
|
noise = invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
assert torch.equal(noise, expected.to(dtype=torch.bfloat16))
|
|
mock_get_noise.assert_not_called()
|
|
|
|
|
|
def test_anima_prepare_noise_rejects_invalid_rank():
|
|
invocation = AnimaDenoiseInvocation.model_construct(
|
|
width=64, height=64, seed=0, noise=MagicMock(latents_name="noise")
|
|
)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = torch.zeros(1, 16, 8, 8)
|
|
|
|
with pytest.raises(ValueError, match="Expected noise with shape"):
|
|
invocation._prepare_noise_tensor(mock_context, torch.bfloat16, torch.device("cpu"))
|
|
|
|
|
|
def test_anima_add_noise_false_ignores_connected_noise():
|
|
invocation = AnimaDenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
noise=MagicMock(latents_name="noise"),
|
|
add_noise=False,
|
|
width=64,
|
|
height=64,
|
|
steps=4,
|
|
denoising_start=0.0,
|
|
denoising_end=1.0,
|
|
positive_conditioning=SimpleNamespace(conditioning_name="positive", mask=None),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
seed=123,
|
|
scheduler="euler",
|
|
)
|
|
init_latents = torch.full((1, 16, 8, 8), 2.0)
|
|
loaded_text_conditioning = [SimpleNamespace(mask=None)]
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
mock_context.models.load.return_value = MagicMock()
|
|
|
|
with (
|
|
patch(
|
|
"invokeai.app.invocations.anima_denoise.TorchDevice.choose_torch_device", return_value=torch.device("cpu")
|
|
),
|
|
patch(
|
|
"invokeai.app.invocations.anima_denoise.TorchDevice.choose_bfloat16_safe_dtype", return_value=torch.bfloat16
|
|
),
|
|
patch.object(invocation, "_load_text_conditionings", return_value=loaded_text_conditioning),
|
|
patch.object(invocation, "_prepare_noise_tensor", side_effect=AssertionError("noise should be ignored")),
|
|
patch.object(invocation, "_get_sigmas", return_value=[0.75]),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
assert torch.equal(result, init_latents)
|
|
|
|
|
|
def test_flux2_add_noise_false_ignores_connected_noise():
|
|
invocation = Flux2DenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
noise=MagicMock(latents_name="noise"),
|
|
add_noise=False,
|
|
width=64,
|
|
height=64,
|
|
num_steps=4,
|
|
denoising_start=0.25,
|
|
denoising_end=0.25,
|
|
positive_text_conditioning=MagicMock(conditioning_name="positive"),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
vae=MagicMock(vae="vae"),
|
|
seed=123,
|
|
)
|
|
init_latents = torch.full((1, 32, 8, 8), 2.0)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
mock_context.conditioning.load.return_value = SimpleNamespace(
|
|
conditionings=[
|
|
SimpleNamespace(
|
|
t5_embeds=torch.zeros(1, 4, 16), to=lambda **_: SimpleNamespace(t5_embeds=torch.zeros(1, 4, 16))
|
|
)
|
|
]
|
|
)
|
|
mock_context.models.get_config.return_value = SimpleNamespace(base=BaseModelType.Flux2, type=ModelType.Main)
|
|
|
|
with (
|
|
patch(
|
|
"invokeai.app.invocations.flux2_denoise.TorchDevice.choose_torch_device", return_value=torch.device("cpu")
|
|
),
|
|
patch("invokeai.app.invocations.flux2_denoise.FLUXConditioningInfo", object),
|
|
patch.object(invocation, "_get_bn_stats", return_value=None),
|
|
patch.object(invocation, "_prepare_noise_tensor", side_effect=AssertionError("noise should be ignored")),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
assert torch.equal(result, init_latents)
|
|
|
|
|
|
def test_flux_metadata_ignores_external_noise_seed_when_noise_not_used():
|
|
invocation = FluxDenoiseLatentsMetaInvocation.model_construct(
|
|
width=64,
|
|
height=64,
|
|
num_steps=4,
|
|
guidance=3.5,
|
|
denoising_start=0.0,
|
|
denoising_end=1.0,
|
|
latents=MagicMock(latents_name="latents"),
|
|
transformer=MagicMock(transformer="transformer", loras=[]),
|
|
noise=MagicMock(seed=123),
|
|
seed=999,
|
|
add_noise=False,
|
|
)
|
|
mock_context = MagicMock()
|
|
output = LatentsOutput.build("latents", torch.zeros(1, 16, 8, 8), seed=None)
|
|
|
|
with patch("invokeai.app.invocations.metadata_linked.FluxDenoiseInvocation.invoke", return_value=output):
|
|
result = invocation.invoke(mock_context)
|
|
|
|
assert result.metadata.root["seed"] == 999
|
|
|
|
|
|
def test_z_image_metadata_ignores_external_noise_seed_when_noise_not_used():
|
|
invocation = ZImageDenoiseMetaInvocation.model_construct(
|
|
width=64,
|
|
height=64,
|
|
steps=8,
|
|
guidance_scale=1.0,
|
|
denoising_start=0.0,
|
|
denoising_end=1.0,
|
|
scheduler="euler",
|
|
latents=MagicMock(latents_name="latents"),
|
|
transformer=MagicMock(transformer="transformer", loras=[]),
|
|
noise=MagicMock(seed=123),
|
|
seed=999,
|
|
add_noise=False,
|
|
)
|
|
mock_context = MagicMock()
|
|
output = LatentsOutput.build("latents", torch.zeros(1, 16, 8, 8), seed=None)
|
|
|
|
with patch("invokeai.app.invocations.metadata_linked.ZImageDenoiseInvocation.invoke", return_value=output):
|
|
result = invocation.invoke(mock_context)
|
|
|
|
assert result.metadata.root["seed"] == 999
|
|
|
|
|
|
def _get_first_scheduler_sigma(
|
|
scheduler, *, scheduler_name: str, sigmas: list[float], mu: float | None = None
|
|
) -> float:
|
|
set_timesteps_signature = inspect.signature(scheduler.set_timesteps)
|
|
if scheduler_name != "lcm" and "sigmas" in set_timesteps_signature.parameters:
|
|
kwargs: dict[str, object] = {"sigmas": sigmas, "device": "cpu"}
|
|
if mu is not None and "mu" in set_timesteps_signature.parameters:
|
|
kwargs["mu"] = mu
|
|
scheduler.set_timesteps(**kwargs)
|
|
else:
|
|
kwargs = {"num_inference_steps": len(sigmas) - 1, "device": "cpu"}
|
|
if mu is not None and "mu" in set_timesteps_signature.parameters:
|
|
kwargs["mu"] = mu
|
|
scheduler.set_timesteps(**kwargs)
|
|
return float(scheduler.sigmas[0])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"scheduler_name",
|
|
[
|
|
"euler",
|
|
pytest.param(
|
|
"heun",
|
|
marks=pytest.mark.xfail(
|
|
reason="Known img2img preblend mismatch for FLUX with scheduler-defined first step.",
|
|
strict=True,
|
|
),
|
|
),
|
|
pytest.param(
|
|
"lcm",
|
|
marks=pytest.mark.xfail(
|
|
reason="Known img2img preblend mismatch for FLUX with scheduler-defined first step.",
|
|
strict=True,
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_flux_img2img_preblend_matches_scheduler_first_sigma(scheduler_name: str):
|
|
sigmas = clip_timestep_schedule_fractional(get_schedule(num_steps=4, image_seq_len=16, shift=True), 0.25, 1.0)
|
|
scheduler_class = FLUX_SCHEDULER_MAP[scheduler_name]
|
|
scheduler = scheduler_class(num_train_timesteps=1000)
|
|
|
|
assert sigmas[0] == pytest.approx(
|
|
_get_first_scheduler_sigma(scheduler, scheduler_name=scheduler_name, sigmas=sigmas)
|
|
)
|
|
|
|
|
|
def test_flux2_partial_denoise_short_circuit_uses_first_clipped_timestep():
|
|
invocation = Flux2DenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
width=64,
|
|
height=64,
|
|
num_steps=4,
|
|
denoising_start=0.25,
|
|
denoising_end=0.25,
|
|
positive_text_conditioning=MagicMock(conditioning_name="positive"),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
vae=MagicMock(vae="vae"),
|
|
seed=0,
|
|
scheduler="lcm",
|
|
)
|
|
init_latents = torch.full((1, 32, 8, 8), 2.0)
|
|
noise = torch.full((1, 32, 8, 8), 10.0)
|
|
dummy_conditioning = SimpleNamespace(t5_embeds=torch.zeros(1, 4, 16))
|
|
dummy_conditioning.to = lambda **_: dummy_conditioning
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
mock_context.conditioning.load.return_value = SimpleNamespace(conditionings=[dummy_conditioning])
|
|
mock_context.models.get_config.return_value = SimpleNamespace(base=BaseModelType.Flux2, type=ModelType.Main)
|
|
|
|
with (
|
|
patch(
|
|
"invokeai.app.invocations.flux2_denoise.TorchDevice.choose_torch_device", return_value=torch.device("cpu")
|
|
),
|
|
patch("invokeai.app.invocations.flux2_denoise.FLUXConditioningInfo", object),
|
|
patch.object(invocation, "_get_bn_stats", return_value=None),
|
|
patch.object(invocation, "_prepare_noise_tensor", return_value=noise),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
timesteps = clip_timestep_schedule_fractional(get_schedule_flux2(num_steps=4, image_seq_len=16), 0.25, 0.25)
|
|
expected = timesteps[0] * noise + (1.0 - timesteps[0]) * init_latents
|
|
assert torch.equal(result, expected)
|
|
|
|
|
|
def test_flux2_lcm_scheduler_setup_passes_mu():
|
|
from invokeai.backend.flux2.denoise import denoise
|
|
|
|
class DummyScheduler:
|
|
def __init__(self) -> None:
|
|
self.received_mu = None
|
|
self.timesteps = torch.tensor([750.0, 500.0], dtype=torch.float32)
|
|
self.sigmas = torch.tensor([0.75, 0.5, 0.0], dtype=torch.float32)
|
|
self.config = SimpleNamespace(num_train_timesteps=1000)
|
|
|
|
def set_timesteps(self, num_inference_steps: int, device: str | torch.device, mu: float | None = None) -> None:
|
|
self.received_mu = mu
|
|
|
|
def step(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor):
|
|
return SimpleNamespace(prev_sample=sample)
|
|
|
|
class DummyModel(torch.nn.Module):
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
img_ids: torch.Tensor,
|
|
txt_ids: torch.Tensor,
|
|
guidance: torch.Tensor,
|
|
return_dict: bool = False,
|
|
):
|
|
return (torch.zeros_like(hidden_states),)
|
|
|
|
scheduler = DummyScheduler()
|
|
denoise(
|
|
model=DummyModel(),
|
|
img=torch.zeros(1, 4, 8),
|
|
img_ids=torch.zeros(1, 4, 4, dtype=torch.long),
|
|
txt=torch.zeros(1, 4, 8),
|
|
txt_ids=torch.zeros(1, 4, 4, dtype=torch.long),
|
|
timesteps=[0.75, 0.5, 0.0],
|
|
step_callback=lambda _: None,
|
|
guidance=1.0,
|
|
cfg_scale=[1.0, 1.0],
|
|
scheduler=scheduler,
|
|
mu=0.42,
|
|
)
|
|
|
|
assert scheduler.received_mu == pytest.approx(0.42)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"scheduler_name",
|
|
[
|
|
"euler",
|
|
pytest.param(
|
|
"heun",
|
|
marks=pytest.mark.xfail(
|
|
reason="Known img2img preblend mismatch for Z-Image with scheduler-defined first step.",
|
|
strict=True,
|
|
),
|
|
),
|
|
pytest.param(
|
|
"lcm",
|
|
marks=pytest.mark.xfail(
|
|
reason="Known img2img preblend mismatch for Z-Image with scheduler-defined first step.",
|
|
strict=True,
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_z_image_img2img_preblend_matches_scheduler_first_sigma(scheduler_name: str):
|
|
invocation = ZImageDenoiseInvocation.model_construct(steps=8, width=1024, height=1024)
|
|
img_seq_len = (invocation.height // 8 // 2) * (invocation.width // 8 // 2)
|
|
shift = invocation._calculate_shift(img_seq_len)
|
|
sigmas = invocation._get_sigmas(shift, invocation.steps)
|
|
sigmas = sigmas[int(0.25 * (len(sigmas) - 1)) :]
|
|
scheduler_class = ZIMAGE_SCHEDULER_MAP[scheduler_name]
|
|
scheduler = scheduler_class(num_train_timesteps=1000, shift=1.0)
|
|
|
|
assert sigmas[0] == pytest.approx(
|
|
_get_first_scheduler_sigma(scheduler, scheduler_name=scheduler_name, sigmas=sigmas)
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"scheduler_name",
|
|
[
|
|
"euler",
|
|
pytest.param(
|
|
"heun",
|
|
marks=pytest.mark.xfail(
|
|
reason="Known img2img preblend mismatch for Anima with scheduler-defined first step.",
|
|
strict=True,
|
|
),
|
|
),
|
|
pytest.param(
|
|
"lcm",
|
|
marks=pytest.mark.xfail(
|
|
reason="Known img2img preblend mismatch for Anima with scheduler-defined first step.",
|
|
strict=True,
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_anima_img2img_preblend_matches_scheduler_first_sigma(scheduler_name: str):
|
|
invocation = AnimaDenoiseInvocation.model_construct(steps=30)
|
|
sigmas = invocation._get_sigmas(invocation.steps)
|
|
sigmas = sigmas[int(0.25 * (len(sigmas) - 1)) :]
|
|
scheduler_class, scheduler_kwargs = ANIMA_SCHEDULER_MAP[scheduler_name]
|
|
scheduler = scheduler_class(num_train_timesteps=1000, **scheduler_kwargs)
|
|
|
|
assert sigmas[0] == pytest.approx(
|
|
_get_first_scheduler_sigma(scheduler, scheduler_name=scheduler_name, sigmas=sigmas)
|
|
)
|
|
|
|
|
|
def test_sd3_partial_denoise_short_circuit_uses_first_clipped_timestep():
|
|
invocation = SD3DenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
width=64,
|
|
height=64,
|
|
steps=4,
|
|
denoising_start=0.25,
|
|
denoising_end=0.25,
|
|
positive_conditioning=MagicMock(conditioning_name="positive"),
|
|
negative_conditioning=MagicMock(conditioning_name="negative"),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
seed=0,
|
|
)
|
|
init_latents = torch.full((1, 16, 8, 8), 2.0)
|
|
noise = torch.full((1, 16, 8, 8), 10.0)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
mock_context.models.load.return_value = MagicMock(
|
|
model=MagicMock(config=MagicMock(in_channels=16, joint_attention_dim=4096))
|
|
)
|
|
|
|
with (
|
|
patch("invokeai.app.invocations.sd3_denoise.TorchDevice.choose_torch_device", return_value=torch.device("cpu")),
|
|
patch("invokeai.app.invocations.sd3_denoise.TorchDevice.choose_torch_dtype", return_value=torch.float32),
|
|
patch.object(invocation, "_prepare_noise_tensor", return_value=noise),
|
|
patch.object(invocation, "_load_text_conditioning", return_value=(torch.zeros(1, 1, 1), torch.zeros(1, 1))),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
timesteps = clip_timestep_schedule_fractional(torch.linspace(1, 0, invocation.steps + 1).tolist(), 0.25, 0.25)
|
|
expected = timesteps[0] * noise + (1.0 - timesteps[0]) * init_latents
|
|
assert torch.equal(result, expected)
|
|
|
|
|
|
def test_cogview4_partial_denoise_short_circuit_uses_first_clipped_sigma():
|
|
invocation = CogView4DenoiseInvocation.model_construct(
|
|
latents=MagicMock(latents_name="latents"),
|
|
width=64,
|
|
height=64,
|
|
steps=4,
|
|
denoising_start=0.25,
|
|
denoising_end=0.25,
|
|
positive_conditioning=MagicMock(conditioning_name="positive"),
|
|
negative_conditioning=MagicMock(conditioning_name="negative"),
|
|
transformer=MagicMock(transformer="transformer"),
|
|
seed=0,
|
|
)
|
|
init_latents = torch.full((1, 16, 8, 8), 2.0)
|
|
noise = torch.full((1, 16, 8, 8), 10.0)
|
|
mock_context = MagicMock()
|
|
mock_context.tensors.load.return_value = init_latents
|
|
transformer_model = MagicMock(config=MagicMock(in_channels=16, patch_size=2))
|
|
mock_context.models.load.return_value = MagicMock(model=transformer_model)
|
|
|
|
with (
|
|
patch("invokeai.app.invocations.cogview4_denoise.CogView4Transformer2DModel", object),
|
|
patch(
|
|
"invokeai.app.invocations.cogview4_denoise.TorchDevice.choose_torch_device",
|
|
return_value=torch.device("cpu"),
|
|
),
|
|
patch.object(invocation, "_prepare_noise_tensor", return_value=noise),
|
|
patch.object(invocation, "_load_text_conditioning", return_value=torch.zeros(1, 1, 1)),
|
|
):
|
|
result = invocation._run_diffusion(mock_context)
|
|
|
|
timesteps = clip_timestep_schedule_fractional(torch.linspace(1, 0, invocation.steps + 1).tolist(), 0.25, 0.25)
|
|
sigmas = invocation._convert_timesteps_to_sigmas(
|
|
image_seq_len=((invocation.height // 8) * (invocation.width // 8)) // (2**2),
|
|
timesteps=torch.tensor(timesteps),
|
|
)
|
|
expected = sigmas[0] * noise + (1.0 - sigmas[0]) * init_latents
|
|
assert torch.allclose(result, expected, atol=2e-3, rtol=0)
|