175 lines
6.3 KiB
Python
175 lines
6.3 KiB
Python
"""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
|