LoRA inference (#14)

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