Compare commits
15
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
39cd4436c4 | ||
|
|
3e6c7e88ef | ||
|
|
911461a4ce | ||
|
|
e71c4fbc9e | ||
|
|
7f615f0146 | ||
|
|
00d3d2551b | ||
|
|
051bba506a | ||
|
|
876be80ce5 | ||
|
|
63dd35814c | ||
|
|
11b2efbb95 | ||
|
|
6a081b4e7b | ||
|
|
f28c87e0d5 | ||
|
|
b25c595dcc | ||
|
|
a902471919 | ||
|
|
ca5575c4ef |
@@ -3,15 +3,15 @@ from __future__ import annotations
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from .tide_core import TIDEConfig, TIDEAttentionOverride, TIDEAttentionPatch, TIDEModelWrapper, install_tide_wan_patch
|
||||
from .tide_core import TIDEConfig, TIDEAttentionOverride, TIDEAttentionPatch, TIDEModelWrapper, aspect_adaptive_base_resolution, install_tide_wan_patch
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name not in {f"{__package__}.tide_core", "tide_core"}:
|
||||
raise
|
||||
from tide_core import TIDEConfig, TIDEAttentionOverride, TIDEAttentionPatch, TIDEModelWrapper, install_tide_wan_patch
|
||||
from tide_core import TIDEConfig, TIDEAttentionOverride, TIDEAttentionPatch, TIDEModelWrapper, aspect_adaptive_base_resolution, install_tide_wan_patch
|
||||
except ImportError as exc:
|
||||
if "attempted relative import with no known parent package" not in str(exc):
|
||||
raise
|
||||
from tide_core import TIDEConfig, TIDEAttentionOverride, TIDEAttentionPatch, TIDEModelWrapper, install_tide_wan_patch
|
||||
from tide_core import TIDEConfig, TIDEAttentionOverride, TIDEAttentionPatch, TIDEModelWrapper, aspect_adaptive_base_resolution, install_tide_wan_patch
|
||||
|
||||
|
||||
class TIDEHighResolutionExtrapolation:
|
||||
@@ -110,6 +110,24 @@ class TIDEHighResolutionExtrapolation:
|
||||
return (patched,)
|
||||
|
||||
|
||||
def _resolve_wan_base_resolution(
|
||||
*,
|
||||
width: int,
|
||||
height: int,
|
||||
base_width: int,
|
||||
base_height: int,
|
||||
base_resolution_mode: str,
|
||||
) -> tuple[int, int]:
|
||||
mode = str(base_resolution_mode)
|
||||
if mode == "manual":
|
||||
return int(base_width), int(base_height)
|
||||
if mode == "aspect_adaptive_720p":
|
||||
return aspect_adaptive_base_resolution(width, height, 1280, 720)
|
||||
if mode == "aspect_adaptive_480p":
|
||||
return aspect_adaptive_base_resolution(width, height, 832, 480)
|
||||
raise ValueError(f"Unsupported WAN base_resolution_mode: {base_resolution_mode!r}")
|
||||
|
||||
|
||||
class TIDEWANHighResolutionExtrapolation:
|
||||
"""Patch a WAN 2.1/2.2-style DiT model with TIDE Dynamic Temperature Control."""
|
||||
|
||||
@@ -126,8 +144,12 @@ class TIDEWANHighResolutionExtrapolation:
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"base_width": ("INT", {"default": 640, "min": 16, "max": 16384, "step": 16}),
|
||||
"base_height": ("INT", {"default": 640, "min": 16, "max": 16384, "step": 16}),
|
||||
"base_resolution_mode": (
|
||||
["aspect_adaptive_720p", "aspect_adaptive_480p", "manual"],
|
||||
{"default": "aspect_adaptive_720p", "tooltip": "WAN native reference. Adaptive modes preserve target aspect and keep a 1280x720 or 832x480 native pixel budget; manual uses base_width/base_height exactly."},
|
||||
),
|
||||
"base_width": ("INT", {"default": 1280, "min": 16, "max": 16384, "step": 16}),
|
||||
"base_height": ("INT", {"default": 720, "min": 16, "max": 16384, "step": 16}),
|
||||
"alpha_low": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 8.0, "step": 0.05}),
|
||||
"alpha_high": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 8.0, "step": 0.05}),
|
||||
"tau_max": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 4.0, "step": 0.01}),
|
||||
@@ -149,8 +171,9 @@ class TIDEWANHighResolutionExtrapolation:
|
||||
width: int,
|
||||
height: int,
|
||||
temperature_strength: float,
|
||||
base_width: int = 640,
|
||||
base_height: int = 640,
|
||||
base_resolution_mode: str = "aspect_adaptive_720p",
|
||||
base_width: int = 1280,
|
||||
base_height: int = 720,
|
||||
alpha_low: float = 0.6,
|
||||
alpha_high: float = 0.2,
|
||||
tau_max: float = 1.0,
|
||||
@@ -159,6 +182,14 @@ class TIDEWANHighResolutionExtrapolation:
|
||||
preserve_existing_wrapper: bool = True,
|
||||
debug: bool = False,
|
||||
):
|
||||
base_width, base_height = _resolve_wan_base_resolution(
|
||||
width=int(width),
|
||||
height=int(height),
|
||||
base_width=int(base_width),
|
||||
base_height=int(base_height),
|
||||
base_resolution_mode=str(base_resolution_mode),
|
||||
)
|
||||
|
||||
config = TIDEConfig(
|
||||
width=int(width),
|
||||
height=int(height),
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "tide"
|
||||
description = "A ComfyUI custom node implementing TIDE (Text-Informed Dynamic Extrapolation with Step-Aware Temperature Control) for diffusion models."
|
||||
version = "1.0.0"
|
||||
version = "1.0.2"
|
||||
license = { file = "LICENSE" }
|
||||
|
||||
[project.urls]
|
||||
|
||||
+31
-1
@@ -2,13 +2,14 @@ import math
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
ROOT = pathlib.Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
from tide_core.config import TIDEConfig
|
||||
from tide_core.math import adaptive_text_bias, get_default_temperature, rope_temperature_scale
|
||||
from tide_core.math import adaptive_text_bias, aspect_adaptive_base_resolution, get_default_temperature, rope_temperature_scale
|
||||
|
||||
|
||||
def test_adaptive_text_bias_matches_paper_and_official_flux_script():
|
||||
@@ -43,3 +44,32 @@ def test_temperature_strength_zero_disables_scaling():
|
||||
cfg = TIDEConfig(width=4096, height=4096, temperature_strength=0.0)
|
||||
scale = rope_temperature_scale(cfg, timestep=1.0, device=torch.device("cpu"))
|
||||
assert torch.allclose(scale, torch.ones_like(scale), atol=1e-6)
|
||||
|
||||
|
||||
def test_aspect_adaptive_base_resolution_preserves_720p_area_budget_for_square_targets():
|
||||
assert aspect_adaptive_base_resolution(960, 960, 1280, 720) == (960, 960)
|
||||
|
||||
|
||||
def test_aspect_adaptive_base_resolution_preserves_native_landscape_and_portrait():
|
||||
assert aspect_adaptive_base_resolution(1280, 720, 1280, 720) == (1280, 720)
|
||||
assert aspect_adaptive_base_resolution(720, 1280, 1280, 720) == (720, 1280)
|
||||
|
||||
|
||||
def test_aspect_adaptive_base_resolution_uses_target_aspect_for_nonstandard_i2v_size():
|
||||
assert aspect_adaptive_base_resolution(896, 656, 1280, 720) == (1120, 816)
|
||||
|
||||
|
||||
def test_aspect_adaptive_base_resolution_keeps_snapped_result_within_native_budget():
|
||||
base_width, base_height = aspect_adaptive_base_resolution(1000, 777, 1280, 720)
|
||||
assert base_width * base_height <= 1280 * 720
|
||||
|
||||
|
||||
def test_aspect_adaptive_base_resolution_preserves_wan_480p_budget():
|
||||
assert aspect_adaptive_base_resolution(832, 480, 832, 480) == (832, 480)
|
||||
base_width, base_height = aspect_adaptive_base_resolution(480, 832, 832, 480)
|
||||
assert base_width * base_height <= 832 * 480
|
||||
|
||||
|
||||
def test_aspect_adaptive_base_resolution_rejects_unsnappable_budget():
|
||||
with pytest.raises(ValueError, match="No grid-aligned base resolution"):
|
||||
aspect_adaptive_base_resolution(100_000_000, 1, 1280, 720)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from .config import TIDEConfig
|
||||
from .math import (
|
||||
adaptive_text_bias,
|
||||
aspect_adaptive_base_resolution,
|
||||
get_default_temperature,
|
||||
get_mscale,
|
||||
rope_temperature_scale,
|
||||
@@ -11,6 +12,7 @@ from .wan import TIDEWanDiffusionWrapper, install_tide_wan_patch
|
||||
__all__ = [
|
||||
"TIDEConfig",
|
||||
"adaptive_text_bias",
|
||||
"aspect_adaptive_base_resolution",
|
||||
"get_default_temperature",
|
||||
"get_mscale",
|
||||
"rope_temperature_scale",
|
||||
|
||||
@@ -26,6 +26,72 @@ def get_default_temperature(scale: float) -> float:
|
||||
return 1.0 / (mscale * mscale)
|
||||
|
||||
|
||||
def aspect_adaptive_base_resolution(
|
||||
width: int,
|
||||
height: int,
|
||||
native_width: int,
|
||||
native_height: int,
|
||||
*,
|
||||
grid: int = 16,
|
||||
) -> tuple[int, int]:
|
||||
"""Return native-budget base dimensions matched to the target aspect.
|
||||
|
||||
WAN 480p/720p checkpoints are better represented as a native pixel budget
|
||||
than as one fixed landscape rectangle. For an arbitrary I2V aspect, preserve
|
||||
the target aspect while keeping native_width * native_height as the maximum
|
||||
native reference area used by TIDE's extrapolation gates and RoPE-axis scale
|
||||
factors.
|
||||
"""
|
||||
|
||||
width = int(width)
|
||||
height = int(height)
|
||||
native_width = int(native_width)
|
||||
native_height = int(native_height)
|
||||
grid = int(grid)
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError(f"width and height must be > 0, got {width}x{height}")
|
||||
if native_width <= 0 or native_height <= 0:
|
||||
raise ValueError(f"native base must be > 0, got {native_width}x{native_height}")
|
||||
if grid <= 0:
|
||||
raise ValueError(f"grid must be > 0, got {grid}")
|
||||
|
||||
native_area = float(native_width) * float(native_height)
|
||||
aspect = float(width) / float(height)
|
||||
raw_width = math.sqrt(native_area * aspect)
|
||||
raw_height = math.sqrt(native_area / aspect)
|
||||
|
||||
def snapped_neighbors(value: float) -> tuple[int, int]:
|
||||
scaled = value / grid
|
||||
lower = max(grid, int(math.floor(scaled)) * grid)
|
||||
upper = max(grid, int(math.ceil(scaled)) * grid)
|
||||
return lower, upper
|
||||
|
||||
candidates = {
|
||||
(candidate_width, candidate_height)
|
||||
for candidate_width in snapped_neighbors(raw_width)
|
||||
for candidate_height in snapped_neighbors(raw_height)
|
||||
}
|
||||
under_budget = [
|
||||
candidate
|
||||
for candidate in candidates
|
||||
if candidate[0] * candidate[1] <= native_width * native_height
|
||||
]
|
||||
if not under_budget:
|
||||
raise ValueError(
|
||||
"No grid-aligned base resolution satisfies the native area budget "
|
||||
f"for {width}x{height} against {native_width}x{native_height} on grid {grid}"
|
||||
)
|
||||
candidates = set(under_budget)
|
||||
|
||||
def score(candidate: tuple[int, int]) -> tuple[float, float]:
|
||||
candidate_width, candidate_height = candidate
|
||||
area_error = abs((candidate_width * candidate_height) - native_area) / native_area
|
||||
aspect_error = abs((candidate_width / candidate_height) - aspect) / aspect
|
||||
return area_error, aspect_error
|
||||
|
||||
return min(candidates, key=score)
|
||||
|
||||
|
||||
def adaptive_text_bias(config: TIDEConfig) -> float:
|
||||
"""Paper Eq. 17/18: beta = log(lambda), with lambda = pixel ratio.
|
||||
|
||||
|
||||
@@ -235,8 +235,39 @@ class TIDEModelWrapper:
|
||||
tide_opts["width"] = self.config.width
|
||||
tide_opts["height"] = self.config.height
|
||||
transformer_options["tide"] = tide_opts
|
||||
|
||||
# WAN 2.1/Spectrum may replace BaseModel.apply_model with a closure and
|
||||
# may also bypass some lower wrapper paths. Put the WAN TIDE config into
|
||||
# the live transformer_options before any downstream wrapper sees them;
|
||||
# this is skipped for Flux/SDXL because only install_tide_wan_patch marks
|
||||
# model options with the WAN runtime contract.
|
||||
try:
|
||||
from .wan import has_tide_wan_options, inject_tide_wan_options
|
||||
|
||||
if has_tide_wan_options(transformer_options):
|
||||
inject_tide_wan_options(transformer_options, self.config, timestep=args.get("timestep"), source="model_function_wrapper")
|
||||
except (ImportError, AttributeError, KeyError, TypeError, ValueError) as exc:
|
||||
if not getattr(self, "_wan_injection_warned", False):
|
||||
_LOG.warning("TIDE WAN transformer_options injection failed and was skipped: %s", exc, exc_info=True)
|
||||
self._wan_injection_warned = True
|
||||
elif self.config.debug:
|
||||
_LOG.exception("TIDE WAN transformer_options injection failed and was skipped: %s", exc)
|
||||
|
||||
c["transformer_options"] = transformer_options
|
||||
|
||||
# WAN 2.1 Spectrum can complete from an APPLY_MODEL wrapper without
|
||||
# entering the lower DIFFUSION_MODEL wrapper chain where TIDE WAN normally
|
||||
# binds lazily. If this is a WAN-looking model, bind the live inner here
|
||||
# before apply_model runs. The helper is intentionally best-effort and is
|
||||
# a no-op for Flux/SDXL/non-WAN models.
|
||||
try:
|
||||
from .wan import prepare_tide_wan_apply_model
|
||||
|
||||
prepare_tide_wan_apply_model(apply_model, self.config, transformer_options)
|
||||
except Exception as exc:
|
||||
if self.config.debug:
|
||||
_LOG.exception("TIDE WAN apply_model preparation failed and was skipped: %s", exc)
|
||||
|
||||
if self.config.preserve_existing_wrapper and self.old_wrapper is not None:
|
||||
return self.old_wrapper(apply_model, args | {"c": c})
|
||||
return apply_model(args["input"], args["timestep"], **c)
|
||||
|
||||
+338
-17
@@ -17,8 +17,18 @@ _CONFIG_KEY = "tide_wan_config"
|
||||
_ENABLED_KEY = "tide_wan_enabled"
|
||||
_WRAPPER_KEY = "tide_wan_rope_temperature"
|
||||
_SCALED_FREQS_ID_KEY = "_tide_wan_scaled_freqs_id"
|
||||
_CONFIG_SOURCE_KEY = "_tide_wan_config_source"
|
||||
_INNER_KEY = "_tide_wan_inner"
|
||||
_BLOCK_PATCH_KEY = "_tide_wan_block_dtc_patch"
|
||||
_BLOCK_PATCH_CONFIG_KEY = "_tide_wan_block_dtc_config"
|
||||
_BLOCK_PATCH_INNER_KEY = "_tide_wan_block_dtc_inner"
|
||||
_BLOCK_PATCH_SOURCE_KEY = "_tide_wan_block_dtc_source"
|
||||
_BLOCK_PATCH_TRANSFORMER_OPTIONS_KEY = "_tide_wan_block_dtc_transformer_options"
|
||||
_BLOCK_SCALED_CACHE_KEY = "_tide_wan_block_scaled_freqs_cache"
|
||||
_TRACE_LOG_COUNT_KEY = "_tide_wan_trace_log_count"
|
||||
_DIFFUSION_MODEL_WRAPPER_TYPE = "diffusion_model"
|
||||
_DEBUG_SCALE_LOG_LIMIT = 12
|
||||
_DEBUG_TRACE_LOG_LIMIT = 24
|
||||
|
||||
|
||||
def _debug_print(config: Optional[TIDEConfig], message: str) -> None:
|
||||
@@ -58,6 +68,56 @@ def _config_dict(config: TIDEConfig) -> dict[str, Any]:
|
||||
return data
|
||||
|
||||
|
||||
def inject_tide_wan_options(
|
||||
transformer_options: Optional[dict[str, Any]],
|
||||
config: TIDEConfig,
|
||||
*,
|
||||
timestep: Any = None,
|
||||
source: str = "unknown",
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""Inject the WAN TIDE runtime contract into the live transformer_options.
|
||||
|
||||
The WAN RoPE/forward hooks run much lower than the Comfy model_function_wrapper.
|
||||
They must therefore be able to resolve the current TIDE config from the exact
|
||||
transformer_options dict that reaches WanModel.rope_encode/forward_orig. This
|
||||
helper intentionally mutates that dict in place.
|
||||
"""
|
||||
|
||||
if not isinstance(transformer_options, dict):
|
||||
return transformer_options
|
||||
|
||||
transformer_options[_CONFIG_KEY] = config
|
||||
transformer_options[_ENABLED_KEY] = config.should_apply_temperature() and config.temperature_strength != 0.0
|
||||
transformer_options[_CONFIG_SOURCE_KEY] = source
|
||||
|
||||
tide_opts = transformer_options.get("tide", {})
|
||||
if not isinstance(tide_opts, dict):
|
||||
tide_opts = {}
|
||||
if timestep is not None:
|
||||
tide_opts["timestep"] = _safe_timestep01(timestep)
|
||||
tide_opts["width"] = config.width
|
||||
tide_opts["height"] = config.height
|
||||
transformer_options["tide"] = tide_opts
|
||||
return transformer_options
|
||||
|
||||
|
||||
def has_tide_wan_options(transformer_options: Optional[dict[str, Any]]) -> bool:
|
||||
return isinstance(transformer_options, dict) and (
|
||||
"tide_wan" in transformer_options or isinstance(transformer_options.get(_CONFIG_KEY), TIDEConfig)
|
||||
)
|
||||
|
||||
|
||||
def _trace_inner(config: Optional[TIDEConfig], inner: Any, message: str) -> None:
|
||||
if config is None or not config.debug:
|
||||
return
|
||||
count = int(getattr(inner, _TRACE_LOG_COUNT_KEY, 0)) if inner is not None else 0
|
||||
if count >= _DEBUG_TRACE_LOG_LIMIT:
|
||||
return
|
||||
if inner is not None:
|
||||
setattr(inner, _TRACE_LOG_COUNT_KEY, count + 1)
|
||||
_debug_print(config, message)
|
||||
|
||||
|
||||
def _looks_like_wan_inner(inner: Any) -> bool:
|
||||
# ComfyUI WAN 2.1/2.2 models expose this narrow runtime contract; avoid
|
||||
# class-name checks so compatible WAN subclasses can use the same path.
|
||||
@@ -149,9 +209,10 @@ def _scale_wan_freqs(
|
||||
out = freqs * scale.reshape(view_shape)
|
||||
|
||||
if config.debug:
|
||||
count = int(getattr(inner, "_tide_wan_scale_log_count", 0))
|
||||
count = int(getattr(inner, "_tide_wan_scale_log_count", 0)) if inner is not None else 0
|
||||
if count < _DEBUG_SCALE_LOG_LIMIT:
|
||||
setattr(inner, "_tide_wan_scale_log_count", count + 1)
|
||||
if inner is not None:
|
||||
setattr(inner, "_tide_wan_scale_log_count", count + 1)
|
||||
_debug_print(
|
||||
config,
|
||||
"[ComfyUI-TIDE] WAN DTC applied: "
|
||||
@@ -198,6 +259,7 @@ def _ensure_wan_rope_encode_wrapped(inner: Any, fallback_config: Optional[TIDECo
|
||||
if not _looks_like_wan_inner(inner):
|
||||
return False
|
||||
if getattr(inner, "_tide_wan_rope_encode_wrapped", False):
|
||||
inner._tide_wan_fallback_config = fallback_config
|
||||
return True
|
||||
|
||||
original_rope_encode = inner.rope_encode
|
||||
@@ -205,7 +267,15 @@ def _ensure_wan_rope_encode_wrapped(inner: Any, fallback_config: Optional[TIDECo
|
||||
def tide_wan_rope_encode(*args: Any, **kwargs: Any):
|
||||
transformer_options = kwargs.get("transformer_options", None)
|
||||
freqs = original_rope_encode(*args, **kwargs)
|
||||
config = _resolve_config(transformer_options, fallback_config)
|
||||
config = _resolve_config(transformer_options, getattr(inner, "_tide_wan_fallback_config", fallback_config))
|
||||
_trace_inner(
|
||||
config,
|
||||
inner,
|
||||
"[ComfyUI-TIDE] WAN rope_encode hook entered: "
|
||||
f"freqs_tensor={torch.is_tensor(freqs)} "
|
||||
f"freqs_shape={tuple(freqs.shape) if torch.is_tensor(freqs) else None} "
|
||||
f"config_source={transformer_options.get(_CONFIG_SOURCE_KEY) if isinstance(transformer_options, dict) else None}",
|
||||
)
|
||||
if config is None:
|
||||
return freqs
|
||||
timestep = _resolve_timestep(transformer_options)
|
||||
@@ -214,6 +284,7 @@ def _ensure_wan_rope_encode_wrapped(inner: Any, fallback_config: Optional[TIDECo
|
||||
return out
|
||||
|
||||
inner._tide_wan_original_rope_encode = original_rope_encode
|
||||
inner._tide_wan_fallback_config = fallback_config
|
||||
inner.rope_encode = tide_wan_rope_encode
|
||||
inner._tide_wan_rope_encode_wrapped = True
|
||||
_debug_print(fallback_config, f"[ComfyUI-TIDE] WAN wrapped rope_encode on {type(inner).__name__} id={id(inner)}")
|
||||
@@ -224,6 +295,7 @@ def _ensure_wan_forward_orig_wrapped(inner: Any, fallback_config: Optional[TIDEC
|
||||
if not _looks_like_wan_inner(inner):
|
||||
return False
|
||||
if getattr(inner, "_tide_wan_forward_orig_wrapped", False):
|
||||
inner._tide_wan_fallback_config = fallback_config
|
||||
return True
|
||||
|
||||
original_forward_orig = inner.forward_orig
|
||||
@@ -237,8 +309,18 @@ def _ensure_wan_forward_orig_wrapped(inner: Any, fallback_config: Optional[TIDEC
|
||||
transformer_options=None,
|
||||
**kwargs,
|
||||
):
|
||||
config = _resolve_config(transformer_options, fallback_config)
|
||||
if config is not None and torch.is_tensor(freqs) and not _freqs_already_scaled(transformer_options, freqs):
|
||||
config = _resolve_config(transformer_options, getattr(inner, "_tide_wan_fallback_config", fallback_config))
|
||||
already_scaled = _freqs_already_scaled(transformer_options, freqs)
|
||||
_trace_inner(
|
||||
config,
|
||||
inner,
|
||||
"[ComfyUI-TIDE] WAN forward_orig hook entered: "
|
||||
f"freqs_tensor={torch.is_tensor(freqs)} "
|
||||
f"freqs_shape={tuple(freqs.shape) if torch.is_tensor(freqs) else None} "
|
||||
f"already_scaled={bool(already_scaled)} "
|
||||
f"config_source={transformer_options.get(_CONFIG_SOURCE_KEY) if isinstance(transformer_options, dict) else None}",
|
||||
)
|
||||
if config is not None and torch.is_tensor(freqs) and not already_scaled:
|
||||
timestep = _resolve_timestep(transformer_options, t)
|
||||
freqs = _scale_wan_freqs(config, inner, freqs, timestep=timestep)
|
||||
_mark_scaled_freqs(transformer_options, freqs)
|
||||
@@ -256,12 +338,205 @@ def _ensure_wan_forward_orig_wrapped(inner: Any, fallback_config: Optional[TIDEC
|
||||
)
|
||||
|
||||
inner._tide_wan_original_forward_orig = original_forward_orig
|
||||
inner._tide_wan_fallback_config = fallback_config
|
||||
inner.forward_orig = tide_wan_forward_orig
|
||||
inner._tide_wan_forward_orig_wrapped = True
|
||||
_debug_print(fallback_config, f"[ComfyUI-TIDE] WAN wrapped forward_orig on {type(inner).__name__} id={id(inner)}")
|
||||
return True
|
||||
|
||||
|
||||
def _wan_candidates_from_outer(outer: Any) -> list[Any]:
|
||||
if outer is None:
|
||||
return []
|
||||
|
||||
candidates = [getattr(outer, "diffusion_model", None)]
|
||||
model = getattr(outer, "model", None)
|
||||
if model is not None:
|
||||
candidates.extend(
|
||||
[
|
||||
getattr(model, "diffusion_model", None),
|
||||
getattr(getattr(model, "model", None), "diffusion_model", None),
|
||||
]
|
||||
)
|
||||
return candidates
|
||||
|
||||
|
||||
def _resolve_apply_model_outer(apply_model: Any) -> Any:
|
||||
outer = getattr(apply_model, "__self__", None)
|
||||
if outer is not None:
|
||||
return outer
|
||||
|
||||
# Spectrum WAN replaces BaseModel.apply_model with a per-instance closure.
|
||||
# Functions assigned to an instance are not descriptors, so __self__ is lost.
|
||||
# Recover the captured BaseModel/outer object from the closure instead of
|
||||
# silently skipping the WAN prepare path.
|
||||
closure = getattr(apply_model, "__closure__", None)
|
||||
if not closure:
|
||||
return None
|
||||
|
||||
matches: list[tuple[Any, Any]] = []
|
||||
for cell in closure:
|
||||
try:
|
||||
value = cell.cell_contents
|
||||
except ValueError:
|
||||
continue
|
||||
bound_self = getattr(value, "__self__", None)
|
||||
if bound_self is not None:
|
||||
value = bound_self
|
||||
for candidate in _wan_candidates_from_outer(value):
|
||||
if _looks_like_wan_inner(candidate) and not any(candidate is seen for _, seen in matches):
|
||||
matches.append((value, candidate))
|
||||
if len(matches) == 1:
|
||||
return matches[0][0]
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_apply_model_wan_inner(outer: Any) -> Any:
|
||||
for candidate in _wan_candidates_from_outer(outer):
|
||||
if _looks_like_wan_inner(candidate):
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
|
||||
def _get_block_scaled_freqs(
|
||||
transformer_options: Optional[dict[str, Any]],
|
||||
config: TIDEConfig,
|
||||
inner: Any,
|
||||
freqs: Any,
|
||||
) -> Any:
|
||||
if not torch.is_tensor(freqs):
|
||||
return freqs
|
||||
if _freqs_already_scaled(transformer_options, freqs):
|
||||
return freqs
|
||||
|
||||
cache = None
|
||||
timestep = _resolve_timestep(transformer_options)
|
||||
cache_key = (id(freqs), timestep, id(config))
|
||||
if isinstance(transformer_options, dict):
|
||||
cache = transformer_options.setdefault(_BLOCK_SCALED_CACHE_KEY, {})
|
||||
if isinstance(cache, dict):
|
||||
cached = cache.get(cache_key)
|
||||
if torch.is_tensor(cached):
|
||||
return cached
|
||||
|
||||
out = _scale_wan_freqs(config, inner, freqs, timestep=timestep)
|
||||
_mark_scaled_freqs(transformer_options, out)
|
||||
if isinstance(cache, dict) and torch.is_tensor(out):
|
||||
cache[cache_key] = out
|
||||
return out
|
||||
|
||||
|
||||
def _install_wan_block_dtc_patches(
|
||||
transformer_options: Optional[dict[str, Any]],
|
||||
config: TIDEConfig,
|
||||
inner: Any,
|
||||
*,
|
||||
source: str,
|
||||
) -> bool:
|
||||
"""Install a block-level WAN PE scaling fallback.
|
||||
|
||||
Spectrum WAN 2.1 can run actual forwards through an APPLY_MODEL driver where
|
||||
neither the native WanModel.rope_encode hook nor the WanModel.forward_orig
|
||||
hook is reached by TIDE, even though transformer block execution still goes
|
||||
through Comfy's patches_replace path. Wrapping the block replacement hooks
|
||||
lets TIDE scale the shared RoPE matrix at the last stable point before each
|
||||
attention block while preserving existing DiffAid/Spectrum patches.
|
||||
"""
|
||||
|
||||
if not isinstance(transformer_options, dict) or not _looks_like_wan_inner(inner):
|
||||
return False
|
||||
|
||||
patches_replace = transformer_options.get("patches_replace")
|
||||
if not isinstance(patches_replace, dict):
|
||||
patches_replace = {}
|
||||
dit_replace = patches_replace.get("dit")
|
||||
if not isinstance(dit_replace, dict):
|
||||
dit_replace = {}
|
||||
|
||||
block_count = len(getattr(inner, "blocks", ()) or ())
|
||||
if block_count <= 0:
|
||||
return False
|
||||
|
||||
def make_tide_wan_block_patch(previous_patch: Any, block_index: int):
|
||||
def tide_wan_block_patch(args, extra_options):
|
||||
local_config = getattr(tide_wan_block_patch, _BLOCK_PATCH_CONFIG_KEY, config)
|
||||
local_inner = getattr(tide_wan_block_patch, _BLOCK_PATCH_INNER_KEY, inner)
|
||||
fallback_transformer_options = getattr(
|
||||
tide_wan_block_patch,
|
||||
_BLOCK_PATCH_TRANSFORMER_OPTIONS_KEY,
|
||||
transformer_options,
|
||||
)
|
||||
block_transformer_options = None
|
||||
if isinstance(args, dict):
|
||||
block_transformer_options = args.get("transformer_options")
|
||||
if not isinstance(block_transformer_options, dict):
|
||||
block_transformer_options = fallback_transformer_options
|
||||
|
||||
resolved_config = _resolve_config(block_transformer_options, local_config)
|
||||
if resolved_config is not None and isinstance(args, dict) and "pe" in args:
|
||||
if isinstance(block_transformer_options, dict):
|
||||
block_transformer_options[_INNER_KEY] = local_inner
|
||||
pe = args.get("pe")
|
||||
scaled_pe = _get_block_scaled_freqs(block_transformer_options, resolved_config, local_inner, pe)
|
||||
if scaled_pe is not pe:
|
||||
args = dict(args)
|
||||
args["pe"] = scaled_pe
|
||||
if resolved_config.debug and not getattr(tide_wan_block_patch, "_tide_wan_block_scaled_logged", False):
|
||||
setattr(tide_wan_block_patch, "_tide_wan_block_scaled_logged", True)
|
||||
_debug_print(
|
||||
resolved_config,
|
||||
"[ComfyUI-TIDE] WAN block DTC fallback scaled PE: "
|
||||
f"block={block_index} source={getattr(tide_wan_block_patch, _BLOCK_PATCH_SOURCE_KEY, source)} "
|
||||
f"pe_shape={tuple(pe.shape) if torch.is_tensor(pe) else None}",
|
||||
)
|
||||
|
||||
if callable(previous_patch):
|
||||
return previous_patch(args, extra_options)
|
||||
original_block = extra_options.get("original_block") if isinstance(extra_options, dict) else None
|
||||
if not callable(original_block):
|
||||
raise RuntimeError("TIDE WAN block DTC patch missing original_block")
|
||||
return original_block(args)
|
||||
|
||||
return tide_wan_block_patch
|
||||
|
||||
installed = 0
|
||||
for block_index in range(block_count):
|
||||
block_key = ("double_block", block_index)
|
||||
previous_patch = dit_replace.get(block_key)
|
||||
|
||||
if callable(previous_patch) and getattr(previous_patch, _BLOCK_PATCH_KEY, False):
|
||||
setattr(previous_patch, _BLOCK_PATCH_CONFIG_KEY, config)
|
||||
setattr(previous_patch, _BLOCK_PATCH_INNER_KEY, inner)
|
||||
setattr(previous_patch, _BLOCK_PATCH_SOURCE_KEY, source)
|
||||
setattr(previous_patch, _BLOCK_PATCH_TRANSFORMER_OPTIONS_KEY, transformer_options)
|
||||
installed += 1
|
||||
continue
|
||||
|
||||
tide_wan_block_patch = make_tide_wan_block_patch(previous_patch, block_index)
|
||||
setattr(tide_wan_block_patch, _BLOCK_PATCH_KEY, True)
|
||||
setattr(tide_wan_block_patch, _BLOCK_PATCH_CONFIG_KEY, config)
|
||||
setattr(tide_wan_block_patch, _BLOCK_PATCH_INNER_KEY, inner)
|
||||
setattr(tide_wan_block_patch, _BLOCK_PATCH_SOURCE_KEY, source)
|
||||
setattr(tide_wan_block_patch, _BLOCK_PATCH_TRANSFORMER_OPTIONS_KEY, transformer_options)
|
||||
dit_replace[block_key] = tide_wan_block_patch
|
||||
installed += 1
|
||||
|
||||
patches_replace["dit"] = dit_replace
|
||||
transformer_options["patches_replace"] = patches_replace
|
||||
|
||||
if config.debug:
|
||||
marker = "_tide_wan_block_dtc_logged"
|
||||
if not getattr(inner, marker, False):
|
||||
setattr(inner, marker, True)
|
||||
_debug_print(
|
||||
config,
|
||||
"[ComfyUI-TIDE] WAN block DTC fallback installed: "
|
||||
f"blocks={installed}/{block_count} source={source}",
|
||||
)
|
||||
return installed > 0
|
||||
|
||||
|
||||
class TIDEWanDiffusionWrapper:
|
||||
"""ComfyUI diffusion_model wrapper that enables WAN RoPE DTC hooks lazily."""
|
||||
|
||||
@@ -284,20 +559,14 @@ class TIDEWanDiffusionWrapper:
|
||||
):
|
||||
if transformer_options is None:
|
||||
transformer_options = {}
|
||||
transformer_options[_CONFIG_KEY] = self.config
|
||||
transformer_options[_ENABLED_KEY] = self.config.should_apply_temperature() and self.config.temperature_strength != 0.0
|
||||
|
||||
tide_opts = transformer_options.get("tide", {})
|
||||
if not isinstance(tide_opts, dict):
|
||||
tide_opts = {}
|
||||
tide_opts["timestep"] = _safe_timestep01(timestep)
|
||||
tide_opts["width"] = self.config.width
|
||||
tide_opts["height"] = self.config.height
|
||||
transformer_options["tide"] = tide_opts
|
||||
inject_tide_wan_options(transformer_options, self.config, timestep=timestep, source="diffusion_wrapper")
|
||||
|
||||
inner = getattr(executor, "class_obj", None)
|
||||
if isinstance(transformer_options, dict):
|
||||
transformer_options[_INNER_KEY] = inner
|
||||
wrapped_rope = _ensure_wan_rope_encode_wrapped(inner, self.config)
|
||||
wrapped_forward_orig = _ensure_wan_forward_orig_wrapped(inner, self.config)
|
||||
installed_block_dtc = _install_wan_block_dtc_patches(transformer_options, self.config, inner, source="diffusion_wrapper")
|
||||
if self.config.debug:
|
||||
if not getattr(self, "_logged_runtime", False):
|
||||
_debug_print(
|
||||
@@ -305,6 +574,7 @@ class TIDEWanDiffusionWrapper:
|
||||
"[ComfyUI-TIDE] WAN diffusion wrapper active: "
|
||||
f"inner={type(inner).__name__ if inner is not None else None} "
|
||||
f"wrapped_rope={bool(wrapped_rope)} wrapped_forward_orig={bool(wrapped_forward_orig)} "
|
||||
f"block_dtc={bool(installed_block_dtc)} "
|
||||
f"enabled={bool(transformer_options[_ENABLED_KEY])}",
|
||||
)
|
||||
self._logged_runtime = True
|
||||
@@ -323,12 +593,60 @@ class TIDEWanDiffusionWrapper:
|
||||
)
|
||||
|
||||
|
||||
def prepare_tide_wan_apply_model(
|
||||
apply_model: Any,
|
||||
config: TIDEConfig,
|
||||
transformer_options: Optional[dict[str, Any]],
|
||||
) -> None:
|
||||
"""Prepare WAN DTC before Comfy's APPLY_MODEL wrappers run.
|
||||
|
||||
Spectrum WAN 2.1 can drive the WAN model from an APPLY_MODEL wrapper and
|
||||
return without entering Comfy's native DIFFUSION_MODEL wrapper chain. The
|
||||
normal TIDE WAN diffusion wrapper therefore cannot be the only lazy binding
|
||||
point. This hook runs from TIDE's model_function_wrapper, mutates the same
|
||||
transformer_options dict the sampler will pass into apply_model, and wraps
|
||||
the live WanModel before Spectrum's APPLY_MODEL wrapper can bypass the lower
|
||||
wrapper chain.
|
||||
"""
|
||||
|
||||
if transformer_options is None:
|
||||
return
|
||||
|
||||
outer = _resolve_apply_model_outer(apply_model)
|
||||
inner = _resolve_apply_model_wan_inner(outer)
|
||||
if not _looks_like_wan_inner(inner):
|
||||
return
|
||||
|
||||
inject_tide_wan_options(transformer_options, config, source="model_function_wrapper")
|
||||
transformer_options[_INNER_KEY] = inner
|
||||
|
||||
wrapped_rope = _ensure_wan_rope_encode_wrapped(inner, config)
|
||||
wrapped_forward_orig = _ensure_wan_forward_orig_wrapped(inner, config)
|
||||
installed_block_dtc = _install_wan_block_dtc_patches(transformer_options, config, inner, source="model_function_wrapper")
|
||||
|
||||
if config.debug:
|
||||
# Log once per live inner. This is intentionally separate from
|
||||
# TIDEWanDiffusionWrapper's runtime log because Spectrum WAN 2.1 may
|
||||
# bypass that path completely.
|
||||
marker = "_tide_wan_apply_model_prepare_logged"
|
||||
if not getattr(inner, marker, False):
|
||||
setattr(inner, marker, True)
|
||||
_debug_print(
|
||||
config,
|
||||
"[ComfyUI-TIDE] WAN apply_model prepare active: "
|
||||
f"outer={type(outer).__name__ if outer is not None else None} "
|
||||
f"inner={type(inner).__name__ if inner is not None else None} "
|
||||
f"wrapped_rope={bool(wrapped_rope)} wrapped_forward_orig={bool(wrapped_forward_orig)} "
|
||||
f"block_dtc={bool(installed_block_dtc)} "
|
||||
f"enabled={bool(transformer_options[_ENABLED_KEY])}",
|
||||
)
|
||||
|
||||
|
||||
def install_tide_wan_patch(model: Any, config: TIDEConfig) -> Any:
|
||||
"""Install the WAN-specific TIDE DTC path on a cloned ComfyUI MODEL."""
|
||||
|
||||
transformer_options = _ensure_transformer_options(model)
|
||||
transformer_options[_CONFIG_KEY] = config
|
||||
transformer_options[_ENABLED_KEY] = config.should_apply_temperature() and config.temperature_strength != 0.0
|
||||
inject_tide_wan_options(transformer_options, config, source="node_install")
|
||||
transformer_options["tide_wan"] = _config_dict(config)
|
||||
|
||||
wrapper = TIDEWanDiffusionWrapper(config)
|
||||
@@ -367,6 +685,9 @@ def install_tide_wan_patch(model: Any, config: TIDEConfig) -> Any:
|
||||
|
||||
__all__ = [
|
||||
"TIDEWanDiffusionWrapper",
|
||||
"has_tide_wan_options",
|
||||
"install_tide_wan_patch",
|
||||
"prepare_tide_wan_apply_model",
|
||||
"inject_tide_wan_options",
|
||||
"_scale_wan_freqs",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user