diff --git a/config_utils.py b/config_utils.py index df529f9..750f854 100644 --- a/config_utils.py +++ b/config_utils.py @@ -361,6 +361,75 @@ def expand_lora_stack(lora_input): return expanded_loras +def _normalize_ltx_audio_mode(entry): + """Convert audio_mode='both' to ['on', 'off'] for cartesian expansion.""" + if entry.get("audio_mode") == "both": + e = dict(entry) + e["audio_mode"] = ["on", "off"] + return e + return entry + + +def _expand_ltx_entry(entry): + """Expand a single ltx_video entry into a list of fully-resolved LTX configs. + + Each LTX field that is a list (except clip_models which has special semantics) + contributes a Cartesian axis. Returns a list of dicts ready for the orchestrator. + """ + e = _normalize_ltx_audio_mode(entry) + + def to_list(x): + return x if isinstance(x, list) else [x] + + # clip_models special case: + # - [str, str] (2 strings) -> single value (a CLIP pair) + # - [[str, str], [str, str], ...] -> sweep + raw_clip = e.get("clip_models", ["", ""]) + if isinstance(raw_clip, list) and len(raw_clip) == 2 and all(isinstance(c, str) for c in raw_clip): + clip_options = [raw_clip] + elif isinstance(raw_clip, list) and all(isinstance(c, list) for c in raw_clip): + clip_options = raw_clip + else: + clip_options = [raw_clip] + + # All other fields: simple to_list + field_axes = { + "model": to_list(e.get("model", "")), + "clip_models": clip_options, + "vae_video": to_list(e.get("vae_video", "")), + "vae_audio": to_list(e.get("vae_audio", "")), + "latent_upscaler": to_list(e.get("latent_upscaler", "")), + "duration_seconds": to_list(e.get("duration_seconds", 5)), + "frame_rate": to_list(e.get("frame_rate", 25)), + "sampler_stage1": to_list(e.get("sampler_stage1", "euler_ancestral_cfg_pp")), + "sampler_stage2": to_list(e.get("sampler_stage2", "euler_cfg_pp")), + "sigmas_stage1": to_list(e.get("sigmas_stage1", "")), + "sigmas_stage2": to_list(e.get("sigmas_stage2", "")), + "cfg": to_list(e.get("cfg", 1.0)), + "seed": to_list(e.get("seed", 0)), + "input_image": to_list(e.get("input_image")) if e.get("input_image") is not None else [None], + "image_strength_stage1": to_list(e.get("image_strength_stage1", 0.8)), + "image_strength_stage2": to_list(e.get("image_strength_stage2", 1.0)), + "img_compression": to_list(e.get("img_compression", 18)), + "audio_mode": to_list(e.get("audio_mode", "on")), + "lora": to_list(e.get("lora", "None")), + "positive": to_list(e.get("positive", "")), + "negative": to_list(e.get("negative", "")), + "width": to_list(e.get("width", 446)), + "height": to_list(e.get("height", 576)), + } + + keys = list(field_axes.keys()) + values = [field_axes[k] for k in keys] + + expanded = [] + for combo in itertools.product(*values): + cfg = dict(zip(keys, combo)) + cfg["model_type"] = "ltx_video" + expanded.append(cfg) + return expanded + + def expand_configs(raw_configs, pos_prompts, neg_prompts, denoise_values, seed, extra_seeds, ckpt_name=None): """ Expand raw config entries into full configuration combinations. @@ -396,6 +465,11 @@ def expand_configs(raw_configs, pos_prompts, neg_prompts, denoise_values, seed, return x if isinstance(x, list) else [x] for entry in raw_configs: + # Branch early for LTX video — different field shape, separate cartesian + if entry.get("model_type") == "ltx_video": + expanded.extend(_expand_ltx_entry(entry)) + continue + # ==== PER-CONFIG PROMPT HANDLING ==== # If this config defines its own prompts, parse them and use instead of defaults if "positive" in entry or "negative" in entry: diff --git a/tests/conftest.py b/tests/conftest.py index 040373d..1f317aa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,7 @@ _COMFY_STUBS = [ "comfy.utils", "comfy.sd", "comfy.model_management", + "comfy.samplers", ] for _mod_name in _COMFY_STUBS: @@ -38,8 +39,22 @@ for _mod_name in _COMFY_STUBS: def put(self, *a, **kw): return lambda fn: fn _stub.RouteTableDef = _RouteTableDef + # comfy.samplers needs KSampler with SCHEDULERS and SAMPLERS lists + if _mod_name == "comfy.samplers": + class _KSampler: # noqa: N801 + SCHEDULERS = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"] + SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral", + "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", + "dpmpp_sde", "dpmpp_2m", "dpmpp_2m_sde", "ddim", "uni_pc", + "uni_pc_bh2"] + _stub.KSampler = _KSampler sys.modules[_mod_name] = _stub +# Wire comfy.samplers as an attribute of the comfy stub so +# `import comfy.samplers` and `comfy.samplers.KSampler` both resolve +if hasattr(sys.modules.get("comfy"), "__name__"): + sys.modules["comfy"].samplers = sys.modules["comfy.samplers"] + # 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: diff --git a/tests/test_ltx_config_expansion.py b/tests/test_ltx_config_expansion.py new file mode 100644 index 0000000..8547e1c --- /dev/null +++ b/tests/test_ltx_config_expansion.py @@ -0,0 +1,120 @@ +"""Tests for LTX field cartesian expansion.""" +import pytest +import sys +import os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from config_utils import expand_configs + + +def _ltx_base(): + 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", + "positive": "a video", + "negative": "ugly", + "width": 446, + "height": 576, + "duration_seconds": 5, + "frame_rate": 25, + "sampler_stage1": "euler_ancestral_cfg_pp", + "sampler_stage2": "euler_cfg_pp", + "sigmas_stage1": "1.0, 0.5, 0.0", + "sigmas_stage2": "0.85, 0.4, 0.0", + "cfg": 1.0, + "seed": 42, + "input_image": None, + "image_strength_stage1": 0.8, + "image_strength_stage2": 1.0, + "img_compression": 18, + "audio_mode": "on", + "lora": "None", + } + + +def _call(base): + """Wrap expand_configs with the dummy outer args it needs for non-LTX paths.""" + return expand_configs( + raw_configs=[base], + pos_prompts=[base.get("positive", "")], + neg_prompts=[base.get("negative", "")], + denoise_values=[1.0], + seed=base.get("seed", 0), + extra_seeds=[], + ckpt_name=None, + ) + + +def test_single_value_returns_one_config(): + result = _call(_ltx_base()) + assert len(result) == 1 + assert result[0]["model_type"] == "ltx_video" + assert result[0]["model"] == "ltx_model.safetensors" + + +def test_sigmas_stage2_array_expands(): + base = _ltx_base() + base["sigmas_stage2"] = ["0.85, 0.4, 0.0", "0.9, 0.5, 0.0"] + result = _call(base) + assert len(result) == 2 + sigmas_seen = sorted(c["sigmas_stage2"] for c in result) + assert sigmas_seen == ["0.85, 0.4, 0.0", "0.9, 0.5, 0.0"] + + +def test_audio_mode_both_expands_to_on_off(): + base = _ltx_base() + base["audio_mode"] = "both" + result = _call(base) + assert len(result) == 2 + modes_seen = sorted(c["audio_mode"] for c in result) + assert modes_seen == ["off", "on"] + + +def test_audio_mode_array_accepted(): + base = _ltx_base() + base["audio_mode"] = ["on", "off"] + result = _call(base) + assert len(result) == 2 + + +def test_duration_array_expands(): + base = _ltx_base() + base["duration_seconds"] = [5, 10] + result = _call(base) + assert len(result) == 2 + + +def test_image_strength_arrays_cross_with_sigmas(): + base = _ltx_base() + base["image_strength_stage1"] = [0.6, 0.8] + base["sigmas_stage2"] = ["0.85, 0.4, 0.0", "0.9, 0.5, 0.0"] + result = _call(base) + assert len(result) == 4 # 2 x 2 + + +def test_input_image_array_expands(): + base = _ltx_base() + base["input_image"] = ["a.png", "b.png", "c.png"] + result = _call(base) + assert len(result) == 3 + + +def test_clip_models_2tuple_does_not_expand(): + """clip_models is a [name1, name2] tuple, not 2 sweep values.""" + base = _ltx_base() + base["clip_models"] = ["g.safetensors", "tp.safetensors"] + result = _call(base) + assert len(result) == 1 + assert result[0]["clip_models"] == ["g.safetensors", "tp.safetensors"] + + +def test_clip_models_list_of_lists_expands(): + """[[c1a, c2a], [c1b, c2b]] sweeps 2 CLIP pairs.""" + base = _ltx_base() + base["clip_models"] = [["g.safetensors", "tp1.safetensors"], ["g.safetensors", "tp2.safetensors"]] + result = _call(base) + assert len(result) == 2