diff options
Diffstat (limited to 'modules/scripts.py')
| -rw-r--r-- | modules/scripts.py | 82 |
1 files changed, 71 insertions, 11 deletions
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
|