Add LoRA fine-tuning scripts (#5)

vibevoice/modular/modeling_vibevoice.py

@@ -368,11 +368,18 @@ class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
                 )
             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]
+                    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.unsqueeze(-1) & speech_masks] # only part audio need diffuse
-                speech_connect_features = speech_all_connect_features[speeches_loss_input.unsqueeze(-1) & speech_masks]
+
+                # Select only the target segments' latents for diffusion loss.
+                # Both masks are [num_segments, max_latent_len]; using 2D mask on [B,T,D] selects [N_true, D].
+                target_latent_mask = speeches_loss_input & speech_masks
+                speech_features = speech_all_features[target_latent_mask]
+                speech_connect_features = speech_all_connect_features[target_latent_mask]
         else:
             speech_features, speech_connect_features = self.forward_speech_features(
                     speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
@@ -485,4 +492,4 @@ __all__ = [
     "VibeVoiceForConditionalGeneration",
     "VibeVoiceCausalLMOutputWithPast",
     "VibeVoiceGenerationOutput",
-]
+]