From 187cf6203600a45cac0b9527977cb0e933b30af4 Mon Sep 17 00:00:00 2001 From: mrfakename Date: Mon, 27 Oct 2025 08:51:49 -0700 Subject: [PATCH] Add support for loading Safetensors LoRAs --- vibevoice/modular/lora_loading.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vibevoice/modular/lora_loading.py b/vibevoice/modular/lora_loading.py index 8bd157b..8929fd6 100644 --- a/vibevoice/modular/lora_loading.py +++ b/vibevoice/modular/lora_loading.py @@ -74,6 +74,7 @@ def _load_diffusion_head( diff_dir = adapter_root / "diffusion_head" adapter_config = diff_dir / "adapter_config.json" adapter_model = diff_dir / "adapter_model.bin" + adapter_model_safetensors = diff_dir / "adapter_model.safetensors" try: from peft import PeftModel @@ -82,7 +83,7 @@ def _load_diffusion_head( "peft is required to load diffusion head adapters but is not installed" ) from exc - if adapter_config.exists() and adapter_model.exists(): + if adapter_config.exists() and (adapter_model.exists() or adapter_model_safetensors.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) @@ -113,7 +114,8 @@ def _load_language_model( ) -> None: config_file = adapter_root / "adapter_config.json" bin_file = adapter_root / "adapter_model.bin" - if not (config_file.exists() and bin_file.exists()): + safe_tensors_file = adapter_root / "adapter_model.safetensors" + if not (config_file.exists() and (bin_file.exists() or safe_tensors_file.exists())): return try: