LoRA.py 598 B

123456789101112131415161718
  1. from pathlib import Path
  2. from peft import PeftModel
  3. import modules.shared as shared
  4. from modules.models import load_model
  5. def add_lora_to_model(lora_name):
  6. # Is there a more efficient way of returning to the base model?
  7. if lora_name == "None":
  8. print("Reloading the model to remove the LoRA...")
  9. shared.model, shared.tokenizer = load_model(shared.model_name)
  10. else:
  11. # Why doesn't this work in 16-bit mode?
  12. print(f"Adding the LoRA {lora_name} to the model...")
  13. shared.model = PeftModel.from_pretrained(shared.model, Path(f"loras/{lora_name}"))