added poor man's inpainting script

1 files changed, 62 insertions, 38 deletions

diff --git a/modules/scripts.py b/modules/scripts.py
index 99502857..89a0618d 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -18,6 +18,9 @@ class Script:
     def ui(self, is_img2img):

         pass

 

+    def show(self, is_img2img):

+        return True

+

     def run(self, *args):

         raise NotImplementedError()

 

@@ -25,7 +28,7 @@ class Script:
         return ""

 

 

-scripts = []

+scripts_data = []

 

 

 def load_scripts(basedir):

@@ -49,10 +52,8 @@ def load_scripts(basedir):
 

             for key, script_class in module.__dict__.items():

                 if type(script_class) == type and issubclass(script_class, Script):

-                    obj = script_class()

-                    obj.filename = path

+                    scripts_data.append((script_class, path))

 

-                    scripts.append(obj)

         except Exception:

             print(f"Error loading script: {filename}", file=sys.stderr)

             print(traceback.format_exc(), file=sys.stderr)

@@ -69,52 +70,75 @@ def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
     return default

 

 

-def setup_ui(is_img2img):

-    titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in scripts]

+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)

+

+            controls = wrap_call(script.ui, script.filename, "ui", is_img2img)

+

+            if controls is None:

+                continue

 

-    dropdown = gr.Dropdown(label="Script", choices=["None"] + titles, value="None", type="index")

+            for control in controls:

+                control.visible = False

 

-    inputs = [dropdown]

+            inputs += controls

+            script.args_to = len(inputs)

 

-    for script in scripts:

-        script.args_from = len(inputs)

-        controls = script.ui(is_img2img)

+        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

 

-        for control in controls:

-            control.visible = False

+            return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]

 

-        inputs += controls

-        script.args_to = len(inputs)

+        dropdown.change(

+            fn=select_script,

+            inputs=[dropdown],

+            outputs=inputs

+        )

 

-    def select_script(index):

-        if index > 0:

-            script = scripts[index-1]

-            args_from = script.args_from

-            args_to = script.args_to

-        else:

-            args_from = 0

-            args_to = 0

+        return inputs

 

-        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

-    )

+    def run(self, p: StableDiffusionProcessing, *args):

+        script_index = args[0]

 

-    return inputs

+        if script_index == 0:

+            return None

 

+        script = self.scripts[script_index-1]

 

-def run(p: StableDiffusionProcessing, *args):

-    script_index = args[0] - 1

+        if script is None:

+            return None

 

-    if script_index < 0 or script_index >= len(scripts):

-        return None

+        script_args = args[script.args_from:script.args_to]

+        processed = script.run(p, *script_args)

 

-    script = scripts[script_index]

+        return processed

 

-    script_args = args[script.args_from:script.args_to]

-    processed = script.run(p, *script_args)

 

-    return processed

+scripts_txt2img = ScriptRunner()

+scripts_img2img = ScriptRunner()