cymic/Waifu_Diffusion_Webui

1import os2import sys3import traceback4 5import modules.ui as ui6import gradio as gr7 8from modules.processing import StableDiffusionProcessing9from modules import shared10 11class Script:12    filename = None13    args_from = None14    args_to = None15 16    # The title of the script. This is what will be displayed in the dropdown menu.17    def title(self):18        raise NotImplementedError()19 20    # How the script is displayed in the UI. See https://gradio.app/docs/#components21    # for the different UI components you can use and how to create them.22    # Most UI components can return a value, such as a boolean for a checkbox.23    # The returned values are passed to the run method as parameters.24    def ui(self, is_img2img):25        pass26 27    # Determines when the script should be shown in the dropdown menu via the 28    # returned value. As an example:29    # is_img2img is True if the current tab is img2img, and False if it is txt2img.30    # Thus, return is_img2img to only show the script on the img2img tab.31    def show(self, is_img2img):32        return True33 34    # This is where the additional processing is implemented. The parameters include35    # self, the model object "p" (a StableDiffusionProcessing class, see36    # processing.py), and the parameters returned by the ui method.37    # Custom functions can be defined here, and additional libraries can be imported 38    # to be used in processing. The return value should be a Processed object, which is39    # what is returned by the process_images method.40    def run(self, *args):41        raise NotImplementedError()42 43    # The description method is currently unused.44    # To add a description that appears when hovering over the title, amend the "titles" 45    # dict in script.js to include the script title (returned by title) as a key, and 46    # your description as the value.47    def describe(self):48        return ""49 50 51scripts_data = []52 53 54def load_scripts(basedir):55    if not os.path.exists(basedir):56        return57 58    for filename in sorted(os.listdir(basedir)):59        path = os.path.join(basedir, filename)60 61        if not os.path.isfile(path):62            continue63 64        try:65            with open(path, "r", encoding="utf8") as file:66                text = file.read()67 68            from types import ModuleType69            compiled = compile(text, path, 'exec')70            module = ModuleType(filename)71            exec(compiled, module.__dict__)72 73            for key, script_class in module.__dict__.items():74                if type(script_class) == type and issubclass(script_class, Script):75                    scripts_data.append((script_class, path))76 77        except Exception:78            print(f"Error loading script: {filename}", file=sys.stderr)79            print(traceback.format_exc(), file=sys.stderr)80 81 82def wrap_call(func, filename, funcname, *args, default=None, **kwargs):83    try:84        res = func(*args, **kwargs)85        return res86    except Exception:87        print(f"Error calling: {filename}/{funcname}", file=sys.stderr)88        print(traceback.format_exc(), file=sys.stderr)89 90    return default91 92 93class ScriptRunner:94    def __init__(self):95        self.scripts = []96 97    def setup_ui(self, is_img2img):98        for script_class, path in scripts_data:99            script = script_class()100            script.filename = path101 102            if not script.show(is_img2img):103                continue104 105            self.scripts.append(script)106 107        titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.scripts]108 109        dropdown = gr.Dropdown(label="Script", choices=["None"] + titles, value="None", type="index")110        inputs = [dropdown]111 112        for script in self.scripts:113            script.args_from = len(inputs)114            script.args_to = len(inputs)115 116            controls = wrap_call(script.ui, script.filename, "ui", is_img2img)117 118            if controls is None:119                continue120 121            for control in controls:122                control.custom_script_source = os.path.basename(script.filename)123                control.visible = False124 125            inputs += controls126            script.args_to = len(inputs)127 128        def select_script(script_index):129            if 0 < script_index <= len(self.scripts):130                script = self.scripts[script_index-1]131                args_from = script.args_from132                args_to = script.args_to133            else:134                args_from = 0135                args_to = 0136 137            return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]138 139        dropdown.change(140            fn=select_script,141            inputs=[dropdown],142            outputs=inputs143        )144 145        return inputs146 147    def run(self, p: StableDiffusionProcessing, *args):148        script_index = args[0]149 150        if script_index == 0:151            return None152 153        script = self.scripts[script_index-1]154 155        if script is None:156            return None157 158        script_args = args[script.args_from:script.args_to]159        processed = script.run(p, *script_args)160 161        shared.total_tqdm.clear()162 163        return processed164 165    def reload_sources(self):166        for si, script in list(enumerate(self.scripts)):167            with open(script.filename, "r", encoding="utf8") as file:168                args_from = script.args_from169                args_to = script.args_to170                filename = script.filename171                text = file.read()172 173                from types import ModuleType174 175                compiled = compile(text, filename, 'exec')176                module = ModuleType(script.filename)177                exec(compiled, module.__dict__)178 179                for key, script_class in module.__dict__.items():180                    if type(script_class) == type and issubclass(script_class, Script):181                        self.scripts[si] = script_class()182                        self.scripts[si].filename = filename183                        self.scripts[si].args_from = args_from184                        self.scripts[si].args_to = args_to185 186scripts_txt2img = ScriptRunner()187scripts_img2img = ScriptRunner()188 189def reload_script_body_only():190    scripts_txt2img.reload_sources()191    scripts_img2img.reload_sources()192 193 194def reload_scripts(basedir):195    global scripts_txt2img, scripts_img2img196 197    scripts_data.clear()198    load_scripts(basedir)199 200    scripts_txt2img = ScriptRunner()201    scripts_img2img = ScriptRunner()202