| import os | |
| import gradio as gr | |
| from modules import errors, shared | |
| class PostprocessedImage: | |
| def __init__(self, image): | |
| self.image = image | |
| self.info = {} | |
| class ScriptPostprocessing: | |
| filename = None | |
| controls = None | |
| args_from = None | |
| args_to = None | |
| order = 1000 | |
| """scripts will be ordred by this value in postprocessing UI""" | |
| name = None | |
| """this function should return the title of the script.""" | |
| group = None | |
| """A gr.Group component that has all script's UI inside it""" | |
| def ui(self): | |
| """ | |
| This function should create gradio UI elements. See https://gradio.app/docs/#components | |
| The return value should be a dictionary that maps parameter names to components used in processing. | |
| Values of those components will be passed to process() function. | |
| """ | |
| pass | |
| def process(self, pp: PostprocessedImage, **args): | |
| """ | |
| This function is called to postprocess the image. | |
| args contains a dictionary with all values returned by components from ui() | |
| """ | |
| pass | |
| def image_changed(self): | |
| pass | |
| def wrap_call(func, filename, funcname, *args, default=None, **kwargs): | |
| try: | |
| res = func(*args, **kwargs) | |
| return res | |
| except Exception as e: | |
| errors.display(e, f"calling {filename}/{funcname}") | |
| return default | |
| class ScriptPostprocessingRunner: | |
| def __init__(self): | |
| self.scripts = None | |
| self.ui_created = False | |
| def initialize_scripts(self, scripts_data): | |
| self.scripts = [] | |
| for script_class, path, basedir, script_module in scripts_data: | |
| script: ScriptPostprocessing = script_class() | |
| script.filename = path | |
| if script.name == "Simple Upscale": | |
| continue | |
| self.scripts.append(script) | |
| def create_script_ui(self, script, inputs): | |
| script.args_from = len(inputs) | |
| script.args_to = len(inputs) | |
| script.controls = wrap_call(script.ui, script.filename, "ui") | |
| for control in script.controls.values(): | |
| control.custom_script_source = os.path.basename(script.filename) | |
| inputs += list(script.controls.values()) | |
| script.args_to = len(inputs) | |
| def scripts_in_preferred_order(self): | |
| if self.scripts is None: | |
| import modules.scripts | |
| self.initialize_scripts(modules.scripts.postprocessing_scripts_data) | |
| scripts_order = shared.opts.postprocessing_operation_order | |
| def script_score(name): | |
| for i, possible_match in enumerate(scripts_order): | |
| if possible_match == name: | |
| return i | |
| return len(self.scripts) | |
| script_scores = {script.name: (script_score(script.name), script.order, script.name, original_index) for original_index, script in enumerate(self.scripts)} | |
| return sorted(self.scripts, key=lambda x: script_scores[x.name]) | |
| def setup_ui(self): | |
| inputs = [] | |
| for script in self.scripts_in_preferred_order(): | |
| with gr.Box() as group: | |
| self.create_script_ui(script, inputs) | |
| script.group = group | |
| self.ui_created = True | |
| return inputs | |
| def run(self, pp: PostprocessedImage, args): | |
| for script in self.scripts_in_preferred_order(): | |
| shared.state.job = script.name | |
| script_args = args[script.args_from:script.args_to] | |
| process_args = {} | |
| for (name, component), value in zip(script.controls.items(), script_args): | |
| process_args[name] = value | |
| script.process(pp, **process_args) | |
| def create_args_for_run(self, scripts_args): | |
| if not self.ui_created: | |
| with gr.Blocks(analytics_enabled=False): | |
| self.setup_ui() | |
| scripts = self.scripts_in_preferred_order() | |
| args = [None] * max([x.args_to for x in scripts]) | |
| for script in scripts: | |
| script_args_dict = scripts_args.get(script.name, None) | |
| if script_args_dict is not None: | |
| for i, name in enumerate(script.controls): | |
| args[script.args_from + i] = script_args_dict.get(name, None) | |
| return args | |
| def image_changed(self): | |
| for script in self.scripts_in_preferred_order(): | |
| script.image_changed() | |