Coverage for tests/test_wrapper_vibevoice.py: 98%
318 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3from __future__ import annotations
5import os
6import sys
7import gc
8import base64
9import inspect
10import tempfile
11import pytest
13import torch
15from typing import Any
17from types import ModuleType
19from unittest.mock import patch
20from unittest.mock import MagicMock
22sys.path.append("wrapper")
23sys.path.append("wrapper/vibevoice")
26class DummySchedulerMixin:
27 pass
30class _FrozenDict(dict):
31 """Dict-like config that also supports attribute access (mimics diffusers FrozenDict)."""
32 def __getattr__(self, name: str) -> object:
33 try:
34 return self[name]
35 except KeyError:
36 raise AttributeError(name)
39class DummyConfigMixin:
40 """Mimics diffusers ConfigMixin: register_to_config stores __init__ kwargs in self.config."""
41 pass
44class DummySchedulerOutput:
45 pass
48class DummyModelOutput:
49 """Minimal stand-in for transformers.ModelOutput so @dataclass can resolve __mro__."""
50 pass
53class DummyPretrainedConfig:
54 """Minimal stand-in for transformers.PretrainedConfig."""
55 model_type = ""
56 is_composition = False
57 sub_configs: dict[str, object] = {}
59 def __init__(self, **kwargs: object) -> None:
60 for k, v in kwargs.items():
61 setattr(self, k, v)
64class DummyBaseModelOutputWithPast(DummyModelOutput):
65 pass
68class DummyPreTrainedModel(torch.nn.Module):
69 """Minimal stand-in for transformers.PreTrainedModel."""
70 config_class = None
71 base_model_prefix = ""
73 def __init__(self, config: object = None, *args: object, **kwargs: object) -> None:
74 super().__init__()
75 self.config = config
77 def _init_weights(self, module: object) -> None:
78 pass
80 def post_init(self) -> None:
81 pass
84class DummyLlamaRMSNorm(torch.nn.Module):
85 """Minimal stand-in for LlamaRMSNorm."""
86 def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
87 super().__init__()
89 def forward(self, x: object) -> object:
90 return x
93def passthrough_register_to_config(func: Any) -> Any:
94 """Mimics @register_to_config: wraps __init__ to store kwargs as self.config."""
95 sig = inspect.signature(func)
97 def wrapper(
98 self: Any,
99 *args: Any,
100 **kwargs: Any,
101 ) -> Any:
102 bound = sig.bind(self, *args, **kwargs)
103 bound.apply_defaults()
104 cfg = {k: v for k, v in bound.arguments.items() if k != "self"}
105 self.config = _FrozenDict(cfg)
106 return func(self, *args, **kwargs)
108 return wrapper
111diffusers_sched = ModuleType("diffusers.schedulers.scheduling_utils")
112diffusers_sched.SchedulerMixin = DummySchedulerMixin # type: ignore[attr-defined]
113diffusers_sched.SchedulerOutput = DummySchedulerOutput # type: ignore[attr-defined]
114diffusers_sched.KarrasDiffusionSchedulers = MagicMock() # type: ignore[attr-defined]
116diffusers_conf = ModuleType("diffusers.configuration_utils")
117diffusers_conf.ConfigMixin = DummyConfigMixin # type: ignore[attr-defined]
118diffusers_conf.register_to_config = passthrough_register_to_config # type: ignore[attr-defined]
120# Build transformers mock modules with real classes where needed for inheritance
121mock_transformers = MagicMock()
123mock_modeling_outputs = MagicMock()
124mock_modeling_outputs.ModelOutput = DummyModelOutput
125mock_modeling_outputs.BaseModelOutputWithPast = DummyBaseModelOutputWithPast
127mock_modeling_utils = MagicMock()
128mock_modeling_utils.PreTrainedModel = DummyPreTrainedModel
130mock_llama_modeling = MagicMock()
131mock_llama_modeling.LlamaRMSNorm = DummyLlamaRMSNorm
133mock_transformers_config = ModuleType("transformers.configuration_utils")
134mock_transformers_config.PretrainedConfig = DummyPretrainedConfig # type: ignore[attr-defined]
136mock_modules = {
137 "modeling_vibevoice_inference": MagicMock(),
138 # "modeling_vibevoice": MagicMock(),
139 "transformers": mock_transformers,
140 "transformers.utils": MagicMock(),
141 "transformers.modeling_utils": mock_modeling_utils,
142 "transformers.modeling_outputs": mock_modeling_outputs,
143 "transformers.generation": MagicMock(),
144 "transformers.models": MagicMock(),
145 "transformers.models.llama": MagicMock(),
146 "transformers.models.llama.modeling_llama": mock_llama_modeling,
147 "transformers.models.qwen2": MagicMock(),
148 "transformers.models.qwen2.tokenization_qwen2": MagicMock(),
149 "transformers.models.qwen2.tokenization_qwen2_fast": MagicMock(),
150 "transformers.models.qwen2.configuration_qwen2": MagicMock(),
151 "transformers.tokenization_utils_base": MagicMock(),
152 "transformers.feature_extraction_utils": MagicMock(),
153 "transformers.modeling_flash_attention_utils": MagicMock(),
154 "transformers.configuration_utils": mock_transformers_config,
155 "transformers.activations": MagicMock(),
156 "diffusers": MagicMock(),
157 "diffusers.schedulers": ModuleType("diffusers.schedulers"),
158 "diffusers.schedulers.scheduling_utils": diffusers_sched,
159 "diffusers.configuration_utils": diffusers_conf,
160 "diffusers.utils": MagicMock(),
161 "diffusers.utils.torch_utils": MagicMock(),
162}
165with patch.dict(sys.modules, mock_modules):
166 from vibevoice.wrapper_vibevoice import VibeVoiceGeneration
167 from vibevoice.wrapper_vibevoice import VoiceMapper
169 from configuration_vibevoice import VibeVoiceConfig
170 from configuration_vibevoice import VibeVoiceSemanticTokenizerConfig
171 from modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
172 from audio_streamer import AudioStreamer
174 # from modeling_vibevoice import VibeVoicePreTrainedModel
175 from modeling_vibevoice import VibeVoiceCausalLMOutputWithPast
176 from modeling_vibevoice import SpeechConnector
177 from modeling_vibevoice import VibeVoiceGenerationOutput
178 from modeling_vibevoice import VibeVoicePreTrainedModel
179 from modular_vibevoice_diffusion_head import RMSNorm
180 from modular_vibevoice_tokenizer import NormConvTranspose1d
181 from modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache
182 from modular_vibevoice_tokenizer import VibeVoiceSemanticTokenizerModel
183 from schedule.timestep_sampler import UniformSampler, LogitNormalSampler
184 from schedule.dpm_solver import DPMSolverMultistepScheduler
185 from schedule.dpm_solver import rescale_zero_terminal_snr
186 from schedule.dpm_solver import betas_for_alpha_bar
187 from modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizer
190@pytest.mark.asyncio
191async def test_vibevoice() -> None:
192 model = VibeVoiceGeneration()
193 assert model is not None
194 assert model.model_name == "vibevoice"
195 assert model.status == "initializing"
197 with pytest.raises(AttributeError): # TODO
198 model.init()
199 assert model.status == "failed"
201 health = model.get_health()
202 assert health is not None
203 assert len(health) > 1
204 timestamps = model.get_timestamps()
205 assert timestamps is not None
206 assert len(timestamps) >= 1
208 with pytest.raises(ValueError):
209 await model.get_rest_args(None)
210 with pytest.raises(ValueError):
211 await model.get_rest_args({})
212 await model.get_rest_args({
213 "text": "Test text"
214 })
216 with pytest.raises(ValueError, match="Model not initialized"):
217 await model.warmup()
219 with pytest.raises(ValueError, match="Model not initialized"):
220 await model.generate(
221 text="Test text",
222 output_type="audio_path")
224 del model
225 gc.collect()
228def test_vibevoice_libs() -> None:
229 assert VibeVoiceCausalLMOutputWithPast is not None
230 assert VibeVoiceGenerationOutput is not None
231 assert SpeechConnector is not None
232 assert VibeVoicePreTrainedModel is not None
233 assert RMSNorm is not None
236def test_tokenizer() -> None:
237 assert NormConvTranspose1d is not None
238 assert VibeVoiceTokenizerStreamingCache is not None
239 tokenizer = VibeVoiceSemanticTokenizerModel(VibeVoiceSemanticTokenizerConfig())
240 assert tokenizer is not None
243def test_text_tokenizer() -> None:
244 tokenizer = VibeVoiceTextTokenizer(
245 vocab_file="missing_vocab_file",
246 merges_file="missing_merges_file")
247 assert tokenizer is not None
250def test_model_inference() -> None:
251 # Configuration
252 vibevoice_config = VibeVoiceConfig()
253 assert vibevoice_config is not None
255 # Inference
256 inference = VibeVoiceForConditionalGenerationInference(vibevoice_config)
257 assert inference is not None
258 assert inference.forward() is not None
259 assert inference.generate() is not None
262def test_audio_streamer() -> None:
263 audio_streamer = AudioStreamer(batch_size=8)
264 assert audio_streamer is not None
265 audio_streamer.put(
266 audio_chunks=MagicMock(),
267 sample_indices=MagicMock(),
268 )
269 audio_streamer.end(sample_indices=MagicMock())
272def test_voice_mapper() -> None:
273 voice_mapper = VoiceMapper()
274 voice_mapper.setup_voice_presets()
276 # voice_presets is a dict (may be empty or populated depending on environment)
277 assert isinstance(voice_mapper.voice_presets, dict)
279 if not voice_mapper.voice_presets:
280 # No voices available: get_voice_path raises ValueError
281 with pytest.raises(ValueError, match="No voice presets available"):
282 voice_mapper.get_voice_path("any_speaker")
283 else:
284 # Voices available: get_voice_path returns a path for any speaker name
285 path = voice_mapper.get_voice_path("any_speaker")
286 assert isinstance(path, str)
289@pytest.mark.asyncio
290async def test_vibevoice_get_rest_args_voice() -> None:
291 model = VibeVoiceGeneration()
293 # Custom voice parameter is returned in args
294 result = await model.get_rest_args({"text": "Hello world", "voice": "custom_voice"})
295 assert result["args"]["voice"] == "custom_voice"
297 # Default voice returned when voice not specified
298 result_default = await model.get_rest_args({"text": "Hello world"})
299 assert "voice" in result_default["args"]
300 assert result_default["args"]["voice"] == "af_heart"
303@pytest.mark.asyncio
304async def test_vibevoice_get_rest_args_voice_sample() -> None:
305 """get_rest_args forwards voice_sample when present and omits it when absent."""
306 model = VibeVoiceGeneration()
308 with open("tests/data/sample.wav", "rb") as f:
309 wav_bytes = f.read()
310 dummy_audio = base64.b64encode(wav_bytes).decode()
312 # voice_sample present -> included in args
313 result = await model.get_rest_args({
314 "text": "Hello world",
315 "voice_sample": dummy_audio,
316 })
317 assert result["args"].get("voice_sample") == dummy_audio
319 # voice_sample absent -> not included in args
320 result_no_sample = await model.get_rest_args({"text": "Hello world"})
321 assert "voice_sample" not in result_no_sample["args"]
324def test_decode_voice_sample_to_tmp_file() -> None:
325 """_decode_voice_sample_to_tmp_file writes the decoded bytes to a temp WAV file."""
326 model = VibeVoiceGeneration()
328 with open("tests/data/sample.wav", "rb") as f:
329 audio_content = f.read()
330 voice_sample_b64 = base64.b64encode(audio_content).decode()
332 tmp_path = model._decode_voice_sample_to_tmp_file(voice_sample_b64)
333 try:
334 assert os.path.exists(tmp_path), "Temp file must be created"
335 assert tmp_path.endswith(".wav"), "Temp file must have .wav extension"
336 with open(tmp_path, "rb") as f:
337 written = f.read()
338 assert written == audio_content, "Written bytes must match original audio content"
339 finally:
340 if os.path.exists(tmp_path):
341 os.unlink(tmp_path)
344@pytest.mark.asyncio
345async def test_vibevoice_voice_sample_debug_file_saved() -> None:
346 """When voice_sample + job_id are provided, the decoded audio is persisted to
347 /tmp/{job_id}_voice_sample.wav before processing begins."""
348 from unittest.mock import patch, MagicMock
350 model = VibeVoiceGeneration()
352 with open("tests/data/sample.wav", "rb") as f:
353 wav_bytes = f.read()
354 voice_sample_b64 = base64.b64encode(wav_bytes).decode()
356 job_id = "vibevoice_debug_test_001"
357 debug_path = f"/tmp/{os.path.basename(job_id)}_voice_sample.wav"
359 if os.path.exists(debug_path):
360 os.unlink(debug_path)
362 # Bypass model-init guard; leave processor=None so generate() raises *after*
363 # it has already written the debug file.
364 model.voice_mapper = MagicMock()
365 model.processor = None
367 with patch.object(model, "_assert_model_init"):
368 with pytest.raises(ValueError, match="Processor not initialized"):
369 await model.generate(text="Hello", voice_sample=voice_sample_b64, job_id=job_id)
371 try:
372 assert os.path.exists(debug_path), "Debug voice sample file must be created"
373 with open(debug_path, "rb") as f:
374 assert f.read() == wav_bytes, "Debug file content must match original audio"
375 finally:
376 if os.path.exists(debug_path):
377 os.unlink(debug_path)
380@pytest.mark.asyncio
381async def test_vibevoice_voice_sample_no_job_id_uses_tmp_file() -> None:
382 """When voice_sample is provided but job_id is None, a temporary file is used
383 (and will be cleaned up), not a named debug file."""
384 from unittest.mock import patch, MagicMock
386 model = VibeVoiceGeneration()
388 with open("tests/data/sample.wav", "rb") as f:
389 wav_bytes = f.read()
390 voice_sample_b64 = base64.b64encode(wav_bytes).decode()
392 recorded: list[str] = []
393 original_decode = model._decode_voice_sample_to_tmp_file
395 def capturing_decode(vs: str) -> str:
396 path = original_decode(vs)
397 recorded.append(path)
398 return path
400 model.voice_mapper = MagicMock()
401 model.processor = None
403 with patch.object(model, "_assert_model_init"):
404 with patch.object(model, "_decode_voice_sample_to_tmp_file", side_effect=capturing_decode):
405 with pytest.raises(ValueError, match="Processor not initialized"):
406 # job_id=None → should use temp file, not a named debug path
407 await model.generate(text="Hello", voice_sample=voice_sample_b64, job_id=None)
409 assert len(recorded) == 1, "_decode_voice_sample_to_tmp_file must be called once"
410 # Temp file should have been cleaned up by the finally block inside generate()
411 assert not os.path.exists(recorded[0]), "Temp voice file must be deleted after generate()"
414def test_cleanup_tmp_voice_file() -> None:
415 """_cleanup_tmp_voice_file removes the file and handles None / missing paths gracefully."""
416 model = VibeVoiceGeneration()
418 # None input is a no-op (must not raise)
419 model._cleanup_tmp_voice_file(None)
421 # Existing file is deleted
422 with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
423 tmp_path = f.name
424 assert os.path.exists(tmp_path)
425 model._cleanup_tmp_voice_file(tmp_path)
426 assert not os.path.exists(tmp_path), "File must be deleted by _cleanup_tmp_voice_file"
428 # Already-deleted path must not raise
429 model._cleanup_tmp_voice_file(tmp_path)
432def test_timestep_samplers() -> None:
433 uniform = UniformSampler(timesteps=100)
434 result_u = uniform.sample(batch_size=4, device=torch.device('cpu'))
435 assert result_u is not None
437 logit = LogitNormalSampler(timesteps=100)
438 result_l = logit.sample(batch_size=4, device=torch.device('cpu'))
439 assert result_l is not None
442def test_dpm_solver_scheduler() -> None:
443 scheduler = DPMSolverMultistepScheduler(num_train_timesteps=100)
444 assert scheduler is not None
445 assert hasattr(scheduler, "betas")
446 assert hasattr(scheduler, "alphas_cumprod")
448 # Test other beta schedules to cover more branches
449 for schedule in ["scaled_linear", "squaredcos_cap_v2", "cosine"]:
450 s = DPMSolverMultistepScheduler(num_train_timesteps=50, beta_schedule=schedule)
451 assert hasattr(s, "betas")
453 # Unknown schedule raises NotImplementedError
454 with pytest.raises(NotImplementedError):
455 DPMSolverMultistepScheduler(beta_schedule="unknown_schedule")
457 # Test betas_for_alpha_bar directly with different alpha_transform_types
458 for alpha_type in ["cosine", "exp", "cauchy", "laplace"]:
459 betas = betas_for_alpha_bar(10, alpha_transform_type=alpha_type)
460 assert betas is not None
462 with pytest.raises(ValueError):
463 betas_for_alpha_bar(10, alpha_transform_type="unknown")
466def test_dpm_solver_set_timesteps() -> None:
468 # set_timesteps with different timestep_spacing values
469 for spacing in ["linspace", "leading", "trailing"]:
470 s = DPMSolverMultistepScheduler(num_train_timesteps=100, timestep_spacing=spacing)
471 s.set_timesteps(10)
472 assert s.num_inference_steps == 10
473 assert len(s.timesteps) == 10
475 # set_timesteps with use_karras_sigmas covers _convert_to_karras + _sigma_to_t
476 s_karras = DPMSolverMultistepScheduler(num_train_timesteps=100, use_karras_sigmas=True)
477 s_karras.set_timesteps(10)
478 assert s_karras.num_inference_steps == 10
480 # set_timesteps with use_lu_lambdas covers _convert_to_lu + _sigma_to_t
481 s_lu = DPMSolverMultistepScheduler(num_train_timesteps=100, use_lu_lambdas=True)
482 s_lu.set_timesteps(10)
483 assert s_lu.num_inference_steps == 10
485 # set_timesteps with custom timesteps list
486 s_custom = DPMSolverMultistepScheduler(num_train_timesteps=100)
487 s_custom.set_timesteps(timesteps=[80, 60, 40, 20])
488 assert s_custom.num_inference_steps == 4
490 # Error: neither argument provided
491 with pytest.raises(ValueError, match="Must pass exactly one of"):
492 s_custom.set_timesteps()
494 # Error: both provided
495 with pytest.raises(ValueError, match="Can only pass one"):
496 s_custom.set_timesteps(num_inference_steps=10, timesteps=[80, 60])
498 # Error: custom timesteps with karras_sigmas
499 with pytest.raises(ValueError, match="use_karras_sigmas"):
500 s_karras.set_timesteps(timesteps=[80, 60])
502 # __len__, set_begin_index, step_index and begin_index properties
503 s = DPMSolverMultistepScheduler(num_train_timesteps=200)
504 assert len(s) == 200
505 s.set_begin_index(5)
506 assert s.begin_index == 5
507 assert s.step_index is None # before any step
509 # rescale_zero_terminal_snr
510 betas = torch.linspace(0.0001, 0.02, 100)
511 rescaled = rescale_zero_terminal_snr(betas)
512 assert rescaled is not None
513 assert rescaled.shape == betas.shape
515 # add_noise and get_velocity
516 s.set_timesteps(10)
517 original = torch.randn(2, 4)
518 noise = torch.randn(2, 4)
519 timesteps = torch.IntTensor([50, 80])
520 noisy = s.add_noise(original, noise, timesteps)
521 assert noisy.shape == original.shape
522 velocity = s.get_velocity(original, noise, timesteps)
523 assert velocity.shape == original.shape
525 # _sigma_to_alpha_sigma_t and _sigma_to_t are exercised by set_timesteps
526 # Make sure sigmas attribute exists after set_timesteps
527 assert hasattr(s, "sigmas")
528 assert s.sigmas is not None
531def test_vibevoice_model_methods() -> None:
532 """Test _assert_model_init, init_parallelism, init_model_parallelism, model_compile."""
533 model = VibeVoiceGeneration()
535 # _assert_model_init raises before model is initialized (status != "ok")
536 with pytest.raises(ValueError, match="Model not initialized"):
537 model._assert_model_init()
539 # init_parallelism runs without error (no CUDA in test env, uses CPU path)
540 model.init_parallelism()
541 assert model.rank == 0
542 assert model.world_size == 1
544 # init_model_parallelism: world_size=1, just logs a warning if > 1
545 model.init_model_parallelism()
547 # model_compile: short-circuits when torch_compile=False
548 model.torch_compile = False
549 model.model_compile() # should return immediately without error
551 del model