From f0bbd8ecbc4f87efd826d2734124209175471caf Mon Sep 17 00:00:00 2001 From: capitan01R Date: Tue, 18 Aug 2026 06:32:45 -0400 Subject: [PATCH] Add files via upload --- __init__.py | 12 + common.py | 117 ++++++++ middle_cache.py | 704 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 833 insertions(+) create mode 100644 __init__.py create mode 100644 common.py create mode 100644 middle_cache.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6742e9d --- /dev/null +++ b/__init__.py @@ -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"] diff --git a/common.py b/common.py new file mode 100644 index 0000000..c7d98e6 --- /dev/null +++ b/common.py @@ -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) diff --git a/middle_cache.py b/middle_cache.py new file mode 100644 index 0000000..7b4d7a6 --- /dev/null +++ b/middle_cache.py @@ -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", +]