From 7e2b3248016c6fdc93ab8ec40422cd63e203107d Mon Sep 17 00:00:00 2001 From: mrfakename Date: Thu, 2 Oct 2025 16:15:25 -0700 Subject: [PATCH] Allow passing a custom tqdm class --- vibevoice/modular/modeling_vibevoice_inference.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vibevoice/modular/modeling_vibevoice_inference.py b/vibevoice/modular/modeling_vibevoice_inference.py index 2076795..a47b273 100644 --- a/vibevoice/modular/modeling_vibevoice_inference.py +++ b/vibevoice/modular/modeling_vibevoice_inference.py @@ -343,6 +343,7 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener return_speech: bool = True, cfg_scale: float = 1.0, stop_check_fn: Optional[Callable[[], bool]] = None, + tqdm_class: Optional[type] = None, **kwargs, ) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]: """ @@ -423,7 +424,8 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener # Create progress iterator if verbose if kwargs.get("show_progress_bar", True): - progress_bar = tqdm(range(max_steps), desc="Generating", leave=False) + tqdm_fn = tqdm_class if tqdm_class is not None else tqdm + progress_bar = tqdm_fn(range(max_steps), desc="Generating", leave=False) else: progress_bar = range(max_steps)