Make is_prefill configurable (#11)

2025-09-23 11:39:53 -07:00
parent 74d57a72a9
commit ae6cb88889
1 changed files with 1 additions and 1 deletions
@@ -339,6 +339,7 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
        speech_tensors: Optional[torch.FloatTensor] = None,
        speech_masks: Optional[torch.BoolTensor] = None,
        speech_input_mask: Optional[torch.BoolTensor] = None,
+        is_prefill: bool = True,
        return_speech: bool = True,
        cfg_scale: float = 1.0,
        stop_check_fn: Optional[Callable[[], bool]] = None,
@@ -390,7 +391,6 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
        device = input_ids.device
        finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device)
        correct_cnt = torch.zeros(batch_size, dtype=torch.long, device=device)
-        is_prefill = True
        inputs_embeds = None
        verbose = kwargs.get("verbose", False)