Coverage for wrapper/vibevoice/modeling_vibevoice.py: 30%
234 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1# mypy: ignore-errors
2# Copy from: https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/modeling_vibevoice.py
4from dataclasses import dataclass
6from typing import Any
7from typing import List
8from typing import Optional
9from typing import Tuple
10from typing import Union
12import torch
13import torch.nn as nn
14import torch.nn.functional as F
15import torch.distributed as dist
17from transformers import AutoModel
18from transformers import AutoModelForCausalLM
20from transformers.modeling_outputs import BaseModelOutputWithPast
21from transformers.modeling_outputs import ModelOutput
22from transformers.models.llama.modeling_llama import LlamaRMSNorm
23from transformers.modeling_utils import PreTrainedModel
24from transformers.utils import logging
26from configuration_vibevoice import VibeVoiceConfig
27from modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
28from schedule.dpm_solver import DPMSolverMultistepScheduler
31logger = logging.get_logger(__name__)
34@dataclass
35class VibeVoiceCausalLMOutputWithPast(ModelOutput):
36 loss: Optional[torch.FloatTensor] = None
37 diffusion_loss: Optional[torch.FloatTensor] = None
38 speech_token_num: Optional[int] = None
39 logits: torch.FloatTensor = None
40 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
41 hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
42 attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
45@dataclass
46class VibeVoiceGenerationOutput(ModelOutput):
47 """
48 Output type for VibeVoice generation.
50 Args:
51 sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
52 The generated sequences.
53 speech_outputs (`List[torch.FloatTensor]`, *optional*):
54 List of generated speech waveforms or latents for each speech segment.
55 """
56 sequences: torch.LongTensor = None
57 speech_outputs: Optional[List[torch.FloatTensor]] = None
60class SpeechConnector(nn.Module):
61 def __init__(self, input_dim: int, output_dim: int) -> None:
62 super().__init__()
63 self.fc1 = nn.Linear(input_dim, output_dim)
64 self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
65 self.fc2 = nn.Linear(output_dim, output_dim)
67 def forward(self, features: Any, **kwargs: Any) -> Any:
68 x = self.fc1(features)
69 x = self.norm(x)
70 x = self.fc2(x)
71 return x
74# @auto_docstring
75class VibeVoicePreTrainedModel(PreTrainedModel):
76 config_class = VibeVoiceConfig
77 base_model_prefix = "model"
78 supports_gradient_checkpointing = True
79 _skip_keys_device_placement = "past_key_values"
80 _supports_cache_class = True
81 _supports_flash_attn_2 = True
82 _supports_sdpa = True
83 _supports_quantized_cache = True
84 _supports_static_cache = True
85 _supports_attention_backend = True
87 def _init_weights(self, module: nn.Module) -> None:
88 if isinstance(module, VibeVoiceDiffusionHead):
89 module.initialize_weights()
90 return
92 # Use the language model's initializer_range if available
93 if hasattr(
94 self.config, 'language_model_config'
95 ) and hasattr(
96 self.config.language_model_config, 'initializer_range'
97 ):
98 std = self.config.language_model_config.initializer_range
99 elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
100 std = self.config.decoder_config.initializer_range
101 else:
102 std = 0.02 # Default value
104 if isinstance(module, nn.Linear):
105 module.weight.data.normal_(mean=0.0, std=std)
106 if module.bias is not None:
107 module.bias.data.zero_()
108 elif isinstance(module, nn.LayerNorm):
109 module.weight.data.fill_(1.0)
110 module.bias.data.zero_()
113# @auto_docstring
114class VibeVoiceModel(VibeVoicePreTrainedModel):
115 def __init__(self, config: VibeVoiceConfig) -> None:
116 super().__init__(config)
118 if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
119 if isinstance(config.torch_dtype, str):
120 dtype = getattr(torch, config.torch_dtype)
121 else:
122 dtype = config.torch_dtype
123 else:
124 dtype = torch.float32
126 # Initialize Qwen2 model for language modeling
127 lm_config = config.decoder_config
128 self.language_model = AutoModel.from_config(lm_config)
130 # Initialize speech components if needed
131 self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
132 self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
134 self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
135 self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
137 # Register scaling factors as buffers - use 1D tensors for FSDP compatibility
138 self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))
139 self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
141 # Initialize prediction head for speech generation
142 self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
144 # Initialize noise scheduler
145 self.noise_scheduler = DPMSolverMultistepScheduler(
146 num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
147 beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
148 prediction_type=config.diffusion_head_config.prediction_type
149 )
151 def get_input_embeddings(self) -> nn.Module:
152 if hasattr(self.language_model, 'embed_tokens'):
153 # If the language model has an embed_tokens attribute, return it
154 return self.language_model.embed_tokens
156 for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
157 if attr.orig_name == 'embed_tokens.weight':
158 return getattr(self.language_model, name)
159 assert False, 'should not arrive here'
161 def set_input_embeddings(self, value: nn.Module) -> None:
162 self.language_model.embed_tokens = value
164 def set_speech_tokenizers(
165 self, acoustic_tokenizer: Optional[Any] = None, semantic_tokenizer: Optional[Any] = None
166 ) -> None:
167 """Set the speech tokenizers used for encoding and decoding speech."""
168 self.acoustic_tokenizer = acoustic_tokenizer
169 self.semantic_tokenizer = semantic_tokenizer
171 # Reset the encoder to evaluation mode
172 if self.acoustic_tokenizer is not None:
173 self.acoustic_tokenizer.eval()
175 if self.semantic_tokenizer is not None:
176 self.semantic_tokenizer.eval()
178 def forward(
179 self,
180 input_ids: torch.LongTensor = None,
181 attention_mask: Optional[torch.Tensor] = None,
182 position_ids: Optional[torch.LongTensor] = None,
183 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
184 inputs_embeds: Optional[torch.FloatTensor] = None,
185 use_cache: Optional[bool] = None,
186 output_attentions: Optional[bool] = None,
187 output_hidden_states: Optional[bool] = None,
188 return_dict: Optional[bool] = None,
189 cache_position: Optional[torch.LongTensor] = None,
190 **kwargs: Any,
191 ) -> Union[Tuple, BaseModelOutputWithPast]:
193 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
195 # Forward through language model
196 outputs = self.language_model(
197 input_ids=input_ids,
198 attention_mask=attention_mask,
199 position_ids=position_ids,
200 past_key_values=past_key_values,
201 inputs_embeds=inputs_embeds,
202 use_cache=use_cache,
203 output_attentions=output_attentions,
204 output_hidden_states=output_hidden_states,
205 return_dict=return_dict,
206 cache_position=cache_position,
207 **kwargs,
208 )
210 if not return_dict:
211 return outputs
213 return BaseModelOutputWithPast(
214 last_hidden_state=outputs.last_hidden_state,
215 past_key_values=outputs.past_key_values,
216 hidden_states=outputs.hidden_states,
217 attentions=outputs.attentions,
218 )
221class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
222 _tied_weights_keys = ["lm_head.weight"]
223 _tp_plan = {"lm_head": "colwise_rep"}
225 def __init__(
226 self,
227 config: VibeVoiceConfig,
228 ) -> None:
229 super().__init__(config)
230 self.model = VibeVoiceModel(config)
231 self.vocab_size = config.decoder_config.vocab_size
232 self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False)
234 self.post_init()
236 def get_input_embeddings(self) -> nn.Module:
237 return self.model.get_input_embeddings()
239 def set_input_embeddings(
240 self,
241 value: nn.Module,
242 ) -> None:
243 self.model.set_input_embeddings(value)
245 def get_output_embeddings(self) -> nn.Module:
246 return self.lm_head
248 def set_decoder(self, decoder: nn.Module) -> None:
249 self.model.language_model = decoder
251 def get_decoder(self) -> nn.Module:
252 return self.model.language_model
254 def tie_weights(self) -> None:
255 """
256 Tie the weights between the input embeddings and the output embeddings.
257 """
258 if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
259 # The standard PreTrainedModel method will handle the tying.
260 # It typically does a simple parameter object assignment, which is
261 # CORRECT to do BEFORE FSDP wraps the model.
262 output_embeddings = self.get_output_embeddings()
263 input_embeddings = self.get_input_embeddings()
264 if hasattr(input_embeddings, 'weight'):
265 output_embeddings.weight = input_embeddings.weight
266 else:
267 # maybe returned input_embeddings a tensor directly
268 output_embeddings.weight = input_embeddings
270 if getattr(output_embeddings, "bias", None) is not None:
271 output_embeddings.bias.data = nn.functional.pad(
272 output_embeddings.bias.data,
273 (0, output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0]),
274 "constant",
275 0,
276 )
277 logger.info("✅ Tied input and output embeddings using standard assignment.")
278 else:
279 logger.info("ℹ️ tie_word_embeddings is False, not tying weights.")
281 # Also, ensure set_output_embeddings is safe, though your implementation looks okay.
282 # The key is to avoid calling it after accelerator.prepare().
283 def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
284 # Your current implementation using data.copy_ is good practice,
285 # but the best way is to not call this after prepare().
286 self.lm_head = new_embeddings
288 def forward_speech_features(
289 self,
290 speech_tensors: Optional[torch.Tensor] = None,
291 speech_masks: Optional[torch.Tensor] = None,
292 speech_type: str = "audio",
293 return_unmask: bool = False
294 ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
295 if speech_tensors is None:
296 # Use config to get vae_dim instead of non-existent self.args
297 vae_dim = self.config.acoustic_tokenizer_config.vae_dim
298 audio_features = torch.zeros(1, 1, vae_dim).to(self.get_input_embeddings().weight)
299 connect_features = self.model.acoustic_connector(audio_features)
300 return audio_features, connect_features
301 else:
302 with torch.no_grad():
303 if speech_type == "audio":
304 with torch.no_grad():
305 frames = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))[0][0]
306 audio_tokens = frames.sample(self.model.acoustic_tokenizer.std_dist_type)[0]
308 elif speech_type == "vae":
309 # Use config to get vae_dim instead of non-existent self.args
310 vae_dim = self.config.acoustic_tokenizer_config.vae_dim
311 speech_mode = speech_tensors.reshape(speech_tensors.size(0), -1, vae_dim)
313 # gaussian sample from the speech_mode
314 batch_size = speech_mode.size(0)
315 value = self.model.acoustic_tokenizer.fix_std / 0.8
316 std = torch.randn(batch_size, dtype=speech_mode.dtype, device=speech_mode.device) * value
317 std = std.view(-1, *[1] * (speech_mode.dim() - 1))
318 audio_tokens = speech_mode + std * torch.randn(speech_mode.shape).to(speech_mode)
319 else:
320 raise NotImplementedError(f"Speech type {speech_type} not implemented")
322 if torch.isnan(self.model.speech_scaling_factor) or torch.isnan(self.model.speech_bias_factor):
323 scaling_factor = 1. / audio_tokens[speech_masks].flatten().std()
324 bias_factor = -audio_tokens[speech_masks].flatten().mean()
326 # Only use distributed operations if the process group is initialized
327 if dist.is_available() and dist.is_initialized():
328 dist.all_reduce(scaling_factor, op=dist.ReduceOp.SUM)
329 dist.all_reduce(bias_factor, op=dist.ReduceOp.SUM)
330 world_size = dist.get_world_size()
331 self.model.speech_scaling_factor.copy_(scaling_factor / world_size)
332 self.model.speech_bias_factor.copy_(bias_factor / world_size)
333 logger.info(
334 f"Speech scaling factor (distributed): {self.model.speech_scaling_factor}, "
335 f"bias factor: {self.model.speech_bias_factor}", flush=True)
336 else:
337 # Single process case
338 self.model.speech_scaling_factor.copy_(scaling_factor)
339 self.model.speech_bias_factor.copy_(bias_factor)
340 logger.info(
341 f"Speech scaling factor (single process): {self.model.speech_scaling_factor}, "
342 f"bias factor: {self.model.speech_bias_factor}", flush=True)
344 audio_features = (audio_tokens + self.model.speech_bias_factor) * self.model.speech_scaling_factor
346 connect_features = self.model.acoustic_connector(audio_features)
347 if return_unmask:
348 return audio_features, connect_features
349 return audio_features[speech_masks], connect_features[speech_masks]
351 def forward(
352 self,
353 input_ids: torch.LongTensor = None,
354 attention_mask: Optional[torch.Tensor] = None,
355 position_ids: Optional[torch.LongTensor] = None,
356 past_key_values: Optional[List[torch.FloatTensor]] = None,
357 inputs_embeds: Optional[torch.FloatTensor] = None,
358 labels: Optional[torch.LongTensor] = None,
359 use_cache: Optional[bool] = False,
360 output_attentions: Optional[bool] = None,
361 output_hidden_states: Optional[bool] = None,
362 return_dict: Optional[bool] = None,
363 cache_position: Optional[torch.LongTensor] = None,
364 # New arguments for speech processing and loss calculation
365 speech_tensors: Optional[torch.FloatTensor] = None,
366 speech_masks: Optional[torch.BoolTensor] = None,
367 speeches_loss_input: Optional[torch.FloatTensor] = None,
368 speech_semantic_tensors: Optional[torch.FloatTensor] = None,
369 acoustic_input_mask: Optional[torch.BoolTensor] = None,
370 acoustic_loss_mask: Optional[torch.BoolTensor] = None,
371 ddpm_batch_mul: int = 1,
372 **kwargs: Any,
373 ) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
375 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
377 x = self.get_input_embeddings()(input_ids)
379 semantic_speech_all_connect_features = self.model.semantic_connector(speech_semantic_tensors)
380 if speeches_loss_input is not None:
381 # only part audio need diffuse
382 speech_all_features, speech_all_connect_features = self.forward_speech_features(
383 speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
384 speech_masks=speech_masks,
385 speech_type=kwargs.get("speech_type", "audio"),
386 return_unmask=True
387 )
388 if speech_tensors is not None:
389 if semantic_speech_all_connect_features is not None:
390 x[acoustic_input_mask] = speech_all_connect_features[speech_masks] + \
391 semantic_speech_all_connect_features[speech_masks]
392 else:
393 x[acoustic_input_mask] = speech_all_connect_features[speech_masks]
394 mask = speeches_loss_input.unsqueeze(-1) & speech_masks
395 speech_features = speech_all_features[mask] # only part audio need diffuse
396 speech_connect_features = speech_all_connect_features[mask]
397 else:
398 speech_features, speech_connect_features = self.forward_speech_features(
399 speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
400 speech_masks=speech_masks,
401 speech_type=kwargs.get("speech_type", "audio"),
402 )
403 if speech_tensors is not None:
404 x[acoustic_input_mask] = speech_connect_features
406 outputs = self.model(
407 input_ids=None,
408 attention_mask=attention_mask,
409 position_ids=position_ids,
410 past_key_values=past_key_values,
411 inputs_embeds=x,
412 use_cache=use_cache,
413 output_attentions=output_attentions,
414 output_hidden_states=False,
415 return_dict=return_dict,
416 cache_position=cache_position,
417 )
419 hidden_states = outputs.last_hidden_state
420 logits = self.lm_head(hidden_states)
421 # logits = logits.float()
423 loss = None
424 if labels is not None:
425 # The custom CE loss with masking is calculated in the training script.
426 # We leave the standard loss calculation here as None.
427 pass
429 # --- Diffusion Loss Calculation ---
430 diffusion_loss = None
431 # This block is executed only if we are in a context that involves speech.
432 if speech_tensors is not None and acoustic_loss_mask is not None and acoustic_loss_mask.sum().item() > 0:
433 condition_features = hidden_states[acoustic_loss_mask]
435 speech_len, latent_size = speech_features.shape
437 noise = torch.randn(
438 (speech_len * ddpm_batch_mul, latent_size),
439 device=hidden_states.device,
440 dtype=hidden_states.dtype
441 )
443 timesteps = torch.multinomial(
444 torch.ones(self.config.diffusion_head_config.ddpm_num_steps),
445 speech_len * ddpm_batch_mul,
446 replacement=True,
447 ).to(hidden_states.device)
449 speech_features_repeated = speech_features.repeat_interleave(ddpm_batch_mul, dim=0)
450 condition_features_repeated = condition_features.repeat_interleave(ddpm_batch_mul, dim=0)
452 noisy_speech_features = self.model.noise_scheduler.add_noise(
453 speech_features_repeated, noise, timesteps
454 )
456 model_output = self.model.prediction_head(
457 noisy_speech_features,
458 timesteps.type_as(x),
459 condition_features_repeated
460 )
462 prediction_type = self.config.diffusion_head_config.prediction_type
463 if prediction_type == "epsilon":
464 target_for_loss = noise
465 elif prediction_type == "v_prediction":
466 target_for_loss = self.model.noise_scheduler.get_velocity(
467 speech_features_repeated, noise, timesteps
468 )
469 else:
470 raise NotImplementedError(f"Prediction type {prediction_type} not implemented")
472 diffusion_loss = F.mse_loss(model_output.float(), target_for_loss.float(), reduction='sum')
473 if latent_size > 0 and ddpm_batch_mul > 0:
474 diffusion_loss = diffusion_loss / latent_size / ddpm_batch_mul
475 else:
476 diffusion_loss = torch.tensor(0.0, device=diffusion_loss.device)
478 else:
479 # Dummy loss for DDP to work when there are no speech samples in a batch,
480 # but we are in a speech context.
481 diffusion_loss = sum(p.sum() for p in self.model.prediction_head.parameters()) * 0.0
482 diffusion_loss += sum(p.sum() for p in self.model.acoustic_connector.parameters()) * 0.0
483 diffusion_loss += sum(p.sum() for p in self.model.semantic_connector.parameters()) * 0.0
484 # --- End Diffusion Loss Calculation ---
486 if not return_dict:
487 output = (logits, speech_len) + outputs.to_tuple()[1:]
488 return (loss, diffusion_loss) + output
490 return VibeVoiceCausalLMOutputWithPast(
491 loss=loss,
492 diffusion_loss=diffusion_loss,
493 speech_token_num=speech_len if speech_tensors is not None else 0,
494 logits=logits,
495 past_key_values=outputs.past_key_values,
496 hidden_states=outputs.hidden_states,
497 attentions=outputs.attentions,
498 )
501AutoModel.register(VibeVoiceConfig, VibeVoiceModel)
502AutoModelForCausalLM.register(VibeVoiceConfig, VibeVoiceForConditionalGeneration)
504__all__ = [
505 "VibeVoiceModel",
506 "VibeVoicePreTrainedModel",
507 "VibeVoiceForConditionalGeneration",
508 "VibeVoiceCausalLMOutputWithPast",
509 "VibeVoiceGenerationOutput",
510]