Files
2026-09-11 17:55:08 +08:00

686 lines
27 KiB
Python

"""UniBlockSwap: DynamicVRAM bypass switch for ComfyUI >= 0.35.
"""
import contextlib
import logging
import os
log = logging.getLogger("UniBlockSwap.LegacyVRAM")
__all__ = [
"SWITCH", "MODE", "ENV", "NO_PIN", "NO_PIN_ENV", "NODE_TYPES",
"apply", "set_bypass", "restore", "bypass_scope", "mode",
"status", "runtime_report", "description_suffix",
]
# ---------------------------------------------------------------------------
# 开关
# ---------------------------------------------------------------------------
SWITCH = True # 总开关 (False = 完全不做任何事)
MODE = "auto" # "auto" = 只在工作流用到本插件节点时才绕过 (默认)
# "always" = 老行为: 进程启动就绕过 (整机生效)
ENV = "UNIBLOCKSWAP_LEGACY_VRAM" # 环境变量覆盖 (可选)
MODE_ENV = ENV # 同上, 语义化别名
# 本插件的节点 class_type (auto 模式靠它扫图; 新增节点记得加进来)
NODE_TYPES = ("UniBlockSwap", "UniBlockSwapTE", "UniBlockSwapCacheControl")
# 产出模型对象的节点 (auto 模式需要强制它们重新执行, 见 _bust_model_reload)
MODEL_OUTPUT_TYPES = ("MODEL", "CLIP", "VAE")
BUST_KEY = "__uniblockswap_reload__"
# 可选: 同时清零 pinned memory 预算 (Windows 上"共享 GPU 内存"增长的来源)。
# 会让"共享显存"数字变干净, 但 H2D 传输会慢一点。
NO_PIN = False
NO_PIN_ENV = "UNIBLOCKSWAP_LEGACY_VRAM_NO_PIN"
MIN_VERSION = (0, 35) # 从 0.35 起 DynamicVRAM 默认开启
MIN_VERSION_STR = "0.35"
_ON = ("1", "true", "yes", "on", "enable", "enabled", "open", "always", "force")
_OFF = ("0", "false", "no", "off", "disable", "disabled", "close", "never")
_DRY = ("soft", "dry", "dry-run", "dryrun", "test", "check")
_AUTO = ("auto", "smart", "graph", "on-demand", "ondemand", "prompt", "node")
_state = {"applied": False, "reason": "not evaluated", "version": "unknown", "changes": []}
# /prompt 钩子的状态
_gate = {"installed": False, "last": None, "bannered": False, "seq": 0, "busted": 0}
# 首次 flip 之前的全局量快照 (回退目标 = 启动时的原值)
_original = None
# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------
def _env(name, default=""):
return str(os.environ.get(name, default)).strip().lower()
def mode():
"""当前模式: "off" / "soft" / "auto" / "always" (环境变量优先于代码里的 SWITCH/MODE)。"""
env = _env(ENV)
if env in _OFF:
return "off"
if env in _DRY:
return "soft"
if env in _AUTO:
return "auto"
if env in _ON:
return "always"
if not SWITCH:
return "off"
text = str(MODE).strip().lower()
return text if text in ("auto", "always") else "auto"
def switch_enabled():
"""开关是否开启 (环境变量优先于代码里的 SWITCH)。"""
return mode() != "off"
def dry_run():
return _env(ENV) in _DRY
def no_pin_enabled():
env = _env(NO_PIN_ENV)
if env in _ON:
return True
if env in _OFF:
return False
return bool(NO_PIN)
def comfyui_version():
try:
import comfyui_version
return str(getattr(comfyui_version, "__version__", "unknown"))
except Exception:
return "unknown"
def _parse_version(text):
parts = []
for chunk in str(text).split("."):
digits = ""
for ch in chunk:
if ch.isdigit():
digits += ch
else:
break
if not digits:
break
parts.append(int(digits))
return tuple(parts)
def version_ok():
"""ComfyUI 版本是否 >= 0.35 (拿不到版本号时退化为能力探测)。"""
parsed = _parse_version(comfyui_version())
if len(parsed) >= 2:
return parsed >= MIN_VERSION
import comfy.model_patcher as model_patcher
return hasattr(model_patcher, "ModelPatcherDynamic")
def dynamic_active():
"""当前进程是否还处于 DynamicVRAM 模式。"""
import comfy.memory_management as memory_management
import comfy.model_patcher as model_patcher
core = getattr(model_patcher, "CoreModelPatcher", None)
if core is not None and core is not getattr(model_patcher, "ModelPatcher", None):
return True
return bool(getattr(memory_management, "aimdo_enabled", False))
# ---------------------------------------------------------------------------
# auto 模式: /prompt 钩子 (只在工作流真的用到本插件节点时才绕过)
# ---------------------------------------------------------------------------
def _graph_has_node(prompt):
"""递归扫描工作流图, 判断有没有本插件的节点 (含嵌套 subgraph 定义)。"""
if not isinstance(prompt, dict):
return False
stack = [prompt]
while stack:
obj = stack.pop()
if isinstance(obj, dict):
cls = obj.get("class_type")
if isinstance(cls, str) and cls in NODE_TYPES:
return True
stack.extend(obj.values())
elif isinstance(obj, list):
stack.extend(obj)
return False
def _pending_uses_node():
"""队列里正在执行 / 排队等待的图有没有用到本插件节点。"""
try:
import server
srv = getattr(server.PromptServer, "instance", None)
queue = getattr(srv, "prompt_queue", None)
if queue is None:
return False
running, queued = queue.get_current_queue_volatile()
for item in list(running) + list(queued):
try:
if _graph_has_node(item[2]):
return True
except Exception:
continue
return False
except Exception:
return False
def _loader_class_types():
"""产出 MODEL/CLIP/VAE 的 class_type 集合 = 模型对象的来源节点。"""
try:
import nodes
mapping = getattr(nodes, "NODE_CLASS_MAPPINGS", {}) or {}
except Exception:
return frozenset()
found = set()
for cls_name, cls in mapping.items():
try:
types = getattr(cls, "RETURN_TYPES", ()) or ()
except Exception:
continue
if isinstance(types, str):
types = (types,)
if any(t in MODEL_OUTPUT_TYPES for t in types):
found.add(cls_name)
return frozenset(found)
def _bust_model_reload(prompt, token):
"""让"会产出模型"的节点缓存失效, 迫使它们重新执行 (返回命中的节点数)。
为什么必须这么做 —— ComfyUI 的节点输出缓存是**跨提交保留**的
(execution.py:670-673, reset() 只在 __init__ 调用), 缓存键 = 输入签名。
而 CheckpointLoaderSimple 这类 loader **没有 IS_CHANGED** ->
IsChangedCache.get() 直接返回 False (execution.py:82-84) -> 只要 ckpt_name
没变, 签名恒定 (caching.py:109-127) => 缓存命中 => **loader 根本不执行**,
上一轮那个 dynamic patcher 被原样复用, 我们翻转的全局量对它完全无效。
手段: 给它的 inputs 塞一个多余的常量键 -> 签名变化 -> 缓存不命中 -> 重跑
-> 重新构造 patcher (此时全局量已是 legacy)。下游节点因为签名包含 ancestors
(caching.py:125) 会自动一起失效。
为什么安全 (已对真模块实测):
- validate_inputs 只遍历节点自己声明的输入 -> 多余键不参与校验 (execution.py:896-898)
- get_input_data 对未声明的输入 key 直接忽略 -> 不会传给节点函数 (execution.py:174-190)
- 值必须是**常量**: 长度 2 的 list 会被当成连线 link 去解析上游输出
"""
loaders = _loader_class_types()
if not loaders:
return 0
hits = 0
stack = [prompt]
while stack:
obj = stack.pop()
if isinstance(obj, dict):
cls = obj.get("class_type")
if isinstance(cls, str) and cls in loaders:
inputs = obj.get("inputs")
if isinstance(inputs, dict):
inputs[BUST_KEY] = token
hits += 1
stack.extend(obj.values())
elif isinstance(obj, list):
stack.extend(obj)
return hits
def _install_gate():
"""装上 /prompt 钩子 (官方 add_on_prompt_handler)。失败只记日志, 不抛。"""
if _gate["installed"]:
return True
try:
import server
srv = getattr(server.PromptServer, "instance", None)
if srv is None or not hasattr(srv, "add_on_prompt_handler"):
return False
handlers = getattr(srv, "on_prompt_handlers", None)
if isinstance(handlers, list) and _on_prompt in handlers:
_gate["installed"] = True
return True
srv.add_on_prompt_handler(_on_prompt)
_gate["installed"] = True
return True
except Exception:
log.warning("[UniBlockSwap] could not install the /prompt gate", exc_info=True)
return False
def _on_prompt(json_data):
"""每次提交工作流时决定要不要绕过。
⚠️ server 的 trigger_on_prompt 会把返回值当成新的 json_data, 必须原样返回。
"""
try:
if mode() != "auto":
return json_data
prompt = json_data.get("prompt") if isinstance(json_data, dict) else None
used = _graph_has_node(prompt) or _pending_uses_node()
_gate["last"] = bool(used)
if used:
changed = set_bypass(True, quiet=True,
reason="auto: workflow uses %s" % "/".join(NODE_TYPES))
if changed:
# 刚翻过去 -> 换一个 token, 让上一轮的模型缓存彻底失效
_gate["seq"] += 1
if _graph_has_node(prompt):
# 本轮图真的用到本插件节点: 必须保证模型**重新构造**。模型通常在上游
# loader 里早就建好了(而且很可能命中跨提交的输出缓存), 光翻全局量对它
# 没用。token 只在翻转时递增 -> 未翻转的后续提交用同一个 token, 缓存键
# 稳定, 可以正常命中"已经按 legacy 构造好"的那份。
forced = _bust_model_reload(prompt, str(_gate["seq"]))
_gate["busted"] = forced
if changed and forced:
log.info("[UniBlockSwap] forced %d model loader(s) to re-run so the model "
"is rebuilt under the legacy path", forced)
if changed:
if not _gate["bannered"]:
_gate["bannered"] = True
_banner(
[_rule()]
+ [" UniBlockSwap: DynamicVRAM BYPASSED (mode=auto)"]
+ [" " + c for c in changed]
+ [" why : this workflow (or one still in the queue) uses "
+ "/".join(NODE_TYPES),
" scope: models built from now on. To undo WITHOUT restarting a",
" workflow: submit a workflow that does not use this node",
" (the bypass is restored automatically before loading).",
" Already-built models cannot be converted - but any loader",
" that still has to run will rebuild them the right way."]
+ [_rule()]
)
elif _state["applied"]:
changed = set_bypass(False, quiet=True,
reason="auto: no UniBlockSwap node in this workflow")
if changed:
_gate["busted"] = 0
log.info("[UniBlockSwap] workflow does not use this node -> DynamicVRAM "
"restored for models loaded from now on")
except Exception:
log.warning("[UniBlockSwap] /prompt gate failed", exc_info=True)
return json_data
# ---------------------------------------------------------------------------
# 回退机制: 快照 / 应用 / 还原
# ---------------------------------------------------------------------------
def _snapshot():
"""记下当前(通常是启动时)的全局量, 作为 restore() 的目标。"""
import comfy.memory_management as memory_management
import comfy.model_patcher as model_patcher
from comfy.cli_args import args
snap = {
"CoreModelPatcher": model_patcher.CoreModelPatcher,
"aimdo_enabled": bool(getattr(memory_management, "aimdo_enabled", False)),
"disable_dynamic_vram": bool(getattr(args, "disable_dynamic_vram", False)),
"disable_pinned_memory": bool(getattr(args, "disable_pinned_memory", False)),
"MAX_PINNED_MEMORY": None,
}
try:
import comfy.model_management as model_management
snap["MAX_PINNED_MEMORY"] = int(getattr(model_management, "MAX_PINNED_MEMORY", 0))
except Exception:
pass
return snap
def _apply_legacy(no_pin):
"""把全局量翻成 legacy 口径, 返回实际改动项。"""
import comfy.memory_management as memory_management
import comfy.model_patcher as model_patcher
from comfy.cli_args import args
changed = []
if getattr(memory_management, "aimdo_enabled", False):
memory_management.aimdo_enabled = False
changed.append("comfy.memory_management.aimdo_enabled = False")
if model_patcher.CoreModelPatcher is not model_patcher.ModelPatcher:
model_patcher.CoreModelPatcher = model_patcher.ModelPatcher
changed.append("comfy.model_patcher.CoreModelPatcher = ModelPatcher")
if not getattr(args, "disable_dynamic_vram", False):
args.disable_dynamic_vram = True
changed.append("args.disable_dynamic_vram = True")
if no_pin:
try:
import comfy.model_management as model_management
if getattr(model_management, "MAX_PINNED_MEMORY", 0) > 0:
model_management.MAX_PINNED_MEMORY = 0
changed.append("comfy.model_management.MAX_PINNED_MEMORY = 0")
except Exception:
pass
if not getattr(args, "disable_pinned_memory", False):
args.disable_pinned_memory = True
changed.append("args.disable_pinned_memory = True")
return changed
def _restore_state(snap):
"""按快照还原全局量, 返回实际改动项。"""
import comfy.memory_management as memory_management
import comfy.model_patcher as model_patcher
from comfy.cli_args import args
changed = []
if bool(getattr(memory_management, "aimdo_enabled", False)) != snap["aimdo_enabled"]:
memory_management.aimdo_enabled = snap["aimdo_enabled"]
changed.append("comfy.memory_management.aimdo_enabled = %s" % snap["aimdo_enabled"])
if model_patcher.CoreModelPatcher is not snap["CoreModelPatcher"]:
model_patcher.CoreModelPatcher = snap["CoreModelPatcher"]
changed.append("comfy.model_patcher.CoreModelPatcher = %s"
% getattr(snap["CoreModelPatcher"], "__name__", "?"))
if bool(getattr(args, "disable_dynamic_vram", False)) != snap["disable_dynamic_vram"]:
args.disable_dynamic_vram = snap["disable_dynamic_vram"]
changed.append("args.disable_dynamic_vram = %s" % snap["disable_dynamic_vram"])
if bool(getattr(args, "disable_pinned_memory", False)) != snap["disable_pinned_memory"]:
args.disable_pinned_memory = snap["disable_pinned_memory"]
changed.append("args.disable_pinned_memory = %s" % snap["disable_pinned_memory"])
if snap["MAX_PINNED_MEMORY"] is not None:
try:
import comfy.model_management as model_management
if int(getattr(model_management, "MAX_PINNED_MEMORY", 0)) != snap["MAX_PINNED_MEMORY"]:
model_management.MAX_PINNED_MEMORY = snap["MAX_PINNED_MEMORY"]
changed.append("comfy.model_management.MAX_PINNED_MEMORY = %d"
% snap["MAX_PINNED_MEMORY"])
except Exception:
pass
return changed
def set_bypass(enabled=True, no_pin=None, quiet=False, reason=""):
"""运行时切换绕过态 (幂等)。返回实际改动项 list, 空 list = 已经是目标态。
对**之后构造**的模型立即生效; 已构造的模型不受影响 (想让它们也变,
需要让它们重新加载: comfy 的 /free 卸载模型、换 checkpoint、或重启)。
"""
global _original
if _original is None:
_original = _snapshot()
if enabled:
if no_pin is None:
no_pin = no_pin_enabled()
changed = _apply_legacy(no_pin)
else:
changed = _restore_state(_original)
_state["applied"] = not dynamic_active()
_state["changes"] = changed
_state["reason"] = reason or ("runtime: bypass on" if enabled else "runtime: bypass off")
if changed and not quiet:
title = ("DynamicVRAM BYPASSED (runtime)" if enabled
else "DynamicVRAM bypass DISABLED (runtime revert)")
extra = ([" note : this only affects models constructed from now on;",
" already-built models stay on the legacy path - unload them",
" (comfy /free, or switch checkpoint) or restart to get back."]
if not enabled else
[" note : only models constructed from now on are affected;"])
_banner([_rule()]
+ [" UniBlockSwap: " + title]
+ [" " + c for c in changed]
+ extra
+ [_rule()])
return changed
def restore():
"""运行时回退: 回到启动时的原值 (== set_bypass(False))。"""
return set_bypass(False)
@contextlib.contextmanager
def bypass_scope(enabled=True, no_pin=None, quiet=True):
"""上下文管理器: 域内**新构造**的模型走 (或不走) legacy, 退出时还原。
with legacy_vram.bypass_scope(): # eval / 脚本化加载时用
model = comfy.sd.load_diffusion_model(path)
退出时按"进入前"的快照还原 (内部 try/finally, 体内抛异常也照样还原),
嵌套安全。注意它只能影响域内新建的模型。
"""
snap = _snapshot()
applied_before = _state["applied"]
reason_before = _state["reason"]
try:
set_bypass(enabled, no_pin=no_pin, quiet=quiet,
reason="bypass_scope(%s)" % bool(enabled))
yield
finally:
try:
_restore_state(snap)
_state["applied"] = applied_before
_state["reason"] = reason_before
except Exception:
log.warning("[UniBlockSwap] bypass_scope restore failed", exc_info=True)
def status():
return {
"switch": switch_enabled(),
"mode": mode(),
"dry_run": dry_run(),
"version": comfyui_version(),
"version_ok": version_ok(),
"dynamic_active": dynamic_active(),
"applied": _state["applied"],
"reason": _state["reason"],
"changes": list(_state["changes"]),
"revertable": _original is not None,
"gate_installed": bool(_gate["installed"]),
"last_prompt_used_node": _gate["last"],
"last_forced_reload_nodes": _gate["busted"],
}
def _banner(lines):
"""醒目的控制台提示 (绝不抛异常, 免得 custom node 导入失败)。"""
for line in lines:
try:
log.warning("%s", line)
except Exception:
pass
try:
print("\n".join(lines), flush=True)
except Exception:
pass
def _rule(char="="):
return char * 74
def description_suffix():
"""给节点 DESCRIPTION 用的状态后缀 (UI 上可见)。"""
try:
if _state["applied"]:
return " [DynamicVRAM bypass: ON]"
current = mode()
if current == "auto" and dynamic_active():
return " [DynamicVRAM bypass: AUTO - 只在本节点被用到时才生效]"
if current == "always" and dynamic_active():
return " [DynamicVRAM bypass: 未生效, 需重启]"
return ""
except Exception:
return ""
# ---------------------------------------------------------------------------
# 生效
# ---------------------------------------------------------------------------
def _apply_impl():
version = comfyui_version()
_state["version"] = version
current = mode()
if current == "off":
_state["reason"] = "switch off (env %s=%s)" % (ENV, _env(ENV) or "unset")
log.info("[UniBlockSwap] DynamicVRAM bypass switch is OFF; keeping DynamicVRAM")
return False
if not version_ok():
_state["reason"] = "ComfyUI %s < %s, bypass not needed" % (version, MIN_VERSION_STR)
log.info("[UniBlockSwap] ComfyUI %s: no DynamicVRAM to bypass", version)
return False
if not dynamic_active():
_state["reason"] = "DynamicVRAM already off (startup args)"
log.info("[UniBlockSwap] DynamicVRAM is already off at startup; nothing to bypass")
return False
from comfy.cli_args import args
if getattr(args, "enable_dynamic_vram", False):
_state["reason"] = "--enable-dynamic-vram given explicitly, bypass skipped"
_banner([
_rule(),
" UniBlockSwap: bypass switch is ON but NOT applied",
" reason: ComfyUI was started with --enable-dynamic-vram (explicit CLI flag wins)",
" fix : drop that flag, or set UNIBLOCKSWAP_LEGACY_VRAM=0 to silence this",
_rule(),
])
return False
if current == "soft":
_state["reason"] = "dry run"
_banner([
_rule(),
" UniBlockSwap: DRY RUN (UNIBLOCKSWAP_LEGACY_VRAM=%s)" % _env(ENV),
" would set: comfy.memory_management.aimdo_enabled = False",
" would set: comfy.model_patcher.CoreModelPatcher = ModelPatcher",
" nothing was changed, DynamicVRAM stays active",
_rule(),
])
return False
if current == "auto":
armed = _install_gate()
_state["reason"] = ("auto: armed, engages only when a workflow uses %s"
% "/".join(NODE_TYPES))
_banner(
[_rule()]
+ [" UniBlockSwap: DynamicVRAM bypass ARMED (mode=auto)"]
+ [" DynamicVRAM is NOT touched right now - it stays as configured at startup."
if armed else
" ⚠️ could not install the /prompt gate; bypass stays OFF (see warnings above)."]
+ [" rule : 只有当提交的工作流真的用到 %s 时," % "/".join(NODE_TYPES),
" 才会在模型构造之前把 DynamicVRAM 翻回 legacy 口径;",
" 队列里没有这个节点时自动保持/还原 dynamic。",
" cost : while it is engaged: no aimdo demand paging / prefetch /",
" cuda-graph weight load + conservative legacy VRAM estimates.",
" manual: legacy_vram.set_bypass(True/False) / restore() /",
" with legacy_vram.bypass_scope(): (脚本化加载用)",
" force : UNIBLOCKSWAP_LEGACY_VRAM=1 -> 老行为, 进程启动就整机绕过",
_rule()]
)
return armed
changed = set_bypass(True, quiet=True, reason="import-time bypass (mode=always)")
_state["reason"] = "bypassed DynamicVRAM (mode=always)"
_banner(
[_rule()]
+ [" UniBlockSwap: DynamicVRAM BYPASSED (mode=always, ComfyUI %s >= %s)"
% (version, MIN_VERSION_STR)]
+ [" " + c for c in changed]
+ [
" effect: this whole process uses the pre-0.35 legacy ModelPatcher path",
" (block swap keeps its _load_list filtering + per-block sync copy)",
" cost : no aimdo demand paging / prefetch / cuda-graph weight load,",
" legacy conservative VRAM estimates (big models may OOM on small cards)",
" note : EVERY workflow is affected, whether or not it uses this node.",
" If you only want it for workflows using this node, use mode=auto",
" (unset UNIBLOCKSWAP_LEGACY_VRAM / set it to 'auto').",
" revert: legacy_vram.set_bypass(False) (or restore()) - runtime, applies to",
" models loaded from then on; or UNIBLOCKSWAP_LEGACY_VRAM=0 + restart",
_rule(),
]
)
return True
def apply():
"""在 __init__.py 里 import 时调用一次 (auto 模式只装钩子, 不翻转全局量)。
任何异常都不会阻断 custom node 加载。
"""
try:
return _apply_impl()
except Exception:
_state["reason"] = "apply() failed"
log.warning("[UniBlockSwap] DynamicVRAM bypass failed, keeping defaults", exc_info=True)
return False
# ---------------------------------------------------------------------------
# 节点执行时的提示
# ---------------------------------------------------------------------------
_HINTED = set()
def runtime_report(node="UniBlockSwap", patcher=None):
"""节点执行时调用: 把当前生效状态提示给用户, 返回状态字符串。"""
try:
dynamic = dynamic_active()
patcher_dynamic = False
try:
patcher_dynamic = bool(patcher is not None and patcher.is_dynamic())
except Exception:
pass
if _state["applied"] and not patcher_dynamic:
if "ok" not in _HINTED:
_HINTED.add("ok")
log.info("[%s] DynamicVRAM bypass active (legacy ModelPatcher), "
"block swap runs on its native code path", node)
return "bypass"
if patcher_dynamic or dynamic:
if "warn" not in _HINTED:
_HINTED.add("warn")
lines = [
_rule("!"),
" UniBlockSwap: still running under DynamicVRAM",
" reason: %s" % _state["reason"],
" hint : block swap degrades into the old synchronous cast path and",
" saves much less VRAM.",
]
if mode() == "auto":
lines += [
" note : mode=auto only helps models that get CONSTRUCTED after the",
" prompt is submitted, and it forces every MODEL/CLIP/VAE",
" producing node to re-run exactly so that happens.",
" Seeing this warning means the model STILL came from a cache:",
" its loader did not re-run (custom loader with its own cache,",
" e.g. GGUF) or it was built outside the /prompt path.",
" fix : restart ComfyUI, or pick another checkpoint and switch back,",
" or start with --cache-none.",
]
else:
lines += [
" hint : to bypass: legacy_vram.set_bypass(True) then RELOAD the",
" model (already-built models cannot be converted);",
" UNIBLOCKSWAP_LEGACY_VRAM=1 + restart also works.",
]
lines.append(_rule("!"))
_banner(lines)
return "dynamic"
if "off" not in _HINTED:
_HINTED.add("off")
log.info("[%s] DynamicVRAM bypass switch off; legacy path used anyway (%s)",
node, _state["reason"])
return "off"
except Exception:
return "unknown"