text-generation-webui/server.py

258 lines
9.2 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
2023-12-04 16:27:44 +01:00
import accelerate # This early import makes Intel GPUs happy
import modules.one_click_installer_check
2023-07-07 07:24:52 +02:00
from modules.block_requests import OpenMonkeyPatch, RequestBlocker
2023-09-11 23:49:30 +02:00
from modules.logging_colors import logger
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')
2023-10-11 03:20:49 +02:00
warnings.filterwarnings('ignore', category=UserWarning, message='Using the update method is deprecated')
warnings.filterwarnings('ignore', category=UserWarning, message='Field "model_name" has conflict')
with RequestBlocker():
import gradio as gr
2023-04-05 19:05:50 +02:00
import matplotlib
2023-09-11 23:49:30 +02:00
2023-05-04 02:43:17 +02:00
matplotlib.use('Agg') # This fixes LaTeX rendering on some systems
import json
import os
2023-12-05 06:16:16 +01:00
import signal
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,
2023-09-11 23:49:30 +02:00
utils
2023-08-07 02:49:27 +02:00
)
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 (
2023-09-11 23:49:30 +02:00
get_fallback_settings,
get_model_metadata,
2023-06-25 06:44:36 +02:00
update_model_parameters
)
from modules.utils import gradio
2023-03-17 01:31:39 +01:00
2023-12-05 06:16:16 +01:00
def signal_handler(sig, frame):
2023-12-05 06:17:05 +01:00
logger.info("Received Ctrl+C. Shutting down Text generation web UI gracefully.")
2023-12-05 06:16:16 +01:00
if 'interface' in shared.gradio:
shared.gradio['interface'].close()
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
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',
2023-08-13 06:12:15 +02:00
'mode': shared.settings['mode'],
'character_menu': shared.args.character or shared.settings['character'],
'instruction_template': shared.settings['instruction_template'],
'prompt_menu-default': shared.settings['prompt-default'],
'prompt_menu-notebook': shared.settings['prompt-notebook'],
'filter_by_loader': shared.args.loader or 'All'
})
2023-08-13 06:12:15 +02:00
if Path("cache/pfp_character.png").exists():
Path("cache/pfp_character.png").unlink()
# css/js strings
2023-08-13 06:12:15 +02:00
css = ui.css
js = ui.js
css += apply_extensions('css')
js += apply_extensions('js')
2023-08-13 06:12:15 +02:00
# Interface state elements
2023-08-07 02:49:27 +02:00
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
2023-08-13 06:12:15 +02:00
# Interface state
shared.gradio['interface_state'] = gr.State({k: None for k in shared.input_elements})
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
2023-08-13 06:12:15 +02:00
ui_chat.create_ui()
ui_default.create_ui()
ui_notebook.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
2023-08-13 06:12:15 +02:00
ui_chat.create_event_handlers()
ui_default.create_event_handlers()
ui_notebook.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')")
2023-08-13 06:12:15 +02:00
shared.gradio['interface'].load(lambda: None, None, None, _js=f"() => {{{js}}}")
shared.gradio['interface'].load(None, gradio('show_controls'), None, _js=f'(x) => {{{ui.show_controls_js}; toggle_controls(x)}}')
shared.gradio['interface'].load(partial(ui.apply_interface_values, {}, use_persistent=True), None, gradio(ui.list_interface_input_elements()), show_progress=False)
2023-08-13 06:12:15 +02:00
shared.gradio['interface'].load(chat.redraw_html, gradio(ui_chat.reload_arr), 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
2023-09-11 23:49:30 +02:00
shared.model_config['.*'] = get_fallback_settings()
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()
2023-08-13 06:12:15 +02:00
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':
2023-09-11 23:49:30 +02:00
p = Path(shared.model_name)
if p.exists():
model_name = p.parts[-1]
2023-09-19 22:13:24 +02:00
shared.model_name = model_name
2023-09-11 23:49:30 +02:00
else:
model_name = shared.model_name
model_settings = get_model_metadata(model_name)
update_model_parameters(model_settings, initial=True) # hijack the command-line arguments
# Load the model
2023-09-11 23:49:30 +02:00
shared.model, shared.tokenizer = load_model(model_name)
if shared.args.lora:
add_lora_to_model(shared.args.lora)
shared.generation_lock = Lock()
if shared.args.nowebui:
# Start the API in standalone mode
shared.args.extensions = [x for x in shared.args.extensions if x != 'gallery']
if shared.args.extensions is not None and len(shared.args.extensions) > 0:
extensions_module.load_extensions()
else:
# 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()