Fix LTXAVTextEncoderLoader and LTXVAudioVAELoader for standalone gemma4 layout
LTXAVTextEncoderLoader required a diffusion checkpoint to be present even though the official LTX-2.5 gemma4 text encoder ships as a standalone file with no matching checkpoint. Add a "none" option so ckpt_name can be skipped and only the text encoder file is loaded. LTXVAudioVAELoader read its file list from the checkpoints folder, but the official Lightricks release ships the audio VAE under models/vae/, leaving the combo box empty for that layout. Read from the vae folder instead, matching VAELoader. Fixes #15595
This commit is contained in:
chelsealong
parent
a779de4d89
commit
fc5a5e895c
|
|
@ -16,7 +16,7 @@ class LTXVAudioVAELoader(io.ComfyNode):
|
|||
inputs=[
|
||||
io.Combo.Input(
|
||||
"ckpt_name",
|
||||
options=folder_paths.get_filename_list("checkpoints"),
|
||||
options=folder_paths.get_filename_list("vae"),
|
||||
tooltip="Audio VAE checkpoint to load.",
|
||||
)
|
||||
],
|
||||
|
|
@ -25,7 +25,7 @@ class LTXVAudioVAELoader(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, ckpt_name: str) -> io.NodeOutput:
|
||||
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
||||
ckpt_path = folder_paths.get_full_path_or_raise("vae", ckpt_name)
|
||||
sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
|
||||
sd = comfy.utils.state_dict_prefix_replace(sd, {"audio_vae.": "autoencoder.", "vocoder.": "vocoder."}, filter_keys=True)
|
||||
vae = comfy.sd.VAE(sd=sd, metadata=metadata)
|
||||
|
|
@ -181,7 +181,8 @@ class LTXAVTextEncoderLoader(io.ComfyNode):
|
|||
),
|
||||
io.Combo.Input(
|
||||
"ckpt_name",
|
||||
options=folder_paths.get_filename_list("checkpoints"),
|
||||
options=["none"] + folder_paths.get_filename_list("checkpoints"),
|
||||
tooltip="Optional diffusion checkpoint the text encoder was split from. Set to 'none' when the text encoder file is standalone.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"device",
|
||||
|
|
@ -196,14 +197,15 @@ class LTXAVTextEncoderLoader(io.ComfyNode):
|
|||
def execute(cls, text_encoder, ckpt_name, device="default"):
|
||||
clip_type = comfy.sd.CLIPType.LTXV
|
||||
|
||||
clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", text_encoder)
|
||||
clip_path2 = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
||||
clip_paths = [folder_paths.get_full_path_or_raise("text_encoders", text_encoder)]
|
||||
if ckpt_name != "none":
|
||||
clip_paths.append(folder_paths.get_full_path_or_raise("checkpoints", ckpt_name))
|
||||
|
||||
model_options = {}
|
||||
if device == "cpu":
|
||||
model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
|
||||
|
||||
clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)
|
||||
clip = comfy.sd.load_clip(ckpt_paths=clip_paths, embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)
|
||||
return io.NodeOutput(clip)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,153 @@
|
|||
import sys
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
class FakeComfyNode:
|
||||
pass
|
||||
|
||||
|
||||
class FakeNodeOutput:
|
||||
def __init__(self, *args):
|
||||
self.args = args
|
||||
|
||||
|
||||
class FakeInput:
|
||||
def __init__(self, id, options=None, **kwargs):
|
||||
self.id = id
|
||||
self.options = options
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
class FakeCombo:
|
||||
Input = staticmethod(lambda id, **kwargs: FakeInput(id, **kwargs))
|
||||
|
||||
|
||||
class FakeSchema:
|
||||
def __init__(self, inputs=None, outputs=None, **kwargs):
|
||||
self.inputs = inputs or []
|
||||
self.outputs = outputs or []
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
def _fake_io_module():
|
||||
io_mock = MagicMock()
|
||||
io_mock.ComfyNode = FakeComfyNode
|
||||
io_mock.NodeOutput = FakeNodeOutput
|
||||
io_mock.Schema = FakeSchema
|
||||
io_mock.Combo = FakeCombo
|
||||
return io_mock
|
||||
|
||||
|
||||
def _fake_folder_paths(checkpoints=None, vae=None, text_encoders=None):
|
||||
checkpoints = checkpoints or []
|
||||
vae = vae or []
|
||||
text_encoders = text_encoders if text_encoders is not None else ["gemma4-12b-with-proj.safetensors"]
|
||||
lists_by_folder = {"checkpoints": checkpoints, "vae": vae, "text_encoders": text_encoders}
|
||||
|
||||
def get_filename_list(folder_name):
|
||||
return list(lists_by_folder.get(folder_name, []))
|
||||
|
||||
def get_full_path_or_raise(folder_name, filename):
|
||||
if filename not in lists_by_folder.get(folder_name, []):
|
||||
raise FileNotFoundError(f"no file '{filename}' registered under '{folder_name}'")
|
||||
return f"/models/{folder_name}/{filename}"
|
||||
|
||||
fp = MagicMock()
|
||||
fp.get_filename_list.side_effect = get_filename_list
|
||||
fp.get_full_path_or_raise.side_effect = get_full_path_or_raise
|
||||
fp.get_folder_paths.return_value = ["/models/embeddings"]
|
||||
return fp
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _nodes_lt_audio(folder_paths_mock):
|
||||
comfy_api_latest_mock = MagicMock()
|
||||
comfy_api_latest_mock.ComfyExtension = object
|
||||
comfy_api_latest_mock.io = _fake_io_module()
|
||||
|
||||
nodes_audio_mock = MagicMock()
|
||||
nodes_audio_mock.VAEEncodeAudio = type("VAEEncodeAudio", (), {})
|
||||
|
||||
sd_mock = MagicMock()
|
||||
utils_mock = MagicMock()
|
||||
utils_mock.load_torch_file.return_value = ({}, {})
|
||||
model_management_mock = MagicMock()
|
||||
|
||||
modules = {
|
||||
"torch": MagicMock(),
|
||||
"folder_paths": folder_paths_mock,
|
||||
"comfy.utils": utils_mock,
|
||||
"comfy.model_management": model_management_mock,
|
||||
"comfy.sd": sd_mock,
|
||||
"comfy_api.latest": comfy_api_latest_mock,
|
||||
"comfy_extras.nodes_audio": nodes_audio_mock,
|
||||
}
|
||||
comfy_submodule_attrs = {"utils": utils_mock, "model_management": model_management_mock, "sd": sd_mock}
|
||||
|
||||
sentinel = object()
|
||||
import comfy as comfy_pkg
|
||||
prior_attrs = {name: getattr(comfy_pkg, name, sentinel) for name in comfy_submodule_attrs}
|
||||
prior_module = sys.modules.get("comfy_extras.nodes_lt_audio", sentinel)
|
||||
sys.modules.pop("comfy_extras.nodes_lt_audio", None)
|
||||
|
||||
try:
|
||||
with patch.dict(sys.modules, modules):
|
||||
for name, mock in comfy_submodule_attrs.items():
|
||||
setattr(comfy_pkg, name, mock)
|
||||
module = __import__("comfy_extras.nodes_lt_audio", fromlist=["dummy"])
|
||||
yield module, sd_mock
|
||||
finally:
|
||||
for name, prior in prior_attrs.items():
|
||||
if prior is sentinel:
|
||||
if hasattr(comfy_pkg, name):
|
||||
delattr(comfy_pkg, name)
|
||||
else:
|
||||
setattr(comfy_pkg, name, prior)
|
||||
sys.modules.pop("comfy_extras.nodes_lt_audio", None)
|
||||
if prior_module is not sentinel:
|
||||
sys.modules["comfy_extras.nodes_lt_audio"] = prior_module
|
||||
|
||||
|
||||
class TestLTXVAudioVAELoader:
|
||||
def test_reads_from_vae_folder_not_checkpoints(self):
|
||||
fp = _fake_folder_paths(
|
||||
checkpoints=["some_checkpoint.safetensors"],
|
||||
vae=["ltx-2.5-audio-vae-bf16.safetensors"],
|
||||
)
|
||||
with _nodes_lt_audio(fp) as (module, _sd_mock):
|
||||
schema = module.LTXVAudioVAELoader.define_schema()
|
||||
ckpt_input = next(i for i in schema.inputs if i.id == "ckpt_name")
|
||||
assert ckpt_input.options == ["ltx-2.5-audio-vae-bf16.safetensors"]
|
||||
|
||||
module.LTXVAudioVAELoader.execute("ltx-2.5-audio-vae-bf16.safetensors")
|
||||
fp.get_full_path_or_raise.assert_called_with("vae", "ltx-2.5-audio-vae-bf16.safetensors")
|
||||
|
||||
|
||||
class TestLTXAVTextEncoderLoader:
|
||||
def test_none_ckpt_loads_standalone_text_encoder(self):
|
||||
fp = _fake_folder_paths(checkpoints=["some_checkpoint.safetensors"])
|
||||
with _nodes_lt_audio(fp) as (module, sd_mock):
|
||||
module.LTXAVTextEncoderLoader.execute("gemma4-12b-with-proj.safetensors", "none")
|
||||
|
||||
sd_mock.load_clip.assert_called_once()
|
||||
ckpt_paths = sd_mock.load_clip.call_args.kwargs["ckpt_paths"]
|
||||
assert ckpt_paths == ["/models/text_encoders/gemma4-12b-with-proj.safetensors"]
|
||||
|
||||
def test_with_ckpt_still_loads_both_files(self):
|
||||
fp = _fake_folder_paths(checkpoints=["some_checkpoint.safetensors"])
|
||||
with _nodes_lt_audio(fp) as (module, sd_mock):
|
||||
module.LTXAVTextEncoderLoader.execute("gemma4-12b-with-proj.safetensors", "some_checkpoint.safetensors")
|
||||
|
||||
ckpt_paths = sd_mock.load_clip.call_args.kwargs["ckpt_paths"]
|
||||
assert ckpt_paths == [
|
||||
"/models/text_encoders/gemma4-12b-with-proj.safetensors",
|
||||
"/models/checkpoints/some_checkpoint.safetensors",
|
||||
]
|
||||
|
||||
def test_none_is_a_valid_schema_option(self):
|
||||
fp = _fake_folder_paths(checkpoints=["some_checkpoint.safetensors"])
|
||||
with _nodes_lt_audio(fp) as (module, _sd_mock):
|
||||
schema = module.LTXAVTextEncoderLoader.define_schema()
|
||||
ckpt_input = next(i for i in schema.inputs if i.id == "ckpt_name")
|
||||
assert ckpt_input.options[0] == "none"
|
||||
Loading…
Reference in New Issue