133 lines
3.8 KiB
Python
133 lines
3.8 KiB
Python
# tests/test_generator.py
|
|
from unittest.mock import patch, MagicMock
|
|
import torch
|
|
import pytest
|
|
from nodes.generator import OmniVoiceGenerate
|
|
|
|
|
|
def make_mock_model(return_tensors=None):
|
|
mock = MagicMock()
|
|
if return_tensors is None:
|
|
return_tensors = [torch.zeros(1, 24000)] # 1 second of silence
|
|
mock.generate.return_value = return_tensors
|
|
return mock
|
|
|
|
|
|
def test_input_types_structure():
|
|
inputs = OmniVoiceGenerate.INPUT_TYPES()
|
|
required = inputs["required"]
|
|
assert "model" in required
|
|
assert "text" in required
|
|
assert "mode" in required
|
|
optional = inputs.get("optional", {})
|
|
assert "ref_audio" in optional
|
|
assert "ref_text" in optional
|
|
assert "instruct" in optional
|
|
assert "speed" in optional
|
|
assert "num_step" in optional
|
|
|
|
|
|
def test_return_type():
|
|
assert OmniVoiceGenerate.RETURN_TYPES == ("AUDIO",)
|
|
|
|
|
|
def test_generate_auto_voice():
|
|
node = OmniVoiceGenerate()
|
|
mock_model = make_mock_model()
|
|
result = node.generate(
|
|
model=mock_model,
|
|
text="Hello world",
|
|
mode="auto_voice",
|
|
speed=1.0,
|
|
num_step=32,
|
|
)
|
|
audio = result[0]
|
|
assert "waveform" in audio
|
|
assert "sample_rate" in audio
|
|
assert audio["sample_rate"] == 24000
|
|
mock_model.generate.assert_called_once_with(
|
|
text="Hello world", speed=1.0, num_step=32
|
|
)
|
|
|
|
|
|
def test_generate_voice_design():
|
|
node = OmniVoiceGenerate()
|
|
mock_model = make_mock_model()
|
|
result = node.generate(
|
|
model=mock_model,
|
|
text="Hello world",
|
|
mode="voice_design",
|
|
instruct="female, low pitch",
|
|
speed=1.0,
|
|
num_step=32,
|
|
)
|
|
audio = result[0]
|
|
assert audio["sample_rate"] == 24000
|
|
mock_model.generate.assert_called_once_with(
|
|
text="Hello world", instruct="female, low pitch", speed=1.0, num_step=32
|
|
)
|
|
|
|
|
|
def test_generate_voice_cloning():
|
|
node = OmniVoiceGenerate()
|
|
mock_model = make_mock_model()
|
|
# Simulate ComfyUI AUDIO input: waveform shape (batch, channels, samples)
|
|
ref_waveform = torch.zeros(1, 1, 24000)
|
|
ref_audio_input = {"waveform": ref_waveform, "sample_rate": 24000}
|
|
|
|
with patch("nodes.generator.torchaudio.save") as mock_save:
|
|
result = node.generate(
|
|
model=mock_model,
|
|
text="Hello world",
|
|
mode="voice_cloning",
|
|
ref_audio=ref_audio_input,
|
|
ref_text="reference text",
|
|
speed=1.0,
|
|
num_step=32,
|
|
)
|
|
|
|
assert mock_save.called
|
|
call_kwargs = mock_model.generate.call_args[1]
|
|
assert call_kwargs["ref_text"] == "reference text"
|
|
assert "ref_audio" in call_kwargs
|
|
|
|
|
|
def test_voice_cloning_without_ref_audio_raises():
|
|
node = OmniVoiceGenerate()
|
|
mock_model = make_mock_model()
|
|
with pytest.raises(ValueError, match="ref_audio"):
|
|
node.generate(
|
|
model=mock_model,
|
|
text="Hello",
|
|
mode="voice_cloning",
|
|
speed=1.0,
|
|
num_step=32,
|
|
)
|
|
|
|
|
|
def test_voice_design_without_instruct_raises():
|
|
node = OmniVoiceGenerate()
|
|
mock_model = make_mock_model()
|
|
with pytest.raises(ValueError, match="instruct"):
|
|
node.generate(
|
|
model=mock_model,
|
|
text="Hello",
|
|
mode="voice_design",
|
|
speed=1.0,
|
|
num_step=32,
|
|
)
|
|
|
|
|
|
def test_output_waveform_shape():
|
|
node = OmniVoiceGenerate()
|
|
# Simulate two chunks returned by OmniVoice
|
|
chunk1 = torch.zeros(1, 24000)
|
|
chunk2 = torch.zeros(1, 12000)
|
|
mock_model = make_mock_model(return_tensors=[chunk1, chunk2])
|
|
result = node.generate(
|
|
model=mock_model, text="Long text", mode="auto_voice", speed=1.0, num_step=32
|
|
)
|
|
waveform = result[0]["waveform"]
|
|
# Shape must be (batch=1, channels=1, samples=36000)
|
|
assert waveform.shape == (1, 1, 36000)
|