removed aesthetic gradients as built-in

added support for extensions

1 files changed, 156 insertions, 54 deletions

diff --git a/modules/scripts.py b/modules/scripts.py
index 1039fa9c..65f25f49 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -1,86 +1,153 @@
 import os

 import sys

 import traceback

+from collections import namedtuple

 

 import modules.ui as ui

 import gradio as gr

 

 from modules.processing import StableDiffusionProcessing

-from modules import shared

+from modules import shared, paths, script_callbacks

+

+AlwaysVisible = object()

+

 

 class Script:

     filename = None

     args_from = None

     args_to = None

+    alwayson = False

+

+    infotext_fields = None

+    """if set in ui(), this is a list of pairs of gradio component + text; the text will be used when

+    parsing infotext to set the value for the component; see ui.py's txt2img_paste_fields for an example

+    """

 

-    # The title of the script. This is what will be displayed in the dropdown menu.

     def title(self):

+        """this function should return the title of the script. This is what will be displayed in the dropdown menu."""

+

         raise NotImplementedError()

 

-    # How the script is displayed in the UI. See https://gradio.app/docs/#components

-    # for the different UI components you can use and how to create them.

-    # Most UI components can return a value, such as a boolean for a checkbox.

-    # The returned values are passed to the run method as parameters.

     def ui(self, is_img2img):

+        """this function should create gradio UI elements. See https://gradio.app/docs/#components

+        The return value should be an array of all components that are used in processing.

+        Values of those returned componenbts will be passed to run() and process() functions.

+        """

+

         pass

 

-    # Determines when the script should be shown in the dropdown menu via the 

-    # returned value. As an example:

-    # is_img2img is True if the current tab is img2img, and False if it is txt2img.

-    # Thus, return is_img2img to only show the script on the img2img tab.

     def show(self, is_img2img):

+        """

+        is_img2img is True if this function is called for the img2img interface, and Fasle otherwise

+

+        This function should return:

+         - False if the script should not be shown in UI at all

+         - True if the script should be shown in UI if it's scelected in the scripts drowpdown

+         - script.AlwaysVisible if the script should be shown in UI at all times

+         """

+

         return True

 

-    # This is where the additional processing is implemented. The parameters include

-    # self, the model object "p" (a StableDiffusionProcessing class, see

-    # processing.py), and the parameters returned by the ui method.

-    # Custom functions can be defined here, and additional libraries can be imported 

-    # to be used in processing. The return value should be a Processed object, which is

-    # what is returned by the process_images method.

-    def run(self, *args):

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

+        """

+        This function is called if the script has been selected in the script dropdown.

+        It must do all processing and return the Processed object with results, same as

+        one returned by processing.process_images.

+

+        Usually the processing is done by calling the processing.process_images function.

+

+        args contains all values returned by components from ui()

+        """

+

         raise NotImplementedError()

 

-    # The description method is currently unused.

-    # To add a description that appears when hovering over the title, amend the "titles" 

-    # dict in script.js to include the script title (returned by title) as a key, and 

-    # your description as the value.

+    def process(self, p, *args):

+        """

+        This function is called before processing begins for AlwaysVisible scripts.

+        scripts. You can modify the processing object (p) here, inject hooks, etc.

+        """

+

+        pass

+

     def describe(self):

+        """unused"""

         return ""

 

 

+current_basedir = paths.script_path

+

+

+def basedir():

+    """returns the base directory for the current script. For scripts in the main scripts directory,

+    this is the main directory (where webui.py resides), and for scripts in extensions directory

+    (ie extensions/aesthetic/script/aesthetic.py), this is extension's directory (extensions/aesthetic)

+    """

+    return current_basedir

+

+

 scripts_data = []

+ScriptFile = namedtuple("ScriptFile", ["basedir", "filename", "path"])

+ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir"])

+

+

+def list_scripts(scriptdirname, extension):

+    scripts_list = []

+

+    basedir = os.path.join(paths.script_path, scriptdirname)

+    if os.path.exists(basedir):

+        for filename in sorted(os.listdir(basedir)):

+            scripts_list.append(ScriptFile(paths.script_path, filename, os.path.join(basedir, filename)))

+

+    extdir = os.path.join(paths.script_path, "extensions")

+    if os.path.exists(extdir):

+        for dirname in sorted(os.listdir(extdir)):

+            dirpath = os.path.join(extdir, dirname)

+            if not os.path.isdir(dirpath):

+                continue

 

+            for filename in sorted(os.listdir(os.path.join(dirpath, scriptdirname))):

+                scripts_list.append(ScriptFile(dirpath, filename, os.path.join(dirpath, scriptdirname, filename)))

 

-def load_scripts(basedir):

-    if not os.path.exists(basedir):

-        return

+    scripts_list = [x for x in scripts_list if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)]

 

-    for filename in sorted(os.listdir(basedir)):

-        path = os.path.join(basedir, filename)

+    return scripts_list

 

-        if os.path.splitext(path)[1].lower() != '.py':

-            continue

 

-        if not os.path.isfile(path):

-            continue

+def load_scripts():

+    global current_basedir

+    scripts_data.clear()

+    script_callbacks.clear_callbacks()

+

+    scripts_list = list_scripts("scripts", ".py")

+

+    syspath = sys.path

 

+    for scriptfile in sorted(scripts_list):

         try:

-            with open(path, "r", encoding="utf8") as file:

+            if scriptfile.basedir != paths.script_path:

+                sys.path = [scriptfile.basedir] + sys.path

+            current_basedir = scriptfile.basedir

+

+            with open(scriptfile.path, "r", encoding="utf8") as file:

                 text = file.read()

 

             from types import ModuleType

-            compiled = compile(text, path, 'exec')

-            module = ModuleType(filename)

+            compiled = compile(text, scriptfile.path, 'exec')

+            module = ModuleType(scriptfile.filename)

             exec(compiled, module.__dict__)

 

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

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

-                    scripts_data.append((script_class, path))

+                    scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir))

 

         except Exception:

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

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

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

 

+        finally:

+            sys.path = syspath

+            current_basedir = paths.script_path

+

 

 def wrap_call(func, filename, funcname, *args, default=None, **kwargs):

     try:

@@ -96,56 +163,80 @@ def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
 class ScriptRunner:

     def __init__(self):

         self.scripts = []

+        self.selectable_scripts = []

+        self.alwayson_scripts = []

         self.titles = []

+        self.infotext_fields = []

 

     def setup_ui(self, is_img2img):

-        for script_class, path in scripts_data:

+        for script_class, path, basedir in scripts_data:

             script = script_class()

             script.filename = path

 

-            if not script.show(is_img2img):

-                continue

+            visibility = script.show(is_img2img)

 

-            self.scripts.append(script)

+            if visibility == AlwaysVisible:

+                self.scripts.append(script)

+                self.alwayson_scripts.append(script)

+                script.alwayson = True

 

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

+            elif visibility:

+                self.scripts.append(script)

+                self.selectable_scripts.append(script)

 

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

-        dropdown.save_to_config = True

-        inputs = [dropdown]

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

+

+        inputs = [None]

+        inputs_alwayson = [True]

 

-        for script in self.scripts:

+        def create_script_ui(script, inputs, inputs_alwayson):

             script.args_from = len(inputs)

             script.args_to = len(inputs)

 

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

 

             if controls is None:

-                continue

+                return

 

             for control in controls:

                 control.custom_script_source = os.path.basename(script.filename)

-                control.visible = False

+                if not script.alwayson:

+                    control.visible = False

+

+            if script.infotext_fields is not None:

+                self.infotext_fields += script.infotext_fields

 

             inputs += controls

+            inputs_alwayson += [script.alwayson for _ in controls]

             script.args_to = len(inputs)

 

+        for script in self.alwayson_scripts:

+            with gr.Group():

+                create_script_ui(script, inputs, inputs_alwayson)

+

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

+        dropdown.save_to_config = True

+        inputs[0] = dropdown

+

+        for script in self.selectable_scripts:

+            create_script_ui(script, inputs, inputs_alwayson)

+

         def select_script(script_index):

-            if 0 < script_index <= len(self.scripts):

-                script = self.scripts[script_index-1]

+            if 0 < script_index <= len(self.selectable_scripts):

+                script = self.selectable_scripts[script_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))]

+            return [ui.gr_show(True if i == 0 else args_from <= i < args_to or is_alwayson) for i, is_alwayson in enumerate(inputs_alwayson)]

 

         def init_field(title):

             if title == 'None':

                 return

             script_index = self.titles.index(title)

-            script = self.scripts[script_index]

+            script = self.selectable_scripts[script_index]

             for i in range(script.args_from, script.args_to):

                 inputs[i].visible = True

 

@@ -164,7 +255,7 @@ class ScriptRunner:
         if script_index == 0:

             return None

 

-        script = self.scripts[script_index-1]

+        script = self.selectable_scripts[script_index-1]

 

         if script is None:

             return None

@@ -176,6 +267,15 @@ class ScriptRunner:
 

         return processed

 

+    def run_alwayson_scripts(self, p):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.process(p, *script_args)

+            except Exception:

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

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

+

     def reload_sources(self):

         for si, script in list(enumerate(self.scripts)):

             with open(script.filename, "r", encoding="utf8") as file:

@@ -197,19 +297,21 @@ class ScriptRunner:
                         self.scripts[si].args_from = args_from

                         self.scripts[si].args_to = args_to

 

+

 scripts_txt2img = ScriptRunner()

 scripts_img2img = ScriptRunner()

 

+

 def reload_script_body_only():

     scripts_txt2img.reload_sources()

     scripts_img2img.reload_sources()

 

 

-def reload_scripts(basedir):

+def reload_scripts():

     global scripts_txt2img, scripts_img2img

 

-    scripts_data.clear()

-    load_scripts(basedir)

+    load_scripts()

 

     scripts_txt2img = ScriptRunner()

     scripts_img2img = ScriptRunner()

+