938 lines
33 KiB
Python
938 lines
33 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from comfyui_spectrum_wan.config import SpectrumWanConfig
|
|
from comfyui_spectrum_wan.handlers import resolve_handler
|
|
from comfyui_spectrum_wan.runtime import SpectrumWanRuntime
|
|
import comfyui_spectrum_wan.wan as wan
|
|
from comfyui_spectrum_wan.wan import _RUNTIME_KEY, _run_spectrum_forward, WanSpectrumPatcher
|
|
|
|
|
|
class DummyInner:
|
|
def __init__(self) -> None:
|
|
self.blocks = []
|
|
self.ref_conv = None
|
|
self.original_calls = 0
|
|
|
|
def forward_orig(self, *args, **kwargs):
|
|
self.original_calls += 1
|
|
return "orig"
|
|
|
|
def patch_embedding(self, x):
|
|
return x
|
|
|
|
def condition_embedder(self, t, context, clip_fea):
|
|
batch = t.shape[0]
|
|
return (
|
|
torch.zeros((batch, 1), dtype=t.dtype),
|
|
torch.zeros((batch, 6), dtype=t.dtype),
|
|
torch.zeros((batch, 1, 1), dtype=t.dtype),
|
|
None,
|
|
)
|
|
|
|
def head(self, x, temb):
|
|
return x
|
|
|
|
def unpatchify(self, x, grid_sizes):
|
|
return x
|
|
|
|
|
|
class NewApiInner(DummyInner):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.condition_embedder = None
|
|
self.freq_dim = 8
|
|
self.img_emb = None
|
|
|
|
def time_embedding(self, x):
|
|
return torch.zeros((*x.shape[:-1], 1), dtype=x.dtype, device=x.device)
|
|
|
|
def time_projection(self, e):
|
|
return torch.zeros((*e.shape[:-1], 6), dtype=e.dtype, device=e.device)
|
|
|
|
def text_embedding(self, context):
|
|
return context
|
|
|
|
|
|
class DummyOuter:
|
|
def __init__(self, inner: DummyInner) -> None:
|
|
self.diffusion_model = inner
|
|
|
|
def apply_model(
|
|
self,
|
|
x,
|
|
timestep,
|
|
context,
|
|
clip_fea=None,
|
|
time_dim_concat=None,
|
|
transformer_options=None,
|
|
**kwargs,
|
|
):
|
|
if hasattr(self.diffusion_model, "_forward") and callable(getattr(self.diffusion_model, "_forward")):
|
|
return self.diffusion_model._forward(
|
|
x,
|
|
timestep,
|
|
context,
|
|
clip_fea=clip_fea,
|
|
time_dim_concat=time_dim_concat,
|
|
transformer_options=transformer_options,
|
|
**kwargs,
|
|
)
|
|
if callable(self.diffusion_model):
|
|
return self.diffusion_model(
|
|
x,
|
|
timestep,
|
|
context,
|
|
clip_fea=clip_fea,
|
|
transformer_options=transformer_options,
|
|
**kwargs,
|
|
)
|
|
return self.diffusion_model.forward_orig(
|
|
x,
|
|
timestep,
|
|
context,
|
|
clip_fea=clip_fea,
|
|
transformer_options=transformer_options,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
class NonWanInner:
|
|
pass
|
|
|
|
|
|
class LegacyOnlyInner(DummyInner):
|
|
pass
|
|
|
|
|
|
class ForwardOrigOnlyInner:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def forward_orig(self, *args, **kwargs):
|
|
self.calls += 1
|
|
return "orig-only"
|
|
|
|
|
|
class ForwardOrigProxyWrapper:
|
|
def __init__(self, inner) -> None:
|
|
self.model = inner
|
|
|
|
def forward_orig(self, *args, **kwargs):
|
|
return self.model.forward_orig(*args, **kwargs)
|
|
|
|
|
|
class DummyInnerWrapper:
|
|
def __init__(self, inner) -> None:
|
|
self.model = inner
|
|
self.forward_calls = 0
|
|
self.seen_transformer_options = None
|
|
|
|
def __call__(self, x, timestep, context, clip_fea=None, transformer_options=None, **kwargs):
|
|
self.forward_calls += 1
|
|
self.seen_transformer_options = transformer_options
|
|
return self.model.forward_orig(
|
|
x,
|
|
timestep,
|
|
context,
|
|
clip_fea=clip_fea,
|
|
transformer_options=transformer_options,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
class DummyModel:
|
|
model_name = "wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors"
|
|
|
|
def __init__(self) -> None:
|
|
self.model = DummyOuter(DummyInner())
|
|
self.model_options = None
|
|
|
|
def clone(self):
|
|
return self
|
|
|
|
|
|
class DummyModelForwardOrigOnly(DummyModel):
|
|
def __init__(self) -> None:
|
|
self.model = DummyOuter(ForwardOrigOnlyInner())
|
|
self.model_options = None
|
|
|
|
|
|
def _cfg() -> SpectrumWanConfig:
|
|
return SpectrumWanConfig(
|
|
backend="auto",
|
|
blend_weight=1.0,
|
|
degree=4,
|
|
ridge_lambda=0.1,
|
|
window_size=2.0,
|
|
flex_window=0.75,
|
|
warmup_steps=2,
|
|
history_size=8,
|
|
debug=True,
|
|
).validated()
|
|
|
|
|
|
def test_direct_forward_orig_requires_runtime_in_current_transformer_options() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
out = inner.forward_orig(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={
|
|
_RUNTIME_KEY: runtime,
|
|
"sample_sigmas": sample_sigmas,
|
|
"cond_or_uncond": [0, 1],
|
|
},
|
|
)
|
|
|
|
assert torch.is_tensor(out)
|
|
assert inner.original_calls == 0
|
|
assert inner._spectrum_wan_runtime.last_info["last_sigma"] == 1.0
|
|
|
|
|
|
def test_patched_wan_does_not_leak_runtime_into_later_bypassed_calls() -> None:
|
|
model = DummyModel()
|
|
patched = WanSpectrumPatcher.patch(model, _cfg())
|
|
inner = patched.model.diffusion_model
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
|
|
first_out = inner.forward_orig(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={
|
|
_RUNTIME_KEY: runtime,
|
|
"sample_sigmas": sample_sigmas,
|
|
"cond_or_uncond": [0, 1],
|
|
},
|
|
)
|
|
|
|
bypass_out = model.model.diffusion_model.forward_orig(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[1]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]},
|
|
)
|
|
|
|
assert torch.is_tensor(first_out)
|
|
assert bypass_out == "orig"
|
|
assert inner.original_calls == 1
|
|
|
|
|
|
class DummyInnerWithForward(DummyInner):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.live_calls = 0
|
|
self.seen_time_dim_concat = None
|
|
self.seen_transformer_options = None
|
|
|
|
def _forward(
|
|
self,
|
|
x,
|
|
timestep,
|
|
context,
|
|
clip_fea=None,
|
|
time_dim_concat=None,
|
|
transformer_options=None,
|
|
**kwargs,
|
|
):
|
|
self.live_calls += 1
|
|
self.seen_time_dim_concat = time_dim_concat
|
|
self.seen_transformer_options = transformer_options
|
|
return "live"
|
|
|
|
|
|
class DummyModelWithForward(DummyModel):
|
|
def __init__(self) -> None:
|
|
self.model = DummyOuter(DummyInnerWithForward())
|
|
self.model_options = None
|
|
|
|
|
|
def test_patcher_wraps__forward_as_runtime_passthrough_when_available() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
transformer_options = {
|
|
_RUNTIME_KEY: patched.model_options["transformer_options"][_RUNTIME_KEY],
|
|
"sample_sigmas": sample_sigmas,
|
|
"cond_or_uncond": [0, 1],
|
|
}
|
|
time_dim_concat = torch.ones((1, 1, 1), dtype=torch.float32)
|
|
out = inner._forward(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
time_dim_concat=time_dim_concat,
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
assert out == "live"
|
|
assert inner.live_calls == 1
|
|
assert inner.original_calls == 0
|
|
assert inner.seen_time_dim_concat is time_dim_concat
|
|
assert inner.seen_transformer_options is transformer_options
|
|
assert inner.seen_transformer_options[_RUNTIME_KEY] is patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
assert inner._spectrum_wan__forward_wrapped is True
|
|
assert inner._spectrum_wan_wrapped_attr == "forward_orig"
|
|
assert inner._spectrum_wan_runtime.last_info["hook_target"] == "model.diffusion_model.forward_orig"
|
|
|
|
|
|
def test_patcher_rebinds_runtime_when_live_inner_changes_before_apply_model() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
live_inner = DummyInner()
|
|
patched.model.diffusion_model = live_inner
|
|
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
out = patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]},
|
|
)
|
|
|
|
assert torch.is_tensor(out)
|
|
assert live_inner.original_calls == 0
|
|
assert live_inner._spectrum_wan_runtime is runtime
|
|
assert patched.model._spectrum_wan_bound_inner_id == id(live_inner)
|
|
assert runtime.last_info["live_inner_id"] == id(live_inner)
|
|
assert runtime.last_info["live_inner_type"] == "DummyInner"
|
|
|
|
|
|
def test_patcher_rebinds__forward_passthrough_when_live_inner_changes_before_apply_model() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
live_inner = DummyInnerWithForward()
|
|
patched.model.diffusion_model = live_inner
|
|
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
transformer_options = {"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]}
|
|
out = patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
assert out == "live"
|
|
assert live_inner.live_calls == 1
|
|
assert live_inner.seen_transformer_options is transformer_options
|
|
assert live_inner.seen_transformer_options[_RUNTIME_KEY] is runtime
|
|
assert live_inner._spectrum_wan_runtime is runtime
|
|
assert live_inner._spectrum_wan__forward_wrapped is True
|
|
|
|
|
|
def test_patcher_rebinds_runtime_when_live_inner_is_wrapped_one_level_deeper() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
live_inner = DummyInner()
|
|
wrapped_live_inner = DummyInnerWrapper(live_inner)
|
|
patched.model.diffusion_model = wrapped_live_inner
|
|
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
out = patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={"sample_sigmas": sample_sigmas, "cond_or_uncond": [0, 1]},
|
|
)
|
|
|
|
assert torch.is_tensor(out)
|
|
assert wrapped_live_inner.forward_calls == 1
|
|
assert wrapped_live_inner.seen_transformer_options[_RUNTIME_KEY] is runtime
|
|
assert live_inner._spectrum_wan_runtime is runtime
|
|
assert patched.model._spectrum_wan_bound_inner_id == id(live_inner)
|
|
assert runtime.last_info["live_inner_type"] == "DummyInner"
|
|
assert runtime.last_info["hook_target"] == "model.diffusion_model.model.forward_orig"
|
|
|
|
|
|
def test_patcher_binds_live_inner_when_forward_orig_exists_even_if_runtime_attrs_are_missing() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelForwardOrigOnly(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
inner = patched.model.diffusion_model
|
|
assert inner._spectrum_wan_runtime is runtime
|
|
assert runtime.last_info["patched"] is True
|
|
assert runtime.last_info["hook_target"] == "model.diffusion_model.forward_orig"
|
|
assert "blocks" in runtime.last_info["runtime_missing_attrs"]
|
|
|
|
|
|
def test_runtime_missing_attrs_accepts_legacy_conditioning_api() -> None:
|
|
inner = LegacyOnlyInner()
|
|
|
|
missing = wan._spectrum_runtime_missing_attrs(inner)
|
|
|
|
assert "condition_embedder|time_embedding" not in missing
|
|
assert "condition_embedder|time_projection" not in missing
|
|
assert "condition_embedder|text_embedding" not in missing
|
|
|
|
|
|
def test_runtime_missing_attrs_accepts_new_conditioning_api() -> None:
|
|
inner = NewApiInner()
|
|
|
|
missing = wan._spectrum_runtime_missing_attrs(inner)
|
|
|
|
assert "condition_embedder|time_embedding" not in missing
|
|
assert "condition_embedder|time_projection" not in missing
|
|
assert "condition_embedder|text_embedding" not in missing
|
|
|
|
|
|
def test_locate_prefers_full_new_api_wan_descendant_over_forward_orig_only_proxy() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
live_inner = NewApiInner()
|
|
patched.model.diffusion_model = ForwardOrigProxyWrapper(live_inner)
|
|
original = wan._upstream_sinusoidal_embedding_1d
|
|
wan._upstream_sinusoidal_embedding_1d = (
|
|
lambda dim, t: torch.zeros((t.shape[0], dim), dtype=torch.float32, device=t.device)
|
|
)
|
|
try:
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={"sample_sigmas": torch.tensor([1.0, 0.5, 0.0]), "cond_or_uncond": [0, 1]},
|
|
)
|
|
finally:
|
|
wan._upstream_sinusoidal_embedding_1d = original
|
|
|
|
assert live_inner.original_calls == 0
|
|
assert patched.model._spectrum_wan_bound_inner_id == id(live_inner)
|
|
assert runtime.last_info["live_inner_type"] == "NewApiInner"
|
|
|
|
|
|
def test_run_spectrum_forward_new_api_preserves_multitoken_timestep_shape() -> None:
|
|
class CaptureBlock:
|
|
def __init__(self) -> None:
|
|
self.seen_e_shape = None
|
|
|
|
def __call__(
|
|
self,
|
|
x,
|
|
e,
|
|
freqs=None,
|
|
context=None,
|
|
context_img_len=None,
|
|
transformer_options=None,
|
|
):
|
|
self.seen_e_shape = tuple(e.shape)
|
|
return x
|
|
|
|
class RuntimeStub:
|
|
def __init__(self) -> None:
|
|
self.last_info = {}
|
|
|
|
def begin_step(self, transformer_options, timesteps):
|
|
return {"step_idx": 0, "actual_forward": True, "global_step": 0}
|
|
|
|
def can_forecast(self, transformer_options):
|
|
raise AssertionError("can_forecast should not run when actual_forward is True")
|
|
|
|
def _debug_log(self, message: str) -> None:
|
|
pass
|
|
|
|
def observe_feature(self, transformer_options, step_idx, feature, global_step=None):
|
|
pass
|
|
|
|
inner = NewApiInner()
|
|
block = CaptureBlock()
|
|
inner.blocks = [block]
|
|
runtime = RuntimeStub()
|
|
|
|
original = wan._upstream_sinusoidal_embedding_1d
|
|
wan._upstream_sinusoidal_embedding_1d = (
|
|
lambda dim, t: torch.zeros((t.shape[0], dim), dtype=torch.float32, device=t.device)
|
|
)
|
|
try:
|
|
_run_spectrum_forward(
|
|
inner,
|
|
runtime,
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([[1.0, 0.5]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={},
|
|
)
|
|
finally:
|
|
wan._upstream_sinusoidal_embedding_1d = original
|
|
|
|
assert block.seen_e_shape == (1, 2, 6, 1)
|
|
|
|
|
|
def test_locate_prefers_full_wan_descendant_over_forward_orig_only_proxy() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
live_inner = DummyInner()
|
|
patched.model.diffusion_model = ForwardOrigProxyWrapper(live_inner)
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={"sample_sigmas": torch.tensor([1.0, 0.5, 0.0]), "cond_or_uncond": [0, 1]},
|
|
)
|
|
|
|
assert patched.model._spectrum_wan_bound_inner_id == id(live_inner)
|
|
assert runtime.last_info["live_inner_type"] == "DummyInner"
|
|
|
|
|
|
def test_forward_orig_wrapper_falls_back_and_logs_when_runtime_attrs_are_missing(capsys) -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelForwardOrigOnly(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
out = inner.forward_orig(
|
|
torch.tensor([1.0]),
|
|
torch.tensor([1.0]),
|
|
torch.tensor([1.0]),
|
|
transformer_options={_RUNTIME_KEY: runtime},
|
|
)
|
|
captured = capsys.readouterr()
|
|
|
|
assert out == "orig-only"
|
|
assert "runtime path unavailable" in captured.err
|
|
|
|
|
|
def test_patcher_clears_live_binding_state_when_current_inner_is_not_wan_like() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
patched.model.diffusion_model = NonWanInner()
|
|
|
|
try:
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={},
|
|
)
|
|
except AttributeError:
|
|
pass
|
|
|
|
assert patched.model._spectrum_wan_bound_inner_id is None
|
|
assert runtime.last_info["patched"] is False
|
|
assert runtime.last_info["hook_target"] == "model.diffusion_model"
|
|
assert "live_inner_id" not in runtime.last_info
|
|
assert "live_inner_type" not in runtime.last_info
|
|
|
|
|
|
def test_successful_rebind_clears_stale_live_inner_root_type() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModel(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
patched.model.diffusion_model = NonWanInner()
|
|
try:
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={},
|
|
)
|
|
except AttributeError:
|
|
pass
|
|
|
|
assert runtime.last_info["live_inner_root_type"] == "NonWanInner"
|
|
|
|
live_inner = DummyInner()
|
|
patched.model.diffusion_model = DummyInnerWrapper(live_inner)
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={"sample_sigmas": torch.tensor([1.0, 0.5, 0.0]), "cond_or_uncond": [0, 1]},
|
|
)
|
|
|
|
assert "live_inner_root_type" not in runtime.last_info
|
|
|
|
|
|
def test_patcher__forward_passthrough_preserves_current_transformer_runtime() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
|
|
stale_runtime = SpectrumWanRuntime(_cfg(), resolve_handler("auto", DummyModel()))
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
transformer_options = {
|
|
_RUNTIME_KEY: stale_runtime,
|
|
"sample_sigmas": sample_sigmas,
|
|
"cond_or_uncond": [0, 1],
|
|
}
|
|
|
|
inner._forward(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
assert inner.seen_transformer_options is transformer_options
|
|
assert inner.seen_transformer_options[_RUNTIME_KEY] is stale_runtime
|
|
|
|
|
|
def test_patcher__forward_passthrough_injects_bound_runtime_when_missing() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
current_runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
|
|
transformer_options = {"cond_or_uncond": [0, 1]}
|
|
|
|
inner._forward(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
assert inner.seen_transformer_options is transformer_options
|
|
assert inner.seen_transformer_options[_RUNTIME_KEY] is current_runtime
|
|
|
|
|
|
def test_runtime_accumulates_history_when_sample_sigmas_are_missing() -> None:
|
|
cfg = SpectrumWanConfig(
|
|
backend="wan21",
|
|
warmup_steps=0,
|
|
window_size=2.0,
|
|
flex_window=0.75,
|
|
).validated()
|
|
runtime = SpectrumWanRuntime(cfg, resolve_handler("wan21", DummyModel()))
|
|
transformer_options = {"cond_or_uncond": [0, 1]}
|
|
|
|
first = runtime.begin_step(transformer_options, torch.tensor([1.0], dtype=torch.float32))
|
|
assert first["step_idx"] == 0
|
|
assert first["actual_forward"] is True
|
|
runtime.observe_feature(transformer_options, first["step_idx"], torch.ones((1, 2), dtype=torch.float32))
|
|
runtime.end_step(transformer_options, first["step_idx"])
|
|
|
|
second = runtime.begin_step(transformer_options, torch.tensor([0.8], dtype=torch.float32))
|
|
assert second["step_idx"] == 1
|
|
assert second["actual_forward"] is True
|
|
runtime.observe_feature(transformer_options, second["step_idx"], torch.full((1, 2), 2.0, dtype=torch.float32))
|
|
runtime.end_step(transformer_options, second["step_idx"])
|
|
|
|
third = runtime.begin_step(transformer_options, torch.tensor([0.6], dtype=torch.float32))
|
|
assert third["step_idx"] == 2
|
|
assert third["actual_forward"] is False
|
|
|
|
|
|
def test_runtime_missing_sample_sigmas_does_not_reuse_previous_schedule_length() -> None:
|
|
cfg = SpectrumWanConfig(
|
|
backend="wan21",
|
|
warmup_steps=0,
|
|
window_size=2.0,
|
|
flex_window=0.75,
|
|
).validated()
|
|
runtime = SpectrumWanRuntime(cfg, resolve_handler("wan21", DummyModel()))
|
|
|
|
scheduled_options = {
|
|
"sample_sigmas": torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32),
|
|
"cond_or_uncond": [0, 1],
|
|
}
|
|
first_scheduled = runtime.begin_step(scheduled_options, torch.tensor([1.0], dtype=torch.float32))
|
|
runtime.observe_feature(
|
|
scheduled_options,
|
|
first_scheduled["step_idx"],
|
|
torch.ones((1, 2), dtype=torch.float32),
|
|
)
|
|
runtime.end_step(scheduled_options, first_scheduled["step_idx"])
|
|
second_scheduled = runtime.begin_step(scheduled_options, torch.tensor([0.5], dtype=torch.float32))
|
|
runtime.observe_feature(
|
|
scheduled_options,
|
|
second_scheduled["step_idx"],
|
|
torch.full((1, 2), 2.0, dtype=torch.float32),
|
|
)
|
|
runtime.end_step(scheduled_options, second_scheduled["step_idx"])
|
|
assert runtime.num_steps() == 2
|
|
|
|
unscheduled_options = scheduled_options
|
|
unscheduled_options.pop("sample_sigmas")
|
|
first = runtime.begin_step(unscheduled_options, torch.tensor([1.0], dtype=torch.float32))
|
|
assert first["step_idx"] == 0
|
|
assert "spectrum_wan_active_num_steps" not in unscheduled_options
|
|
assert runtime.last_info["num_steps"] == 0
|
|
runtime.observe_feature(unscheduled_options, first["step_idx"], torch.ones((1, 2), dtype=torch.float32))
|
|
runtime.end_step(unscheduled_options, first["step_idx"])
|
|
|
|
second = runtime.begin_step(unscheduled_options, torch.tensor([0.8], dtype=torch.float32))
|
|
assert second["step_idx"] == 1
|
|
runtime.observe_feature(unscheduled_options, second["step_idx"], torch.full((1, 2), 2.0, dtype=torch.float32))
|
|
runtime.end_step(unscheduled_options, second["step_idx"])
|
|
|
|
third = runtime.begin_step(unscheduled_options, torch.tensor([0.6], dtype=torch.float32))
|
|
assert third["step_idx"] == 2
|
|
assert third["actual_forward"] is False
|
|
|
|
|
|
def test_patcher_apply_model_overwrites_stale_runtime_with_current_outer_runtime() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
current_runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
inner = patched.model.diffusion_model
|
|
|
|
stale_runtime = SpectrumWanRuntime(_cfg(), resolve_handler("auto", DummyModel()))
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
transformer_options = {
|
|
_RUNTIME_KEY: stale_runtime,
|
|
"sample_sigmas": sample_sigmas,
|
|
"cond_or_uncond": [0, 1],
|
|
}
|
|
|
|
out = patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
assert out == "live"
|
|
assert inner.seen_transformer_options is transformer_options
|
|
assert inner.seen_transformer_options[_RUNTIME_KEY] is current_runtime
|
|
assert inner.seen_transformer_options[_RUNTIME_KEY] is not stale_runtime
|
|
|
|
|
|
def test_apply_model_installs_diffusion_wrapper_into_live_transformer_options() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
transformer_options = {"cond_or_uncond": [0, 1]}
|
|
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
wrapper_slot = transformer_options["wrappers"]["diffusion_model"]["spectrum_wan_runtime"]
|
|
assert transformer_options[_RUNTIME_KEY] is runtime
|
|
assert wan._spectrum_wan_diffusion_model_wrapper in wrapper_slot
|
|
assert runtime.last_info["live_diffusion_wrapper_installed"] is True
|
|
|
|
|
|
def test_apply_model_installs_live_diffusion_wrapper_once() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
transformer_options = {"cond_or_uncond": [0, 1]}
|
|
|
|
for _ in range(2):
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options=transformer_options,
|
|
)
|
|
|
|
wrapper_slot = transformer_options["wrappers"]["diffusion_model"]["spectrum_wan_runtime"]
|
|
assert wrapper_slot.count(wan._spectrum_wan_diffusion_model_wrapper) == 1
|
|
|
|
|
|
def test_apply_model_finds_positional_transformer_options() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelWithForward(), _cfg())
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
transformer_options = {"cond_or_uncond": [0, 1]}
|
|
|
|
patched.model.apply_model(
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
None,
|
|
None,
|
|
transformer_options,
|
|
)
|
|
|
|
wrapper_slot = transformer_options["wrappers"]["diffusion_model"]["spectrum_wan_runtime"]
|
|
assert transformer_options[_RUNTIME_KEY] is runtime
|
|
assert wan._spectrum_wan_diffusion_model_wrapper in wrapper_slot
|
|
|
|
|
|
class DummyInnerLegacyForward(DummyInner):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.forward_calls = 0
|
|
|
|
def forward(self, *args, **kwargs):
|
|
self.forward_calls += 1
|
|
return "legacy-forward"
|
|
|
|
|
|
class DummyModelLegacyForward(DummyModel):
|
|
def __init__(self) -> None:
|
|
self.model = DummyOuter(DummyInnerLegacyForward())
|
|
self.model_options = None
|
|
|
|
|
|
def test_patcher_prefers_forward_orig_over_legacy_forward_when_no__forward() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelLegacyForward(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
|
|
assert inner._spectrum_wan_wrapped_attr == "forward_orig"
|
|
assert inner._spectrum_wan_runtime.last_info["hook_target"] == "model.diffusion_model.forward_orig"
|
|
|
|
|
|
def test_run_spectrum_forward_records_forecast_errors_and_falls_back(capsys) -> None:
|
|
inner = DummyInner()
|
|
|
|
class RuntimeStub:
|
|
def __init__(self) -> None:
|
|
self.last_info = {}
|
|
self.observed = []
|
|
|
|
def begin_step(self, transformer_options, timesteps):
|
|
return {"step_idx": 2, "actual_forward": False, "global_step": 7}
|
|
|
|
def can_forecast(self, transformer_options):
|
|
return True
|
|
|
|
def predict_feature(self, transformer_options, step_idx, global_step=None):
|
|
raise RuntimeError("forecast broke")
|
|
|
|
def _debug_log(self, message: str) -> None:
|
|
print(message, file=sys.stderr, flush=True)
|
|
|
|
def observe_feature(self, transformer_options, step_idx, feature, global_step=None):
|
|
self.observed.append((step_idx, global_step, tuple(feature.shape)))
|
|
|
|
runtime = RuntimeStub()
|
|
out = _run_spectrum_forward(
|
|
inner,
|
|
runtime,
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={},
|
|
)
|
|
|
|
captured = capsys.readouterr()
|
|
assert torch.is_tensor(out)
|
|
assert runtime.observed == [(2, 7, (1, 4, 1))]
|
|
assert runtime.last_info["forecast_error"] == "RuntimeError: forecast broke"
|
|
assert "forecast_error step=2 global_step=7 RuntimeError: forecast broke" in captured.err
|
|
|
|
|
|
def test_run_spectrum_forward_clears_stale_forecast_error_before_non_forecast_step() -> None:
|
|
inner = DummyInner()
|
|
|
|
class RuntimeStub:
|
|
def __init__(self) -> None:
|
|
self.last_info = {"forecast_error": "stale"}
|
|
self.observed = []
|
|
|
|
def begin_step(self, transformer_options, timesteps):
|
|
return {"step_idx": 3, "actual_forward": True, "global_step": 8}
|
|
|
|
def can_forecast(self, transformer_options):
|
|
raise AssertionError("can_forecast should not run when actual_forward is True")
|
|
|
|
def _debug_log(self, message: str) -> None:
|
|
pass
|
|
|
|
def observe_feature(self, transformer_options, step_idx, feature, global_step=None):
|
|
self.observed.append((step_idx, global_step, tuple(feature.shape)))
|
|
|
|
runtime = RuntimeStub()
|
|
out = _run_spectrum_forward(
|
|
inner,
|
|
runtime,
|
|
torch.ones((1, 1, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
transformer_options={},
|
|
)
|
|
|
|
assert torch.is_tensor(out)
|
|
assert "forecast_error" not in runtime.last_info
|
|
assert runtime.observed == [(3, 8, (1, 4, 1))]
|
|
|
|
|
|
class DummyInnerForDiffusionWrapper(DummyInner):
|
|
patch_size = (1, 1, 1)
|
|
|
|
def rope_encode(self, t, h, w, device=None, dtype=None, transformer_options=None):
|
|
return torch.zeros((1,), device=device, dtype=dtype)
|
|
|
|
def unpatchify(self, x, grid_sizes):
|
|
return torch.zeros((1, 1, *grid_sizes), dtype=x.dtype, device=x.device)
|
|
|
|
|
|
class DummyModelForDiffusionWrapper(DummyModel):
|
|
def __init__(self) -> None:
|
|
self.model = DummyOuter(DummyInnerForDiffusionWrapper())
|
|
self.model_options = None
|
|
|
|
|
|
class DummyDiffusionExecutor:
|
|
def __init__(self, class_obj) -> None:
|
|
self.class_obj = class_obj
|
|
self.calls = 0
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
self.calls += 1
|
|
return torch.full((1, 1, 1, 1, 1), -1.0)
|
|
|
|
|
|
def test_patcher_installs_diffusion_model_wrapper() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelForDiffusionWrapper(), _cfg())
|
|
wrapper_slot = patched.model_options["transformer_options"]["wrappers"]["diffusion_model"]["spectrum_wan_runtime"]
|
|
|
|
assert wan._spectrum_wan_diffusion_model_wrapper in wrapper_slot
|
|
|
|
|
|
def test_diffusion_model_wrapper_runs_spectrum_without_calling_captured_forward_executor() -> None:
|
|
patched = WanSpectrumPatcher.patch(DummyModelForDiffusionWrapper(), _cfg())
|
|
inner = patched.model.diffusion_model
|
|
runtime = patched.model_options["transformer_options"][_RUNTIME_KEY]
|
|
executor = DummyDiffusionExecutor(inner)
|
|
sample_sigmas = torch.tensor([1.0, 0.5, 0.0], dtype=torch.float32)
|
|
transformer_options = {
|
|
_RUNTIME_KEY: runtime,
|
|
"sample_sigmas": sample_sigmas,
|
|
"cond_or_uncond": [0, 1],
|
|
}
|
|
|
|
out = wan._spectrum_wan_diffusion_model_wrapper(
|
|
executor,
|
|
torch.ones((1, 1, 2, 2, 2), dtype=torch.float32),
|
|
torch.tensor([sample_sigmas[0]], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
None,
|
|
None,
|
|
transformer_options,
|
|
)
|
|
|
|
assert executor.calls == 0
|
|
assert torch.is_tensor(out)
|
|
assert tuple(out.shape) == (1, 1, 2, 2, 2)
|
|
assert inner.original_calls == 0
|
|
assert runtime.last_info["last_sigma"] == 1.0
|
|
|
|
|
|
def test_diffusion_model_wrapper_falls_back_for_non_wan_executor_target() -> None:
|
|
cfg = _cfg()
|
|
runtime = SpectrumWanRuntime(cfg, resolve_handler("auto", DummyModelForDiffusionWrapper()))
|
|
executor = DummyDiffusionExecutor(NonWanInner())
|
|
|
|
out = wan._spectrum_wan_diffusion_model_wrapper(
|
|
executor,
|
|
torch.ones((1, 1, 2, 2, 2), dtype=torch.float32),
|
|
torch.tensor([1.0], dtype=torch.float32),
|
|
torch.zeros((1, 1, 1), dtype=torch.float32),
|
|
None,
|
|
None,
|
|
{_RUNTIME_KEY: runtime},
|
|
)
|
|
|
|
assert executor.calls == 1
|
|
assert torch.equal(out, torch.full((1, 1, 1, 1, 1), -1.0))
|