diff options
Diffstat (limited to 'modules/scripts.py')
| -rw-r--r-- | modules/scripts.py | 207 |
1 files changed, 188 insertions, 19 deletions
diff --git a/modules/scripts.py b/modules/scripts.py index e8518ad0..b6fcf96e 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -11,11 +11,31 @@ from modules import shared, paths, script_callbacks, extensions, script_loading, AlwaysVisible = object()
+class MaskBlendArgs:
+ def __init__(self, current_latent, nmask, init_latent, mask, blended_latent, denoiser=None, sigma=None):
+ self.current_latent = current_latent
+ self.nmask = nmask
+ self.init_latent = init_latent
+ self.mask = mask
+ self.blended_latent = blended_latent
+
+ self.denoiser = denoiser
+ self.is_final_blend = denoiser is None
+ self.sigma = sigma
+
+class PostSampleArgs:
+ def __init__(self, samples):
+ self.samples = samples
class PostprocessImageArgs:
def __init__(self, image):
self.image = image
+class PostProcessMaskOverlayArgs:
+ def __init__(self, index, mask_for_overlay, overlay_image):
+ self.index = index
+ self.mask_for_overlay = mask_for_overlay
+ self.overlay_image = overlay_image
class PostprocessBatchListArgs:
def __init__(self, images):
@@ -206,6 +226,25 @@ class Script: pass
+ def on_mask_blend(self, p, mba: MaskBlendArgs, *args):
+ """
+ Called in inpainting mode when the original content is blended with the inpainted content.
+ This is called at every step in the denoising process and once at the end.
+ If is_final_blend is true, this is called for the final blending stage.
+ Otherwise, denoiser and sigma are defined and may be used to inform the procedure.
+ """
+
+ pass
+
+ def post_sample(self, p, ps: PostSampleArgs, *args):
+ """
+ Called after the samples have been generated,
+ but before they have been decoded by the VAE, if applicable.
+ Check getattr(samples, 'already_decoded', False) to test if the images are decoded.
+ """
+
+ pass
+
def postprocess_image(self, p, pp: PostprocessImageArgs, *args):
"""
Called for every image after it has been generated.
@@ -213,6 +252,13 @@ class Script: pass
+ def postprocess_maskoverlay(self, p, ppmo: PostProcessMaskOverlayArgs, *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.
@@ -311,20 +357,113 @@ scripts_data = [] postprocessing_scripts_data = []
ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir", "module"])
+def topological_sort(dependencies):
+ """Accepts a dictionary mapping name to its dependencies, returns a list of names ordered according to dependencies.
+ Ignores errors relating to missing dependeencies or circular dependencies
+ """
+
+ visited = {}
+ result = []
+
+ def inner(name):
+ visited[name] = True
+
+ for dep in dependencies.get(name, []):
+ if dep in dependencies and dep not in visited:
+ inner(dep)
+
+ result.append(name)
+
+ for depname in dependencies:
+ if depname not in visited:
+ inner(depname)
+
+ return result
+
+
+@dataclass
+class ScriptWithDependencies:
+ script_canonical_name: str
+ file: ScriptFile
+ requires: list
+ load_before: list
+ load_after: list
+
def list_scripts(scriptdirname, extension, *, include_extensions=True):
- scripts_list = []
+ scripts = {}
+
+ loaded_extensions = {ext.canonical_name: ext for ext in extensions.active()}
+ loaded_extensions_scripts = {ext.canonical_name: [] for ext in extensions.active()}
+
+ # build script dependency map
+ root_script_basedir = os.path.join(paths.script_path, scriptdirname)
+ if os.path.exists(root_script_basedir):
+ for filename in sorted(os.listdir(root_script_basedir)):
+ if not os.path.isfile(os.path.join(root_script_basedir, filename)):
+ continue
+
+ if os.path.splitext(filename)[1].lower() != extension:
+ continue
- 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)))
+ script_file = ScriptFile(paths.script_path, filename, os.path.join(root_script_basedir, filename))
+ scripts[filename] = ScriptWithDependencies(filename, script_file, [], [], [])
if include_extensions:
for ext in extensions.active():
- scripts_list += ext.list_files(scriptdirname, extension)
-
- scripts_list = [x for x in scripts_list if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)]
+ extension_scripts_list = ext.list_files(scriptdirname, extension)
+ for extension_script in extension_scripts_list:
+ if not os.path.isfile(extension_script.path):
+ continue
+
+ script_canonical_name = ("builtin/" if ext.is_builtin else "") + ext.canonical_name + "/" + extension_script.filename
+ relative_path = scriptdirname + "/" + extension_script.filename
+
+ script = ScriptWithDependencies(
+ script_canonical_name=script_canonical_name,
+ file=extension_script,
+ requires=ext.metadata.get_script_requirements("Requires", relative_path, scriptdirname),
+ load_before=ext.metadata.get_script_requirements("Before", relative_path, scriptdirname),
+ load_after=ext.metadata.get_script_requirements("After", relative_path, scriptdirname),
+ )
+
+ scripts[script_canonical_name] = script
+ loaded_extensions_scripts[ext.canonical_name].append(script)
+
+ for script_canonical_name, script in scripts.items():
+ # load before requires inverse dependency
+ # in this case, append the script name into the load_after list of the specified script
+ for load_before in script.load_before:
+ # if this requires an individual script to be loaded before
+ other_script = scripts.get(load_before)
+ if other_script:
+ other_script.load_after.append(script_canonical_name)
+
+ # if this requires an extension
+ other_extension_scripts = loaded_extensions_scripts.get(load_before)
+ if other_extension_scripts:
+ for other_script in other_extension_scripts:
+ other_script.load_after.append(script_canonical_name)
+
+ # if After mentions an extension, remove it and instead add all of its scripts
+ for load_after in list(script.load_after):
+ if load_after not in scripts and load_after in loaded_extensions_scripts:
+ script.load_after.remove(load_after)
+
+ for other_script in loaded_extensions_scripts.get(load_after, []):
+ script.load_after.append(other_script.script_canonical_name)
+
+ dependencies = {}
+
+ for script_canonical_name, script in scripts.items():
+ for required_script in script.requires:
+ if required_script not in scripts and required_script not in loaded_extensions:
+ errors.report(f'Script "{script_canonical_name}" requires "{required_script}" to be loaded, but it is not.', exc_info=False)
+
+ dependencies[script_canonical_name] = script.load_after
+
+ ordered_scripts = topological_sort(dependencies)
+ scripts_list = [scripts[script_canonical_name].file for script_canonical_name in ordered_scripts]
return scripts_list
@@ -365,15 +504,9 @@ def load_scripts(): elif issubclass(script_class, scripts_postprocessing.ScriptPostprocessing):
postprocessing_scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))
- def orderby(basedir):
- # 1st webui, 2nd extensions-builtin, 3rd extensions
- priority = {os.path.join(paths.script_path, "extensions-builtin"):1, paths.script_path:0}
- for key in priority:
- if basedir.startswith(key):
- return priority[key]
- return 9999
-
- for scriptfile in sorted(scripts_list, key=lambda x: [orderby(x.basedir), x]):
+ # here the scripts_list is already ordered
+ # processing_script is not considered though
+ for scriptfile in scripts_list:
try:
if scriptfile.basedir != paths.script_path:
sys.path = [scriptfile.basedir] + sys.path
@@ -473,17 +606,25 @@ class ScriptRunner: on_after.clear()
def create_script_ui(self, script):
- import modules.api.models as api_models
script.args_from = len(self.inputs)
script.args_to = len(self.inputs)
+ try:
+ self.create_script_ui_inner(script)
+ except Exception:
+ errors.report(f"Error creating UI for {script.name}: ", exc_info=True)
+
+ def create_script_ui_inner(self, script):
+ import modules.api.models as api_models
+
controls = wrap_call(script.ui, script.filename, "ui", script.is_img2img)
if controls is None:
return
script.name = wrap_call(script.title, script.filename, "title", default=script.filename).lower()
+
api_args = []
for control in controls:
@@ -491,11 +632,15 @@ class ScriptRunner: arg_info = api_models.ScriptArg(label=control.label or "")
- for field in ("value", "minimum", "maximum", "step", "choices"):
+ for field in ("value", "minimum", "maximum", "step"):
v = getattr(control, field, None)
if v is not None:
setattr(arg_info, field, v)
+ choices = getattr(control, 'choices', None) # as of gradio 3.41, some items in choices are strings, and some are tuples where the first elem is the string
+ if choices is not None:
+ arg_info.choices = [x[0] if isinstance(x, tuple) else x for x in choices]
+
api_args.append(arg_info)
script.api_info = api_models.ScriptInfo(
@@ -668,6 +813,22 @@ class ScriptRunner: except Exception:
errors.report(f"Error running postprocess_batch_list: {script.filename}", exc_info=True)
+ def post_sample(self, p, ps: PostSampleArgs):
+ for script in self.alwayson_scripts:
+ try:
+ script_args = p.script_args[script.args_from:script.args_to]
+ script.post_sample(p, ps, *script_args)
+ except Exception:
+ errors.report(f"Error running post_sample: {script.filename}", exc_info=True)
+
+ def on_mask_blend(self, p, mba: MaskBlendArgs):
+ for script in self.alwayson_scripts:
+ try:
+ script_args = p.script_args[script.args_from:script.args_to]
+ script.on_mask_blend(p, mba, *script_args)
+ except Exception:
+ errors.report(f"Error running post_sample: {script.filename}", exc_info=True)
+
def postprocess_image(self, p, pp: PostprocessImageArgs):
for script in self.alwayson_scripts:
try:
@@ -676,6 +837,14 @@ class ScriptRunner: except Exception:
errors.report(f"Error running postprocess_image: {script.filename}", exc_info=True)
+ def postprocess_maskoverlay(self, p, ppmo: PostProcessMaskOverlayArgs):
+ for script in self.alwayson_scripts:
+ try:
+ script_args = p.script_args[script.args_from:script.args_to]
+ script.postprocess_maskoverlay(p, ppmo, *script_args)
+ except Exception:
+ errors.report(f"Error running postprocess_image: {script.filename}", exc_info=True)
+
def before_component(self, component, **kwargs):
for callback, script in self.on_before_component_elem_id.get(kwargs.get("elem_id"), []):
try:
|