Fix UI error caused by --idle-timeout

This commit is contained in:
oobabooga 2024-07-28 18:30:06 -07:00
parent 3aa646c1d0
commit 514fb2e451

View file

@ -368,14 +368,16 @@ def clear_torch_cache():
torch.cuda.empty_cache() torch.cuda.empty_cache()
def unload_model(): def unload_model(keep_model_name=False):
shared.model = shared.tokenizer = None shared.model = shared.tokenizer = None
shared.previous_model_name = shared.model_name shared.previous_model_name = shared.model_name
shared.model_name = 'None'
shared.lora_names = [] shared.lora_names = []
shared.model_dirty_from_training = False shared.model_dirty_from_training = False
clear_torch_cache() clear_torch_cache()
if not keep_model_name:
shared.model_name = 'None'
def reload_model(): def reload_model():
unload_model() unload_model()
@ -393,7 +395,7 @@ def unload_model_if_idle():
if time.time() - last_generation_time > shared.args.idle_timeout * 60: if time.time() - last_generation_time > shared.args.idle_timeout * 60:
if shared.model is not None: if shared.model is not None:
logger.info("Unloading the model for inactivity.") logger.info("Unloading the model for inactivity.")
unload_model() unload_model(keep_model_name=True)
finally: finally:
shared.generation_lock.release() shared.generation_lock.release()