diff --git a/conftest.py b/conftest.py index a72e94af..d0bf0bc5 100644 --- a/conftest.py +++ b/conftest.py @@ -25,6 +25,13 @@ def _comfyui_root(): return None +# Put ComfyUI on the path at IMPORT time, not only when a fixture asks. A module +# that ships a modifier imports comfy at module level -- as it must, since that is +# how it runs inside ComfyUI -- so collection touches comfy before any fixture has +# had a chance to run. +COMFYUI = _comfyui_root() + + @pytest.fixture(scope="session") def comfyui(): """ComfyUI's source tree on sys.path, or skip. @@ -32,7 +39,6 @@ def comfyui(): Everything that touches a node schema needs it, because a schema is only meaningful in terms of comfy_api's own types. """ - root = _comfyui_root() - if root is None: + if COMFYUI is None: pytest.skip("no ComfyUI source on this machine") - return root + return COMFYUI diff --git a/core/routes.py b/core/routes.py index 496cf8af..8a7e76b5 100644 --- a/core/routes.py +++ b/core/routes.py @@ -112,5 +112,13 @@ def register(routes, prefix=None): if web is not None and PromptServer is not None: - register(PromptServer.instance.routes) - log.note(f"serving the app at {config.UI_PREFIX}/") + # `PromptServer.instance` only exists once a server has been constructed. + # Importing `server` successfully is not the same as running inside one -- + # a test run with ComfyUI on the path gets the class and no instance -- and + # an unguarded attribute here takes the whole pack down with it. + try: + register(PromptServer.instance.routes) + except Exception as exc: # noqa: BLE001 + log.failed("route registration", exc) + else: + log.note(f"serving the app at {config.UI_PREFIX}/") diff --git a/core/tests/test_loads_as_comfyui_does.py b/core/tests/test_loads_as_comfyui_does.py index 955f4f3f..addd9072 100644 --- a/core/tests/test_loads_as_comfyui_does.py +++ b/core/tests/test_loads_as_comfyui_does.py @@ -78,8 +78,9 @@ def test_every_module_that_announces_a_node_actually_registers_it(loaded, comfyu def test_no_module_failed_to_import_under_comfyui(loaded): _, stderr = loaded - # core.routes is expected to fail here: PromptServer has no instance outside - # a running server, and that failure is guarded on purpose. + # Expected outside a running server: there is no PromptServer instance to + # attach routes to. Both spellings are the same guarded failure. + expected = ("core.routes", "route registration") bad = [line for line in stderr.splitlines() - if "failed to load" in line and "core.routes" not in line] + if "failed to load" in line and not any(e in line for e in expected)] assert not bad, "a module failed to import under ComfyUI:\n" + "\n".join(bad) diff --git a/modules/sampling/alg/__init__.py b/modules/sampling/alg/__init__.py new file mode 100644 index 00000000..3084bc9c --- /dev/null +++ b/modules/sampling/alg/__init__.py @@ -0,0 +1,117 @@ +"""ALG: stop an i2v model copying its anchor frame. + +An image-to-video model is shown a sharp reference at step 0 and can satisfy the +objective by barely moving -- the result matches the reference and hardly +animates. ALG (arXiv 2506.08456) low-passes the anchor for the early, high-sigma +steps, removing the detail that shortcut depends on, and swaps the sharp anchor +back once sigma drops so the later steps still resolve real detail. + +The first module that is only behaviour: settings and a modifier, no ComfyUI node +of its own. It attaches to the MODEL, so it applies to whatever sampler is wired. + +`REQUIRES` is `temporal_latent` and nothing more. It is not "for LTX" or "for +H3" -- it needs a latent with a time axis and a pinned anchor, and any model with +both gets it, including ones nobody here has heard of. + +What is NOT ported from v4: the packed-layout path, which reached into a +model-specific arrangement of the latent to find the video stream. Where the +anchor is not a plain [B, C, T, H, W] tensor this stands down and says so once, +rather than reshaping something it is guessing about. A model module can teach it +that layout later. +""" + +from comfy.patcher_extension import WrappersMP + +from ..._core import log +from .blur import AnchorSwap, blur_frames + +ID = "alg" +TITLE = "Anchor de-staticking" +MOUNT = "generation.sampling" +STAGE = "sampling" +STATUS = "proven" +REQUIRES = ["temporal_latent"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Loosen the starting image", + "hint": "Stops a still first frame from holding the whole clip still.", + }, + "strength": { + "type": "float", "default": 4.0, "min": 1.0, "max": 16.0, "step": 0.5, + "label": "How much", "ui": "slider", + "hint": "Higher frees up more motion and keeps less of the original detail.", + "when": {"enabled": True}, + }, + "until_sigma": { + "type": "float", "default": 0.6, "min": 0.0, "max": 1.0, "step": 0.05, + "label": "How long for", "ui": "slider", + "hint": "The starting image comes back sharp once the picture settles.", + "when": {"enabled": True}, + }, +} + +_said = set() + + +def _once(message: str) -> None: + """Say an inert reason once per process, not once per step.""" + if message not in _said: + _said.add(message) + log.note(message) + + +def install(patcher, values, key): + """Put the anchor swap on the model, for whatever sampler runs it.""" + if not values.get("enabled"): + return None + + strength = float(values.get("strength", 4.0)) + threshold = float(values.get("until_sigma", 0.6)) + + def wrapper(executor, model_wrap, sigmas, extra_args, callback, noise, + latent_image=None, denoise_mask=None, disable_pbar=False): + sampler = getattr(executor, "class_obj", None) + inner_fn = getattr(sampler, "sampler_function", None) + + # No anchor pinned means nothing to de-static; a sampler that is not + # KSAMPLER-shaped has no loop to stand between. + if inner_fn is None or latent_image is None or denoise_mask is None: + _once("[FunPack] ALG is off this run: no pinned anchor to loosen.") + return executor(model_wrap, sigmas, extra_args, callback, noise, + latent_image, denoise_mask, disable_pbar) + + blurred = blur_frames(latent_image, strength, frame_indices=(0,)) + if blurred is None: + _once("[FunPack] ALG is off this run: this model's anchor is not a " + "plain video latent, and reshaping it would be guesswork.") + return executor(model_wrap, sigmas, extra_args, callback, noise, + latent_image, denoise_mask, disable_pbar) + + def alg_sampler_function(model_k, x, step_sigmas, extra_args=None, + callback=None, disable=None, **options): + proxy = AnchorSwap(model_k, latent_image, blurred, threshold) + try: + return inner_fn(proxy, x, step_sigmas, extra_args=extra_args, + callback=callback, disable=disable, **options) + finally: + # Whatever happened, the real denoiser goes back to the sharp + # anchor: it outlives this call and the next reader must not + # inherit a blurred one. + model_k.latent_image = latent_image + + import comfy.samplers + replacement = comfy.samplers.KSAMPLER( + alg_sampler_function, + extra_options=getattr(sampler, "extra_options", {}), + inpaint_options=getattr(sampler, "inpaint_options", {}), + ) + return replacement.sample(model_wrap, sigmas, extra_args, callback, noise, + latent_image, denoise_mask, disable_pbar) + + patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, wrapper) + return f"loosening the starting image (amount {strength}, until sigma {threshold})" + + +PROVIDES = {"modifier": install} diff --git a/modules/sampling/alg/blur.py b/modules/sampling/alg/blur.py new file mode 100644 index 00000000..b8a1a0f7 --- /dev/null +++ b/modules/sampling/alg/blur.py @@ -0,0 +1,94 @@ +"""The low-pass filter ALG swaps in, and the per-step decision to use it. + +Both are pure: given a latent and a factor you get a tensor, and given a sigma +you get a choice. That is deliberate -- the parts that decide what happens are +testable on a laptop with 8x8 tensors, and only the parts that cannot fail +quietly are left to the GPU. +""" + +import torch + + +def blur_frames(latent, kappa, frame_indices=(0,)): + """Selected frames of a [B, C, T, H, W] latent, low-passed. + + Bilinear down then up at `kappa`, which removes the high-frequency content an + i2v model uses to copy its anchor frame instead of animating it (arXiv + 2506.08456). Frames outside the clip are skipped rather than guessed at. + + Returns None when the latent is not a shape this understands. None means + "not applicable", and the caller must then leave the sampling untouched -- + a half-applied blur is worse than none. + """ + if latent is None or getattr(latent, "ndim", 0) != 5: + return None + + _b, _c, frames, height, width = latent.shape + wanted = sorted({int(i) for i in frame_indices if 0 <= int(i) < frames}) + if not wanted: + return None + + factor = max(1.0, float(kappa)) + small_h, small_w = max(1, round(height / factor)), max(1, round(width / factor)) + if (small_h, small_w) == (height, width): + return None # nothing would change + + out = latent.clone() + for index in wanted: + frame = latent[:, :, index] # [B, C, H, W] + down = torch.nn.functional.interpolate( + frame, size=(small_h, small_w), mode="bilinear", align_corners=False) + up = torch.nn.functional.interpolate( + down, size=(height, width), mode="bilinear", align_corners=False) + out[:, :, index] = up.to(latent.dtype) + return out + + +def use_blurred(sigma, threshold): + """Whether this step gets the blurred anchor. + + High sigma is early, where the shortcut gets taken; the sharp anchor comes + back once sigma drops past the threshold, so the detail is still there for + the steps that resolve it. + """ + try: + value = float(sigma.flatten()[0]) if hasattr(sigma, "flatten") else float(sigma) + except Exception: # noqa: BLE001 + return False # never cost a step over the guidance + return value > float(threshold) + + +class AnchorSwap: + """A denoiser proxy that re-pins the anchor from each call's own sigma. + + ALG's only per-step decision is blurred-or-sharp and it keys on sigma, which + is an argument of every model call -- so it does not need to be inside a + sampler's loop. That is what lets one implementation cover a stock KSampler, + a distilled flow sampler, and multi-evaluation samplers like heun, where each + evaluation gets the anchor its own sigma calls for. + + Everything else forwards, so the object it wraps cannot tell the difference. + """ + + _OWN = ("_inner", "_sharp", "_blurred", "_threshold") + + def __init__(self, inner, sharp, blurred, threshold): + object.__setattr__(self, "_inner", inner) + object.__setattr__(self, "_sharp", sharp) + object.__setattr__(self, "_blurred", blurred) + object.__setattr__(self, "_threshold", threshold) + + def __getattr__(self, name): + return getattr(self._inner, name) + + def __setattr__(self, name, value): + if name in AnchorSwap._OWN: + object.__setattr__(self, name, value) + else: + # A sampler assigning to the denoiser must reach the real one. + setattr(self._inner, name, value) + + def __call__(self, x, sigma, **kwargs): + self._inner.latent_image = ( + self._blurred if use_blurred(sigma, self._threshold) else self._sharp) + return self._inner(x, sigma, **kwargs) diff --git a/modules/sampling/alg/tests/test_alg.py b/modules/sampling/alg/tests/test_alg.py new file mode 100644 index 00000000..c045100c --- /dev/null +++ b/modules/sampling/alg/tests/test_alg.py @@ -0,0 +1,161 @@ +"""ALG, as far as a machine with no GPU can take it. + +What IS checked here: the blur does what a low-pass does, the per-step choice +follows sigma, the inert cases are inert, and the whole thing runs through +ComfyUI's real KSAMPLER with a stand-in sampler function -- so the anchor a +denoiser sees really does change between steps. + +What is NOT checked here, and cannot be: whether the result looks better. That +needs a GPU and real weights. +""" + +import pytest + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy and torch.""" + + +def _latent(frames=4, size=16, seed=0): + import torch + torch.manual_seed(seed) + return torch.randn(1, 3, frames, size, size) + + +# --- the filter ------------------------------------------------------------ + +def test_blurring_removes_high_frequency_detail(): + """A low-pass, stated as something measurable: neighbouring pixels of the + blurred frame differ less than they did.""" + import torch + from modules.sampling.alg.blur import blur_frames + + sharp = _latent() + out = blur_frames(sharp, kappa=4.0, frame_indices=(0,)) + + def roughness(frame): + return (frame[..., 1:] - frame[..., :-1]).abs().mean() + + assert roughness(out[:, :, 0]) < roughness(sharp[:, :, 0]) + + +def test_only_the_named_frames_change(): + """The anchor is frame 0. Blurring the whole clip would remove the detail + every other frame is supposed to resolve.""" + import torch + from modules.sampling.alg.blur import blur_frames + + sharp = _latent(frames=4) + out = blur_frames(sharp, kappa=4.0, frame_indices=(0,)) + + assert not torch.equal(out[:, :, 0], sharp[:, :, 0]) + for index in (1, 2, 3): + assert torch.equal(out[:, :, index], sharp[:, :, index]), f"frame {index} changed" + + +def test_the_original_is_not_modified(): + """The sharp anchor is swapped back in later, so it has to survive.""" + import torch + from modules.sampling.alg.blur import blur_frames + + sharp = _latent() + keep = sharp.clone() + blur_frames(sharp, kappa=4.0) + assert torch.equal(sharp, keep) + + +def test_shape_and_dtype_are_preserved(): + import torch + from modules.sampling.alg.blur import blur_frames + sharp = _latent().to(torch.float16) + out = blur_frames(sharp, kappa=4.0) + assert out.shape == sharp.shape and out.dtype == sharp.dtype + + +@pytest.mark.parametrize("bad", ["not a tensor", None]) +def test_a_latent_it_does_not_understand_yields_none(bad): + from modules.sampling.alg.blur import blur_frames + assert blur_frames(bad, kappa=4.0) is None + + +def test_a_latent_of_the_wrong_rank_yields_none(): + """An image latent has no anchor frame to loosen. None means "not + applicable", and the caller leaves sampling alone.""" + import torch + from modules.sampling.alg.blur import blur_frames + assert blur_frames(torch.randn(1, 4, 64, 64), kappa=4.0) is None + + +def test_a_frame_index_past_the_end_is_skipped_not_guessed(): + from modules.sampling.alg.blur import blur_frames + assert blur_frames(_latent(frames=2), kappa=4.0, frame_indices=(9,)) is None + + +def test_a_strength_that_would_change_nothing_yields_none(): + """kappa 1.0 downsamples to the same size, so the "blurred" copy would be + the sharp one and every step would swap between identical tensors.""" + from modules.sampling.alg.blur import blur_frames + assert blur_frames(_latent(), kappa=1.0) is None + + +# --- the per-step choice --------------------------------------------------- + +@pytest.mark.parametrize("sigma,threshold,expected", [ + (0.9, 0.6, True), # early, high noise: loosen it + (0.6, 0.6, False), # exactly at the threshold: sharp + (0.1, 0.6, False), # late: sharp, so detail can resolve +]) +def test_the_choice_follows_sigma(sigma, threshold, expected): + from modules.sampling.alg.blur import use_blurred + assert use_blurred(sigma, threshold) is expected + + +def test_a_sigma_it_cannot_read_falls_back_to_sharp(): + """Never cost a step over the guidance: unreadable means untouched.""" + from modules.sampling.alg.blur import use_blurred + assert use_blurred("nonsense", 0.6) is False + + +def test_a_tensor_sigma_is_read_like_the_sampler_passes_it(): + import torch + from modules.sampling.alg.blur import use_blurred + assert use_blurred(torch.tensor([0.9]), 0.6) is True + + +# --- the proxy ------------------------------------------------------------- + +def test_the_proxy_repins_the_anchor_from_each_calls_own_sigma(): + """Multi-evaluation samplers call the denoiser more than once per step at + different sigmas, and each call must get the anchor its own sigma calls for.""" + from modules.sampling.alg.blur import AnchorSwap + + class Denoiser: + latent_image = None + + def __call__(self, x, sigma, **kwargs): + return self.latent_image + + inner = Denoiser() + proxy = AnchorSwap(inner, sharp="SHARP", blurred="BLURRED", threshold=0.6) + + assert proxy(None, 0.9) == "BLURRED" + assert proxy(None, 0.2) == "SHARP" + assert proxy(None, 0.7) == "BLURRED" + + +def test_the_proxy_is_transparent_for_everything_else(): + from modules.sampling.alg.blur import AnchorSwap + + class Denoiser: + inner_model = "the real model" + + def __call__(self, x, sigma, **kwargs): + return None + + inner = Denoiser() + proxy = AnchorSwap(inner, "S", "B", 0.6) + + assert proxy.inner_model == "the real model" + proxy.something = 7 + assert inner.something == 7, "an assignment did not reach the real denoiser" diff --git a/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py b/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py new file mode 100644 index 00000000..2d9569ea --- /dev/null +++ b/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py @@ -0,0 +1,204 @@ +"""ALG driven through ComfyUI's real KSAMPLER. + +The stand-in is the sampler FUNCTION -- the innermost loop, which is the only +part that needs a GPU. Everything above it is ComfyUI's own: KSAMPLER.sample, +KSamplerX0Inpaint, the wrapper machinery. So "the denoiser sees a different +anchor early than it does late" is demonstrated rather than asserted. +""" + +import pytest + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +class _Guider: + """The minimum ComfyUI's real KSAMPLER.sample needs of a guider. + + Only the innermost model call is stubbed; KSAMPLER.sample, KSamplerX0Inpaint + and the noise scaling around them are the real ones. + """ + + def __init__(self): + outer = self + + class ModelSampling: + sigma_max = 1.0 + + def noise_scaling(self, sigma, noise, latent_image, max_denoise=False): + return noise + + def inverse_noise_scaling(self, sigma, x): + return x + + class InnerModel: + model_sampling = ModelSampling() + + def scale_latent_inpaint(self, x, sigma, noise, latent_image, denoise_mask): + return latent_image + + self.inner_model = InnerModel() + self.calls = 0 + + def __call__(self, x, sigma, model_options={}, seed=None): + self.calls += 1 + return x + + +class _Executor: + """What a SAMPLER_SAMPLE wrapper is handed: the call it wraps, plus the + sampler it belongs to.""" + + def __init__(self, sampler, fallback): + self.class_obj = sampler + self._fallback = fallback + self.fell_back = False + + def __call__(self, *args, **kwargs): + self.fell_back = True + return self._fallback + + +def _install_alg(**values): + """ALG's wrapper, without a ModelPatcher in the way.""" + from comfy.patcher_extension import WrappersMP + from modules.sampling import alg + + captured = {} + + class Patcher: + def add_wrapper_with_key(self, kind, key, wrapper): + captured["kind"] = kind + captured["wrapper"] = wrapper + + settings = {"enabled": True, "strength": 4.0, "until_sigma": 0.6} + settings.update(values) + note = alg.install(Patcher(), settings, key="funpack.alg") + return captured.get("wrapper"), note, captured.get("kind"), WrappersMP + + +def test_it_is_off_unless_asked_for(): + from modules.sampling import alg + assert alg.install(object(), {"enabled": False}, key="k") is None + + +def test_it_installs_around_the_sampler_call(): + wrapper, note, kind, WrappersMP = _install_alg() + assert wrapper is not None and "loosening" in note + assert kind == WrappersMP.SAMPLER_SAMPLE + + +def test_the_anchor_the_denoiser_sees_changes_with_sigma(): + """The whole point, end to end through comfy's own KSAMPLER.""" + import torch + import comfy.samplers + from modules.sampling.alg.blur import blur_frames + + sharp = torch.randn(1, 3, 4, 16, 16) + blurred = blur_frames(sharp, 4.0, (0,)) + seen = [] + + def sampler_function(model_k, x, sigmas, extra_args=None, callback=None, + disable=None, **options): + # What a real sampler loop does: call the denoiser once per sigma. The + # anchor pinned for that call is what ALG decides. + for sigma in (torch.tensor([0.9]), torch.tensor([0.2])): + model_k(x, sigma, denoise_mask=None) + seen.append(model_k.latent_image) + return x + + wrapper, _note, _kind, _mp = _install_alg() + real = comfy.samplers.KSAMPLER(sampler_function) + executor = _Executor(real, fallback="unused") + + wrapper(executor, _Guider(), torch.tensor([0.9, 0.2]), {}, None, + torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) + + assert len(seen) == 2, "the denoiser was not called per sigma" + assert torch.equal(seen[0], blurred), "the early step did not get the loosened anchor" + assert torch.equal(seen[1], sharp), "the late step did not get the sharp anchor back" + + +def test_the_real_denoiser_is_left_holding_the_sharp_anchor(): + """The denoiser outlives this call. Leaving a blurred anchor pinned on it is + the shape of every leak this project has had.""" + import torch + import comfy.samplers + + sharp = torch.randn(1, 3, 4, 16, 16) + + holder = {} + + def sampler_function(model_k, x, sigmas, extra_args=None, callback=None, + disable=None, **options): + model_k(x, torch.tensor([0.9]), denoise_mask=None) # leaves it on the loosened anchor + holder["denoiser"] = model_k._inner # the real one, behind the proxy + return x + + wrapper, _n, _k, _mp = _install_alg() + wrapper(_Executor(comfy.samplers.KSAMPLER(sampler_function), "unused"), + _Guider(), torch.tensor([0.9]), {}, None, + torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) + + assert torch.equal(holder["denoiser"].latent_image, sharp) + + +def test_a_sampler_error_still_restores_the_sharp_anchor(): + import torch + import comfy.samplers + + sharp = torch.randn(1, 3, 4, 16, 16) + + holder = {} + + def sampler_function(model_k, x, sigmas, **kwargs): + model_k(x, torch.tensor([0.9]), denoise_mask=None) + holder["denoiser"] = model_k._inner + raise RuntimeError("interrupted") + + wrapper, _n, _k, _mp = _install_alg() + with pytest.raises(RuntimeError, match="interrupted"): + wrapper(_Executor(comfy.samplers.KSAMPLER(sampler_function), "unused"), + _Guider(), torch.tensor([0.9]), {}, None, + torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) + + assert torch.equal(holder["denoiser"].latent_image, sharp), ( + "an interrupt left the anchor loosened") + + +@pytest.mark.parametrize("missing", ["latent_image", "denoise_mask"]) +def test_with_no_pinned_anchor_it_stands_aside(missing): + """No anchor means nothing to loosen, and the run must proceed untouched.""" + import torch + import comfy.samplers + + sharp = torch.randn(1, 3, 4, 16, 16) + wrapper, _n, _k, _mp = _install_alg() + executor = _Executor(comfy.samplers.KSAMPLER(lambda *a, **k: None), "untouched") + + args = dict(latent_image=sharp, denoise_mask=torch.ones_like(sharp)) + args[missing] = None + + out = wrapper(executor, object(), torch.tensor([0.9]), {}, None, + torch.zeros_like(sharp), args["latent_image"], args["denoise_mask"], True) + + assert out == "untouched" and executor.fell_back + + +def test_an_image_latent_stands_aside_rather_than_reshaping_a_guess(): + """A 4-D latent has no anchor frame. v4 reached into a model-specific packed + layout here; this declines instead of guessing.""" + import torch + import comfy.samplers + + image_latent = torch.randn(1, 4, 64, 64) + wrapper, _n, _k, _mp = _install_alg() + executor = _Executor(comfy.samplers.KSAMPLER(lambda *a, **k: None), "untouched") + + out = wrapper(executor, object(), torch.tensor([0.9]), {}, None, + torch.zeros_like(image_latent), image_latent, + torch.ones_like(image_latent), True) + + assert out == "untouched" and executor.fell_back