Merge fc5a5e895c into 37ac9ff44f
This commit is contained in:
chelsealong
GitHub
commit
6f6958ba16
|
|
@ -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