import logging import traceback from functools import partial import gradio as gr import extensions import modules.shared as shared state = {} available_extensions = [] setup_called = set() def apply_settings(extension, name): if not hasattr(extension, 'params'): return for param in extension.params: _id = f"{name}-{param}" if _id not in shared.settings: continue extension.params[param] = shared.settings[_id] def load_extensions(): global state, setup_called for i, name in enumerate(shared.args.extensions): if name in available_extensions: if name != 'api': logging.info(f'Loading the extension "{name}"...') try: exec(f"import extensions.{name}.script") extension = getattr(extensions, name).script apply_settings(extension, name) if extension not in setup_called and hasattr(extension, "setup"): setup_called.add(extension) extension.setup() state[name] = [True, i] except: logging.error(f'Failed to load the extension "{name}".') traceback.print_exc() # This iterator returns the extensions in the order specified in the command-line def iterator(): for name in sorted(state, key=lambda x: state[x][1]): if state[name][0]: yield getattr(extensions, name).script, name # Extension functions that map string -> string def _apply_string_extensions(function_name, text): for extension, _ in iterator(): if hasattr(extension, function_name): text = getattr(extension, function_name)(text) return text # Input hijack of extensions def _apply_input_hijack(text, visible_text): for extension, _ in iterator(): if hasattr(extension, 'input_hijack') and extension.input_hijack['state']: extension.input_hijack['state'] = False if callable(extension.input_hijack['value']): text, visible_text = extension.input_hijack['value'](text, visible_text) else: text, visible_text = extension.input_hijack['value'] return text, visible_text # custom_generate_chat_prompt handling - currently only the first one will work def _apply_custom_generate_chat_prompt(text, state, **kwargs): for extension, _ in iterator(): if hasattr(extension, 'custom_generate_chat_prompt'): return extension.custom_generate_chat_prompt(text, state, **kwargs) return None # Extension that modifies the input parameters before they are used def _apply_state_modifier_extensions(state): for extension, _ in iterator(): if hasattr(extension, "state_modifier"): state = getattr(extension, "state_modifier")(state) return state # Extension functions that override the default tokenizer output - currently only the first one will work def _apply_tokenizer_extensions(function_name, state, prompt, input_ids, input_embeds): for extension, _ in iterator(): if hasattr(extension, function_name): return getattr(extension, function_name)(state, prompt, input_ids, input_embeds) return prompt, input_ids, input_embeds # Get prompt length in tokens after applying extension functions which override the default tokenizer output # currently only the first one will work def _apply_custom_tokenized_length(prompt): for extension, _ in iterator(): if hasattr(extension, 'custom_tokenized_length'): return getattr(extension, 'custom_tokenized_length')(prompt) return None # Custom generate reply handling - currently only the first one will work def _apply_custom_generate_reply(): for extension, _ in iterator(): if hasattr(extension, 'custom_generate_reply'): return getattr(extension, 'custom_generate_reply') return None def _apply_custom_css(): all_css = '' for extension, _ in iterator(): if hasattr(extension, 'custom_css'): all_css += getattr(extension, 'custom_css')() return all_css def _apply_custom_js(): all_js = '' for extension, _ in iterator(): if hasattr(extension, 'custom_js'): all_js += getattr(extension, 'custom_js')() return all_js def create_extensions_block(): to_display = [] for extension, name in iterator(): if hasattr(extension, "ui") and not (hasattr(extension, 'params') and extension.params.get('is_tab', False)): to_display.append((extension, name)) # Creating the extension ui elements if len(to_display) > 0: with gr.Column(elem_id="extensions"): for row in to_display: extension, name = row display_name = getattr(extension, 'params', {}).get('display_name', name) gr.Markdown(f"\n### {display_name}") extension.ui() def create_extensions_tabs(): for extension, name in iterator(): if hasattr(extension, "ui") and (hasattr(extension, 'params') and extension.params.get('is_tab', False)): display_name = getattr(extension, 'params', {}).get('display_name', name) with gr.Tab(display_name, elem_classes="extension-tab"): extension.ui() EXTENSION_MAP = { "input": partial(_apply_string_extensions, "input_modifier"), "output": partial(_apply_string_extensions, "output_modifier"), "state": _apply_state_modifier_extensions, "bot_prefix": partial(_apply_string_extensions, "bot_prefix_modifier"), "tokenizer": partial(_apply_tokenizer_extensions, "tokenizer_modifier"), "input_hijack": _apply_input_hijack, "custom_generate_chat_prompt": _apply_custom_generate_chat_prompt, "custom_generate_reply": _apply_custom_generate_reply, "tokenized_length": _apply_custom_tokenized_length, "css": _apply_custom_css, "js": _apply_custom_js } def apply_extensions(typ, *args, **kwargs): if typ not in EXTENSION_MAP: raise ValueError(f"Invalid extension type {typ}") return EXTENSION_MAP[typ](*args, **kwargs)