Unload models on request (#471 from Brawlence/main)

This commit is contained in:
oobabooga 2023-03-26 23:53:39 -03:00 committed by GitHub
commit af603a142a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -50,12 +50,17 @@ def get_available_softprompts():
def get_available_loras(): def get_available_loras():
return ['None'] + sorted([item.name for item in list(Path('loras/').glob('*')) if not item.name.endswith(('.txt', '-np', '.pt', '.json'))], key=str.lower) return ['None'] + sorted([item.name for item in list(Path('loras/').glob('*')) if not item.name.endswith(('.txt', '-np', '.pt', '.json'))], key=str.lower)
def unload_model():
shared.model = shared.tokenizer = None
clear_torch_cache()
def load_model_wrapper(selected_model): def load_model_wrapper(selected_model):
if selected_model != shared.model_name: if selected_model != shared.model_name:
shared.model_name = selected_model shared.model_name = selected_model
shared.model = shared.tokenizer = None
clear_torch_cache() unload_model()
shared.model, shared.tokenizer = load_model(shared.model_name) if selected_model != '':
shared.model, shared.tokenizer = load_model(shared.model_name)
return selected_model return selected_model