alg: the first real modifier, and the first module that is only behaviour
An i2v model shown a sharp reference can satisfy the objective by barely moving. ALG low-passes the anchor while sigma is high, removing the detail that shortcut needs, and hands the sharp one back once the picture settles. It attaches to the model, so it applies to whatever sampler is wired, and it requires temporal_latent and nothing else -- not "for LTX", not "for H3". Any model with a time axis and a pinned anchor gets it, including ones nobody here has heard of. A module with no ComfyUI node of its own: settings and behaviour, which is what most modules should be. The packed-layout path is NOT ported. v4 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 aside and says so once, rather than reshaping something it is guessing about. Tested as far as a machine with no GPU can: the filter really low-passes, only the anchor frame changes, the choice follows sigma, and the whole thing runs through ComfyUI's real KSAMPLER with only the innermost loop stubbed -- so the anchor a denoiser sees differs between an early step and a late one. Whether it looks better needs the rental. Also guarded route registration. PromptServer.instance only exists once a server has been built, and importing `server` successfully is not the same as running inside one -- with ComfyUI on the path the unguarded attribute took the whole pack down. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Opus 5
Claude Opus 5
parent
b3afba43e7
commit
749b1b5db9
+9
-3
@@ -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
|
||||
|
||||
+10
-2
@@ -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}/")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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}
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user