LoRA inference (#14)

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