add an option to enable sections from extras tab in txt2img/img2img

fix some style inconsistenices

1 files changed, 28 insertions, 4 deletions

diff --git a/modules/scripts.py b/modules/scripts.py
index 03907a63..6e9dc0c0 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -6,12 +6,16 @@ from collections import namedtuple
 

 import gradio as gr

 

-from modules.processing import StableDiffusionProcessing

 from modules import shared, paths, script_callbacks, extensions, script_loading, scripts_postprocessing

 

 AlwaysVisible = object()

 

 

+class PostprocessImageArgs:

+    def __init__(self, image):

+        self.image = image

+

+

 class Script:

     filename = None

     args_from = None

@@ -65,7 +69,7 @@ class Script:
         args contains all values returned by components from ui()

         """

 

-        raise NotImplementedError()

+        pass

 

     def process(self, p, *args):

         """

@@ -100,6 +104,13 @@ class Script:
 

         pass

 

+    def postprocess_image(self, p, pp: PostprocessImageArgs, *args):

+        """

+        Called for every image after it has been generated.

+        """

+

+        pass

+

     def postprocess(self, p, processed, *args):

         """

         This function is called after processing ends for AlwaysVisible scripts.

@@ -247,11 +258,15 @@ class ScriptRunner:
         self.infotext_fields = []

 

     def initialize_scripts(self, is_img2img):

+        from modules import scripts_auto_postprocessing

+

         self.scripts.clear()

         self.alwayson_scripts.clear()

         self.selectable_scripts.clear()

 

-        for script_class, path, basedir, script_module in scripts_data:

+        auto_processing_scripts = scripts_auto_postprocessing.create_auto_preprocessing_script_data()

+

+        for script_class, path, basedir, script_module in auto_processing_scripts + scripts_data:

             script = script_class()

             script.filename = path

             script.is_txt2img = not is_img2img

@@ -332,7 +347,7 @@ class ScriptRunner:
 

         return inputs

 

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

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

         script_index = args[0]

 

         if script_index == 0:

@@ -386,6 +401,15 @@ class ScriptRunner:
                 print(f"Error running postprocess_batch: {script.filename}", file=sys.stderr)

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

 

+    def postprocess_image(self, p, pp: PostprocessImageArgs):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.postprocess_image(p, pp, *script_args)

+            except Exception:

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

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

+

     def before_component(self, component, **kwargs):

         for script in self.scripts:

             try: