LoRA inference (#14)

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
 ```
 
 

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,7 +254,9 @@ class VibeVoiceDemo:
                 yield None, "🛑 Generation stopped by user", gr.update(visible=False)
                 return
             
-            # Load voice samples
+            # 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]
@@ -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,12 +1027,10 @@ 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)
@@ -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

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")

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