From 63a21e2b45e908be63765bf312a9ecfb3a588315 Mon Sep 17 00:00:00 2001
From: cseti007 <132441526+cseti007@users.noreply.github.com>
Date: Sun, 5 Oct 2025 04:25:35 +0200
Subject: [PATCH] add ---seed parameter for reproducibility (#18)

---
 demo/inference_from_file.py                 |  17 +-
 vibevoice/scripts/merge_vibevoice_models.py | 490 ++++++++++++++++++++
 2 files changed, 505 insertions(+), 2 deletions(-)
 create mode 100644 vibevoice/scripts/merge_vibevoice_models.py

diff --git a/demo/inference_from_file.py b/demo/inference_from_file.py
index 17d5fa9..e6386f7 100644
--- a/demo/inference_from_file.py
+++ b/demo/inference_from_file.py
@@ -188,7 +188,12 @@ def parse_args():
         default=1.3,
         help="CFG (Classifier-Free Guidance) scale for generation (default: 1.3)",
     )
-    
+    parser.add_argument(
+    "--seed",
+    type=int,
+    default=None,
+    help="Random seed for reproducibility (optional)",
+)
     return parser.parse_args()
 
 def main():
@@ -206,6 +211,12 @@ def main():
 
     print(f"Using device: {args.device}")
 
+    if args.seed is not None:
+        print(f"Setting seed: {args.seed}")
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(args.seed)
+
     # Initialize voice mapper
     voice_mapper = VoiceMapper()
     
@@ -435,7 +446,9 @@ def main():
     print(f"Generation time: {generation_time:.2f} seconds")
     print(f"Audio duration: {audio_duration:.2f} seconds")
     print(f"RTF (Real Time Factor): {rtf:.2f}x")
-    
+    if args.seed is not None:
+        print(f"Seed used: {args.seed}")
+
     print("="*50)
 
 if __name__ == "__main__":
diff --git a/vibevoice/scripts/merge_vibevoice_models.py b/vibevoice/scripts/merge_vibevoice_models.py
new file mode 100644
index 0000000..96f22f3
--- /dev/null
+++ b/vibevoice/scripts/merge_vibevoice_models.py
@@ -0,0 +1,490 @@
+#!/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()
\ No newline at end of file