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()