diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..42f3c8f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,52 @@ +import sys +import types +import torch + +class _DummyProgressBar: + def __init__(self, total): + self.total = total + self.updated = 0 + + def update(self, amount): + self.updated += amount + + +def _ensure_comfy_stubs(): + if "comfy" in sys.modules: + return + + comfy_module = types.ModuleType("comfy") + + model_management_module = types.ModuleType("comfy.model_management") + model_management_module.get_torch_device = lambda: torch.device("cpu") + model_management_module.unet_dtype = lambda: torch.float32 + model_management_module.load_models_gpu = lambda models: None + + utils_module = types.ModuleType("comfy.utils") + utils_module.ProgressBar = _DummyProgressBar + + samplers_module = types.ModuleType("comfy.samplers") + + class _DummyKSampler: + def __init__(self, model, steps, device=None): + self.model = model + self.steps = steps + self.device = device or torch.device("cpu") + self.sigmas = torch.linspace(1.0, 0.0, steps) + + def sample(self, *args, **kwargs): + return torch.zeros((1, 4, 8, 8), device=self.device) + + samplers_module.KSampler = _DummyKSampler + + comfy_module.model_management = model_management_module + comfy_module.utils = utils_module + comfy_module.samplers = samplers_module + + sys.modules["comfy"] = comfy_module + sys.modules["comfy.model_management"] = model_management_module + sys.modules["comfy.utils"] = utils_module + sys.modules["comfy.samplers"] = samplers_module + + +_ensure_comfy_stubs() diff --git a/tests/test_nodes.py b/tests/test_nodes.py new file mode 100644 index 0000000..b6fbc92 --- /dev/null +++ b/tests/test_nodes.py @@ -0,0 +1,102 @@ +import importlib.util +from pathlib import Path + +import torch + + +MODULE_PATH = Path(__file__).resolve().parents[1] / "__init__.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("flowmatching_inverter", MODULE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_latent_gaussian_blur_modifies_values(): + module = load_module() + node = module.LatentGaussianBlur() + latent = {"samples": torch.randn(1, 4, 8, 8)} + + (result,) = node.blur_latent(latent, sigma=1.5, blur_mode="Spatial Only") + + original = latent["samples"] + blurred = result["samples"] + + assert blurred.shape == original.shape + assert not torch.allclose(blurred, original) + + +def test_latent_add_noise_reproducible_with_seed(): + module = load_module() + node = module.LatentAddNoise() + latent = {"samples": torch.randn(1, 4, 8, 8)} + + (first,) = node.add_noise(latent, seed=123, strength=0.8) + (second,) = node.add_noise(latent, seed=123, strength=0.8) + (third,) = node.add_noise(latent, seed=321, strength=0.8) + + assert torch.allclose(first["samples"], second["samples"]) + assert not torch.allclose(first["samples"], third["samples"]) + + +def test_perlin_noise_strength_affects_latent(): + module = load_module() + node = module.LatentPerlinFractalNoise() + latent = {"samples": torch.randn(1, 4, 8, 8)} + + (result,) = node.add_perlin_noise( + latent, + seed=7, + frequency=1.5, + octaves=2, + persistence=0.5, + lacunarity=2.0, + strength=0.5, + channel_mode="shared", + ) + + assert not torch.allclose(result["samples"], latent["samples"]) + + +def test_simplex_noise_zero_strength_is_noop(): + module = load_module() + node = module.LatentSimplexNoise() + latent = {"samples": torch.randn(1, 4, 8, 8)} + + (result,) = node.add_simplex_noise( + latent, + seed=42, + frequency=1.0, + octaves=2, + persistence=0.5, + lacunarity=2.0, + strength=0.0, + channel_mode="shared", + temporal_mode="locked", + ) + + assert torch.allclose(result["samples"], latent["samples"]) + + +def test_conditioning_nodes_modify_embeddings_and_metadata(): + module = load_module() + noise_node = module.ConditioningAddNoise() + blur_node = module.ConditioningGaussianBlur() + + embedding = torch.randn(2, 4, 6) + pooled = torch.randn(4) + conditioning = [[embedding.clone(), {"pooled_output": pooled.clone()}]] + + (noised,) = noise_node.add_noise(conditioning, seed=0, strength=0.5) + (blurred,) = blur_node.blur(noised, sigma=1.0) + + noised_embedding = noised[0][0] + noised_pooled = noised[0][1]["pooled_output"] + + assert not torch.allclose(noised_embedding, embedding) + assert not torch.allclose(noised_pooled, pooled) + + blurred_embedding = blurred[0][0] + assert blurred_embedding.shape == embedding.shape