text-generation-webui/server.py

249 lines
8.8 KiB
Python
Raw Normal View History

2023-04-05 19:05:50 +02:00
import os
2023-04-21 22:20:59 +02:00
import warnings
from modules.logging_colors import logger
2023-07-07 07:24:52 +02:00
from modules.block_requests import OpenMonkeyPatch, RequestBlocker
2023-04-21 22:18:34 +02:00
2023-04-05 19:05:50 +02:00
os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
os.environ['BITSANDBYTES_NOWELCOME'] = '1'
warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
with RequestBlocker():
import gradio as gr
2023-04-05 19:05:50 +02:00
import matplotlib
2023-05-04 02:43:17 +02:00
matplotlib.use('Agg') # This fixes LaTeX rendering on some systems
import json
import os
2023-04-13 02:24:26 +02:00
import sys
import time
from functools import partial
from pathlib import Path
from threading import Lock
import yaml
import modules.extensions as extensions_module
2023-08-07 02:49:27 +02:00
from modules import (
chat,
shared,
training,
ui,
ui_chat,
ui_default,
ui_file_saving,
ui_model_menu,
ui_notebook,
ui_parameters,
ui_session,
utils,
)
from modules.extensions import apply_extensions
2023-03-17 15:42:25 +01:00
from modules.LoRA import add_lora_to_model
2023-08-07 02:49:27 +02:00
from modules.models import load_model
2023-06-25 06:44:36 +02:00
from modules.models_settings import (
get_model_settings_from_yamls,
update_model_parameters
)
from modules.utils import gradio
2023-03-17 01:31:39 +01:00
def create_interface():
title = 'Text generation web UI'
2023-08-07 02:49:27 +02:00
# Password authentication
auth = []
if shared.args.gradio_auth:
2023-08-07 02:49:27 +02:00
auth.extend(x.strip() for x in shared.args.gradio_auth.strip('"').replace('\n', '').split(',') if x.strip())
if shared.args.gradio_auth_path:
with open(shared.args.gradio_auth_path, 'r', encoding="utf8") as file:
2023-08-07 02:49:27 +02:00
auth.extend(x.strip() for line in file for x in line.split(',') if x.strip())
auth = [tuple(cred.split(':')) for cred in auth]
2023-08-07 02:49:27 +02:00
# Import the extensions and execute their setup() functions
2023-03-16 03:29:56 +01:00
if shared.args.extensions is not None and len(shared.args.extensions) > 0:
extensions_module.load_extensions()
2023-08-07 02:49:27 +02:00
# Force some events to be triggered on page load
shared.persistent_interface_state.update({
'loader': shared.args.loader or 'Transformers',
})
if shared.is_chat():
shared.persistent_interface_state.update({
'mode': shared.settings['mode'],
'character_menu': shared.args.character or shared.settings['character'],
'instruction_template': shared.settings['instruction_template']
})
if Path("cache/pfp_character.png").exists():
Path("cache/pfp_character.png").unlink()
# css/js strings
css = ui.css if not shared.is_chat() else ui.css + ui.chat_css
2023-08-04 06:58:37 +02:00
js = ui.main_js
css += apply_extensions('css')
js += apply_extensions('js')
2023-08-07 02:49:27 +02:00
# The input elements for the generation functions
shared.input_elements = ui.list_interface_input_elements()
with gr.Blocks(css=css, analytics_enabled=False, title=title, theme=ui.theme) as shared.gradio['interface']:
2023-08-07 02:49:27 +02:00
# Audio notification
if Path("notification.mp3").exists():
shared.gradio['audio_notification'] = gr.Audio(interactive=False, value="notification.mp3", elem_id="audio_notification", visible=False)
# Floating menus for saving/deleting files
2023-08-07 02:49:27 +02:00
ui_file_saving.create_ui()
2023-08-07 02:49:27 +02:00
# Temporary clipboard for saving files
shared.gradio['temporary_text'] = gr.Textbox(visible=False)
2023-08-07 02:49:27 +02:00
# Text Generation tab
if shared.is_chat():
2023-08-07 02:49:27 +02:00
ui_chat.create_ui()
2023-03-16 03:29:56 +01:00
elif shared.args.notebook:
2023-08-07 02:49:27 +02:00
ui_notebook.create_ui()
2023-03-16 03:29:56 +01:00
else:
2023-08-07 02:49:27 +02:00
ui_default.create_ui()
2023-08-07 02:49:27 +02:00
ui_parameters.create_ui(shared.settings['preset']) # Parameters tab
ui_model_menu.create_ui() # Model tab
training.create_ui() # Training tab
ui_session.create_ui() # Session tab
2023-03-16 03:29:56 +01:00
2023-08-07 02:49:27 +02:00
# Generation events
if shared.is_chat():
2023-08-07 02:49:27 +02:00
ui_chat.create_event_handlers()
elif shared.args.notebook:
ui_notebook.create_event_handlers()
else:
2023-08-07 02:49:27 +02:00
ui_default.create_event_handlers()
2023-08-07 02:49:27 +02:00
# Other events
ui_file_saving.create_event_handlers()
ui_parameters.create_event_handlers()
ui_model_menu.create_event_handlers()
2023-08-07 02:49:27 +02:00
# Interface launch events
2023-08-08 04:03:09 +02:00
if shared.settings['dark_theme']:
shared.gradio['interface'].load(lambda: None, None, None, _js="() => document.getElementsByTagName('body')[0].classList.add('dark')")
shared.gradio['interface'].load(lambda: None, None, None, _js=f"() => {{{js}}}")
shared.gradio['interface'].load(partial(ui.apply_interface_values, {}, use_persistent=True), None, gradio(ui.list_interface_input_elements()), show_progress=False)
if shared.is_chat():
shared.gradio['interface'].load(chat.redraw_html, shared.reload_inputs, gradio('display'))
2023-08-07 02:49:27 +02:00
extensions_module.create_extensions_tabs() # Extensions tabs
extensions_module.create_extensions_block() # Extensions block
2023-03-16 03:29:56 +01:00
# Launch the interface
shared.gradio['interface'].queue(concurrency_count=64)
2023-07-07 07:24:52 +02:00
with OpenMonkeyPatch():
2023-08-04 18:57:31 +02:00
shared.gradio['interface'].launch(
prevent_thread_lock=True,
share=shared.args.share,
2023-08-07 02:49:27 +02:00
server_name=None if not shared.args.listen else (shared.args.listen_host or '0.0.0.0'),
2023-08-04 18:57:31 +02:00
server_port=shared.args.listen_port,
inbrowser=shared.args.auto_launch,
2023-08-07 02:49:27 +02:00
auth=auth or None,
2023-08-04 18:57:31 +02:00
ssl_verify=False if (shared.args.ssl_keyfile or shared.args.ssl_certfile) else True,
ssl_keyfile=shared.args.ssl_keyfile,
ssl_certfile=shared.args.ssl_certfile
)
2023-08-07 02:49:27 +02:00
if __name__ == "__main__":
2023-08-07 02:49:27 +02:00
# Load custom settings
settings_file = None
if shared.args.settings is not None and Path(shared.args.settings).exists():
settings_file = Path(shared.args.settings)
2023-05-29 03:34:12 +02:00
elif Path('settings.yaml').exists():
settings_file = Path('settings.yaml')
elif Path('settings.json').exists():
settings_file = Path('settings.json')
if settings_file is not None:
logger.info(f"Loading settings from {settings_file}...")
2023-05-29 03:34:12 +02:00
file_contents = open(settings_file, 'r', encoding='utf-8').read()
new_settings = json.loads(file_contents) if settings_file.suffix == "json" else yaml.safe_load(file_contents)
2023-08-07 02:49:27 +02:00
shared.settings.update(new_settings)
2023-08-07 02:49:27 +02:00
# Fallback settings for models
shared.model_config['.*'] = {
'wbits': 'None',
'model_type': 'None',
'groupsize': 'None',
'pre_layer': 0,
'mode': shared.settings['mode'],
'skip_special_tokens': shared.settings['skip_special_tokens'],
'custom_stopping_strings': shared.settings['custom_stopping_strings'],
2023-06-05 17:07:52 +02:00
'truncation_length': shared.settings['truncation_length'],
'n_gqa': 0,
'rms_norm_eps': 0,
}
shared.model_config.move_to_end('.*', last=False) # Move to the beginning
2023-08-07 02:49:27 +02:00
# Activate the extensions listed on settings.yaml
2023-05-06 04:14:56 +02:00
extensions_module.available_extensions = utils.get_available_extensions()
if shared.is_chat():
for extension in shared.settings['chat_default_extensions']:
shared.args.extensions = shared.args.extensions or []
if extension not in shared.args.extensions:
shared.args.extensions.append(extension)
else:
for extension in shared.settings['default_extensions']:
shared.args.extensions = shared.args.extensions or []
if extension not in shared.args.extensions:
shared.args.extensions.append(extension)
2023-05-06 04:14:56 +02:00
available_models = utils.get_available_models()
# Model defined through --model
if shared.args.model is not None:
shared.model_name = shared.args.model
# Select the model from a command-line menu
elif shared.args.model_menu:
if len(available_models) == 0:
logger.error('No models are available! Please download at least one.')
sys.exit(0)
else:
print('The following models are available:\n')
for i, model in enumerate(available_models):
print(f'{i+1}. {model}')
print(f'\nWhich one do you want to load? 1-{len(available_models)}\n')
i = int(input()) - 1
print()
shared.model_name = available_models[i]
# If any model has been selected, load it
if shared.model_name != 'None':
model_settings = get_model_settings_from_yamls(shared.model_name)
shared.settings.update(model_settings) # hijacking the interface defaults
update_model_parameters(model_settings, initial=True) # hijacking the command-line arguments
# Load the model
shared.model, shared.tokenizer = load_model(shared.model_name)
if shared.args.lora:
add_lora_to_model(shared.args.lora)
shared.generation_lock = Lock()
# Launch the web UI
create_interface()
while True:
time.sleep(0.5)
if shared.need_restart:
shared.need_restart = False
time.sleep(0.5)
shared.gradio['interface'].close()
time.sleep(0.5)
create_interface()