diff --git a/README.md b/README.md index 869c336..ff7615b 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,21 @@ Qwen Image, and SeedVR2. The retired `*_distill` values remain selectable as compatibility aliases for saved ComfyUI workflows and are normalized before LightX2V's runner is built. +The wrapper keeps one LightX2V GPU runner resident across these node groups. +Switching between Wan/SekoTalk and SeedVR2/SwiftVR releases the previous +runner before loading the next one. Cached restoration loader outputs remain +usable; their runner is rebuilt when selected again. + +When LightX2V loads a runner, it asks ComfyUI to unload its resident models +first. When ComfyUI next frees GPU memory for one of its own models, the wrapper +evicts the resident LightX2V runner. This also applies to cached restoration +loader outputs; they rebuild their runner on the next sampler call. An active +LightX2V request is protected from eviction until it finishes. Switching +models can therefore incur a reload, but resident model allocations do not +accumulate across the two model managers. This hook uses ComfyUI's +`model_management.free_memory` entry point; changes to that internal API +require revalidation. + ## Model files Place each model in its own directory under `ComfyUI/models/lightx2v/`. LoRA diff --git a/README_CN.md b/README_CN.md index 0ec281c..bdc80a8 100644 --- a/README_CN.md +++ b/README_CN.md @@ -61,6 +61,17 @@ git submodule update --init --recursive 变体、SekoTalk、Qwen Image 和 SeedVR2。旧 workflow 使用的 `*_distill` 值仍作为兼容别名保留，构建 LightX2V runner 前会自动归一化。 +这些节点组共用一个常驻的 LightX2V GPU runner。Wan/SekoTalk 与 +SeedVR2/SwiftVR 切换时，插件先释放旧 runner，再加载新模型；ComfyUI +缓存的超分加载器输出仍可复用，再次使用时会重建 runner。 + +LightX2V 加载 runner 前会先让 ComfyUI 卸载它管理的常驻模型；ComfyUI +随后为原生模型清理 GPU 显存时，插件会释放 LightX2V 的常驻 runner。 +缓存的超分加载器输出会在下次采样时重建 runner。正在运行的 LightX2V +推理不会被中途卸载。切换模型需要重新加载，但两套模型管理器的常驻模型 +显存不会因连续切换而一直累积。该机制接入 ComfyUI 的 +`model_management.free_memory`；此内部接口变动时需要重新验证。 + ## 模型目录 每个模型放在 `ComfyUI/models/lightx2v/` 下的独立目录中。LoRA 可以放在 diff --git a/__init__.py b/__init__.py index 7bd2ed6..1b72bb1 100644 --- a/__init__.py +++ b/__init__.py @@ -44,5 +44,8 @@ _setup_env() _register_lightx2v_submodule() from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS # noqa: E402 +from .nodes._runner_memory import install_comfyui_memory_hook # noqa: E402 + +install_comfyui_memory_hook() __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/_lightx2v_runner.py b/nodes/_lightx2v_runner.py index 03c11f3..c7460ad 100644 --- a/nodes/_lightx2v_runner.py +++ b/nodes/_lightx2v_runner.py @@ -48,6 +48,8 @@ def temporary_runner_config(runner, values): def run_request(runner, request, *, runtime_config=None): """Validate one modern request and execute it through the runner API.""" - with temporary_runner_config(runner, runtime_config or {}): + from ._runner_memory import active_lightx2v_request + + with active_lightx2v_request(), temporary_runner_config(runner, runtime_config or {}): input_info = runner.prepare_request(dict(request)) return runner.run_request(input_info) diff --git a/nodes/_runner_memory.py b/nodes/_runner_memory.py new file mode 100644 index 0000000..8947bf6 --- /dev/null +++ b/nodes/_runner_memory.py @@ -0,0 +1,181 @@ +"""Keep at most one resident LightX2V runner across ComfyUI node caches. + +ComfyUI can retain loader outputs after a workflow stops using them. Holding a +runner directly in each output therefore keeps old models on the GPU when the +next workflow loads another model. Handles remain cacheable, but their runner +slot can be cleared and rebuilt when that workflow is used again. +""" + +import functools +import gc +import logging +import sys +import threading +import weakref +from contextlib import contextmanager + +import torch + +from ._lightx2v_runner import build_runner + +_LOCK = threading.RLock() +_HANDLES = weakref.WeakValueDictionary() +_ACTIVE_REQUESTS = 0 +_RELEASING = False + + +@contextmanager +def active_lightx2v_request(): + """Keep ComfyUI's memory hook from evicting a runner during inference.""" + global _ACTIVE_REQUESTS + with _LOCK: + _ACTIVE_REQUESTS += 1 + try: + yield + finally: + _ACTIVE_REQUESTS -= 1 + + +def dispose_runner(runner): + """Drop GPU modules even when an upstream callback still references runner.""" + if runner is None: + return + children = getattr(runner, "clip_generators", None) + if isinstance(children, dict): + for child in children.values(): + dispose_runner(child) + del runner.clip_generators + # DefaultRunner.__del__ covers most of these, but only runs once *all* + # external references disappear. SwiftVR's restorer also holds its own + # transformer and VAE references and is not covered by that destructor. + for name in ("restorer", "model", "text_encoders", "image_encoder", "vae_encoder", "vae_decoder", "inputs"): + if name in getattr(runner, "__dict__", {}): + delattr(runner, name) + + +def _collect_cuda(): + # LightX2V freezes the GC generation after a successful request. Unfreeze + # before dropping runners so cycles containing CUDA tensors can be collected. + gc.unfreeze() + gc.collect() + torch.cuda.empty_cache() + + +def release_restoration_runners(keep=None): + """Drop SwiftVR/SeedVR runners, including cached ComfyUI loader outputs.""" + with _LOCK: + gc.unfreeze() + released_seedvr = False + released_swiftvr = False + for handle in list(_HANDLES.values()): + if handle is not keep: + runner = handle["runner"] + released_seedvr |= runner is not None and handle.startup_config.get("model_cls") == "seedvr2" + released_swiftvr |= runner is not None and handle.startup_config.get("model_cls") == "swiftvr" + dispose_runner(runner) + handle["runner"] = None + del runner + if released_seedvr: + # This @lru_cache is keyed by `self`: without clearing it, every + # SeedVR RoPE instance (and its GPU model graph) remains reachable. + rope_module = sys.modules.get("lightx2v.models.networks.seedvr.utils.rope") + if rope_module is not None: + rope_module.RotaryEmbeddingBase.get_axial_freqs.cache_clear() + if released_swiftvr: + attention_module = sys.modules.get("lightx2v.models.networks.swiftvr.attention") + if attention_module is not None: + attention_module.WindowLayout.clear() + _collect_cuda() + + +def _release_modular_runner(): + from .inference import LightX2VModularInferenceV2 + + if LightX2VModularInferenceV2._current_runner is not None: + LightX2VModularInferenceV2._release_runner() + + +def release_all_runners(): + """Evict every resident LightX2V model before ComfyUI allocates VRAM.""" + global _RELEASING + with _LOCK: + if _ACTIVE_REQUESTS or _RELEASING: + return False + from .inference import LightX2VModularInferenceV2 + + if LightX2VModularInferenceV2._current_runner is None and not any(handle["runner"] is not None for handle in _HANDLES.values()): + return False + _RELEASING = True + try: + logging.info("Releasing resident LightX2V models for ComfyUI memory management") + _release_modular_runner() + release_restoration_runners() + return True + finally: + _RELEASING = False + + +def unload_comfyui_models(): + """Release ComfyUI-owned models before LightX2V constructs a GPU runner.""" + import comfy.model_management as model_management + + model_management.unload_all_models() + model_management.soft_empty_cache() + + +def install_comfyui_memory_hook(): + """Evict LightX2V when ComfyUI's native model manager frees GPU memory. + + Native loads call free_memory through the model_management module, so this + also covers call sites holding an early-bound reference to load_models_gpu. + """ + import comfy.model_management as model_management + + original = model_management.free_memory + if getattr(original, "_lightx2v_memory_hook", False) is True: + return + + @functools.wraps(original) + def free_memory_with_lightx2v(*args, **kwargs): + device = args[1] if len(args) > 1 else kwargs.get("device") + if device is None or getattr(device, "type", None) != "cpu": + release_all_runners() + return original(*args, **kwargs) + + free_memory_with_lightx2v._lightx2v_memory_hook = True + model_management.free_memory = free_memory_with_lightx2v + + +def make_runner_handle(config, **metadata): + """Build a runner after releasing every other LightX2V model.""" + with _LOCK: + release_restoration_runners() + _release_modular_runner() + unload_comfyui_models() + runner = build_runner(config) + handle = RunnerHandle(config, runner, **metadata) + _HANDLES[id(handle)] = handle + return handle + + +def get_runner(model): + """Rebuild an evicted cached loader output before a sampler uses it.""" + if not isinstance(model, RunnerHandle): + return model["runner"] + with _LOCK: + runner = model["runner"] + if runner is None: + release_restoration_runners(keep=model) + _release_modular_runner() + unload_comfyui_models() + runner = build_runner(model.startup_config) + model["runner"] = runner + return runner + + +class RunnerHandle(dict): + """ComfyUI-compatible model output whose GPU state can be evicted.""" + + def __init__(self, startup_config, runner, **metadata): + super().__init__(runner=runner, **metadata) + self.startup_config = dict(startup_config) diff --git a/nodes/inference.py b/nodes/inference.py index 633d4bd..dfd998d 100644 --- a/nodes/inference.py +++ b/nodes/inference.py @@ -63,12 +63,16 @@ class LightX2VModularInferenceV2: @classmethod def _release_runner(cls): - """Drop the singleton runner + force VRAM teardown via DefaultRunner.__del__. + """Drop the singleton runner and its GPU modules before another load. Callers MUST drop their own local refs to the old runner *before* invoking this — otherwise the refcount stays > 0, __del__ doesn't fire, and the next model load OOMs (model_a + model_b alive on GPU at the same time). """ + from ._runner_memory import dispose_runner + + gc.unfreeze() + dispose_runner(cls._current_runner) cls._current_runner = None cls._current_config_hash = None gc.collect() @@ -211,6 +215,13 @@ class LightX2VModularInferenceV2: # switching v2.5 s2v -> v2.7 rs2v). current_runner = None self._release_runner() + # ComfyUI may still cache SwiftVR/SeedVR loader outputs from a + # prior workflow. Evict those GPU runners before loading a + # Wan/SekoTalk model, while retaining rebuildable handles. + from ._runner_memory import release_restoration_runners, unload_comfyui_models + + release_restoration_runners() + unload_comfyui_models() if config.get("task") == "rs2v": from ..lightx2v.lightx2v.shot_runner.rs2v_infer import ShotRS2VPipeline @@ -252,30 +263,29 @@ class LightX2VModularInferenceV2: config["return_result_tensor"] = True config["save_result_path"] = "" config["negative_prompt"] = config.get("negative_prompt", "") - if config.get("task") == "rs2v": - # ShotRS2VPipeline delegates request validation to its clip - # runner. Passing the prepared config wholesale worked with - # the old InputInfo path, but LightX2V 0.6 rejects startup - # fields (model paths, quantization flags, etc.) as request - # fields. Filter the request against the modern runner schema - # before handing it to the shot pipeline. - if _modern_build_runner is not None and _modern_build_startup_config is not None: - shot_runner = current_runner.clip_generators["rs2v_clip"] - request_data = self._build_modern_request(config, shot_runner) - result_dict = current_runner.run_pipeline(request_data) + from ._runner_memory import active_lightx2v_request + + with active_lightx2v_request(): + if config.get("task") == "rs2v": + # ShotRS2VPipeline delegates request validation to its clip + # runner. Filter prepared config against the request schema. + if _modern_build_runner is not None and _modern_build_startup_config is not None: + shot_runner = current_runner.clip_generators["rs2v_clip"] + request_data = self._build_modern_request(config, shot_runner) + result_dict = current_runner.run_pipeline(request_data) + else: + result_dict = current_runner.run_pipeline(config) + elif _modern_build_runner is not None and _modern_build_startup_config is not None: + request_data = self._build_modern_request(config, current_runner) + input_info = current_runner.prepare_request(request_data) + result_dict = current_runner.run_request(input_info) else: - result_dict = current_runner.run_pipeline(config) - elif _modern_build_runner is not None and _modern_build_startup_config is not None: - request_data = self._build_modern_request(config, current_runner) - input_info = current_runner.prepare_request(request_data) - result_dict = current_runner.run_request(input_info) - else: - if _legacy_init_empty_input_info is None or _legacy_update_input_info_from_dict is None: - raise RuntimeError("No compatible LightX2V input-info API found in the installed LightX2V submodule") - input_data = _legacy_init_empty_input_info(config.task) - _legacy_update_input_info_from_dict(input_data, config) - current_runner.set_config(config) - result_dict = current_runner.run_pipeline(input_data) + if _legacy_init_empty_input_info is None or _legacy_update_input_info_from_dict is None: + raise RuntimeError("No compatible LightX2V input-info API found in the installed LightX2V submodule") + input_data = _legacy_init_empty_input_info(config.task) + _legacy_update_input_info_from_dict(input_data, config) + current_runner.set_config(config) + result_dict = current_runner.run_pipeline(input_data) images = result_dict.get("video", None) audio = result_dict.get("audio", None) diff --git a/nodes/seedvr.py b/nodes/seedvr.py index 25a3eeb..af13faf 100644 --- a/nodes/seedvr.py +++ b/nodes/seedvr.py @@ -29,7 +29,8 @@ import folder_paths import torch from comfy.utils import ProgressBar -from ._lightx2v_runner import build_runner, run_request +from ._lightx2v_runner import run_request +from ._runner_memory import get_runner, make_runner_handle from .file_input import probe_video_file, resolve_input_video_path logger = logging.getLogger(__name__) @@ -373,7 +374,7 @@ class LightX2VSeedVR2Loader: else: config["dit_original_ckpt"] = str(ckpt_path) - runner = build_runner(config) + model = make_runner_handle(config, precision=precision, ckpt=ckpt_name) logger.info( "[SeedVR2Loader] loaded %s (%s); cpu_offload=%s, tile_vae=%s, tile=%s, overlap=%s, slice=%s, mem_limit=%sGiB", ckpt_name, @@ -385,7 +386,7 @@ class LightX2VSeedVR2Loader: vae_causal_slice_size, vae_memory_limit_gb, ) - return ({"runner": runner, "precision": precision, "ckpt": ckpt_name},) + return (model,) class LightX2VSeedVR2Sampler: @@ -463,7 +464,7 @@ class LightX2VSeedVR2Sampler: filename_prefix, color_fix, ): - runner = model["runner"] + runner = get_runner(model) if int(infer_steps) != 1: raise ValueError("SeedVR2 requires infer_steps=1") @@ -644,7 +645,7 @@ class LightX2VSeedVR2FileSampler: "color_fix": str(color_fix), } - runner = model["runner"] + runner = get_runner(model) progress = ProgressBar(100) logger.info( "[SeedVR2FileSampler] input=%s (%sx%s @ %.3f fps), target=%sx%s, segment=%s/%s", diff --git a/nodes/swiftvr.py b/nodes/swiftvr.py index f7e6774..15b27aa 100644 --- a/nodes/swiftvr.py +++ b/nodes/swiftvr.py @@ -11,7 +11,8 @@ import folder_paths import torch from comfy.utils import ProgressBar -from ._lightx2v_runner import build_runner, run_request +from ._lightx2v_runner import run_request +from ._runner_memory import get_runner, make_runner_handle from .file_input import probe_video_file, resolve_input_video_path from .seedvr import _prepare_output_video @@ -211,7 +212,7 @@ class LightX2VSwiftVRLoader: "parallel": False, "use_compile": bool(use_compile), } - runner = build_runner(config) + model = make_runner_handle(config, model_name=model_name) logger.info( "[SwiftVRLoader] loaded %s; attention=%s, rope=%s, clip=%s, overlap=%s, reae_batch=%s, compile=%s", model_name, @@ -222,7 +223,7 @@ class LightX2VSwiftVRLoader: reae_frame_batch_size, use_compile, ) - return ({"runner": runner, "model_name": model_name},) + return (model,) class LightX2VSwiftVRSampler: @@ -310,7 +311,7 @@ class LightX2VSwiftVRSampler: sr_ratio, ) - runner = model["runner"] + runner = get_runner(model) output_file = "" output_subfolder = "" temp_dir = None @@ -485,7 +486,7 @@ class LightX2VSwiftVRFileSampler: "ffmpeg_preset": str(ffmpeg_preset), } - runner = model["runner"] + runner = get_runner(model) progress = ProgressBar(100) logger.info( "[SwiftVRFileSampler] input=%s (%sx%s @ %.3f fps), output=%sx%s", diff --git a/tests/test_runner_memory.py b/tests/test_runner_memory.py new file mode 100644 index 0000000..9d9ce77 --- /dev/null +++ b/tests/test_runner_memory.py @@ -0,0 +1,176 @@ +"""Regression checks for switching cached LightX2V model outputs.""" + +import importlib.util +import sys +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import Mock + + +class RunnerMemoryTests(unittest.TestCase): + def setUp(self): + package = "_runner_memory_test_package" + nodes_package = package + ".nodes" + self.build_runner = Mock(side_effect=lambda config: {"model_cls": config["model_cls"]}) + self.release_modular = Mock() + self.empty_cache = Mock() + self.unload_native = Mock() + self.soft_empty_cache = Mock() + self.native_free_memory = Mock(return_value=["native-unloaded"]) + self.modular_node = types.SimpleNamespace(_current_runner=None, _release_runner=self.release_modular) + + modules = { + package: types.ModuleType(package), + nodes_package: types.ModuleType(nodes_package), + nodes_package + "._lightx2v_runner": types.ModuleType(nodes_package + "._lightx2v_runner"), + nodes_package + ".inference": types.ModuleType(nodes_package + ".inference"), + "torch": types.ModuleType("torch"), + "comfy": types.ModuleType("comfy"), + "comfy.model_management": types.ModuleType("comfy.model_management"), + } + modules[nodes_package + "._lightx2v_runner"].build_runner = self.build_runner + modules[nodes_package + ".inference"].LightX2VModularInferenceV2 = self.modular_node + modules["torch"].cuda = types.SimpleNamespace(empty_cache=self.empty_cache) + modules["comfy"].model_management = modules["comfy.model_management"] + modules["comfy.model_management"].unload_all_models = self.unload_native + modules["comfy.model_management"].soft_empty_cache = self.soft_empty_cache + modules["comfy.model_management"].free_memory = self.native_free_memory + self.saved_modules = {name: sys.modules.get(name) for name in modules} + sys.modules.update(modules) + + path = Path(__file__).resolve().parents[1] / "nodes" / "_runner_memory.py" + name = nodes_package + "._runner_memory" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + self.module = module + self.module_name = name + + def tearDown(self): + sys.modules.pop(self.module_name, None) + for name, previous in self.saved_modules.items(): + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + def test_switch_evicts_runner_from_cached_loader_output(self): + first = self.module.make_runner_handle({"model_cls": "seedvr2"}) + second = self.module.make_runner_handle({"model_cls": "swiftvr"}) + + self.assertIsNone(first["runner"]) + self.assertEqual(second["runner"]["model_cls"], "swiftvr") + self.assertEqual(self.unload_native.call_count, 2) + self.assertGreaterEqual(self.empty_cache.call_count, 2) + + restored = self.module.get_runner(first) + self.assertEqual(restored["model_cls"], "seedvr2") + self.assertIsNone(second["runner"]) + self.assertEqual(self.build_runner.call_count, 3) + self.assertEqual(self.unload_native.call_count, 3) + + def test_native_memory_request_evicts_lightx2v_runner(self): + handle = self.module.make_runner_handle({"model_cls": "swiftvr"}) + self.module.install_comfyui_memory_hook() + self.module.install_comfyui_memory_hook() + result = sys.modules["comfy.model_management"].free_memory(100, types.SimpleNamespace(type="cuda")) + + self.assertEqual(result, ["native-unloaded"]) + self.assertIsNone(handle["runner"]) + self.native_free_memory.assert_called_once() + + def test_active_lightx2v_request_cannot_be_evicted(self): + handle = self.module.make_runner_handle({"model_cls": "seedvr2"}) + self.module.install_comfyui_memory_hook() + with self.module.active_lightx2v_request(): + sys.modules["comfy.model_management"].free_memory(100, types.SimpleNamespace(type="cuda")) + self.assertIsNotNone(handle["runner"]) + sys.modules["comfy.model_management"].free_memory(100, types.SimpleNamespace(type="cuda")) + self.assertIsNone(handle["runner"]) + + def test_native_memory_request_evicts_modular_runner(self): + self.modular_node._current_runner = types.SimpleNamespace(model=object()) + self.module.install_comfyui_memory_hook() + sys.modules["comfy.model_management"].free_memory(100, types.SimpleNamespace(type="cuda")) + self.release_modular.assert_called_once_with() + + def test_native_memory_hook_does_not_reenter_during_teardown(self): + self.modular_node._current_runner = types.SimpleNamespace(model=object()) + self.module.install_comfyui_memory_hook() + model_management = sys.modules["comfy.model_management"] + self.release_modular.side_effect = lambda: model_management.free_memory(100, types.SimpleNamespace(type="cuda")) + model_management.free_memory(100, types.SimpleNamespace(type="cuda")) + self.release_modular.assert_called_once_with() + + def test_concurrent_native_load_waits_for_lightx2v_request(self): + handle = self.module.make_runner_handle({"model_cls": "swiftvr"}) + self.module.install_comfyui_memory_hook() + started = threading.Event() + finished = threading.Event() + + def load_native(): + started.set() + sys.modules["comfy.model_management"].free_memory(100, types.SimpleNamespace(type="cuda")) + finished.set() + + with self.module.active_lightx2v_request(): + thread = threading.Thread(target=load_native) + thread.start() + self.assertTrue(started.wait(1)) + self.assertFalse(finished.wait(0.05)) + self.assertIsNotNone(handle["runner"]) + thread.join(timeout=1) + self.assertTrue(finished.is_set()) + self.assertIsNone(handle["runner"]) + + def test_cpu_memory_request_keeps_lightx2v_runner(self): + handle = self.module.make_runner_handle({"model_cls": "swiftvr"}) + self.module.install_comfyui_memory_hook() + sys.modules["comfy.model_management"].free_memory(100, types.SimpleNamespace(type="cpu")) + self.assertIsNotNone(handle["runner"]) + + def test_modular_switch_releases_restoration_runner(self): + cached_output = self.module.make_runner_handle({"model_cls": "swiftvr"}) + self.module.release_restoration_runners() + + self.assertIsNone(cached_output["runner"]) + self.assertEqual(cached_output.startup_config, {"model_cls": "swiftvr"}) + + def test_seedvr_switch_clears_bound_method_cache(self): + clear = Mock() + name = "lightx2v.models.networks.seedvr.utils.rope" + previous = sys.modules.get(name) + rope_module = types.ModuleType(name) + rope_module.RotaryEmbeddingBase = types.SimpleNamespace(get_axial_freqs=types.SimpleNamespace(cache_clear=clear)) + sys.modules[name] = rope_module + try: + cached_output = self.module.make_runner_handle({"model_cls": "seedvr2"}) + self.module.make_runner_handle({"model_cls": "swiftvr"}) + self.assertIsNone(cached_output["runner"]) + clear.assert_called_once_with() + finally: + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + def test_dispose_runner_drops_swiftvr_restorer(self): + runner = types.SimpleNamespace(model=object(), restorer=object(), vae_decoder=object()) + self.module.dispose_runner(runner) + self.assertFalse(hasattr(runner, "model")) + self.assertFalse(hasattr(runner, "restorer")) + self.assertFalse(hasattr(runner, "vae_decoder")) + + def test_dispose_runner_recurses_into_rs2v_clips(self): + child = types.SimpleNamespace(model=object(), vae_decoder=object()) + shot = types.SimpleNamespace(clip_generators={"rs2v_clip": child}) + self.module.dispose_runner(shot) + self.assertFalse(hasattr(shot, "clip_generators")) + self.assertFalse(hasattr(child, "model")) + + +if __name__ == "__main__": + unittest.main()