add 'before_hr callback' script callback

2 files changed, 17 insertions, 0 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 8da73884..35463c37 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -1074,6 +1074,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))

 

+        if self.scripts is not None:

+            self.scripts.before_hr(self)

+

         samples = self.sampler.sample_img2img(self, samples, noise, self.hr_c, self.hr_uc, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)

 

         sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())

diff --git a/modules/scripts.py b/modules/scripts.py
index 99bf836a..6485f398 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -186,6 +186,11 @@ class Script:
 

         return f'script_{tabname}{title}_{item_id}'

 

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

+        """

+        This function is called before hires fix start.

+        """

+        pass

 

 current_basedir = paths.script_path

 

@@ -548,6 +553,15 @@ class ScriptRunner:
                     self.scripts[si].args_to = args_to

 

 

+    def before_hr(self, p):

+        for script in self.alwayson_scripts:

+            try:

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

+                script.before_hr(p, *script_args)

+            except Exception:

+                errors.report(f"Error running before_hr: {script.filename}", exc_info=True)

+

+

 scripts_txt2img: ScriptRunner = None

 scripts_img2img: ScriptRunner = None

 scripts_postproc: scripts_postprocessing.ScriptPostprocessingRunner = None