valenti/spark-vibevoice
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow passing a custom tqdm class

Browse Source
This commit is contained in:
mrfakename
2025-10-02 16:15:25 -07:00
parent c25ae58bb2
commit 7e2b324801
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vibevoice/modular/modeling_vibevoice_inference.py
+3 -1
View File
@@ -343,6 +343,7 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
return_speech: bool = True, return_speech: bool = True,
cfg_scale: float = 1.0, cfg_scale: float = 1.0,
stop_check_fn: Optional[Callable[[], bool]] = None, stop_check_fn: Optional[Callable[[], bool]] = None,
tqdm_class: Optional[type] = None,
**kwargs, **kwargs,
) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]: ) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]:
""" """
@@ -423,7 +424,8 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
# Create progress iterator if verbose # Create progress iterator if verbose
if kwargs.get("show_progress_bar", True): 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: else:
progress_bar = range(max_steps) progress_bar = range(max_steps)