diff --git a/vibevoice/finetune/train_vibevoice.py b/vibevoice/finetune/train_vibevoice.py index f45f3bc..35a297e 100644 --- a/vibevoice/finetune/train_vibevoice.py +++ b/vibevoice/finetune/train_vibevoice.py @@ -603,6 +603,164 @@ def main() -> None: logger.warning(f"LoRA debug (on_step_end) failed: {e}") class VibeVoiceTrainer(Trainer): + def training_forward(self, model: VibeVoiceForConditionalGeneration, inputs: Dict[str, Any]): + """Custom forward pass for training with new diffusion loss calculation.""" + # Extract inputs + input_ids = inputs.get("input_ids") + attention_mask = inputs.get("attention_mask") + position_ids = inputs.get("position_ids") + past_key_values = inputs.get("past_key_values") + inputs_embeds = inputs.get("inputs_embeds") + use_cache = inputs.get("use_cache", False) + output_attentions = inputs.get("output_attentions") + output_hidden_states = inputs.get("output_hidden_states") + return_dict = inputs.get("return_dict", True) + cache_position = inputs.get("cache_position") + + # Speech-related inputs + speech_tensors = inputs.get("speech_tensors") + speech_masks = inputs.get("speech_masks") + speeches_loss_input = inputs.get("speeches_loss_input") + speech_semantic_tensors = inputs.get("speech_semantic_tensors") + acoustic_input_mask = inputs.get("acoustic_input_mask") + acoustic_loss_mask = inputs.get("acoustic_loss_mask") + ddmp_batch_mul = training_args.ddpm_batch_mul + kwargs = {} + + # --- START: Copy of model forward logic with new diffusion loss --- + x = model.get_input_embeddings()(input_ids) + + semantic_speech_all_connect_features = model.model.semantic_connector(speech_semantic_tensors) + if speeches_loss_input is not None: + # only part audio need diffuse + speech_all_features, speech_all_connect_features = model.forward_speech_features( + speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None, + speech_masks=speech_masks, + speech_type=kwargs.get("speech_type", "audio"), + return_unmask=True + ) + if speech_tensors is not None: + if semantic_speech_all_connect_features is not None: + x[acoustic_input_mask] = speech_all_connect_features[speech_masks] + semantic_speech_all_connect_features[speech_masks] + else: + x[acoustic_input_mask] = speech_all_connect_features[speech_masks] + speech_features = speech_all_features[speeches_loss_input & speech_masks] # only part audio need diffuse + speech_connect_features = speech_all_connect_features[speeches_loss_input & speech_masks] + # Forward-time consistency check: selected latent count should match number of acoustic placeholders + try: + if acoustic_input_mask is not None: + assert speech_connect_features.shape[0] == int(acoustic_input_mask.sum().item()), ( + f"Mismatch between selected speech connectors ({speech_connect_features.shape[0]}) and acoustic_input_mask sum ({int(acoustic_input_mask.sum().item())})" + ) + except Exception: + pass + else: + speech_features, speech_connect_features = model.forward_speech_features( + speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None, + speech_masks=speech_masks, + speech_type=kwargs.get("speech_type", "audio"), + ) + if speech_tensors is not None: + x[acoustic_input_mask] = speech_connect_features + + outputs = model.model( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=x, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=False, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs.last_hidden_state + logits = model.lm_head(hidden_states) + + loss = None + + # --- NEW Diffusion Loss Calculation --- + diffusion_loss = None + # This block is executed only if we are in a context that involves speech. + if speech_tensors is not None and acoustic_loss_mask.sum().item() > 0: + # Build conditioning mask from positions whose NEXT token is a speech latent (shift left by 1) + cond_mask = torch.zeros_like(acoustic_loss_mask, dtype=torch.bool) + cond_mask[:, :-1] = acoustic_loss_mask[:, 1:] + cond_mask[:, 0] = False + condition_features = hidden_states[cond_mask] + + speech_len, latent_size = speech_features.shape + # Sanity check: ensure 1:1 alignment between selected conditions and latents + try: + assert condition_features.shape[0] == speech_len, ( + f"Mismatch: condition_features={condition_features.shape[0]} vs speech_features={speech_len}" + ) + except Exception: + pass + + noise = torch.randn( + (speech_len * ddmp_batch_mul, latent_size), + device=hidden_states.device, + dtype=hidden_states.dtype + ) + + timesteps = torch.multinomial( + torch.ones(model.config.diffusion_head_config.ddpm_num_steps), + speech_len * ddmp_batch_mul, + replacement=True, + ).to(hidden_states.device) + + speech_features_repeated = speech_features.repeat_interleave(ddmp_batch_mul, dim=0) + condition_features_repeated = condition_features.repeat_interleave(ddmp_batch_mul, dim=0) + + noisy_speech_features = model.model.noise_scheduler.add_noise( + speech_features_repeated, noise, timesteps + ) + + model_output = model.model.prediction_head( + noisy_speech_features, + timesteps.type_as(x), + condition_features_repeated + ) + + prediction_type = model.config.diffusion_head_config.prediction_type + if prediction_type == "epsilon": + target_for_loss = noise + elif prediction_type == "v_prediction": + target_for_loss = model.model.noise_scheduler.get_velocity( + speech_features_repeated, noise, timesteps + ) + else: + raise NotImplementedError(f"Prediction type {prediction_type} not implemented") + + diffusion_loss = F.mse_loss(model_output.float(), target_for_loss.float(), reduction='sum') + if latent_size > 0 and ddmp_batch_mul > 0: + # Normalize by latent dim, number of sampled diffusion steps per latent, and number of speech tokens + diffusion_loss = diffusion_loss / latent_size / ddmp_batch_mul / max(speech_len, 1) + else: + diffusion_loss = torch.tensor(0.0, device=diffusion_loss.device) + + else: + # Dummy loss for DDP to work when there are no speech samples in a batch, + # but we are in a speech context. + diffusion_loss = sum(p.sum() for p in model.model.prediction_head.parameters()) * 0.0 + diffusion_loss += sum(p.sum() for p in model.model.acoustic_connector.parameters()) * 0.0 + diffusion_loss += sum(p.sum() for p in model.model.semantic_connector.parameters()) * 0.0 + # --- End NEW Diffusion Loss Calculation --- + + from vibevoice.modular.modeling_vibevoice import VibeVoiceCausalLMOutputWithPast + return VibeVoiceCausalLMOutputWithPast( + loss=loss, + diffusion_loss=diffusion_loss, + speech_token_num=speech_len if speech_tensors is not None else 0, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + def compute_loss(self, model: VibeVoiceForConditionalGeneration, inputs: Dict[str, Any], return_outputs=False, num_items_in_batch: Optional[int] = None): labels = inputs.get("input_ids") attention_mask = inputs.get("attention_mask") @@ -629,17 +787,8 @@ def main() -> None: if isinstance(sem, torch.Tensor): inputs["speech_semantic_tensors"] = sem.to(dtype=target_dtype) - outputs = model( - input_ids=inputs.get("input_ids"), - attention_mask=attention_mask, - speech_tensors=inputs.get("speech_tensors"), - speech_masks=inputs.get("speech_masks"), - speech_semantic_tensors=inputs.get("speech_semantic_tensors"), - acoustic_input_mask=acoustic_input_mask, - acoustic_loss_mask=inputs.get("acoustic_loss_mask"), - speeches_loss_input=inputs.get("speeches_loss_input"), - ddpm_batch_mul=training_args.ddpm_batch_mul, - ) + # Use custom training forward pass with new diffusion loss + outputs = self.training_forward(model, inputs) # Invariants: token/latent selection equality across views (warn, don't assert) try: