diff --git a/tide_core/patches.py b/tide_core/patches.py index b9c5f86..2b09107 100644 --- a/tide_core/patches.py +++ b/tide_core/patches.py @@ -237,6 +237,19 @@ class TIDEModelWrapper: transformer_options["tide"] = tide_opts 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) diff --git a/tide_core/wan.py b/tide_core/wan.py index 14a7540..643bf2a 100644 --- a/tide_core/wan.py +++ b/tide_core/wan.py @@ -262,6 +262,30 @@ def _ensure_wan_forward_orig_wrapped(inner: Any, fallback_config: Optional[TIDEC return True +def _resolve_apply_model_wan_inner(apply_model: Any) -> Any: + outer = getattr(apply_model, "__self__", None) + if outer is None: + return None + + 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), + ] + ) + + for candidate in candidates: + if _looks_like_wan_inner(candidate): + return candidate + return None + + class TIDEWanDiffusionWrapper: """ComfyUI diffusion_model wrapper that enables WAN RoPE DTC hooks lazily.""" @@ -323,6 +347,53 @@ 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 = getattr(apply_model, "__self__", None) + inner = _resolve_apply_model_wan_inner(apply_model) + if not _looks_like_wan_inner(inner): + return + + transformer_options[_CONFIG_KEY] = config + transformer_options[_ENABLED_KEY] = config.should_apply_temperature() and config.temperature_strength != 0.0 + + wrapped_rope = _ensure_wan_rope_encode_wrapped(inner, config) + wrapped_forward_orig = _ensure_wan_forward_orig_wrapped(inner, config) + + 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"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.""" @@ -368,5 +439,6 @@ def install_tide_wan_patch(model: Any, config: TIDEConfig) -> Any: __all__ = [ "TIDEWanDiffusionWrapper", "install_tide_wan_patch", + "prepare_tide_wan_apply_model", "_scale_wan_freqs", ]