Make is_prefill configurable (#11)

This commit is contained in:
Etherll
2025-09-23 11:39:53 -07:00
committed by GitHub
parent 74d57a72a9
commit ae6cb88889
@@ -339,6 +339,7 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
speech_tensors: Optional[torch.FloatTensor] = None, speech_tensors: Optional[torch.FloatTensor] = None,
speech_masks: Optional[torch.BoolTensor] = None, speech_masks: Optional[torch.BoolTensor] = None,
speech_input_mask: Optional[torch.BoolTensor] = None, speech_input_mask: Optional[torch.BoolTensor] = None,
is_prefill: bool = True,
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,
@@ -390,7 +391,6 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
device = input_ids.device device = input_ids.device
finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device) finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device)
correct_cnt = torch.zeros(batch_size, dtype=torch.long, device=device) correct_cnt = torch.zeros(batch_size, dtype=torch.long, device=device)
is_prefill = True
inputs_embeds = None inputs_embeds = None
verbose = kwargs.get("verbose", False) verbose = kwargs.get("verbose", False)