From f40617d6c4e366773677baa8d7f4114ba2893282 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 3 Sep 2022 17:21:15 +0300
Subject: support for scripts

---
 modules/scripts.py | 82 ++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 71 insertions(+), 11 deletions(-)

(limited to 'modules/scripts.py')

diff --git a/modules/scripts.py b/modules/scripts.py
index 20f489ce..be348a70 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -2,19 +2,33 @@ import os
 import sys
 import traceback
 
+import modules.ui as ui
 import gradio as gr
 
+from modules.processing import StableDiffusionProcessing
+
 class Script:
     filename = None
+    args_from = None
+    args_to = None
 
     def title(self):
         raise NotImplementedError()
 
+    def ui(self, is_img2img):
+        pass
+
+    def run(self, *args):
+        raise NotImplementedError()
+
+    def describe(self):
+        return ""
+
 
 scripts = []
 
 
-def load_scripts(basedir, globs):
+def load_scripts(basedir):
     for filename in os.listdir(basedir):
         path = os.path.join(basedir, filename)
 
@@ -27,27 +41,73 @@ def load_scripts(basedir, globs):
         from types import ModuleType
         compiled = compile(text, path, 'exec')
         module = ModuleType(filename)
-        module.__dict__.update(globs)
         exec(compiled, module.__dict__)
 
-        for key, item in module.__dict__.items():
-            if type(item) == type and issubclass(item, Script):
-                item.filename = path
+        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.append(item)
+                scripts.append(obj)
 
 
 def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
     try:
-        res = func()
+        res = func(*args, **kwargs)
         return res
     except Exception:
-        print(f"Error calling: {filename/funcname}", file=sys.stderr)
+        print(f"Error calling: {filename}/{funcname}", file=sys.stderr)
         print(traceback.format_exc(), file=sys.stderr)
 
     return default
 
-def setup_ui():
-    titles = [wrap_call(script.title, script.filename, "title") for script in scripts]
 
-    gr.Dropdown(options=[""] + titles, value="", type="index")
+def setup_ui(is_img2img):
+    titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in scripts]
+
+    dropdown = gr.Dropdown(label="Script", choices=["None"] + titles, value="None", type="index")
+
+    inputs = [dropdown]
+
+    for script in scripts:
+        script.args_from = len(inputs)
+        controls = script.ui(is_img2img)
+
+        for control in controls:
+            control.visible = False
+
+        inputs += controls
+        script.args_to = len(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 [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(p: StableDiffusionProcessing, *args):
+    script_index = args[0] - 1
+
+    if script_index < 0 or script_index >= len(scripts):
+        return None
+
+    script = scripts[script_index]
+
+    script_args = args[script.args_from:script.args_to]
+    processed = script.run(p, *script_args)
+
+    return processed
-- 
cgit v1.2.1