WildAi
e20b4fd9d8
Add a dedicated loader node that loads VibeVoice checkpoints from user-provided .safetensors/.pt files in models/diffusion_models, since ComfyUI's stock Load Diffusion Model cannot detect VibeVoice architecture. The loader emits a VIBEVOICE_MODEL custom type bundle (model/processor/config/state_dict) consumed by an optional external_model input on the TTS, Realtime, and ASR nodes.
- modules/custom_types.py: VibeVoiceModel = io.Custom("VIBEVOICE_MODEL")
- modules/external_loader.py: sidecar config/preprocessor/tokenizer resolution + load_external_vibevoice_model() (TTS/streaming) and load_external_vibevoice_asr_model() (ASR); CPU-first load, in-memory state-dict injection, dtype cast, optional 4-bit quant (TTS only), SageAttention
- nodes/external_loader_node.py: VibeVoiceLoadExternalModel node
- modules/generation.py: ExternalVibeVoiceModelHandler + load_vibevoice_from_external()
- modules/asr_generation.py: ExternalVibeVoiceASRModelHandler + load_asr_from_external()
- nodes/tts_node.py, realtime_node.py, asr_node.py: optional external_model input with kind guards (streaming/ASR/TTS mismatch rejection)
- README: Loading External Models section + v2.2.0 changelog
- example_workflows/VibeVoice_external_model_example.json
- docs/plans/2026-08-16-external-model-input.md (status: COMPLETED)
Tests: +122 new/extended (test_custom_types, test_external_loader, test_external_loader_node, plus extensions to node_schema, generation, asr_generation, realtime_node, asr_node, patcher_behavioral, integration, docs_consistency, workflow, imports, extension). Full suite: 687 passed, 5 pre-existing failures, 4 skipped.
103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
"""Tests for example workflow JSON."""
|
|
|
|
import os
|
|
import json
|
|
import pytest
|
|
|
|
|
|
class TestWorkflow:
|
|
"""Test the example workflow JSON."""
|
|
|
|
def _get_workflow_path(self):
|
|
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
return os.path.join(root, "example_workflows", "VibeVoice_example.json")
|
|
|
|
def test_workflow_json_valid(self):
|
|
"""The workflow JSON should be valid JSON."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
assert data is not None
|
|
|
|
def test_workflow_has_nodes(self):
|
|
"""The workflow should contain nodes."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
# Workflow format varies; check for common keys
|
|
if "nodes" in data:
|
|
assert len(data["nodes"]) > 0
|
|
elif "prompt" in data:
|
|
assert len(data["prompt"]) > 0
|
|
|
|
|
|
class TestExternalModelWorkflow:
|
|
"""Phase 6.2: the external-model example workflow must be valid."""
|
|
|
|
def _get_workflow_path(self):
|
|
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
return os.path.join(
|
|
root, "example_workflows", "VibeVoice_external_model_example.json"
|
|
)
|
|
|
|
def test_external_model_workflow_json_is_valid(self):
|
|
"""The external-model workflow JSON parses and contains nodes."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
assert data is not None
|
|
assert "nodes" in data
|
|
assert len(data["nodes"]) > 0
|
|
|
|
def test_external_model_workflow_has_loader_node(self):
|
|
"""The workflow contains the Load VibeVoice Model node."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
node_types = [n.get("type") for n in data["nodes"]]
|
|
assert "VibeVoiceLoadExternalModel" in node_types
|
|
|
|
def test_external_model_workflow_has_tts_node(self):
|
|
"""The workflow contains the VibeVoice TTS node."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
node_types = [n.get("type") for n in data["nodes"]]
|
|
assert "VibeVoiceTTS" in node_types
|
|
|
|
def test_external_model_workflow_links_loader_to_tts(self):
|
|
"""A VIBEVOICE_MODEL link connects the loader node to the TTS node."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
# Map node id → type.
|
|
node_by_id = {n["id"]: n.get("type") for n in data["nodes"]}
|
|
|
|
# links format: [link_id, from_node, from_slot, to_node, to_slot, type]
|
|
vibevoice_model_links = [
|
|
link for link in data.get("links", [])
|
|
if link[5] == "VIBEVOICE_MODEL"
|
|
]
|
|
assert len(vibevoice_model_links) >= 1
|
|
link = vibevoice_model_links[0]
|
|
assert node_by_id.get(link[1]) == "VibeVoiceLoadExternalModel"
|
|
assert node_by_id.get(link[3]) == "VibeVoiceTTS"
|
|
|
|
def test_external_model_workflow_tts_has_external_model_input(self):
|
|
"""The TTS node declares an external_model input wired to the link."""
|
|
path = self._get_workflow_path()
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
tts_node = next(
|
|
n for n in data["nodes"] if n.get("type") == "VibeVoiceTTS"
|
|
)
|
|
ext_inputs = [
|
|
i for i in tts_node.get("inputs", [])
|
|
if i.get("name") == "external_model"
|
|
]
|
|
assert len(ext_inputs) == 1
|
|
assert ext_inputs[0].get("type") == "VIBEVOICE_MODEL"
|
|
assert ext_inputs[0].get("link") is not None
|