Implement sessions + add basic multi-user support (#2991)

modules/ui.py

@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 
 import gradio as gr
@@ -5,6 +6,7 @@ import torch
 
 from modules import shared
 
+
 with open(Path(__file__).resolve().parent / '../css/main.css', 'r') as f:
     css = f.read()
 with open(Path(__file__).resolve().parent / '../css/chat.css', 'r') as f:
@@ -14,7 +16,7 @@ with open(Path(__file__).resolve().parent / '../css/main.js', 'r') as f:
 with open(Path(__file__).resolve().parent / '../css/chat.js', 'r') as f:
     chat_js = f.read()
 
-refresh_symbol = '\U0001f504'  # 🔄
+refresh_symbol = '🔄'
 delete_symbol = '🗑️'
 save_symbol = '💾'
 
@@ -30,17 +32,103 @@ theme = gr.themes.Default(
 
 
 def list_model_elements():
-    elements = ['loader', 'cpu_memory', 'auto_devices', 'disk', 'cpu', 'bf16', 'load_in_8bit', 'trust_remote_code', 'load_in_4bit', 'compute_dtype', 'quant_type', 'use_double_quant', 'wbits', 'groupsize', 'model_type', 'pre_layer', 'triton', 'desc_act', 'no_inject_fused_attention', 'no_inject_fused_mlp', 'no_use_cuda_fp16', 'threads', 'n_batch', 'no_mmap', 'mlock', 'n_gpu_layers', 'n_ctx', 'llama_cpp_seed', 'gpu_split', 'max_seq_len', 'compress_pos_emb']
+    elements = [
+        'loader',
+        'cpu_memory',
+        'auto_devices',
+        'disk',
+        'cpu',
+        'bf16',
+        'load_in_8bit',
+        'trust_remote_code',
+        'load_in_4bit',
+        'compute_dtype',
+        'quant_type',
+        'use_double_quant',
+        'wbits',
+        'groupsize',
+        'model_type',
+        'pre_layer',
+        'triton',
+        'desc_act',
+        'no_inject_fused_attention',
+        'no_inject_fused_mlp',
+        'no_use_cuda_fp16',
+        'threads',
+        'n_batch',
+        'no_mmap',
+        'mlock',
+        'n_gpu_layers',
+        'n_ctx',
+        'llama_cpp_seed',
+        'gpu_split',
+        'max_seq_len',
+        'compress_pos_emb'
+    ]
+
     for i in range(torch.cuda.device_count()):
         elements.append(f'gpu_memory_{i}')
 
     return elements
 
 
-def list_interface_input_elements(chat=False):
-    elements = ['max_new_tokens', 'seed', 'temperature', 'top_p', 'top_k', 'typical_p', 'epsilon_cutoff', 'eta_cutoff', 'repetition_penalty', 'repetition_penalty_range', 'encoder_repetition_penalty', 'no_repeat_ngram_size', 'min_length', 'do_sample', 'penalty_alpha', 'num_beams', 'length_penalty', 'early_stopping', 'mirostat_mode', 'mirostat_tau', 'mirostat_eta', 'add_bos_token', 'ban_eos_token', 'truncation_length', 'custom_stopping_strings', 'skip_special_tokens', 'preset_menu', 'stream', 'tfs', 'top_a']
-    if chat:
-        elements += ['name1', 'name2', 'greeting', 'context', 'chat_generation_attempts', 'stop_at_newline', 'mode', 'instruction_template', 'character_menu', 'name1_instruct', 'name2_instruct', 'context_instruct', 'turn_template', 'chat_style', 'chat-instruct_command']
+def list_interface_input_elements():
+    elements = [
+        'preset_menu',
+        'max_new_tokens',
+        'seed',
+        'temperature',
+        'top_p',
+        'top_k',
+        'typical_p',
+        'epsilon_cutoff',
+        'eta_cutoff',
+        'repetition_penalty',
+        'repetition_penalty_range',
+        'encoder_repetition_penalty',
+        'no_repeat_ngram_size',
+        'min_length',
+        'do_sample',
+        'penalty_alpha',
+        'num_beams',
+        'length_penalty',
+        'early_stopping',
+        'mirostat_mode',
+        'mirostat_tau',
+        'mirostat_eta',
+        'add_bos_token',
+        'ban_eos_token',
+        'truncation_length',
+        'custom_stopping_strings',
+        'skip_special_tokens',
+        'stream',
+        'tfs',
+        'top_a',
+    ]
+
+    if shared.args.chat:
+        elements += [
+            'character_menu',
+            'history',
+            'name1',
+            'name2',
+            'greeting',
+            'context',
+            'chat_generation_attempts',
+            'stop_at_newline',
+            'mode',
+            'instruction_template',
+            'name1_instruct',
+            'name2_instruct',
+            'context_instruct',
+            'turn_template',
+            'chat_style',
+            'chat-instruct_command',
+        ]
+    else:
+        elements.append('textbox')
+        if not shared.args.notebook:
+            elements.append('output_textbox')
 
     elements += list_model_elements()
     return elements
@@ -48,10 +136,14 @@ def list_interface_input_elements(chat=False):
 
 def gather_interface_values(*args):
     output = {}
-    for i, element in enumerate(shared.input_elements):
+    for i, element in enumerate(list_interface_input_elements()):
         output[element] = args[i]
 
-    shared.persistent_interface_state = output
+    if not shared.args.multi_user:
+        shared.persistent_interface_state = output
+        with open(Path(f'logs/session_{shared.get_mode()}_autosave.json'), 'w') as f:
+            f.write(json.dumps(output, indent=4))
+
     return output
 
 
@@ -59,11 +151,12 @@ def apply_interface_values(state, use_persistent=False):
     if use_persistent:
         state = shared.persistent_interface_state
 
-    elements = list_interface_input_elements(chat=shared.is_chat())
+    elements = list_interface_input_elements()
     if len(state) == 0:
         return [gr.update() for k in elements]  # Dummy, do nothing
     else:
-        return [state[k] if k in state else gr.update() for k in elements]
+        ans = [state[k] if k in state else gr.update() for k in elements]
+        return ans
 
 
 class ToolButton(gr.Button, gr.components.FormComponent):
@@ -92,6 +185,7 @@ def create_refresh_button(refresh_component, refresh_method, refreshed_args, ele
         inputs=[],
         outputs=[refresh_component]
     )
+
     return refresh_button