diff --git a/generation_orchestrator.py b/generation_orchestrator.py index 670056e..6ac2a7b 100644 --- a/generation_orchestrator.py +++ b/generation_orchestrator.py @@ -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}" diff --git a/tests/conftest.py b/tests/conftest.py index 1f317aa..2c370b3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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) diff --git a/tests/test_ltx_cache_keys.py b/tests/test_ltx_cache_keys.py new file mode 100644 index 0000000..2fdc650 --- /dev/null +++ b/tests/test_ltx_cache_keys.py @@ -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