Spaces:
Runtime error
Runtime error
| import os | |
| import sys | |
| import traceback | |
| import modules.ui as ui | |
| import gradio as gr | |
| from modules.processing import StableDiffusionProcessing | |
| from modules import shared | |
| class Script: | |
| filename = None | |
| args_from = None | |
| args_to = None | |
| # The title of the script. This is what will be displayed in the dropdown menu. | |
| def title(self): | |
| raise NotImplementedError() | |
| # How the script is displayed in the UI. See https://gradio.app/docs/#components | |
| # for the different UI components you can use and how to create them. | |
| # Most UI components can return a value, such as a boolean for a checkbox. | |
| # The returned values are passed to the run method as parameters. | |
| def ui(self, is_img2img): | |
| pass | |
| # Determines when the script should be shown in the dropdown menu via the | |
| # returned value. As an example: | |
| # is_img2img is True if the current tab is img2img, and False if it is txt2img. | |
| # Thus, return is_img2img to only show the script on the img2img tab. | |
| def show(self, is_img2img): | |
| return True | |
| # This is where the additional processing is implemented. The parameters include | |
| # self, the model object "p" (a StableDiffusionProcessing class, see | |
| # processing.py), and the parameters returned by the ui method. | |
| # Custom functions can be defined here, and additional libraries can be imported | |
| # to be used in processing. The return value should be a Processed object, which is | |
| # what is returned by the process_images method. | |
| def run(self, *args): | |
| raise NotImplementedError() | |
| # The description method is currently unused. | |
| # To add a description that appears when hovering over the title, amend the "titles" | |
| # dict in script.js to include the script title (returned by title) as a key, and | |
| # your description as the value. | |
| def describe(self): | |
| return "" | |
| scripts_data = [] | |
| def load_scripts(basedir): | |
| if not os.path.exists(basedir): | |
| return | |
| for filename in sorted(os.listdir(basedir)): | |
| path = os.path.join(basedir, filename) | |
| if not os.path.isfile(path): | |
| continue | |
| try: | |
| with open(path, "r", encoding="utf8") as file: | |
| text = file.read() | |
| from types import ModuleType | |
| compiled = compile(text, path, 'exec') | |
| module = ModuleType(filename) | |
| exec(compiled, module.__dict__) | |
| for key, script_class in module.__dict__.items(): | |
| if type(script_class) == type and issubclass(script_class, Script): | |
| scripts_data.append((script_class, path)) | |
| except Exception: | |
| print(f"Error loading script: {filename}", file=sys.stderr) | |
| print(traceback.format_exc(), file=sys.stderr) | |
| def wrap_call(func, filename, funcname, *args, default=None, **kwargs): | |
| try: | |
| res = func(*args, **kwargs) | |
| return res | |
| except Exception: | |
| print(f"Error calling: {filename}/{funcname}", file=sys.stderr) | |
| print(traceback.format_exc(), file=sys.stderr) | |
| return default | |
| class ScriptRunner: | |
| def __init__(self): | |
| self.scripts = [] | |
| def setup_ui(self, is_img2img): | |
| for script_class, path in scripts_data: | |
| script = script_class() | |
| script.filename = path | |
| if not script.show(is_img2img): | |
| continue | |
| self.scripts.append(script) | |
| titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.scripts] | |
| dropdown = gr.Dropdown(label="Script", choices=["None"] + titles, value="None", type="index") | |
| inputs = [dropdown] | |
| for script in self.scripts: | |
| script.args_from = len(inputs) | |
| script.args_to = len(inputs) | |
| controls = wrap_call(script.ui, script.filename, "ui", is_img2img) | |
| if controls is None: | |
| continue | |
| for control in controls: | |
| control.custom_script_source = os.path.basename(script.filename) | |
| control.visible = False | |
| inputs += controls | |
| script.args_to = len(inputs) | |
| def select_script(script_index): | |
| if 0 < script_index <= len(self.scripts): | |
| script = self.scripts[script_index-1] | |
| args_from = script.args_from | |
| args_to = script.args_to | |
| else: | |
| args_from = 0 | |
| args_to = 0 | |
| return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))] | |
| dropdown.change( | |
| fn=select_script, | |
| inputs=[dropdown], | |
| outputs=inputs | |
| ) | |
| return inputs | |
| def run(self, p: StableDiffusionProcessing, *args): | |
| script_index = args[0] | |
| if script_index == 0: | |
| return None | |
| script = self.scripts[script_index-1] | |
| if script is None: | |
| return None | |
| script_args = args[script.args_from:script.args_to] | |
| processed = script.run(p, *script_args) | |
| shared.total_tqdm.clear() | |
| return processed | |
| def reload_sources(self): | |
| for si, script in list(enumerate(self.scripts)): | |
| with open(script.filename, "r", encoding="utf8") as file: | |
| args_from = script.args_from | |
| args_to = script.args_to | |
| filename = script.filename | |
| text = file.read() | |
| from types import ModuleType | |
| compiled = compile(text, filename, 'exec') | |
| module = ModuleType(script.filename) | |
| exec(compiled, module.__dict__) | |
| for key, script_class in module.__dict__.items(): | |
| if type(script_class) == type and issubclass(script_class, Script): | |
| self.scripts[si] = script_class() | |
| self.scripts[si].filename = filename | |
| self.scripts[si].args_from = args_from | |
| self.scripts[si].args_to = args_to | |
| scripts_txt2img = ScriptRunner() | |
| scripts_img2img = ScriptRunner() | |
| def reload_script_body_only(): | |
| scripts_txt2img.reload_sources() | |
| scripts_img2img.reload_sources() | |
| def reload_scripts(basedir): | |
| global scripts_txt2img, scripts_img2img | |
| scripts_data.clear() | |
| load_scripts(basedir) | |
| scripts_txt2img = ScriptRunner() | |
| scripts_img2img = ScriptRunner() | |