Coverage for wrapper/vibevoice/modeling_vibevoice_inference.py: 0%

1# mypy: ignore-errors 

2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/modeling_vibevoice_streaming_inference.py 

3 

4from dataclasses import dataclass 

5 

6from typing import List 

7from typing import Optional 

8from typing import Tuple 

9from typing import Union 

10from typing import Callable 

11from typing import Any 

12from typing import Dict 

13from typing import cast 

14 

15from tqdm import tqdm 

16 

17import torch 

18import torch.nn as nn 

19 

20from transformers import AutoModelForCausalLM 

21 

22from transformers.generation import GenerationMixin 

23from transformers.generation import GenerationConfig 

24from transformers.generation import LogitsProcessor 

25from transformers.generation import LogitsProcessorList 

26from transformers.generation import StoppingCriteriaList 

27from transformers.modeling_outputs import BaseModelOutputWithPast 

28from transformers.modeling_outputs import ModelOutput 

29from transformers.modeling_utils import PreTrainedModel 

30from transformers.utils import logging 

31 

32from configuration_vibevoice import VibeVoiceConfig 

33from modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache 

34 

35from modeling_vibevoice import VibeVoiceModel 

36from modeling_vibevoice import VibeVoicePreTrainedModel 

37from audio_streamer import AudioStreamer 

38from audio_streamer import AsyncAudioStreamer 

39 

40logger = logging.get_logger(__name__) 

41 

42# if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None: 

43# modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"] 

44 

45 

46@dataclass 

47class VibeVoiceCausalLMOutputWithPast(BaseModelOutputWithPast): 

48 logits: Optional[torch.FloatTensor] = None 

49 

50 

51@dataclass 

52class VibeVoiceGenerationOutput(ModelOutput): 

53 """ 

54 Output type for VibeVoice generation. 

55 

56 Args: 

57 sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`): 

58 The generated sequences. 

59 speech_outputs (`List[torch.FloatTensor]`, *optional*): 

60 List of generated speech waveforms or latents for each speech segment. 

61 """ 

62 sequences: Optional[torch.LongTensor] = None 

63 speech_outputs: Optional[List[torch.FloatTensor]] = None 

64 reach_max_step_sample: Optional[torch.BoolTensor] = None 

65 

66 

67class VibeVoiceTokenConstraintProcessor(LogitsProcessor): 

68 """Constrains token generation to only valid tokens during speech generation.""" 

69 

70 def __init__( 

71 self, 

72 valid_token_ids: List[int], 

73 device: Optional[torch.device] = None 

74 ) -> None: 

75 self.valid_token_ids = torch.tensor(valid_token_ids, dtype=torch.long, device=device) 

76 

77 def __call__( 

78 self, 

79 input_ids: torch.LongTensor, 

80 scores: torch.FloatTensor 

81 ) -> torch.FloatTensor: 

82 # Create a mask for valid tokens 

83 mask = torch.full_like(scores, float('-inf')) 

84 mask[:, self.valid_token_ids] = 0 

85 

86 # Apply mask to scores 

87 scores = scores + mask 

88 return cast(torch.FloatTensor, scores) 

89 

90 

91class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, GenerationMixin): 

92 _tied_weights_keys = ["lm_head.weight"] 

93 _tp_plan = {"lm_head": "colwise_rep"} 

94 

95 def __init__( 

96 self, 

97 config: VibeVoiceConfig, 

98 ) -> None: 

99 super().__init__(config) 

100 

101 # Initialize the base model 

102 self.model = VibeVoiceModel(config) 

103 

104 # LM head for text generation 

105 self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.decoder_config.vocab_size, bias=False) 

106 

107 # inference configuration 

108 self.ddpm_inference_steps = config.diffusion_head_config.ddpm_num_inference_steps 

109 

110 # Initialize weights and apply final processing 

111 self.post_init() 

112 

113 @property 

114 def noise_scheduler(self) -> Any: 

115 return self.model.noise_scheduler 

116 

117 @property 

118 def prediction_head(self) -> nn.Module: 

119 return self.model.prediction_head 

120 

121 @property 

122 def speech_scaling_factor(self) -> torch.Tensor: 

123 return self.model.speech_scaling_factor 

124 

125 @property 

126 def speech_bias_factor(self) -> torch.Tensor: 

127 return self.model.speech_bias_factor 

128 

129 @property 

130 def acoustic_tokenizer(self) -> Any: 

131 return self.model.acoustic_tokenizer 

132 

133 @property 

134 def semantic_tokenizer(self) -> Any: 

135 return self.model.semantic_tokenizer 

136 

137 @property 

138 def acoustic_connector(self) -> nn.Module: 

139 return self.model.acoustic_connector 

140 

141 @property 

142 def semantic_connector(self) -> nn.Module: 

143 return self.model.semantic_connector 

144 

145 def tie_weights(self) -> None: 

146 """ 

147 Tie the weights between the input embeddings and the output embeddings. 

148 """ 

149 # Tie lm_head.weight to language_model.embed_tokens.weight 

150 if not getattr(self.config, 'tie_word_embeddings', False): 

151 return 

152 

153 if hasattr(self, 'lm_head') and hasattr(self.model.language_model, 'embed_tokens'): 

154 self.lm_head.weight = self.model.language_model.embed_tokens.weight 

155 

156 def get_input_embeddings(self) -> nn.Module: 

157 return self.model.get_input_embeddings() 

158 

159 def set_input_embeddings( 

160 self, 

161 value: nn.Module 

162 ) -> None: 

163 self.model.set_input_embeddings(value) 

164 

165 def get_output_embeddings(self) -> nn.Module: 

166 return self.lm_head 

167 

168 def set_output_embeddings( 

169 self, 

170 new_embeddings: nn.Module 

171 ) -> None: 

172 self.lm_head = cast(nn.Linear, new_embeddings) 

173 

174 def set_speech_tokenizers( 

175 self, 

176 acoustic_tokenizer: Optional[Any] = None, 

177 semantic_tokenizer: Optional[Any] = None, 

178 ) -> None: 

179 """Set the speech tokenizers used for encoding and decoding speech.""" 

180 self.model.set_speech_tokenizers(acoustic_tokenizer, semantic_tokenizer) 

181 

182 def set_ddpm_inference_steps( 

183 self, 

184 num_steps: Optional[int] = None 

185 ) -> None: 

186 self.ddpm_inference_steps = num_steps or self.config.diffusion_head_config.ddpm_num_inference_steps 

187 

188 def _process_speech_inputs( 

189 self, 

190 speech_tensors: torch.FloatTensor, 

191 speech_masks: torch.BoolTensor, 

192 speech_type: str = "audio" 

193 ) -> Tuple[torch.Tensor, torch.Tensor]: 

194 """Process speech inputs through tokenizers and connectors.""" 

195 encoder_output = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1)) 

196 acoustic_latents = encoder_output.sample(dist_type=self.model.acoustic_tokenizer.std_dist_type)[0] 

197 

198 # Apply scaling and bias 

199 acoustic_features = (acoustic_latents + self.model.speech_bias_factor.to(acoustic_latents.device)) * \ 

200 self.model.speech_scaling_factor.to(acoustic_latents.device) 

201 

202 # Connect to language model space 

203 acoustic_connected = self.model.acoustic_connector(acoustic_features)[speech_masks.cpu()] 

204 

205 return acoustic_features, acoustic_connected 

206 

207 # @can_return_tuple 

208 def forward( 

209 self, 

210 input_ids: Optional[torch.LongTensor] = None, 

211 attention_mask: Optional[torch.Tensor] = None, 

212 position_ids: Optional[torch.LongTensor] = None, 

213 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 

214 inputs_embeds: Optional[torch.FloatTensor] = None, 

215 labels: Optional[torch.LongTensor] = None, 

216 use_cache: Optional[bool] = None, 

217 output_attentions: Optional[bool] = None, 

218 output_hidden_states: Optional[bool] = None, 

219 return_dict: Optional[bool] = None, 

220 cache_position: Optional[torch.LongTensor] = None, 

221 speech_tensors: Optional[torch.FloatTensor] = None, 

222 speech_masks: Optional[torch.BoolTensor] = None, 

223 speech_input_mask: Optional[torch.BoolTensor] = None, 

224 logits_to_keep: Union[int, slice] = 0, 

225 **kwargs: Any, 

226 ) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]: 

227 """ 

228 Args: 

229 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): 

230 Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., 

231 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored 

232 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. 

233 speech_tensors (`torch.FloatTensor`, *optional*): 

234 Input speech waveforms for voice cloning or speech understanding. 

235 speech_masks (`torch.BoolTensor`, *optional*): 

236 Masks indicating valid speech frames. 

237 speech_input_mask (`torch.BoolTensor`, *optional*): 

238 Positions in the input sequence where speech embeddings should be inserted. 

239 

240 Returns: 

241 `VibeVoiceCausalLMOutputWithPast` or tuple 

242 """ 

243 return_dict = return_dict if return_dict is not None else self.config.use_return_dict 

244 

245 # Get embeddings 

246 if inputs_embeds is None: 

247 inputs_embeds = self.model.get_input_embeddings()(input_ids) 

248 

249 # Process speech inputs if provided 

250 if speech_tensors is not None and speech_masks is not None: 

251 acoustic_features, speech_embeds = self._process_speech_inputs( 

252 cast(torch.FloatTensor, speech_tensors.to(self.dtype)), speech_masks) 

253 if speech_input_mask is not None: 

254 inputs_embeds[speech_input_mask] = speech_embeds 

255 

256 outputs = self.model( 

257 inputs_embeds=inputs_embeds, 

258 attention_mask=attention_mask, 

259 position_ids=position_ids, 

260 past_key_values=past_key_values, 

261 use_cache=use_cache, 

262 output_attentions=output_attentions, 

263 output_hidden_states=output_hidden_states, 

264 return_dict=return_dict, 

265 cache_position=cache_position, 

266 **kwargs, 

267 ) 

268 

269 hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state 

270 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss 

271 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep 

272 logits = self.lm_head(hidden_states[:, slice_indices, :]) 

273 

274 if labels is not None: 

275 raise NotImplementedError("Loss computation is not implemented in this version.") 

276 

277 return VibeVoiceCausalLMOutputWithPast( 

278 logits=logits, 

279 past_key_values=outputs.past_key_values, 

280 last_hidden_state=hidden_states, 

281 attentions=outputs.attentions, 

282 ) 

283 

284 def _build_generate_config_model_kwargs( 

285 self, 

286 generation_config: Optional[GenerationConfig], 

287 inputs: Optional[torch.Tensor], 

288 tokenizer: Optional[PreTrainedModel] = None, 

289 return_processors: bool = False, 

290 **kwargs: Dict[str, Any] 

291 ) -> Any: 

292 if generation_config is None: 

293 assert tokenizer is not None, "tokenizer must be provided when generation_config is None" 

294 generation_config = GenerationConfig( 

295 bos_token_id=tokenizer.bos_token_id, 

296 eos_token_id=tokenizer.eos_token_id, 

297 pad_token_id=tokenizer.pad_token_id 

298 ) 

299 else: 

300 assert tokenizer is not None, "tokenizer must be provided" 

301 if isinstance(generation_config, dict): 

302 gen_config_dict = generation_config 

303 else: 

304 gen_config_dict = generation_config.to_dict() 

305 gen_config_dict.pop('bos_token_id', None) 

306 gen_config_dict.pop('eos_token_id', None) 

307 gen_config_dict.pop('pad_token_id', None) 

308 generation_config = GenerationConfig( 

309 **gen_config_dict, 

310 bos_token_id=tokenizer.bos_token_id, 

311 eos_token_id=tokenizer.eos_token_id, 

312 pad_token_id=tokenizer.pad_token_id 

313 ) 

314 

315 generation_config, model_kwargs = self._prepare_generation_config( # type: ignore[call-arg, unused-ignore] 

316 generation_config, 

317 True, 

318 speech_start_id=tokenizer.speech_start_id, 

319 speech_end_id=tokenizer.speech_end_id, 

320 speech_diffusion_id=tokenizer.speech_diffusion_id, 

321 **kwargs 

322 ) 

323 setattr(generation_config, 'speech_start_id', tokenizer.speech_start_id) 

324 setattr(generation_config, 'speech_end_id', tokenizer.speech_end_id) 

325 setattr(generation_config, 'speech_diffusion_id', tokenizer.speech_diffusion_id) 

326 

327 inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs( 

328 inputs, generation_config.bos_token_id, model_kwargs) 

329 batch_size = inputs_tensor.shape[0] 

330 device = self.device 

331 

332 self._prepare_special_tokens(generation_config, True, device=device) 

333 generation_config.use_cache = True 

334 model_kwargs["use_cache"] = generation_config.use_cache 

335 input_ids = inputs_tensor.to(self.device) 

336 

337 input_ids_length = input_ids.shape[1] 

338 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None 

339 has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None 

340 generation_config = self._prepare_generated_length( 

341 generation_config=generation_config, 

342 has_default_max_length=has_default_max_length, 

343 has_default_min_length=has_default_min_length, 

344 model_input_name=model_input_name, 

345 inputs_tensor=inputs_tensor, 

346 input_ids_length=input_ids_length, 

347 ) 

348 

349 max_cache_length = generation_config.max_length - 1 

350 self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length, device) 

351 model_kwargs['cache_position'] = torch.arange(input_ids_length, device=device, dtype=torch.long) 

352 for k, v in model_kwargs.items(): 

353 if isinstance(v, torch.Tensor): 

354 model_kwargs[k] = v.to(device=device) 

355 

356 if return_processors: 

357 logits_processor = self._get_logits_processor( 

358 generation_config=generation_config, 

359 input_ids_seq_length=input_ids_length, 

360 encoder_input_ids=inputs_tensor, 

361 prefix_allowed_tokens_fn=None, 

362 logits_processor=LogitsProcessorList(), 

363 device=inputs_tensor.device, 

364 model_kwargs=model_kwargs, 

365 ) 

366 

367 stopping_criteria = self._get_stopping_criteria( 

368 generation_config=generation_config, 

369 stopping_criteria=StoppingCriteriaList()) 

370 

371 return generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria 

372 else: 

373 return generation_config, model_kwargs, input_ids 

374 

375 @torch.no_grad() 

376 def generate( 

377 self, 

378 inputs: Optional[torch.Tensor] = None, 

379 generation_config: Optional[GenerationConfig] = None, 

380 logits_processor: Optional[LogitsProcessorList] = None, 

381 stopping_criteria: Optional[StoppingCriteriaList] = None, 

382 prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, 

383 synced_gpus: Optional[bool] = None, 

384 assistant_model: Optional["PreTrainedModel"] = None, 

385 audio_streamer: Optional[Union[AudioStreamer, AsyncAudioStreamer]] = None, 

386 negative_prompt_ids: Optional[torch.Tensor] = None, 

387 negative_prompt_attention_mask: Optional[torch.Tensor] = None, 

388 speech_tensors: Optional[torch.FloatTensor] = None, 

389 speech_masks: Optional[torch.BoolTensor] = None, 

390 speech_input_mask: Optional[torch.BoolTensor] = None, 

391 return_speech: bool = True, 

392 cfg_scale: float = 1.0, 

393 stop_check_fn: Optional[Callable[[], bool]] = None, 

394 **kwargs: Any, 

395 ) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]: 

396 """ 

397 Generates sequences of token ids and optionally speech outputs. 

398 

399 Args: 

400 All standard generation arguments from GenerationMixin 

401 negative_prompt_ids: Negative prompt for CFG in speech generation 

402 negative_prompt_attention_mask: Attention mask for negative prompt 

403 speech_tensors: Input speech for voice cloning 

404 speech_masks: Masks for speech tensors 

405 speech_input_mask: Positions to insert speech embeddings 

406 return_speech: Whether to decode and return speech outputs 

407 cfg_scale: CFG scale for speech generation 

408 stop_check_fn: Optional callable that returns True if generation should stop 

409 

410 Returns: 

411 Generated token sequences and optionally speech outputs 

412 """ 

413 # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call 

414 tokenizer = kwargs.pop("tokenizer", None) # Pull this out first, we only use it for stopping criteria 

415 # parsed_scripts = kwargs.pop("parsed_scripts", None) 

416 # all_speakers_list = kwargs.pop("all_speakers_list", None) 

417 max_length_times = kwargs.pop("max_length_times", 2) 

418 

419 if kwargs.get('max_new_tokens', None) is None: 

420 kwargs['max_new_tokens'] = ( 

421 self.config.decoder_config.max_position_embeddings - kwargs['input_ids'].shape[-1] 

422 ) 

423 

424 generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria = \ 

425 self._build_generate_config_model_kwargs( 

426 generation_config, inputs, tokenizer, return_processors=True, **kwargs 

427 ) 

428 

429 negative_kwargs = { 

430 'input_ids': torch.full( 

431 (kwargs['input_ids'].shape[0], 1), 

432 tokenizer.speech_start_id, 

433 dtype=torch.long, 

434 device=kwargs['input_ids'].device), 

435 'attention_mask': torch.ones( 

436 (kwargs['input_ids'].shape[0], 1), 

437 dtype=torch.long, 

438 device=kwargs['input_ids'].device), 

439 'max_new_tokens': kwargs.get('max_new_tokens', 100) 

440 } 

441 negative_generation_config, negative_model_kwargs, negative_input_ids = \ 

442 self._build_generate_config_model_kwargs( 

443 None, None, tokenizer, return_processors=False, **negative_kwargs 

444 ) 

445 

446 acoustic_cache = VibeVoiceTokenizerStreamingCache() 

447 semantic_cache = VibeVoiceTokenizerStreamingCache() 

448 

449 batch_size = input_ids.shape[0] 

450 device = input_ids.device 

451 finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device) 

452 correct_cnt = torch.zeros(batch_size, dtype=torch.long, device=device) 

453 is_prefill = True 

454 inputs_embeds = None 

455 verbose = kwargs.get("verbose", False) 

456 

457 # Initialize audio chunks storage for each sample 

458 audio_chunks: List[List[Any]] = [[] for _ in range(batch_size)] 

459 

460 initial_length = input_ids.shape[-1] 

461 initial_length_per_sample = model_kwargs['attention_mask'].sum(dim=-1) 

462 

463 # Define all valid tokens that can be generated 

464 valid_tokens = [ 

465 generation_config.speech_start_id, 

466 generation_config.speech_end_id, 

467 generation_config.speech_diffusion_id, 

468 generation_config.eos_token_id 

469 ] 

470 # Add bos_token_id if it exists 

471 if hasattr(generation_config, 'bos_token_id') and generation_config.bos_token_id is not None: 

472 valid_tokens.append(generation_config.bos_token_id) 

473 

474 # Add custom processor to constrain token generation 

475 token_constraint_processor = VibeVoiceTokenConstraintProcessor(valid_tokens, device=device) 

476 if logits_processor is None: 

477 logits_processor = LogitsProcessorList() 

478 logits_processor.append(token_constraint_processor) 

479 

480 max_steps = min( 

481 generation_config.max_length - initial_length, 

482 int(max_length_times * initial_length)) 

483 max_step_per_sample = torch.min( 

484 generation_config.max_length - initial_length_per_sample, 

485 (max_length_times * initial_length_per_sample).long()) 

486 reach_max_step_sample = torch.zeros(batch_size, dtype=torch.bool, device=device) 

487 

488 # Create progress iterator if verbose 

489 if kwargs.get("show_progress_bar", True): 

490 progress_bar = tqdm(range(max_steps), desc="Generating", leave=False) 

491 else: 

492 progress_bar = range(max_steps) 

493 

494 for step in progress_bar: 

495 # Check for external stop signal 

496 if stop_check_fn is not None and stop_check_fn(): 

497 if verbose: 

498 print(f"Generation stopped externally at step {step + 1}") 

499 # End the audio streamer if it exists 

500 if audio_streamer is not None: 

501 audio_streamer.end() 

502 break 

503 

504 # Check if audio_streamer has been ended (stopped externally) 

505 if audio_streamer is not None and hasattr(audio_streamer, 'finished_flags'): 

506 if any(audio_streamer.finished_flags): 

507 if verbose: 

508 print(f"Audio generation stopped externally at step {step + 1}") 

509 break 

510 

511 if finished_tags.all(): 

512 if hasattr(progress_bar, 'set_description'): 

513 progress_bar.set_description("Generation complete") 

514 break 

515 

516 if input_ids.shape[-1] >= generation_config.max_length: 

517 print(f"Reached maximum generation length {generation_config.max_length}, stopped it.") 

518 reached_samples = torch.arange(batch_size, device=device)[~finished_tags] 

519 if reached_samples.numel() > 0: 

520 reach_max_step_sample[reached_samples] = True 

521 break 

522 

523 # Update progress bar description with active samples 

524 if hasattr(progress_bar, 'set_description'): 

525 active_samples = (~finished_tags).sum().item() 

526 progress_bar.set_description(f"Generating (active: {active_samples}/{batch_size})") 

527 

528 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) 

529 if is_prefill: 

530 # we process the speech inputs only during the first generation step 

531 assert speech_tensors is not None 

532 assert speech_masks is not None 

533 assert speech_input_mask is not None 

534 prefill_inputs = { 

535 "speech_tensors": speech_tensors.to(device=device), 

536 "speech_masks": speech_masks.to(device), 

537 "speech_input_mask": speech_input_mask.to(device), 

538 } 

539 is_prefill = False 

540 else: 

541 _ = model_inputs.pop('inputs_embeds', None) 

542 prefill_inputs = {'inputs_embeds': inputs_embeds} 

543 

544 # Forward pass through the model 

545 outputs = self( 

546 **model_inputs, **prefill_inputs, 

547 logits_to_keep=1, return_dict=True, 

548 output_attentions=False, output_hidden_states=False, 

549 ) 

550 model_kwargs = self._update_model_kwargs_for_generation( 

551 outputs, model_kwargs, is_encoder_decoder=False, 

552 ) 

553 

554 # Get logits and apply logits processor 

555 next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device) 

556 # next_token_logits = outputs.logits[:, -1, :].to(copy=True, device=input_ids.device) 

557 next_token_scores = logits_processor(input_ids, next_token_logits) 

558 

559 # token selection 

560 if generation_config.do_sample: 

561 probs = nn.functional.softmax(next_token_scores, dim=-1) 

562 # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution 

563 next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) 

564 else: 

565 next_tokens = torch.argmax(next_token_scores, dim=-1) 

566 

567 next_tokens[finished_tags] = generation_config.eos_token_id 

568 input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) 

569 

570 # reached end of generation 

571 if (next_tokens == generation_config.eos_token_id).any(): 

572 eos_indices = (next_tokens == generation_config.eos_token_id).nonzero(as_tuple=False).squeeze(1) 

573 # Only print for samples that are newly finished (not already marked as finished) 

574 new_eos_indices = eos_indices[~finished_tags[eos_indices]] 

575 if new_eos_indices.numel() > 0: 

576 finished_tags[new_eos_indices] = True 

577 if verbose: 

578 logger.info(f"Samples {new_eos_indices.tolist()} reached EOS token at step {step + 1}.") 

579 if audio_streamer is not None: 

580 audio_streamer.end(new_eos_indices) 

581 

582 # Check if any sample reached its maximum generation length 

583 max_length_reached = step >= max_step_per_sample 

584 new_max_length_indices = torch.nonzero(max_length_reached & ~finished_tags, as_tuple=False).squeeze(1) 

585 if new_max_length_indices.numel() > 0: 

586 finished_tags[new_max_length_indices] = True 

587 reach_max_step_sample[new_max_length_indices] = True 

588 if verbose: 

589 logger.info( 

590 f"Samples {new_max_length_indices.tolist()} reached max generation length at step {step + 1}.") 

591 if audio_streamer is not None: 

592 audio_streamer.end(new_max_length_indices) 

593 

594 # speech_end 

595 diffusion_end_indices = (next_tokens == generation_config.speech_end_id).nonzero(as_tuple=False).squeeze(1) 

596 if diffusion_end_indices.numel() > 0: 

597 # Clear tokenizer caches for samples that reached speech end 

598 acoustic_cache.set_to_zero(diffusion_end_indices) 

599 semantic_cache.set_to_zero(diffusion_end_indices) 

600 

601 # speech_begin 

602 diffusion_start_indices = torch.arange(batch_size, device=device)[ 

603 ~finished_tags & (next_tokens == generation_config.speech_start_id) 

604 ] 

605 if diffusion_start_indices.numel() > 0 and kwargs.get('refresh_negative', True): 

606 # update attention mask 

607 for i, sample_idx in enumerate(diffusion_start_indices.tolist()): 

608 negative_model_kwargs['attention_mask'][sample_idx, :] = 0 

609 negative_model_kwargs['attention_mask'][sample_idx, -1] = 1 

610 # update past key values 

611 for layer_idx, (k_cache, v_cache) in enumerate(zip( 

612 negative_model_kwargs['past_key_values'].key_cache, 

613 negative_model_kwargs['past_key_values'].value_cache) 

614 ): 

615 # Process each non-diffusion sample 

616 for sample_idx in diffusion_start_indices.tolist(): 

617 # Shift cache for this sample 

618 k_cache[sample_idx, :, -1, :] = k_cache[sample_idx, :, 0, :].clone() 

619 v_cache[sample_idx, :, -1, :] = v_cache[sample_idx, :, 0, :].clone() 

620 # update negative_input_ids 

621 for sample_idx in diffusion_start_indices.tolist(): 

622 negative_input_ids[sample_idx, -1] = generation_config.speech_start_id 

623 

624 # Prepare inputs_embeds for next iteration 

625 # Initialize with default embeddings for all tokens 

626 # [batch_size, 1, hidden_size] 

627 next_inputs_embeds = self.model.get_input_embeddings()(next_tokens).unsqueeze(1) 

628 

629 # forward diffusion 

630 # Diffusion indices are those that are not finished and not special tokens 

631 diffusion_indices = torch.arange(batch_size, device=device)[ 

632 ~finished_tags & (next_tokens == generation_config.speech_diffusion_id) 

633 ] 

634 

635 if diffusion_indices.numel() > 0: 

636 negative_model_inputs = self.prepare_inputs_for_generation(negative_input_ids, **negative_model_kwargs) 

637 # Forward negative pass through the model 

638 if negative_model_inputs['inputs_embeds'] is None and inputs_embeds is not None: 

639 negative_model_inputs['inputs_embeds'] = inputs_embeds 

640 negative_model_inputs['input_ids'] = None 

641 

642 negative_outputs = self( 

643 **negative_model_inputs, logits_to_keep=0, 

644 return_dict=True, output_attentions=False, 

645 output_hidden_states=False, 

646 ) 

647 negative_model_kwargs = self._update_model_kwargs_for_generation( 

648 negative_outputs, negative_model_kwargs, is_encoder_decoder=False, 

649 ) 

650 negative_input_ids = torch.cat([negative_input_ids, next_tokens[:, None]], dim=-1) 

651 

652 # correct the non-diffusion indices 

653 # we forward all samples' negative outputs even if 

654 # they are not in diffusion mode to keep the cache consistent 

655 # So we need to correct the kv cache of non-diffusion samples 

656 non_diffusion_mask = ~finished_tags & (next_tokens != generation_config.speech_diffusion_id) 

657 if non_diffusion_mask.any(): 

658 non_diffusion_indices = torch.arange(batch_size, device=device)[non_diffusion_mask] 

659 start_indices = correct_cnt[non_diffusion_indices] 

660 

661 # 1. Update attention_mask - need to handle each sample separately 

662 seq_len = negative_model_kwargs['attention_mask'].shape[1] 

663 for i, (sample_idx, start_idx) in enumerate(zip( 

664 non_diffusion_indices.tolist(), start_indices.tolist() 

665 )): 

666 # Shift the attention mask for this sample 

667 if start_idx + 1 < seq_len - 1: 

668 negative_model_kwargs['attention_mask'][sample_idx, start_idx + 1:] = \ 

669 negative_model_kwargs['attention_mask'][sample_idx, start_idx:-1].clone() 

670 negative_model_kwargs['attention_mask'][sample_idx, start_idx] = 0 

671 

672 # 2. Update past_key_values 

673 for layer_idx, (k_cache, v_cache) in enumerate(zip( 

674 negative_model_kwargs['past_key_values'].key_cache, 

675 negative_model_kwargs['past_key_values'].value_cache) 

676 ): 

677 # Process each non-diffusion sample 

678 for sample_idx, start_idx in zip(non_diffusion_indices.tolist(), start_indices.tolist()): 

679 if start_idx + 1 < k_cache.shape[2] - 1: 

680 # Shift cache for this sample 

681 k_cache[sample_idx, :, start_idx + 1:, :] = k_cache[ 

682 sample_idx, :, start_idx:-1, :].clone() 

683 v_cache[sample_idx, :, start_idx + 1:, :] = v_cache[ 

684 sample_idx, :, start_idx:-1, :].clone() 

685 

686 # 3. Update negative_input_ids 

687 for sample_idx, start_idx in zip(non_diffusion_indices.tolist(), start_indices.tolist()): 

688 if start_idx + 1 < negative_input_ids.shape[1] - 1: 

689 negative_input_ids[sample_idx, start_idx + 1:] = \ 

690 negative_input_ids[sample_idx, start_idx:-1].clone() 

691 

692 correct_cnt[non_diffusion_indices] += 1 

693 

694 positive_condition = outputs.last_hidden_state[diffusion_indices, -1, :] 

695 negative_condition = negative_outputs.last_hidden_state[diffusion_indices, -1, :] 

696 

697 speech_latent = self.sample_speech_tokens( 

698 positive_condition, 

699 negative_condition, 

700 cfg_scale=cfg_scale, 

701 ).unsqueeze(1) 

702 

703 # Decode acoustic latent to audio using acoustic streaming cache 

704 scaled_latent = ( 

705 speech_latent / self.model.speech_scaling_factor.to(speech_latent.device) 

706 - self.model.speech_bias_factor.to(speech_latent.device) 

707 ) 

708 audio_chunk = self.model.acoustic_tokenizer.decode( 

709 scaled_latent.to(self.model.acoustic_tokenizer.device), 

710 cache=acoustic_cache, # Use acoustic-specific cache 

711 sample_indices=diffusion_indices.to(self.model.acoustic_tokenizer.device), 

712 use_cache=True, 

713 debug=False 

714 ) 

715 

716 # Store audio chunks for each sample 

717 for i, sample_idx in enumerate(diffusion_indices): 

718 idx = sample_idx.item() 

719 # Only append audio chunk if the sample is not finished 

720 if not finished_tags[idx]: 

721 audio_chunks[idx].append(audio_chunk[i]) 

722 

723 # Add streaming support here 

724 if audio_streamer is not None: 

725 # Stream the audio chunks immediately 

726 audio_streamer.put(audio_chunk, diffusion_indices) 

727 

728 # Encode audio to semantic features using semantic streaming cache 

729 semantic_features = self.model.semantic_tokenizer.encode( 

730 audio_chunk, 

731 cache=semantic_cache, # Use semantic-specific cache 

732 sample_indices=diffusion_indices, 

733 use_cache=True, 

734 debug=False 

735 ).mean # semantic tokenizer has no VAE. 

736 

737 # Combine acoustic and semantic features for next input 

738 acoustic_embed = self.model.acoustic_connector(speech_latent) 

739 semantic_embed = self.model.semantic_connector(semantic_features) 

740 diffusion_embeds = acoustic_embed + semantic_embed 

741 

742 # Update embeddings for diffusion indices 

743 next_inputs_embeds[diffusion_indices] = diffusion_embeds 

744 

745 # Set inputs_embeds for next iteration 

746 inputs_embeds = next_inputs_embeds 

747 

748 if audio_streamer is not None: 

749 audio_streamer.end() 

750 

751 # Concatenate audio chunks for each sample 

752 final_audio_outputs = [] 

753 for sample_chunks in audio_chunks: 

754 if sample_chunks: 

755 # Concatenate all chunks along the time dimension (assumed to be the last dimension) 

756 concatenated_audio = torch.cat(sample_chunks, dim=-1) 

757 final_audio_outputs.append(concatenated_audio) 

758 else: 

759 # If no audio was generated for this sample, append None 

760 final_audio_outputs.append(None) 

761 

762 return VibeVoiceGenerationOutput( 

763 sequences=input_ids, 

764 speech_outputs=final_audio_outputs if return_speech else None,