490 lines
17 KiB
Python
490 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
VibeVoice Universal Model Merger
|
|
|
|
Automatically detects and merges trained components back into the base model:
|
|
- LLM LoRA adapters
|
|
- Diffusion head (LoRA or full fine-tune)
|
|
- Acoustic/Semantic connectors
|
|
|
|
Supports all training configurations from train_vibevoice.py
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import shutil
|
|
from typing import Dict, Optional
|
|
|
|
import torch
|
|
|
|
from vibevoice.modular.modeling_vibevoice import VibeVoiceForConditionalGeneration
|
|
|
|
logging.basicConfig(
|
|
format="%(asctime)s - %(levelname)s - %(message)s",
|
|
level=logging.INFO
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def detect_trained_components(checkpoint_path: str) -> Dict[str, bool]:
|
|
"""Detect which components were trained."""
|
|
|
|
components = {
|
|
"llm_lora": False,
|
|
"diffusion_head": False,
|
|
"acoustic_connector": False,
|
|
"semantic_connector": False,
|
|
}
|
|
|
|
# Check LLM LoRA (adapter files in root checkpoint/)
|
|
llm_adapter_config = os.path.join(checkpoint_path, "adapter_config.json")
|
|
llm_adapter_model = os.path.join(checkpoint_path, "adapter_model.safetensors")
|
|
if not os.path.exists(llm_adapter_model):
|
|
llm_adapter_model = os.path.join(checkpoint_path, "adapter_model.bin")
|
|
|
|
if os.path.exists(llm_adapter_config) and os.path.exists(llm_adapter_model):
|
|
components["llm_lora"] = True
|
|
|
|
# Check diffusion head
|
|
diffusion_head_dir = os.path.join(checkpoint_path, "diffusion_head")
|
|
if os.path.exists(diffusion_head_dir):
|
|
components["diffusion_head"] = True
|
|
|
|
# Check acoustic connector
|
|
acoustic_conn_path = os.path.join(checkpoint_path, "acoustic_connector", "pytorch_model.bin")
|
|
if os.path.exists(acoustic_conn_path):
|
|
components["acoustic_connector"] = True
|
|
|
|
# Check semantic connector
|
|
semantic_conn_path = os.path.join(checkpoint_path, "semantic_connector", "pytorch_model.bin")
|
|
if os.path.exists(semantic_conn_path):
|
|
components["semantic_connector"] = True
|
|
|
|
return components
|
|
|
|
|
|
def merge_llm_lora(model: VibeVoiceForConditionalGeneration, checkpoint_path: str) -> None:
|
|
"""Merge LLM LoRA adapters into base model."""
|
|
|
|
logger.info("Merging LLM LoRA adapters...")
|
|
|
|
try:
|
|
from peft import PeftModel
|
|
except ImportError:
|
|
raise ImportError("peft library required for LoRA merge. Install: pip install peft")
|
|
|
|
# Load and merge LoRA
|
|
language_model_with_lora = PeftModel.from_pretrained(
|
|
model.model.language_model,
|
|
checkpoint_path
|
|
)
|
|
|
|
merged_lm = language_model_with_lora.merge_and_unload()
|
|
model.model.language_model = merged_lm
|
|
|
|
logger.info("✓ LLM LoRA merge completed")
|
|
|
|
|
|
def merge_diffusion_head(model: VibeVoiceForConditionalGeneration, checkpoint_path: str) -> dict:
|
|
"""Merge diffusion head weights (LoRA or full fine-tune).
|
|
|
|
Returns:
|
|
trained_state_dict for verification
|
|
"""
|
|
|
|
logger.info("Merging diffusion head...")
|
|
|
|
diffusion_head_dir = os.path.join(checkpoint_path, "diffusion_head")
|
|
|
|
# Try multiple possible weight file locations
|
|
possible_files = [
|
|
os.path.join(diffusion_head_dir, "model.safetensors"),
|
|
os.path.join(diffusion_head_dir, "diffusion_head_full.bin"),
|
|
os.path.join(checkpoint_path, "diffusion_head_full.bin"),
|
|
]
|
|
|
|
trained_weights_path = None
|
|
for path in possible_files:
|
|
if os.path.exists(path):
|
|
trained_weights_path = path
|
|
break
|
|
|
|
if trained_weights_path is None:
|
|
raise ValueError(
|
|
f"Diffusion head weights not found. Searched:\n" +
|
|
"\n".join(f" - {p}" for p in possible_files)
|
|
)
|
|
|
|
logger.info(f"Loading from: {trained_weights_path}")
|
|
|
|
# Load weights
|
|
if trained_weights_path.endswith(".safetensors"):
|
|
from safetensors.torch import load_file
|
|
trained_state_dict = load_file(trained_weights_path)
|
|
else:
|
|
trained_state_dict = torch.load(trained_weights_path, map_location="cpu")
|
|
|
|
# Check if LoRA-wrapped (has adapter keys like lora_A, lora_B)
|
|
is_lora = any("lora_" in k for k in trained_state_dict.keys())
|
|
|
|
if is_lora:
|
|
logger.info("Detected LoRA format, performing LoRA merge...")
|
|
try:
|
|
from peft import PeftModel
|
|
pred_head_with_lora = PeftModel.from_pretrained(
|
|
model.model.prediction_head,
|
|
diffusion_head_dir
|
|
)
|
|
model.model.prediction_head = pred_head_with_lora.merge_and_unload()
|
|
except Exception as e:
|
|
logger.warning(f"LoRA merge failed: {e}. Trying direct state_dict load...")
|
|
model.model.prediction_head.load_state_dict(trained_state_dict, strict=True)
|
|
else:
|
|
logger.info("Detected full fine-tune format, replacing weights...")
|
|
model.model.prediction_head.load_state_dict(trained_state_dict, strict=True)
|
|
|
|
logger.info("✓ Diffusion head merge completed")
|
|
return trained_state_dict
|
|
|
|
|
|
def merge_connectors(model: VibeVoiceForConditionalGeneration, checkpoint_path: str,
|
|
merge_acoustic: bool, merge_semantic: bool) -> None:
|
|
"""Merge acoustic and/or semantic connector weights."""
|
|
|
|
if merge_acoustic:
|
|
logger.info("Merging acoustic connector...")
|
|
acoustic_path = os.path.join(checkpoint_path, "acoustic_connector", "pytorch_model.bin")
|
|
state_dict = torch.load(acoustic_path, map_location="cpu")
|
|
model.model.acoustic_connector.load_state_dict(state_dict, strict=True)
|
|
logger.info("✓ Acoustic connector merge completed")
|
|
|
|
if merge_semantic:
|
|
logger.info("Merging semantic connector...")
|
|
semantic_path = os.path.join(checkpoint_path, "semantic_connector", "pytorch_model.bin")
|
|
state_dict = torch.load(semantic_path, map_location="cpu")
|
|
model.model.semantic_connector.load_state_dict(state_dict, strict=True)
|
|
logger.info("✓ Semantic connector merge completed")
|
|
|
|
|
|
def verify_merge(
|
|
base_model: VibeVoiceForConditionalGeneration,
|
|
merged_model: VibeVoiceForConditionalGeneration,
|
|
trained_state_dict: Optional[dict],
|
|
component_name: str
|
|
) -> None:
|
|
"""
|
|
Verify that merge was successful.
|
|
|
|
Args:
|
|
base_model: Original base model
|
|
merged_model: Model after merge
|
|
trained_state_dict: Trained weights that were merged (if available)
|
|
component_name: Name of component being verified (e.g., "diffusion_head")
|
|
"""
|
|
|
|
logger.info(f"\n=== Verifying {component_name} merge ===")
|
|
|
|
# Get component modules
|
|
if component_name == "diffusion_head":
|
|
base_module = base_model.model.prediction_head
|
|
merged_module = merged_model.model.prediction_head
|
|
elif component_name == "acoustic_connector":
|
|
base_module = base_model.model.acoustic_connector
|
|
merged_module = merged_model.model.acoustic_connector
|
|
elif component_name == "semantic_connector":
|
|
base_module = base_model.model.semantic_connector
|
|
merged_module = merged_model.model.semantic_connector
|
|
else:
|
|
logger.warning(f"Unknown component: {component_name}, skipping verification")
|
|
return
|
|
|
|
base_state = base_module.state_dict()
|
|
merged_state = merged_module.state_dict()
|
|
|
|
# 1. Check that weights actually changed
|
|
logger.info("Checking if weights changed from base model...")
|
|
weights_changed = False
|
|
changed_params = []
|
|
|
|
for key in base_state.keys():
|
|
if key not in merged_state:
|
|
continue
|
|
if not torch.allclose(base_state[key], merged_state[key], rtol=1e-5, atol=1e-8):
|
|
weights_changed = True
|
|
changed_params.append(key)
|
|
|
|
if not weights_changed:
|
|
# Weights didn't change - this is OK for connectors (may not have been trained)
|
|
# But it's an error for diffusion_head (should always be trained)
|
|
if component_name == "diffusion_head":
|
|
raise ValueError(f"✗ ERROR: {component_name} weights did not change! Merge may have failed.")
|
|
else:
|
|
logger.info(f"✓ {component_name}: unchanged (was not trained)")
|
|
return
|
|
|
|
logger.info(f"✓ Weights changed: {len(changed_params)}/{len(base_state)} parameters modified")
|
|
|
|
# 2. Check trained weights match merged weights (if available)
|
|
if trained_state_dict is not None:
|
|
logger.info("Verifying trained weights match merged model...")
|
|
mismatches = []
|
|
|
|
for key in trained_state_dict.keys():
|
|
if key not in merged_state:
|
|
mismatches.append(f"{key} (missing in merged)")
|
|
continue
|
|
|
|
# Convert to same dtype for comparison (handle bf16/fp32 mismatch)
|
|
trained_tensor = trained_state_dict[key].float()
|
|
merged_tensor = merged_state[key].float()
|
|
|
|
if not torch.allclose(trained_tensor, merged_tensor, rtol=1e-5, atol=1e-8):
|
|
mismatches.append(f"{key} (values differ)")
|
|
|
|
if mismatches:
|
|
logger.error(f"✗ Weight mismatches found:")
|
|
for mm in mismatches[:5]: # Show first 5
|
|
logger.error(f" - {mm}")
|
|
if len(mismatches) > 5:
|
|
logger.error(f" ... and {len(mismatches) - 5} more")
|
|
raise ValueError(f"✗ ERROR: Trained and merged weights do not match!")
|
|
|
|
logger.info(f"✓ All trained weights correctly merged: {len(trained_state_dict)} parameters verified")
|
|
|
|
# 3. Parameter count and structure
|
|
base_params = sum(p.numel() for p in base_module.parameters())
|
|
merged_params = sum(p.numel() for p in merged_module.parameters())
|
|
|
|
if base_params != merged_params:
|
|
raise ValueError(f"✗ ERROR: Parameter count mismatch: base={base_params:,} vs merged={merged_params:,}")
|
|
|
|
logger.info(f"✓ Parameter count matches: {merged_params:,}")
|
|
logger.info(f"✓✓✓ {component_name} verification PASSED ✓✓✓")
|
|
|
|
|
|
def verify_models_only(
|
|
base_model_path: str,
|
|
merged_model_path: str
|
|
) -> None:
|
|
"""
|
|
Verify-only mode: Compare base and merged models without performing merge.
|
|
|
|
Args:
|
|
base_model_path: Path to base model
|
|
merged_model_path: Path to allegedly merged model
|
|
"""
|
|
|
|
logger.info("=== VERIFY-ONLY MODE ===")
|
|
logger.info(f"Base model: {base_model_path}")
|
|
logger.info(f"Merged model: {merged_model_path}")
|
|
|
|
# Load both models
|
|
logger.info("\nLoading base model...")
|
|
base_model = VibeVoiceForConditionalGeneration.from_pretrained(
|
|
base_model_path,
|
|
torch_dtype=torch.float32
|
|
)
|
|
|
|
logger.info("Loading merged model...")
|
|
merged_model = VibeVoiceForConditionalGeneration.from_pretrained(
|
|
merged_model_path,
|
|
torch_dtype=torch.float32
|
|
)
|
|
|
|
# Verify each component
|
|
components_to_check = ["diffusion_head", "acoustic_connector", "semantic_connector"]
|
|
|
|
for component in components_to_check:
|
|
try:
|
|
verify_merge(base_model, merged_model, None, component)
|
|
except ValueError as e:
|
|
if "did not change" in str(e):
|
|
logger.info(f"✓ {component}: unchanged (likely not trained)")
|
|
else:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"✗ {component} verification failed: {e}")
|
|
raise
|
|
|
|
logger.info("\n✓✓✓ VERIFICATION COMPLETE ✓✓✓")
|
|
|
|
|
|
def merge_vibevoice_model(
|
|
base_model_path: str,
|
|
checkpoint_path: str,
|
|
output_path: str,
|
|
output_format: str = "safetensors"
|
|
) -> None:
|
|
"""
|
|
Universal merge function for VibeVoice models.
|
|
Automatically detects and merges all trained components.
|
|
"""
|
|
|
|
# Detect what was trained
|
|
logger.info(f"Scanning trained components in: {checkpoint_path}")
|
|
components = detect_trained_components(checkpoint_path)
|
|
|
|
logger.info("Detected trained components:")
|
|
for name, trained in components.items():
|
|
status = "✓ Found" if trained else "✗ Not found"
|
|
logger.info(f" {name}: {status}")
|
|
|
|
if not any(components.values()):
|
|
raise ValueError("No trained components found in checkpoint path!")
|
|
|
|
# Load base model
|
|
logger.info(f"\nLoading base model from: {base_model_path}")
|
|
base_model = VibeVoiceForConditionalGeneration.from_pretrained(
|
|
base_model_path,
|
|
torch_dtype=torch.float32
|
|
)
|
|
|
|
# Merge components
|
|
logger.info("\n=== Starting merge process ===")
|
|
|
|
trained_diffusion_state = None
|
|
|
|
if components["llm_lora"]:
|
|
merge_llm_lora(base_model, checkpoint_path)
|
|
|
|
if components["diffusion_head"]:
|
|
trained_diffusion_state = merge_diffusion_head(base_model, checkpoint_path)
|
|
|
|
if components["acoustic_connector"] or components["semantic_connector"]:
|
|
merge_connectors(
|
|
base_model,
|
|
checkpoint_path,
|
|
merge_acoustic=components["acoustic_connector"],
|
|
merge_semantic=components["semantic_connector"]
|
|
)
|
|
|
|
# Save merged model
|
|
logger.info(f"\n=== Saving merged model to: {output_path} ===")
|
|
os.makedirs(output_path, exist_ok=True)
|
|
|
|
if output_format == "safetensors":
|
|
base_model.save_pretrained(output_path, safe_serialization=True)
|
|
elif output_format == "bin":
|
|
base_model.save_pretrained(output_path, safe_serialization=False)
|
|
else:
|
|
raise ValueError(f"Unknown output format: {output_format}. Use 'safetensors' or 'bin'")
|
|
|
|
# Copy config and processor files
|
|
logger.info("Copying config and processor files...")
|
|
files_to_copy = [
|
|
"config.json",
|
|
"preprocessor_config.json",
|
|
"generation_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer_config.json",
|
|
"tokenizer.json",
|
|
"vocab.json",
|
|
"merges.txt"
|
|
]
|
|
|
|
for file in files_to_copy:
|
|
src = os.path.join(base_model_path, file)
|
|
dst = os.path.join(output_path, file)
|
|
if os.path.exists(src):
|
|
shutil.copy2(src, dst)
|
|
|
|
# Verification
|
|
logger.info("\n=== Verifying merged model ===")
|
|
try:
|
|
# Reload original base model for comparison (the base_model variable was modified during merge)
|
|
logger.info("Reloading original base model for verification...")
|
|
original_base_model = VibeVoiceForConditionalGeneration.from_pretrained(
|
|
base_model_path,
|
|
torch_dtype=torch.float32
|
|
)
|
|
|
|
logger.info("Loading merged model for verification...")
|
|
test_model = VibeVoiceForConditionalGeneration.from_pretrained(output_path)
|
|
logger.info("✓ Model loads successfully")
|
|
|
|
# Detailed verification for each merged component
|
|
if components["diffusion_head"]:
|
|
verify_merge(original_base_model, test_model, trained_diffusion_state, "diffusion_head")
|
|
|
|
if components["acoustic_connector"]:
|
|
verify_merge(original_base_model, test_model, None, "acoustic_connector")
|
|
|
|
if components["semantic_connector"]:
|
|
verify_merge(original_base_model, test_model, None, "semantic_connector")
|
|
|
|
logger.info("\n✓✓✓ Merge and verification completed successfully! ✓✓✓")
|
|
|
|
except Exception as e:
|
|
logger.error(f"✗ Verification failed: {e}")
|
|
raise
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Universal merger for VibeVoice trained components",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Merge and verify
|
|
python merge_vibevoice_models.py --base_model_path model --checkpoint_path output/lora --output_path merged
|
|
|
|
# Verify existing merge (no actual merging)
|
|
python merge_vibevoice_models.py --base_model_path model --output_path merged --verify_only
|
|
"""
|
|
)
|
|
parser.add_argument(
|
|
"--base_model_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to base VibeVoice model directory"
|
|
)
|
|
parser.add_argument(
|
|
"--checkpoint_path",
|
|
type=str,
|
|
required=False,
|
|
help="Path to checkpoint directory (usually 'lora/' or 'checkpoint-XXX/lora/'). Not needed with --verify_only"
|
|
)
|
|
parser.add_argument(
|
|
"--output_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to save merged model (or path to verify with --verify_only)"
|
|
)
|
|
parser.add_argument(
|
|
"--output_format",
|
|
type=str,
|
|
default="safetensors",
|
|
choices=["safetensors", "bin"],
|
|
help="Output format: 'safetensors' (recommended) or 'bin'"
|
|
)
|
|
parser.add_argument(
|
|
"--verify_only",
|
|
action="store_true",
|
|
help="Only verify existing merge between base_model_path and output_path (no actual merging)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Verify-only mode
|
|
if args.verify_only:
|
|
verify_models_only(
|
|
base_model_path=args.base_model_path,
|
|
merged_model_path=args.output_path
|
|
)
|
|
return
|
|
|
|
# Normal merge mode
|
|
if not args.checkpoint_path:
|
|
parser.error("--checkpoint_path is required unless using --verify_only")
|
|
|
|
merge_vibevoice_model(
|
|
base_model_path=args.base_model_path,
|
|
checkpoint_path=args.checkpoint_path,
|
|
output_path=args.output_path,
|
|
output_format=args.output_format
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |