feat: cache key for ltx_video model_type

2026-04-25 16:48:32 -07:00
parent ef866b50db
commit dca70cb90f
3 changed files with 200 additions and 0 deletions
@@ -236,6 +236,15 @@ def get_model_cache_key(conf):
    model_type = conf.get("model_type", "checkpoint")
    if model_type == "checkpoint":
        return conf["model"]
+    elif model_type == "ltx_video":
+        clip_models = conf.get("clip_models", ["", ""])
+        clip_a = clip_models[0] if len(clip_models) > 0 else ""
+        clip_b = clip_models[1] if len(clip_models) > 1 else ""
+        return (
+            "ltx::" + conf["model"] + "::" + clip_a + "::" + clip_b + "::" +
+            conf.get("vae_video", "") + "::" + conf.get("vae_audio", "") + "::" +
+            conf.get("latent_upscaler", "")
+        )
    else:
        te_key = "|".join(sorted(conf.get("text_encoders", [])))
        return f"{conf['model']}::{model_type}::{te_key}"
@@ -55,7 +55,104 @@ for _mod_name in _COMFY_STUBS:
 if hasattr(sys.modules.get("comfy"), "__name__"):
    sys.modules["comfy"].samplers = sys.modules["comfy.samplers"]

+# comfy.model_management needs InterruptProcessingException for generation_orchestrator
+if "comfy.model_management" in sys.modules:
+    if not hasattr(sys.modules["comfy.model_management"], "InterruptProcessingException"):
+        sys.modules["comfy.model_management"].InterruptProcessingException = type(
+            "InterruptProcessingException", (Exception,), {}
+        )
+
 # Ensure the custom node root is on sys.path so `ltx_video_generation` is importable
 _NODE_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 if _NODE_ROOT not in sys.path:
    sys.path.insert(0, _NODE_ROOT)
+
+# ── Package stub so generation_orchestrator's relative imports resolve ────────
+# generation_orchestrator.py uses `from .trigger_words import ...` etc.
+# When imported as a top-level module (sys.path.insert approach), relative
+# imports fail unless we set up a minimal package context in sys.modules.
+_PKG_NAME = os.path.basename(_NODE_ROOT)
+if _PKG_NAME not in sys.modules:
+    _pkg = types.ModuleType(_PKG_NAME)
+    _pkg.__path__ = [_NODE_ROOT]
+    _pkg.__package__ = _PKG_NAME
+    sys.modules[_PKG_NAME] = _pkg
+
+# Stub the sub-modules that generation_orchestrator imports via relative paths
+_GEN_ORCH_DEPS = {
+    "trigger_words": {
+        "collect_unique_prompts_with_triggers": lambda *a, **kw: [],
+        "build_prompt_with_triggers": lambda *a, **kw: ("", ""),
+        "clear_trigger_caches": lambda *a, **kw: None,
+    },
+    "batch_encoding": {
+        "batch_encode_prompts": lambda *a, **kw: None,
+        "encode_prompt_with_combinators": lambda *a, **kw: None,
+    },
+    "manifest_utils": {
+        "load_existing_manifest": lambda *a, **kw: {},
+        "save_manifest": lambda *a, **kw: None,
+    },
+    "model_loader": {
+        "load_checkpoint": lambda *a, **kw: None,
+        "load_loras": lambda *a, **kw: None,
+        "cleanup_model_references": lambda *a, **kw: None,
+        "get_latent_channels": lambda *a, **kw: 4,
+        "load_loras_for_preencoding": lambda *a, **kw: None,
+        "print_incompatible_loras_summary": lambda *a, **kw: None,
+        "load_diffusion_model_and_clip": lambda *a, **kw: (None, None),
+        "load_vae_by_name": lambda *a, **kw: None,
+    },
+    "lora_utils": {
+        "expand_lora_folder": lambda *a, **kw: [],
+    },
+    "image_generation": {
+        "generate_image": lambda *a, **kw: None,
+        "flush_batch_with_vae": lambda *a, **kw: None,
+        "flush_batch_with_remote_vae": lambda *a, **kw: None,
+        "create_image_metadata": lambda *a, **kw: {},
+        "decode_latent_with_vae": lambda *a, **kw: None,
+        "calculate_eta": lambda *a, **kw: 0.0,
+        "print_generation_progress": lambda *a, **kw: None,
+    },
+    "html_generator": {
+        "get_html_template": lambda *a, **kw: "",
+    },
+    "conditioning_cache": {
+        "ConditioningCache": type("ConditioningCache", (), {}),
+    },
+    "remote_vae": {
+        "RemoteVAEDecodeWorker": type("RemoteVAEDecodeWorker", (), {}),
+        "HF_ENDPOINTS": {},
+    },
+}
+
+for _bare, _attrs in _GEN_ORCH_DEPS.items():
+    _fq = f"{_PKG_NAME}.{_bare}"
+    if _bare not in sys.modules:
+        _s = types.ModuleType(_bare)
+        sys.modules[_bare] = _s
+    if _fq not in sys.modules:
+        sys.modules[_fq] = sys.modules[_bare]
+    for _k, _v in _attrs.items():
+        setattr(sys.modules[_bare], _k, _v)
+        setattr(sys.modules[_fq], _k, _v)
+    # Also wire as package attribute
+    setattr(sys.modules[_PKG_NAME], _bare, sys.modules[_bare])
+
+# Pre-load generation_orchestrator as part of the package so that relative
+# imports (`from .trigger_words import ...`) resolve correctly.  Once loaded
+# under the package name it is also registered as a top-level alias so
+# `from generation_orchestrator import get_model_cache_key` works.
+import importlib.util as _ilu
+_GO_PATH = os.path.join(_NODE_ROOT, "generation_orchestrator.py")
+_GO_FQ = f"{_PKG_NAME}.generation_orchestrator"
+if "generation_orchestrator" not in sys.modules:
+    _spec = _ilu.spec_from_file_location(_GO_FQ, _GO_PATH,
+                                         submodule_search_locations=[])
+    _spec.submodule_search_locations = None  # it is a module, not a package
+    _go_mod = _ilu.module_from_spec(_spec)
+    _go_mod.__package__ = _PKG_NAME
+    sys.modules[_GO_FQ] = _go_mod
+    sys.modules["generation_orchestrator"] = _go_mod
+    _spec.loader.exec_module(_go_mod)
@@ -0,0 +1,94 @@
+"""Tests for LTX cache key generation."""
+import pytest
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from generation_orchestrator import get_model_cache_key
+
+
+def _ltx_conf():
+    return {
+        "model_type": "ltx_video",
+        "model": "ltx_model.safetensors",
+        "clip_models": ["c1.safetensors", "c2.safetensors"],
+        "vae_video": "vv.safetensors",
+        "vae_audio": "va.safetensors",
+        "latent_upscaler": "lup.safetensors",
+    }
+
+
+def test_ltx_key_includes_all_files():
+    conf = _ltx_conf()
+    key = get_model_cache_key(conf)
+    assert "ltx" in key
+    assert "ltx_model.safetensors" in key
+    assert "c1.safetensors" in key
+    assert "c2.safetensors" in key
+    assert "vv.safetensors" in key
+    assert "va.safetensors" in key
+    assert "lup.safetensors" in key
+
+
+def test_ltx_two_identical_configs_have_same_key():
+    a = _ltx_conf()
+    b = _ltx_conf()
+    assert get_model_cache_key(a) == get_model_cache_key(b)
+
+
+def test_ltx_different_upscaler_different_key():
+    a = _ltx_conf()
+    b = _ltx_conf()
+    b["latent_upscaler"] = "different_upscaler.safetensors"
+    assert get_model_cache_key(a) != get_model_cache_key(b)
+
+
+def test_ltx_different_video_vae_different_key():
+    a = _ltx_conf()
+    b = _ltx_conf()
+    b["vae_video"] = "other_vae.safetensors"
+    assert get_model_cache_key(a) != get_model_cache_key(b)
+
+
+def test_ltx_different_clip_different_key():
+    a = _ltx_conf()
+    b = _ltx_conf()
+    b["clip_models"] = ["c1.safetensors", "alt_c2.safetensors"]
+    assert get_model_cache_key(a) != get_model_cache_key(b)
+
+
+def test_ltx_sigma_change_does_not_change_key():
+    """Sweeping sigmas should NOT trigger a model reload."""
+    a = _ltx_conf()
+    b = _ltx_conf()
+    a["sigmas_stage2"] = "0.85, 0.4, 0.0"
+    b["sigmas_stage2"] = "0.9, 0.5, 0.0"
+    assert get_model_cache_key(a) == get_model_cache_key(b)
+
+
+def test_ltx_seed_change_does_not_change_key():
+    a = _ltx_conf()
+    b = _ltx_conf()
+    a["seed"] = 42
+    b["seed"] = 100
+    assert get_model_cache_key(a) == get_model_cache_key(b)
+
+
+def test_checkpoint_key_unchanged():
+    """Existing checkpoint behavior must not regress."""
+    conf = {"model_type": "checkpoint", "model": "sdxl.safetensors"}
+    assert get_model_cache_key(conf) == "sdxl.safetensors"
+
+
+def test_diffusion_model_key_unchanged():
+    """Existing non-checkpoint (diffusion_model/gguf) behavior must not regress."""
+    conf = {
+        "model_type": "diffusion_model",
+        "model": "flux.safetensors",
+        "text_encoders": ["t5.safetensors", "clip_l.safetensors"],
+    }
+    key = get_model_cache_key(conf)
+    assert "flux.safetensors" in key
+    assert "diffusion_model" in key
+    assert "t5.safetensors" in key
+    assert "clip_l.safetensors" in key