mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Clear cache while switching LoRAs

Browse source
This commit is contained in:
oobabooga 2023-03-23 21:56:26 -03:00 • committed by GitHub
parent 4578e88ffd
commit bf22d16ebc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 13 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

15
modules/LoRA.py
View file

@ -2,19 +2,22 @@ from pathlib import Path
import modules.shared as shared
from modules.models import load_model
from modules.text_generation import clear_torch_cache
def reload_model():
shared.model = shared.tokenizer = None
clear_torch_cache()
shared.model, shared.tokenizer = load_model(shared.model_name)
def add_lora_to_model(lora_name):
from peft import PeftModel
# Is there a more efficient way of returning to the base model?
if lora_name == "None":
print("Reloading the model to remove the LoRA...")
shared.model, shared.tokenizer = load_model(shared.model_name)
else:
reload_model()
if lora_name != "None":
print(f"Adding the LoRA {lora_name} to the model...")
params = {}
if not shared.args.cpu:
params['dtype'] = shared.model.dtype