From 7350c712598b748c3cedc2a224887bd839a27d76 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 4 Sep 2022 01:29:43 +0300
Subject: added poor man's inpainting script

---
 modules/scripts.py | 100 +++++++++++++++++++++++++++++++++--------------------
 1 file changed, 62 insertions(+), 38 deletions(-)

(limited to 'modules/scripts.py')

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()
-- 
cgit v1.2.1