Add files via upload
This commit is contained in:
capitan01R
+12
@@ -0,0 +1,12 @@
|
||||
from .middle_cache import MiniMaxH3EulerMiddleCache
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MiniMaxH3EulerMiddleCache": MiniMaxH3EulerMiddleCache,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MiniMaxH3EulerMiddleCache": "MiniMax H3 FastPath Euler Middle Cache",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,117 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
LORA_SIGNATURE_KEY = "minimax_h3_fastpath_lora_signature"
|
||||
LORA_SIGMA_KEY = "minimax_h3_fastpath_sigma"
|
||||
|
||||
CACHE_BINDING_KEY = "minimax_h3_fastpath_cache_binding"
|
||||
CACHE_CALL_KEY = "minimax_h3_fastpath_cache_call"
|
||||
CACHE_ATTACHMENT_KEY = "minimax_h3_fastpath_cache_attachment"
|
||||
CACHE_WRAPPER_KEY = "minimax_h3_fastpath_middle_cache"
|
||||
|
||||
LORA_ATTACHMENT_KEY = "minimax_h3_fastpath_lora_specs"
|
||||
LORA_INJECTION_KEY = "minimax_h3_fastpath_scheduled_lora"
|
||||
|
||||
|
||||
def sigma_anchors(sigmas: Any) -> tuple[float, ...]:
|
||||
values = torch.as_tensor(sigmas).detach().to(device="cpu", dtype=torch.float64).flatten()
|
||||
if values.numel() < 2:
|
||||
raise ValueError("SIGMAS must contain at least one model-call sigma and the final endpoint")
|
||||
if not bool(torch.isfinite(values).all()):
|
||||
raise ValueError("SIGMAS contains NaN or infinite values")
|
||||
if bool((values < 0.0).any()):
|
||||
raise ValueError("SIGMAS must be non-negative")
|
||||
if bool((values[1:] > values[:-1]).any()):
|
||||
raise ValueError("SIGMAS must be monotonically descending")
|
||||
|
||||
anchors = values[:-1]
|
||||
if anchors.numel() > 1 and bool((anchors[1:] >= anchors[:-1]).any()):
|
||||
raise ValueError(
|
||||
"Model-call SIGMAS must be strictly descending; repeated values cannot identify calls safely"
|
||||
)
|
||||
return tuple(float(value) for value in anchors.tolist())
|
||||
|
||||
|
||||
def explicit_strengths(text: str, count: int) -> tuple[float, ...]:
|
||||
cleaned = re.sub(r"[\[\](){}]", " ", text.strip())
|
||||
tokens = [token for token in re.split(r"[,;\s]+", cleaned) if token]
|
||||
if len(tokens) != count:
|
||||
raise ValueError(
|
||||
f"explicit_strengths needs exactly {count} values for sigmas[:-1]; received {len(tokens)}"
|
||||
)
|
||||
try:
|
||||
values = tuple(float(token) for token in tokens)
|
||||
except ValueError as error:
|
||||
raise ValueError("explicit_strengths contains a non-numeric value") from error
|
||||
if not all(math.isfinite(value) for value in values):
|
||||
raise ValueError("explicit_strengths contains NaN or infinite values")
|
||||
return values
|
||||
|
||||
|
||||
def curve_strengths(
|
||||
count: int,
|
||||
start_strength: float,
|
||||
end_strength: float,
|
||||
curve: str,
|
||||
curve_power: float,
|
||||
explicit: str,
|
||||
) -> tuple[float, ...]:
|
||||
if curve == "explicit":
|
||||
return explicit_strengths(explicit, count)
|
||||
if count == 1:
|
||||
return (float(start_strength),)
|
||||
|
||||
values: list[float] = []
|
||||
for index in range(count):
|
||||
progress = index / (count - 1)
|
||||
if curve == "linear":
|
||||
shaped = progress
|
||||
elif curve == "cosine":
|
||||
shaped = 0.5 - 0.5 * math.cos(math.pi * progress)
|
||||
elif curve == "smoothstep":
|
||||
shaped = progress * progress * (3.0 - 2.0 * progress)
|
||||
elif curve == "power":
|
||||
shaped = progress**curve_power
|
||||
else:
|
||||
raise ValueError(f"Unknown strength curve: {curve}")
|
||||
values.append(start_strength + shaped * (end_strength - start_strength))
|
||||
return tuple(float(value) for value in values)
|
||||
|
||||
|
||||
def strength_at_sigma(group: dict[str, Any], sigma: float) -> float:
|
||||
anchors: tuple[float, ...] = group["sigma_anchors"]
|
||||
strengths: tuple[float, ...] = group["strength_anchors"]
|
||||
if sigma >= anchors[0]:
|
||||
return strengths[0]
|
||||
if sigma <= anchors[-1]:
|
||||
return strengths[-1]
|
||||
|
||||
for index in range(len(anchors) - 1):
|
||||
sigma_high = anchors[index]
|
||||
sigma_low = anchors[index + 1]
|
||||
if sigma_high >= sigma >= sigma_low:
|
||||
span = sigma_high - sigma_low
|
||||
if span == 0.0:
|
||||
return strengths[index]
|
||||
alpha = (sigma_high - sigma) / span
|
||||
return strengths[index] + alpha * (strengths[index + 1] - strengths[index])
|
||||
return strengths[-1]
|
||||
|
||||
|
||||
def normalized_items(values: Iterable[Any] | None) -> tuple[str, ...]:
|
||||
if values is None:
|
||||
return ()
|
||||
return tuple(str(value) for value in values)
|
||||
|
||||
|
||||
def rounded_strength(value: float) -> float:
|
||||
# Exact schedule anchors remain exact while insignificant interpolation noise
|
||||
# cannot spuriously invalidate a cache entry.
|
||||
return round(float(value), 10)
|
||||
+704
@@ -0,0 +1,704 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.patcher_extension
|
||||
|
||||
from .common import (
|
||||
CACHE_ATTACHMENT_KEY,
|
||||
CACHE_BINDING_KEY,
|
||||
CACHE_CALL_KEY,
|
||||
CACHE_WRAPPER_KEY,
|
||||
LORA_SIGNATURE_KEY,
|
||||
normalized_items,
|
||||
)
|
||||
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
RUN_ID_KEY = "minimax_h3_fastpath_cache_run_id"
|
||||
STEP_ID_KEY = "minimax_h3_fastpath_cache_step_id"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheConfig:
|
||||
first_middle_block: int
|
||||
last_middle_block: int
|
||||
reuse_threshold: float
|
||||
max_consecutive_reuses: int
|
||||
cache_device: str
|
||||
require_fastpath_schedule: bool
|
||||
strict_branch_identity: bool
|
||||
suppress_candidate_prefetch: bool
|
||||
verbose: bool
|
||||
|
||||
@property
|
||||
def middle_block_count(self) -> int:
|
||||
return self.last_middle_block - self.first_middle_block + 1
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class BranchState:
|
||||
topology: tuple[Any, ...] | None = None
|
||||
lora_signature: tuple[Any, ...] | None = None
|
||||
residual: torch.Tensor | None = None
|
||||
prefix_probe: torch.Tensor | None = None
|
||||
consecutive_reuses: int = 0
|
||||
last_step_id: int = -1
|
||||
|
||||
def reset(self, topology=None, lora_signature=None) -> None:
|
||||
self.topology = topology
|
||||
self.lora_signature = lora_signature
|
||||
self.residual = None
|
||||
self.prefix_probe = None
|
||||
self.consecutive_reuses = 0
|
||||
self.last_step_id = -1
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CacheCall:
|
||||
runtime: "MiddleCacheRuntime"
|
||||
run_id: int
|
||||
step_id: int
|
||||
branch_key: tuple[Any, ...] | None
|
||||
lora_signature: tuple[Any, ...] | None
|
||||
eligible: bool
|
||||
reason: str
|
||||
prefetch_suppressed: bool = False
|
||||
middle_started: bool = False
|
||||
reuse: bool = False
|
||||
completed: bool = False
|
||||
state: BranchState | None = None
|
||||
start_hidden: torch.Tensor | None = None
|
||||
current_probe: torch.Tensor | None = None
|
||||
feature_delta: float | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CacheBinding:
|
||||
runtime: "MiddleCacheRuntime"
|
||||
|
||||
|
||||
class MiddleCacheRuntime:
|
||||
"""Run-local residual cache for a contiguous range of H3 transformer blocks."""
|
||||
|
||||
def __init__(self, config: CacheConfig):
|
||||
self.config = config
|
||||
self.lock = threading.RLock()
|
||||
self.run_serial = 0
|
||||
self.active_run_id: int | None = None
|
||||
self.active_step_id: int | None = None
|
||||
self.supported_sampler = False
|
||||
self.branch_states: dict[tuple[Any, ...], BranchState] = {}
|
||||
self.step_occurrences: dict[tuple[Any, ...], int] = {}
|
||||
self.disabled_reason: str | None = None
|
||||
self.stats = self._new_stats()
|
||||
|
||||
@staticmethod
|
||||
def _new_stats() -> dict[str, int]:
|
||||
return {
|
||||
"model_calls": 0,
|
||||
"cached_full_calls": 0,
|
||||
"reused_calls": 0,
|
||||
"skipped_blocks": 0,
|
||||
"strength_resets": 0,
|
||||
"gate_rejections": 0,
|
||||
"identity_fallbacks": 0,
|
||||
"prefetch_suppressed_calls": 0,
|
||||
}
|
||||
|
||||
def start_run(self, sampler: Any) -> int:
|
||||
with self.lock:
|
||||
self.clear()
|
||||
self.run_serial += 1
|
||||
self.active_run_id = self.run_serial
|
||||
self.supported_sampler = sampler_name(sampler) == "sample_euler"
|
||||
return self.run_serial
|
||||
|
||||
def clear(self) -> None:
|
||||
self.branch_states.clear()
|
||||
self.step_occurrences.clear()
|
||||
self.active_step_id = None
|
||||
self.disabled_reason = None
|
||||
self.stats = self._new_stats()
|
||||
|
||||
def end_run(self, run_id: int) -> dict[str, int]:
|
||||
with self.lock:
|
||||
snapshot = dict(self.stats)
|
||||
if self.active_run_id == run_id:
|
||||
self.clear()
|
||||
self.active_run_id = None
|
||||
self.supported_sampler = False
|
||||
return snapshot
|
||||
|
||||
def begin_step(self, timestep: Any) -> int:
|
||||
del timestep
|
||||
with self.lock:
|
||||
if self.active_run_id is None:
|
||||
return -1
|
||||
self.active_step_id = 0 if self.active_step_id is None else self.active_step_id + 1
|
||||
self.step_occurrences.clear()
|
||||
return self.active_step_id
|
||||
|
||||
def new_model_call(self, options: dict[str, Any]) -> CacheCall:
|
||||
with self.lock:
|
||||
self.stats["model_calls"] += 1
|
||||
run_id = int(options.get(RUN_ID_KEY, -1))
|
||||
step_id = int(options.get(STEP_ID_KEY, -1))
|
||||
|
||||
cond_labels = normalized_items(options.get("cond_or_uncond"))
|
||||
uuids = normalized_items(options.get("uuids"))
|
||||
base_branch: tuple[Any, ...] | None
|
||||
if self.config.strict_branch_identity and not cond_labels:
|
||||
base_branch = None
|
||||
else:
|
||||
base_branch = (cond_labels or ("single",), uuids)
|
||||
|
||||
if base_branch is None:
|
||||
branch_key = None
|
||||
self.stats["identity_fallbacks"] += 1
|
||||
else:
|
||||
occurrence = self.step_occurrences.get(base_branch, 0)
|
||||
self.step_occurrences[base_branch] = occurrence + 1
|
||||
branch_key = (*base_branch, occurrence)
|
||||
|
||||
raw_signature = options.get(LORA_SIGNATURE_KEY)
|
||||
lora_signature = tuple(raw_signature) if raw_signature is not None else None
|
||||
|
||||
eligible = True
|
||||
reason = "eligible"
|
||||
if self.active_run_id != run_id or self.active_step_id != step_id:
|
||||
eligible, reason = False, "outside active Euler step"
|
||||
elif not self.supported_sampler:
|
||||
eligible, reason = False, "sampler is not stock Euler"
|
||||
elif self.disabled_reason is not None:
|
||||
eligible, reason = False, self.disabled_reason
|
||||
elif options.get("multigpu_thread_device") is not None:
|
||||
eligible, reason = False, "multi-GPU calls use the exact path"
|
||||
elif branch_key is None:
|
||||
eligible, reason = False, "CFG branch identity unavailable"
|
||||
elif self.config.require_fastpath_schedule and lora_signature is None:
|
||||
eligible, reason = False, "FastPath scheduled-LoRA signature unavailable"
|
||||
elif lora_signature is None:
|
||||
lora_signature = (("static-model", 1.0),)
|
||||
|
||||
prefetch_suppressed = False
|
||||
if eligible and self.config.suppress_candidate_prefetch:
|
||||
state = self.branch_states.get(branch_key)
|
||||
prefetch_suppressed = bool(
|
||||
state is not None
|
||||
and state.residual is not None
|
||||
and state.lora_signature == lora_signature
|
||||
and state.consecutive_reuses < self.config.max_consecutive_reuses
|
||||
and state.last_step_id != step_id
|
||||
)
|
||||
if prefetch_suppressed:
|
||||
self.stats["prefetch_suppressed_calls"] += 1
|
||||
|
||||
return CacheCall(
|
||||
runtime=self,
|
||||
run_id=run_id,
|
||||
step_id=step_id,
|
||||
branch_key=branch_key,
|
||||
lora_signature=lora_signature,
|
||||
eligible=eligible,
|
||||
reason=reason,
|
||||
prefetch_suppressed=prefetch_suppressed,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _topology(hidden: torch.Tensor, args: dict[str, Any]) -> tuple[Any, ...]:
|
||||
segments = tuple(tuple(int(value) for value in segment) for segment in args["mod_segments"])
|
||||
return (
|
||||
tuple(hidden.shape),
|
||||
str(hidden.dtype),
|
||||
str(hidden.device),
|
||||
segments,
|
||||
tuple(args["t_emb"].shape),
|
||||
tuple(args["rope_freqs"].shape),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _probe(hidden: torch.Tensor, max_rows: int = 96, max_columns: int = 128) -> torch.Tensor:
|
||||
flat = hidden.detach().reshape(-1, hidden.shape[-1])
|
||||
row_stride = max(1, math.ceil(flat.shape[0] / max_rows))
|
||||
col_stride = max(1, math.ceil(flat.shape[1] / max_columns))
|
||||
return flat[::row_stride, ::col_stride][:max_rows, :max_columns].float().clone()
|
||||
|
||||
@staticmethod
|
||||
def _probe_delta(current: torch.Tensor, previous: torch.Tensor) -> float:
|
||||
denominator = previous.abs().mean().clamp_min_(1.0e-6)
|
||||
return float(((current - previous).abs().mean() / denominator).item())
|
||||
|
||||
def begin_middle(self, call: CacheCall, hidden: torch.Tensor, args: dict[str, Any]) -> None:
|
||||
call.middle_started = True
|
||||
if not call.eligible:
|
||||
return
|
||||
|
||||
topology = self._topology(hidden, args)
|
||||
probe = self._probe(hidden)
|
||||
call.current_probe = probe
|
||||
|
||||
with self.lock:
|
||||
state = self.branch_states.setdefault(call.branch_key, BranchState())
|
||||
call.state = state
|
||||
|
||||
if state.topology != topology:
|
||||
state.reset(topology=topology, lora_signature=call.lora_signature)
|
||||
call.reason = "topology changed or cache empty"
|
||||
elif state.lora_signature != call.lora_signature:
|
||||
state.reset(topology=topology, lora_signature=call.lora_signature)
|
||||
self.stats["strength_resets"] += 1
|
||||
call.reason = "LoRA strength signature changed"
|
||||
|
||||
can_reuse = (
|
||||
state.residual is not None
|
||||
and state.prefix_probe is not None
|
||||
and state.consecutive_reuses < self.config.max_consecutive_reuses
|
||||
and state.last_step_id != call.step_id
|
||||
)
|
||||
if can_reuse:
|
||||
delta = self._probe_delta(probe, state.prefix_probe)
|
||||
call.feature_delta = delta
|
||||
if delta <= self.config.reuse_threshold:
|
||||
call.reuse = True
|
||||
state.consecutive_reuses += 1
|
||||
state.last_step_id = call.step_id
|
||||
self.stats["reused_calls"] += 1
|
||||
self.stats["skipped_blocks"] += self.config.middle_block_count
|
||||
call.reason = "middle residual reused"
|
||||
return
|
||||
self.stats["gate_rejections"] += 1
|
||||
call.reason = f"prefix feature delta {delta:.5f} exceeded gate"
|
||||
elif state.residual is not None and state.consecutive_reuses >= self.config.max_consecutive_reuses:
|
||||
call.reason = "maximum consecutive reuse reached"
|
||||
|
||||
try:
|
||||
call.start_hidden = hidden.detach().clone()
|
||||
except RuntimeError as error:
|
||||
if "out of memory" not in str(error).lower():
|
||||
raise
|
||||
with self.lock:
|
||||
self.disabled_reason = "cache allocation ran out of memory"
|
||||
self.branch_states.clear()
|
||||
call.eligible = False
|
||||
call.reason = self.disabled_reason
|
||||
LOG.warning("MiniMax H3 FastPath cache disabled for this run: %s", error)
|
||||
|
||||
def apply_reuse(self, call: CacheCall, hidden: torch.Tensor) -> torch.Tensor:
|
||||
state = call.state
|
||||
if state is None or state.residual is None:
|
||||
raise RuntimeError("FastPath cache selected reuse without a residual")
|
||||
residual = state.residual
|
||||
if residual.device != hidden.device or residual.dtype != hidden.dtype:
|
||||
residual = residual.to(device=hidden.device, dtype=hidden.dtype, non_blocking=True)
|
||||
hidden.add_(residual)
|
||||
return hidden
|
||||
|
||||
def finish_middle(self, call: CacheCall, output: torch.Tensor) -> None:
|
||||
if not call.eligible or call.reuse or call.start_hidden is None or call.state is None:
|
||||
call.completed = True
|
||||
return
|
||||
try:
|
||||
residual = output.detach().clone()
|
||||
residual.sub_(call.start_hidden)
|
||||
if self.config.cache_device == "cpu":
|
||||
residual = residual.to(device="cpu")
|
||||
except RuntimeError as error:
|
||||
if "out of memory" not in str(error).lower():
|
||||
raise
|
||||
with self.lock:
|
||||
self.disabled_reason = "cache residual allocation ran out of memory"
|
||||
self.branch_states.clear()
|
||||
call.completed = True
|
||||
LOG.warning("MiniMax H3 FastPath cache disabled for this run: %s", error)
|
||||
return
|
||||
|
||||
with self.lock:
|
||||
call.state.residual = residual
|
||||
call.state.prefix_probe = call.current_probe
|
||||
call.state.consecutive_reuses = 0
|
||||
call.state.last_step_id = call.step_id
|
||||
self.stats["cached_full_calls"] += 1
|
||||
call.start_hidden = None
|
||||
call.completed = True
|
||||
|
||||
def finish_model_call(self, call: CacheCall) -> None:
|
||||
call.start_hidden = None
|
||||
call.current_probe = None
|
||||
|
||||
|
||||
def sampler_name(sampler: Any) -> str:
|
||||
function = getattr(sampler, "sampler_function", None)
|
||||
return str(getattr(function, "__name__", type(sampler).__name__))
|
||||
|
||||
|
||||
def _binding_from_model_options(model_options: dict[str, Any] | None) -> CacheBinding | None:
|
||||
binding = (model_options or {}).get(CACHE_BINDING_KEY)
|
||||
return binding if isinstance(binding, CacheBinding) else None
|
||||
|
||||
|
||||
def outer_sample_wrapper(executor, *args, **kwargs):
|
||||
guider = executor.class_obj
|
||||
binding = _binding_from_model_options(getattr(guider, "model_options", None))
|
||||
if binding is None:
|
||||
return executor(*args, **kwargs)
|
||||
|
||||
sampler = args[2]
|
||||
runtime = binding.runtime
|
||||
run_id = runtime.start_run(sampler)
|
||||
if runtime.config.verbose:
|
||||
LOG.warning(
|
||||
"MiniMax H3 FastPath cache run start: sampler=%s supported=%s middle=%d..%d",
|
||||
sampler_name(sampler),
|
||||
runtime.supported_sampler,
|
||||
runtime.config.first_middle_block,
|
||||
runtime.config.last_middle_block,
|
||||
)
|
||||
try:
|
||||
return executor(*args, **kwargs)
|
||||
finally:
|
||||
snapshot = dict(runtime.stats)
|
||||
if runtime.config.verbose or snapshot["reused_calls"]:
|
||||
LOG.warning(
|
||||
"MiniMax H3 FastPath cache summary: reused_calls=%d cached_full_calls=%d "
|
||||
"skipped_blocks=%d gate_rejections=%d strength_resets=%d",
|
||||
snapshot["reused_calls"],
|
||||
snapshot["cached_full_calls"],
|
||||
snapshot["skipped_blocks"],
|
||||
snapshot["gate_rejections"],
|
||||
snapshot["strength_resets"],
|
||||
)
|
||||
runtime.end_run(run_id)
|
||||
|
||||
|
||||
def predict_noise_wrapper(executor, x, timestep, model_options=None, seed=None):
|
||||
guider = executor.class_obj
|
||||
binding = _binding_from_model_options(getattr(guider, "model_options", None))
|
||||
if binding is None or binding.runtime.active_run_id is None:
|
||||
return executor(x, timestep, model_options or {}, seed)
|
||||
|
||||
runtime = binding.runtime
|
||||
step_id = runtime.begin_step(timestep)
|
||||
copied = dict(model_options or {})
|
||||
transformer_options = dict(copied.get("transformer_options") or {})
|
||||
copied["transformer_options"] = transformer_options
|
||||
transformer_options[CACHE_BINDING_KEY] = binding
|
||||
transformer_options[RUN_ID_KEY] = runtime.active_run_id
|
||||
transformer_options[STEP_ID_KEY] = step_id
|
||||
return executor(x, timestep, copied, seed)
|
||||
|
||||
|
||||
def diffusion_model_wrapper(
|
||||
executor,
|
||||
x,
|
||||
timestep,
|
||||
context,
|
||||
transformer_options=None,
|
||||
minimax_payload=None,
|
||||
**kwargs,
|
||||
):
|
||||
options = dict(transformer_options or {})
|
||||
binding = options.get(CACHE_BINDING_KEY)
|
||||
if not isinstance(binding, CacheBinding):
|
||||
return executor(
|
||||
x,
|
||||
timestep,
|
||||
context,
|
||||
options,
|
||||
minimax_payload=minimax_payload,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
call = binding.runtime.new_model_call(options)
|
||||
if call.prefetch_suppressed:
|
||||
# Dynamic ModelPatcher prefetch otherwise streams every middle block even
|
||||
# when its replacement returns without executing the block. Candidate
|
||||
# calls turn it off before MiniMax constructs the local prefetch queue.
|
||||
options["prefetch_dynamic_vbars"] = False
|
||||
options[CACHE_CALL_KEY] = call
|
||||
try:
|
||||
return executor(
|
||||
x,
|
||||
timestep,
|
||||
context,
|
||||
options,
|
||||
minimax_payload=minimax_payload,
|
||||
**kwargs,
|
||||
)
|
||||
finally:
|
||||
binding.runtime.finish_model_call(call)
|
||||
|
||||
|
||||
def _execute_previous(previous_patch, args, extra_options):
|
||||
if previous_patch is None:
|
||||
result = extra_options["original_block"](args)
|
||||
else:
|
||||
result = previous_patch(args, extra_options)
|
||||
if not isinstance(result, dict) or "img" not in result:
|
||||
raise TypeError("A MiniMax H3 double_block replacement must return a dict containing 'img'")
|
||||
return result
|
||||
|
||||
|
||||
def make_middle_block_patch(block_index: int, config: CacheConfig, previous_patch=None):
|
||||
is_first = block_index == config.first_middle_block
|
||||
is_last = block_index == config.last_middle_block
|
||||
|
||||
def middle_block_patch(args, extra_options):
|
||||
call = (args.get("transformer_options") or {}).get(CACHE_CALL_KEY)
|
||||
if not isinstance(call, CacheCall):
|
||||
return _execute_previous(previous_patch, args, extra_options)
|
||||
|
||||
hidden = args["img"]
|
||||
if is_first:
|
||||
call.runtime.begin_middle(call, hidden, args)
|
||||
if call.reuse:
|
||||
return {"img": call.runtime.apply_reuse(call, hidden)}
|
||||
elif call.reuse:
|
||||
return {"img": hidden}
|
||||
|
||||
result = _execute_previous(previous_patch, args, extra_options)
|
||||
if is_last:
|
||||
call.runtime.finish_middle(call, result["img"])
|
||||
return result
|
||||
|
||||
return middle_block_patch
|
||||
|
||||
|
||||
def model_clone_callback(source_model: Any, cloned_model: Any) -> None:
|
||||
source_binding = _binding_from_model_options(getattr(source_model, "model_options", None))
|
||||
if source_binding is None:
|
||||
return
|
||||
cloned_model.model_options[CACHE_BINDING_KEY] = CacheBinding(
|
||||
MiddleCacheRuntime(source_binding.runtime.config)
|
||||
)
|
||||
|
||||
|
||||
def _existing_dit_replacements(model) -> dict[Any, Any]:
|
||||
return (
|
||||
model.model_options.get("transformer_options", {})
|
||||
.get("patches_replace", {})
|
||||
.get("dit", {})
|
||||
)
|
||||
|
||||
|
||||
def _install_runtime(model, runtime: MiddleCacheRuntime) -> None:
|
||||
model.model_options[CACHE_BINDING_KEY] = CacheBinding(runtime)
|
||||
wrappers = comfy.patcher_extension.WrappersMP
|
||||
callbacks = comfy.patcher_extension.CallbacksMP
|
||||
|
||||
if not model.get_wrappers(wrappers.OUTER_SAMPLE, CACHE_WRAPPER_KEY):
|
||||
model.add_wrapper_with_key(wrappers.OUTER_SAMPLE, CACHE_WRAPPER_KEY, outer_sample_wrapper)
|
||||
if not model.get_wrappers(wrappers.PREDICT_NOISE, CACHE_WRAPPER_KEY):
|
||||
model.add_wrapper_with_key(wrappers.PREDICT_NOISE, CACHE_WRAPPER_KEY, predict_noise_wrapper)
|
||||
if not model.get_wrappers(wrappers.DIFFUSION_MODEL, CACHE_WRAPPER_KEY):
|
||||
model.add_wrapper_with_key(
|
||||
wrappers.DIFFUSION_MODEL,
|
||||
CACHE_WRAPPER_KEY,
|
||||
diffusion_model_wrapper,
|
||||
)
|
||||
if not model.get_callbacks(callbacks.ON_CLONE, CACHE_WRAPPER_KEY):
|
||||
model.add_callback_with_key(
|
||||
callbacks.ON_CLONE,
|
||||
CACHE_WRAPPER_KEY,
|
||||
model_clone_callback,
|
||||
)
|
||||
|
||||
|
||||
class MiniMaxH3EulerMiddleCache:
|
||||
"""Keep H3 prefix/suffix exact and cache only a configurable middle block range."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"enabled": ("BOOLEAN", {"default": True}),
|
||||
"prefix_blocks": (
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 48,
|
||||
"step": 1,
|
||||
"tooltip": "Exact blocks run before the cached middle on every model call.",
|
||||
},
|
||||
),
|
||||
"suffix_blocks": (
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 48,
|
||||
"step": 1,
|
||||
"tooltip": "Exact final blocks run after the cached middle on every model call.",
|
||||
},
|
||||
),
|
||||
"reuse_threshold": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.12,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.005,
|
||||
"tooltip": "Maximum relative change in a sampled prefix feature before reuse is rejected.",
|
||||
},
|
||||
),
|
||||
"max_consecutive_reuses": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"step": 1,
|
||||
"tooltip": "One gives FULL/PART/FULL for three calls at the same LoRA strength.",
|
||||
},
|
||||
),
|
||||
"cache_device": (
|
||||
["gpu", "cpu"],
|
||||
{
|
||||
"tooltip": "GPU is fastest. CPU saves VRAM but transfers one full hidden residual on reuse.",
|
||||
},
|
||||
),
|
||||
"require_fastpath_schedule": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Refuse approximation unless the matching FastPath scheduled-LoRA signature is present.",
|
||||
},
|
||||
),
|
||||
"strict_branch_identity": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Use the exact path if ComfyUI does not provide CFG branch labels.",
|
||||
},
|
||||
),
|
||||
"suppress_candidate_prefetch": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "On likely reuse calls, stop dynamic INT8 weight prefetch so skipped blocks are not streamed anyway.",
|
||||
},
|
||||
),
|
||||
"verbose": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "MiniMax H3/FastPath"
|
||||
DESCRIPTION = (
|
||||
"Stock-Euler-only, LoRA-strength-aware H3 cache. Runs a prefix and suffix exactly; on accepted "
|
||||
"calls it reuses only the residual of the middle blocks. Cache state is isolated per CFG branch."
|
||||
)
|
||||
|
||||
def patch(
|
||||
self,
|
||||
model,
|
||||
enabled,
|
||||
prefix_blocks,
|
||||
suffix_blocks,
|
||||
reuse_threshold,
|
||||
max_consecutive_reuses,
|
||||
cache_device,
|
||||
require_fastpath_schedule,
|
||||
strict_branch_identity,
|
||||
suppress_candidate_prefetch,
|
||||
verbose,
|
||||
):
|
||||
patched = model.clone()
|
||||
if not enabled:
|
||||
return (patched,)
|
||||
if patched.get_attachment(CACHE_ATTACHMENT_KEY) is not None:
|
||||
raise RuntimeError("MiniMax H3 FastPath middle cache is already installed on this MODEL")
|
||||
if patched.model_options.get("spectrum_h3_binding") is not None:
|
||||
raise RuntimeError(
|
||||
"Do not stack Spectrum forecasting with FastPath middle caching. Disable Spectrum and use one approximation method."
|
||||
)
|
||||
|
||||
try:
|
||||
diffusion_model = patched.get_model_object("diffusion_model")
|
||||
except Exception as error:
|
||||
raise TypeError("MODEL does not expose a diffusion_model") from error
|
||||
if type(diffusion_model).__name__ != "MiniMaxH3Model" or not hasattr(diffusion_model, "blocks"):
|
||||
raise TypeError("MiniMax H3 FastPath middle cache only supports the native MiniMaxH3Model")
|
||||
|
||||
block_count = len(diffusion_model.blocks)
|
||||
prefix_blocks = int(prefix_blocks)
|
||||
suffix_blocks = int(suffix_blocks)
|
||||
if prefix_blocks + suffix_blocks >= block_count:
|
||||
raise ValueError(
|
||||
f"prefix_blocks + suffix_blocks must be below the model's {block_count} blocks"
|
||||
)
|
||||
first_middle = prefix_blocks
|
||||
last_middle = block_count - suffix_blocks - 1
|
||||
|
||||
existing = dict(_existing_dit_replacements(patched))
|
||||
if ("block_loop", 0) in existing:
|
||||
raise RuntimeError(
|
||||
"A whole H3 block-loop replacement is already installed. Remove the older MiniMax H3 cache before FastPath."
|
||||
)
|
||||
|
||||
config = CacheConfig(
|
||||
first_middle_block=first_middle,
|
||||
last_middle_block=last_middle,
|
||||
reuse_threshold=float(reuse_threshold),
|
||||
max_consecutive_reuses=int(max_consecutive_reuses),
|
||||
cache_device=str(cache_device),
|
||||
require_fastpath_schedule=bool(require_fastpath_schedule),
|
||||
strict_branch_identity=bool(strict_branch_identity),
|
||||
suppress_candidate_prefetch=bool(suppress_candidate_prefetch),
|
||||
verbose=bool(verbose),
|
||||
)
|
||||
runtime = MiddleCacheRuntime(config)
|
||||
_install_runtime(patched, runtime)
|
||||
|
||||
for block_index in range(first_middle, last_middle + 1):
|
||||
previous = existing.get(("double_block", block_index))
|
||||
patched.set_model_patch_replace(
|
||||
make_middle_block_patch(block_index, config, previous),
|
||||
"dit",
|
||||
"double_block",
|
||||
block_index,
|
||||
)
|
||||
|
||||
patched.set_attachments(
|
||||
CACHE_ATTACHMENT_KEY,
|
||||
{
|
||||
"prefix_blocks": prefix_blocks,
|
||||
"middle_blocks": config.middle_block_count,
|
||||
"suffix_blocks": suffix_blocks,
|
||||
"reuse_threshold": config.reuse_threshold,
|
||||
"max_consecutive_reuses": config.max_consecutive_reuses,
|
||||
"cache_device": config.cache_device,
|
||||
},
|
||||
)
|
||||
LOG.info(
|
||||
"MiniMax H3 FastPath middle cache installed: %d exact prefix / %d cached middle / %d exact suffix",
|
||||
prefix_blocks,
|
||||
config.middle_block_count,
|
||||
suffix_blocks,
|
||||
)
|
||||
return (patched,)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BranchState",
|
||||
"CacheBinding",
|
||||
"CacheCall",
|
||||
"CacheConfig",
|
||||
"MiddleCacheRuntime",
|
||||
"MiniMaxH3EulerMiddleCache",
|
||||
"diffusion_model_wrapper",
|
||||
"make_middle_block_patch",
|
||||
"model_clone_callback",
|
||||
"outer_sample_wrapper",
|
||||
"predict_noise_wrapper",
|
||||
]
|
||||
Reference in New Issue
Block a user