From 514fb2e451b88cddfc73df38d5052ce865f0b9a4 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Sun, 28 Jul 2024 18:30:06 -0700 Subject: [PATCH] Fix UI error caused by --idle-timeout --- modules/models.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/modules/models.py b/modules/models.py index ea046e9b..ff990ae6 100644 --- a/modules/models.py +++ b/modules/models.py @@ -368,14 +368,16 @@ def clear_torch_cache(): torch.cuda.empty_cache() -def unload_model(): +def unload_model(keep_model_name=False): shared.model = shared.tokenizer = None shared.previous_model_name = shared.model_name - shared.model_name = 'None' shared.lora_names = [] shared.model_dirty_from_training = False clear_torch_cache() + if not keep_model_name: + shared.model_name = 'None' + def reload_model(): unload_model() @@ -393,7 +395,7 @@ def unload_model_if_idle(): if time.time() - last_generation_time > shared.args.idle_timeout * 60: if shared.model is not None: logger.info("Unloading the model for inactivity.") - unload_model() + unload_model(keep_model_name=True) finally: shared.generation_lock.release()