diff --git a/README.md b/README.md index 6ab259b..049eb6c 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,8 @@ We'd like to thank [PsiPi](https://huggingface.co/PsiPi) for sharing an interest ```bash python demo/gradio_demo.py --model_path vibevoice/VibeVoice-1.5B --share # or python demo/gradio_demo.py --model_path vibevoice/VibeVoice-7B --share +# optionally add --checkpoint_path path/to/checkpoint to load a fine-tuned adapter +# use the in-app "Disable voice cloning" setting (Advanced Settings) to skip speaker conditioning ``` **Option 2: Inference from files directly** @@ -89,6 +91,12 @@ python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_pat # or more speakers python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_path demo/text_examples/2p_music.txt --speaker_names Alice Frank + +# load a fine-tuned LoRA checkpoint +python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice --checkpoint_path path/to/checkpoint + +# disable voice cloning (skip speech prefill) +python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice --disable_prefill ``` diff --git a/demo/gradio_demo.py b/demo/gradio_demo.py index 070dba2..9ac413e 100644 --- a/demo/gradio_demo.py +++ b/demo/gradio_demo.py @@ -9,7 +9,7 @@ import sys import tempfile import time from pathlib import Path -from typing import List, Dict, Any, Iterator +from typing import List, Dict, Any, Iterator, Optional from datetime import datetime import threading import numpy as np @@ -22,6 +22,7 @@ import traceback from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference +from vibevoice.modular.lora_loading import load_lora_assets from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor from vibevoice.modular.streamer import AudioStreamer from transformers.utils import logging @@ -32,11 +33,13 @@ logger = logging.get_logger(__name__) class VibeVoiceDemo: - def __init__(self, model_path: str, device: str = "cuda", inference_steps: int = 5): + def __init__(self, model_path: str, device: str = "cuda", inference_steps: int = 5, adapter_path: Optional[str] = None): """Initialize the VibeVoice demo with model loading.""" self.model_path = model_path self.device = device self.inference_steps = inference_steps + self.adapter_path = adapter_path + self.loaded_adapter_root: Optional[str] = None self.is_generating = False # Track generation state self.stop_generation = False # Flag to stop generation self.current_streamer = None # Track current audio streamer @@ -47,6 +50,7 @@ class VibeVoiceDemo: def load_model(self): """Load the VibeVoice model and processor.""" print(f"Loading processor & model from {self.model_path}") + self.loaded_adapter_root = None # Normalize potential 'mpx' if self.device.lower() == "mpx": print("Note: device 'mpx' detected, treating it as 'mps'.") @@ -108,6 +112,29 @@ class VibeVoiceDemo: self.model.to("mps") else: raise e + if self.adapter_path: + print(f"Loading fine-tuned assets from {self.adapter_path}") + report = load_lora_assets(self.model, self.adapter_path) + loaded_components = [ + name for name, loaded in ( + ("language LoRA", report.language_model), + ("diffusion head LoRA", report.diffusion_head_lora), + ("diffusion head weights", report.diffusion_head_full), + ("acoustic connector", report.acoustic_connector), + ("semantic connector", report.semantic_connector), + ) + if loaded + ] + if loaded_components: + print(f"Loaded components: {', '.join(loaded_components)}") + else: + print("Warning: no adapter components were loaded; check the checkpoint path.") + if report.adapter_root is not None: + self.loaded_adapter_root = str(report.adapter_root) + print(f"Adapter assets resolved to: {self.loaded_adapter_root}") + else: + self.loaded_adapter_root = self.adapter_path + self.model.eval() # Use SDE solver by default @@ -182,7 +209,8 @@ class VibeVoiceDemo: speaker_2: str = None, speaker_3: str = None, speaker_4: str = None, - cfg_scale: float = 1.3) -> Iterator[tuple]: + cfg_scale: float = 1.3, + disable_voice_cloning: bool = False) -> Iterator[tuple]: try: # Reset stop flag and set generating state @@ -210,10 +238,15 @@ class VibeVoiceDemo: self.is_generating = False raise gr.Error(f"Error: Please select a valid speaker for Speaker {i+1}.") + voice_cloning_enabled = not disable_voice_cloning + # Build initial log log = f"🎙️ Generating podcast with {num_speakers} speakers\n" log += f"📊 Parameters: CFG Scale={cfg_scale}, Inference Steps={self.inference_steps}\n" log += f"🎭 Speakers: {', '.join(selected_speakers)}\n" + log += f"🔊 Voice cloning: {'Enabled' if voice_cloning_enabled else 'Disabled'}\n" + if self.loaded_adapter_root: + log += f"🧩 LoRA: {self.loaded_adapter_root}\n" # Check for stop signal if self.stop_generation: @@ -221,15 +254,17 @@ class VibeVoiceDemo: yield None, "🛑 Generation stopped by user", gr.update(visible=False) return - # Load voice samples - voice_samples = [] - for speaker_name in selected_speakers: - audio_path = self.available_voices[speaker_name] - audio_data = self.read_audio(audio_path) - if len(audio_data) == 0: - self.is_generating = False - raise gr.Error(f"Error: Failed to load audio for {speaker_name}") - voice_samples.append(audio_data) + # Load voice samples when voice cloning is enabled + voice_samples = None + if voice_cloning_enabled: + voice_samples = [] + for speaker_name in selected_speakers: + audio_path = self.available_voices[speaker_name] + audio_data = self.read_audio(audio_path) + if len(audio_data) == 0: + self.is_generating = False + raise gr.Error(f"Error: Failed to load audio for {speaker_name}") + voice_samples.append(audio_data) # log += f"✅ Loaded {len(voice_samples)} voice samples\n" @@ -268,13 +303,15 @@ class VibeVoiceDemo: start_time = time.time() - inputs = self.processor( - text=[formatted_script], - voice_samples=[voice_samples], - padding=True, - return_tensors="pt", - return_attention_mask=True, - ) + processor_kwargs = { + "text": [formatted_script], + "padding": True, + "return_tensors": "pt", + "return_attention_mask": True, + } + processor_kwargs["voice_samples"] = [voice_samples] if voice_samples is not None else None + + inputs = self.processor(**processor_kwargs) # Move tensors to device target_device = self.device if self.device in ("cuda", "mps") else "cpu" for k, v in inputs.items(): @@ -294,7 +331,7 @@ class VibeVoiceDemo: # Start generation in a separate thread generation_thread = threading.Thread( target=self._generate_with_streamer, - args=(inputs, cfg_scale, audio_streamer) + args=(inputs, cfg_scale, audio_streamer, voice_cloning_enabled) ) generation_thread.start() @@ -476,7 +513,7 @@ class VibeVoiceDemo: traceback.print_exc() yield None, None, error_msg, gr.update(visible=False) - def _generate_with_streamer(self, inputs, cfg_scale, audio_streamer): + def _generate_with_streamer(self, inputs, cfg_scale, audio_streamer, voice_cloning_enabled: bool): """Helper method to run generation with streamer in a separate thread.""" try: # Check for stop signal before starting generation @@ -500,6 +537,7 @@ class VibeVoiceDemo: stop_check_fn=check_stop_generation, # Pass the stop check function verbose=False, # Disable verbose in streaming mode refresh_negative=True, + is_prefill=voice_cloning_enabled, ) except Exception as e: @@ -866,6 +904,11 @@ def create_demo_interface(demo_instance: VibeVoiceDemo): # info="Higher values increase adherence to text", elem_classes="slider-container" ) + disable_voice_cloning = gr.Checkbox( + value=False, + label="Disable voice cloning (skip conditioning voice prompts)", + info="When enabled, sets is_prefill=False so the model ignores provided speaker audio." + ) # Right column - Generation with gr.Column(scale=2, elem_classes="generation-card"): @@ -984,19 +1027,17 @@ Or paste text directly and it will auto-assign speakers.""", ) # Main generation function with streaming - def generate_podcast_wrapper(num_speakers, script, *speakers_and_params): + def generate_podcast_wrapper(num_speakers, script, speaker_1, speaker_2, speaker_3, speaker_4, cfg_scale, disable_voice_cloning): """Wrapper function to handle the streaming generation call.""" try: - # Extract speakers and parameters - speakers = speakers_and_params[:4] # First 4 are speaker selections - cfg_scale = speakers_and_params[4] # CFG scale - + speakers = [speaker_1, speaker_2, speaker_3, speaker_4] + # Clear outputs and reset visibility at start yield None, gr.update(value=None, visible=False), "🎙️ Starting generation...", gr.update(visible=True), gr.update(visible=False), gr.update(visible=True) - + # The generator will yield multiple times final_log = "Starting generation..." - + for streaming_audio, complete_audio, log, streaming_visible in demo_instance.generate_podcast_streaming( num_speakers=int(num_speakers), script=script, @@ -1004,7 +1045,8 @@ Or paste text directly and it will auto-assign speakers.""", speaker_2=speakers[1], speaker_3=speakers[2], speaker_4=speakers[3], - cfg_scale=cfg_scale + cfg_scale=cfg_scale, + disable_voice_cloning=disable_voice_cloning ): final_log = log @@ -1052,7 +1094,7 @@ Or paste text directly and it will auto-assign speakers.""", queue=False ).then( fn=generate_podcast_wrapper, - inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale], + inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale, disable_voice_cloning], outputs=[audio_output, complete_audio_output, log_output, streaming_status, generate_btn, stop_btn], queue=True # Enable Gradio's built-in queue ) @@ -1195,6 +1237,12 @@ def parse_args(): default=7860, help="Port to run the demo on", ) + parser.add_argument( + "--checkpoint_path", + type=str, + default=None, + help="Path to a fine-tuned checkpoint directory containing LoRA adapters (optional)", + ) return parser.parse_args() @@ -1211,7 +1259,8 @@ def main(): demo_instance = VibeVoiceDemo( model_path=args.model_path, device=args.device, - inference_steps=args.inference_steps + inference_steps=args.inference_steps, + adapter_path=args.checkpoint_path, ) # Create interface diff --git a/demo/inference_from_file.py b/demo/inference_from_file.py index 238ea70..17d5fa9 100644 --- a/demo/inference_from_file.py +++ b/demo/inference_from_file.py @@ -7,6 +7,7 @@ import time import torch from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference +from vibevoice.modular.lora_loading import load_lora_assets from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor from transformers.utils import logging @@ -170,6 +171,17 @@ def parse_args(): default=("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")), help="Device for inference: cuda | mps | cpu", ) + parser.add_argument( + "--checkpoint_path", + type=str, + default=None, + help="Path to a fine-tuned checkpoint directory containing LoRA adapters (optional)", + ) + parser.add_argument( + "--disable_prefill", + action="store_true", + help="Disable speech prefill (voice cloning) by setting is_prefill=False during generation", + ) parser.add_argument( "--cfg_scale", type=float, @@ -309,6 +321,35 @@ def main(): raise e + if args.checkpoint_path: + print(f"Loading fine-tuned assets from {args.checkpoint_path}") + try: + report = load_lora_assets(model, args.checkpoint_path) + loaded_components = [ + name for name, loaded in ( + ("language LoRA", report.language_model), + ("diffusion head LoRA", report.diffusion_head_lora), + ("diffusion head weights", report.diffusion_head_full), + ("acoustic connector", report.acoustic_connector), + ("semantic connector", report.semantic_connector), + ) + if loaded + ] + if loaded_components: + print(f"Loaded components: {', '.join(loaded_components)}") + else: + print("Warning: no adapter components were loaded; check the checkpoint path.") + if report.adapter_root is not None: + print(f"Adapter assets resolved to: {report.adapter_root}") + except Exception as exc: + print(f"Failed to load LoRA assets: {exc}") + raise + + if args.disable_prefill: + print("Voice cloning disabled: running generation with is_prefill=False") + else: + print("Voice cloning enabled: running generation with is_prefill=True") + model.eval() model.set_ddpm_inference_steps(num_steps=10) @@ -341,6 +382,7 @@ def main(): tokenizer=processor.tokenizer, generation_config={'do_sample': False}, verbose=True, + is_prefill=not args.disable_prefill, ) generation_time = time.time() - start_time print(f"Generation time: {generation_time:.2f} seconds") diff --git a/vibevoice/modular/lora_loading.py b/vibevoice/modular/lora_loading.py new file mode 100644 index 0000000..8bd157b --- /dev/null +++ b/vibevoice/modular/lora_loading.py @@ -0,0 +1,174 @@ +"""Utilities for loading fine-tuned LoRA adapters and connector weights.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import torch +import torch.nn as nn + +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +@dataclass +class _LoadReport: + """Simple structure capturing what assets were loaded.""" + + language_model: bool = False + diffusion_head_lora: bool = False + diffusion_head_full: bool = False + acoustic_connector: bool = False + semantic_connector: bool = False + adapter_root: Optional[Path] = None + + +class _DiffusionHeadForwardShim(nn.Module): + """Wraps the diffusion head to expose the signature expected by PEFT.""" + + def __init__(self, base: nn.Module): + super().__init__() + self.base = base + + def forward(self, *args, **kwargs): # type: ignore[override] + if len(args) >= 3: + noisy_images, timesteps, condition = args[:3] + else: + noisy_images = kwargs.get("noisy_images") + timesteps = kwargs.get("timesteps") + condition = kwargs.get("condition") + return self.base(noisy_images, timesteps, condition) + + +def _resolve_adapter_root(checkpoint_path: Path) -> Path: + """Return the directory that actually holds adapter assets.""" + + if checkpoint_path.is_file(): + checkpoint_path = checkpoint_path.parent + + if (checkpoint_path / "lora").exists(): + return checkpoint_path / "lora" + return checkpoint_path + + +def _load_connector(module: Optional[nn.Module], path: Path, device: torch.device) -> bool: + if module is None or not path.exists(): + return False + + state_dict = torch.load(path, map_location=device) + missing, unexpected = module.load_state_dict(state_dict, strict=False) + if missing: + logger.warning(f"Connector load missing keys: {missing}") + if unexpected: + logger.warning(f"Connector load unexpected keys: {unexpected}") + module.to(device) + return True + + +def _load_diffusion_head( + model, adapter_root: Path, device: torch.device, report: _LoadReport +) -> None: + diff_dir = adapter_root / "diffusion_head" + adapter_config = diff_dir / "adapter_config.json" + adapter_model = diff_dir / "adapter_model.bin" + + try: + from peft import PeftModel + except ImportError as exc: # pragma: no cover - dependency guaranteed via pyproject, safeguard anyway + raise RuntimeError( + "peft is required to load diffusion head adapters but is not installed" + ) from exc + + if adapter_config.exists() and adapter_model.exists(): + logger.info(f"Loading diffusion head LoRA from {diff_dir}") + shim = _DiffusionHeadForwardShim(model.model.prediction_head) + peft_head = PeftModel.from_pretrained(shim, diff_dir) + peft_head.to(device) + model.model.prediction_head = peft_head + report.diffusion_head_lora = True + return + + # Fallback to full state dict if provided + full_path = diff_dir / "diffusion_head_full.bin" + if not full_path.exists(): + full_path = adapter_root / "diffusion_head_full.bin" + + if full_path.exists(): + logger.info(f"Loading full diffusion head weights from {full_path}") + state_dict = torch.load(full_path, map_location=device) + missing, unexpected = model.model.prediction_head.load_state_dict(state_dict, strict=False) + if missing: + logger.warning(f"Diffusion head load missing keys: {missing}") + if unexpected: + logger.warning(f"Diffusion head load unexpected keys: {unexpected}") + model.model.prediction_head.to(device) + report.diffusion_head_full = True + + +def _load_language_model( + model, adapter_root: Path, device: torch.device, report: _LoadReport +) -> None: + config_file = adapter_root / "adapter_config.json" + bin_file = adapter_root / "adapter_model.bin" + if not (config_file.exists() and bin_file.exists()): + return + + try: + from peft import PeftModel + except ImportError as exc: # pragma: no cover - safeguard + raise RuntimeError( + "peft is required to load language model adapters but is not installed" + ) from exc + + logger.info(f"Loading language model LoRA from {adapter_root}") + peft_lm = PeftModel.from_pretrained(model.model.language_model, adapter_root) + peft_lm.to(device) + model.model.language_model = peft_lm + if hasattr(model, "tie_weights"): + try: + model.tie_weights() + except Exception as exc: + logger.warning(f"Failed to retie weights after loading language LoRA: {exc}") + report.language_model = True + + +def load_lora_assets(model, checkpoint_dir: str, device: Optional[torch.device] = None) -> _LoadReport: + """Load LoRA adapters and connector weights onto an instantiated model. + + Args: + model: The already instantiated `VibeVoiceForConditionalGenerationInference` model. + checkpoint_dir: Directory produced during fine-tuning containing a `lora/` folder. + device: Optional device to place loaded modules on. When ``None`` the function + infers the device from the model parameters. + + Returns: + `_LoadReport` summarizing which assets were successfully loaded. + """ + + adapter_root = _resolve_adapter_root(Path(checkpoint_dir)) + if not adapter_root.exists(): + raise FileNotFoundError(f"Adapter directory not found: {adapter_root}") + + inferred_device = device or next(model.parameters()).device + report = _LoadReport(adapter_root=adapter_root) + + _load_language_model(model, adapter_root, inferred_device, report) + _load_diffusion_head(model, adapter_root, inferred_device, report) + + ac_path = adapter_root / "acoustic_connector" / "pytorch_model.bin" + if _load_connector(getattr(model.model, "acoustic_connector", None), ac_path, inferred_device): + report.acoustic_connector = True + + se_path = adapter_root / "semantic_connector" / "pytorch_model.bin" + if _load_connector(getattr(model.model, "semantic_connector", None), se_path, inferred_device): + report.semantic_connector = True + + if not any(report.__dict__.values()): + logger.warning( + "No adapter assets were loaded. Ensure the checkpoint directory is correct and contains LoRA weights." + ) + + return report