705 lines
26 KiB
Python
705 lines
26 KiB
Python
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",
|
|
]
|