From 32c564509ed615a9627c7dc71fc55d5246fcfd04 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 2 Aug 2023 21:13:16 -0700
Subject: [PATCH] Fix loading session in chat mode

---
 modules/chat.py   |  4 ++++
 modules/shared.py |  4 ++++
 server.py         | 22 ++++++++++++++--------
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/modules/chat.py b/modules/chat.py
index 57a04606..5e4eb245 100644
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -418,6 +418,10 @@ def save_persistent_history(history, character, mode):
 
 
 def load_persistent_history(state):
+    if shared.session_is_loading:
+        shared.session_is_loading = False
+        return state['history']
+
     if state['mode'] == 'instruct':
         return state['history']
 
diff --git a/modules/shared.py b/modules/shared.py
index a2782e65..bac3fa8c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -30,6 +30,10 @@ reload_inputs = []  # Parameters for reloading the chat interface
 # For restarting the interface
 need_restart = False
 
+# To prevent the persistent chat history from being loaded when
+# a session JSON file is being loaded in chat mode
+session_is_loading = False
+
 settings = {
     'dark_theme': True,
     'autoload_model': False,
diff --git a/server.py b/server.py
index a0229995..6715cc4d 100644
--- a/server.py
+++ b/server.py
@@ -511,21 +511,27 @@ def create_file_saving_event_handlers():
         def load_session(file, state):
             decoded_file = file if type(file) == str else file.decode('utf-8')
             data = json.loads(decoded_file)
+
+            if shared.is_chat() and 'character_menu' in data and state.get('character_menu') != data.get('character_menu'):
+                shared.session_is_loading = True
+
             state.update(data)
-
-            if shared.is_chat():
-                chat.save_persistent_history(state['history'], state['character_menu'], state['mode'])
-
             return state
 
         shared.gradio['save_session'].click(
             ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
             lambda x: json.dumps(x, indent=4), gradio('interface_state'), gradio('temporary_text')).then(
-                None, gradio('temporary_text'), None, _js=f"(contents) => {{{ui.save_files_js}; saveSession(contents, \"{shared.get_mode()}\")}}")
+            None, gradio('temporary_text'), None, _js=f"(contents) => {{{ui.save_files_js}; saveSession(contents, \"{shared.get_mode()}\")}}")
 
-        shared.gradio['load_session'].upload(
-            load_session, gradio('load_session', 'interface_state'), gradio('interface_state')).then(
-            ui.apply_interface_values, gradio('interface_state'), gradio(ui.list_interface_input_elements()), show_progress=False)
+        if shared.is_chat():
+            shared.gradio['load_session'].upload(
+                load_session, gradio('load_session', 'interface_state'), gradio('interface_state')).then(
+                ui.apply_interface_values, gradio('interface_state'), gradio(ui.list_interface_input_elements()), show_progress=False).then(
+                chat.redraw_html, shared.reload_inputs, gradio('display'))
+        else:
+            shared.gradio['load_session'].upload(
+                load_session, gradio('load_session', 'interface_state'), gradio('interface_state')).then(
+                ui.apply_interface_values, gradio('interface_state'), gradio(ui.list_interface_input_elements()), show_progress=False)
 
 
 def set_interface_arguments(interface_mode, extensions, bool_active):