LoRA.py 550 B

1234567891011121314151617
  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. print(f"Adding the LoRA {lora_name} to the model...")
  12. shared.model = PeftModel.from_pretrained(shared.model, Path(f"loras/{lora_name}"))