scripts.py 32.6 KB
Newer Older
1
import os
2
import re
3
import sys
H
huchenlei 已提交
4
import inspect
5
from collections import namedtuple
6
from dataclasses import dataclass
7 8 9

import gradio as gr

10
from modules import shared, paths, script_callbacks, extensions, script_loading, scripts_postprocessing, errors, timer
11 12 13

AlwaysVisible = object()

A
AUTOMATIC 已提交
14

15 16 17 18 19
class PostprocessImageArgs:
    def __init__(self, image):
        self.image = image


L
ljleb 已提交
20 21 22 23 24
class PostprocessBatchListArgs:
    def __init__(self, images):
        self.images = images


25 26 27 28 29
@dataclass
class OnComponent:
    component: gr.blocks.Block


30
class Script:
A
AUTOMATIC 已提交
31 32 33
    name = None
    """script's internal name derived from title"""

34 35 36
    section = None
    """name of UI section that the script's controls will be placed into"""

37
    filename = None
A
AUTOMATIC 已提交
38 39
    args_from = None
    args_to = None
40 41
    alwayson = False

42 43
    is_txt2img = False
    is_img2img = False
44
    tabname = None
45

46
    group = None
47 48 49 50
    """A gr.Group component that has all script's UI inside it."""

    create_group = True
    """If False, for alwayson scripts, a group component will not be created."""
51

52 53 54 55
    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
    """
56

57 58 59 60 61
    paste_field_names = None
    """if set in ui(), this is a list of names of infotext fields; the fields will be sent through the
    various "Send to <X>" buttons when clicked
    """

A
AUTOMATIC 已提交
62 63 64
    api_info = None
    """Generated value of type modules.api.models.ScriptInfo with information about the script for API"""

A
AUTOMATIC1111 已提交
65
    on_before_component_elem_id = None
66 67
    """list of callbacks to be called before a component with an elem_id is created"""

A
AUTOMATIC1111 已提交
68
    on_after_component_elem_id = None
69 70
    """list of callbacks to be called after a component with an elem_id is created"""

A
AUTOMATIC1111 已提交
71 72 73
    setup_for_ui_only = False
    """If true, the script setup will only be run in Gradio UI, not in API"""

74
    def title(self):
75 76
        """this function should return the title of the script. This is what will be displayed in the dropdown menu."""

77 78
        raise NotImplementedError()

A
AUTOMATIC 已提交
79
    def ui(self, is_img2img):
80 81
        """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.
J
Jim Hays 已提交
82
        Values of those returned components will be passed to run() and process() functions.
83 84
        """

A
AUTOMATIC 已提交
85 86
        pass

A
AUTOMATIC 已提交
87
    def show(self, is_img2img):
88 89 90 91 92
        """
        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
J
Jim Hays 已提交
93
         - True if the script should be shown in UI if it's selected in the scripts dropdown
94 95 96
         - script.AlwaysVisible if the script should be shown in UI at all times
         """

A
AUTOMATIC 已提交
97 98
        return True

99 100 101 102 103 104 105 106 107 108 109
    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()
        """

110
        pass
A
AUTOMATIC 已提交
111

A
AUTOMATIC1111 已提交
112 113 114 115 116 117 118
    def setup(self, p, *args):
        """For AlwaysVisible scripts, this function is called when the processing object is set up, before any processing starts.
        args contains all values returned by components from ui().
        """
        pass


119 120
    def before_process(self, p, *args):
        """
A
AUTOMATIC1111 已提交
121
        This function is called very early during processing begins for AlwaysVisible scripts.
122 123 124 125 126 127
        You can modify the processing object (p) here, inject hooks, etc.
        args contains all values returned by components from ui()
        """

        pass

128 129 130
    def process(self, p, *args):
        """
        This function is called before processing begins for AlwaysVisible scripts.
A
AUTOMATIC 已提交
131 132 133 134 135 136
        You can modify the processing object (p) here, inject hooks, etc.
        args contains all values returned by components from ui()
        """

        pass

137 138 139 140 141 142 143 144 145 146 147 148 149 150
    def before_process_batch(self, p, *args, **kwargs):
        """
        Called before extra networks are parsed from the prompt, so you can add
        new extra network keywords to the prompt with this callback.

        **kwargs will have those items:
          - batch_number - index of current batch, from 0 to number of batches-1
          - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
          - seeds - list of seeds for current batch
          - subseeds - list of subseeds for current batch
        """

        pass

151 152
    def after_extra_networks_activate(self, p, *args, **kwargs):
        """
W
w-e-w 已提交
153
        Called after extra networks activation, before conds calculation
154 155 156 157 158 159 160 161 162 163 164 165
        allow modification of the network after extra networks activation been applied
        won't be call if p.disable_extra_networks

        **kwargs will have those items:
          - batch_number - index of current batch, from 0 to number of batches-1
          - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
          - seeds - list of seeds for current batch
          - subseeds - list of subseeds for current batch
          - extra_network_data - list of ExtraNetworkParams for current stage
        """
        pass

166
    def process_batch(self, p, *args, **kwargs):
A
Artem Zagidulin 已提交
167
        """
168 169 170 171 172 173 174
        Same as process(), but called for every batch.

        **kwargs will have those items:
          - batch_number - index of current batch, from 0 to number of batches-1
          - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
          - seeds - list of seeds for current batch
          - subseeds - list of subseeds for current batch
A
Artem Zagidulin 已提交
175 176 177 178
        """

        pass

179 180 181 182 183 184 185 186 187 188 189
    def postprocess_batch(self, p, *args, **kwargs):
        """
        Same as process_batch(), but called for every batch after it has been generated.

        **kwargs will have same items as process_batch, and also:
          - batch_number - index of current batch, from 0 to number of batches-1
          - images - torch tensor with all generated images, with values ranging from 0 to 1;
        """

        pass

L
ljleb 已提交
190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
    def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, *args, **kwargs):
        """
        Same as postprocess_batch(), but receives batch images as a list of 3D tensors instead of a 4D tensor.
        This is useful when you want to update the entire batch instead of individual images.

        You can modify the postprocessing object (pp) to update the images in the batch, remove images, add images, etc.
        If the number of images is different from the batch size when returning,
        then the script has the responsibility to also update the following attributes in the processing object (p):
          - p.prompts
          - p.negative_prompts
          - p.seeds
          - p.subseeds

        **kwargs will have same items as process_batch, and also:
          - batch_number - index of current batch, from 0 to number of batches-1
        """

        pass

209 210 211 212 213 214 215
    def postprocess_image(self, p, pp: PostprocessImageArgs, *args):
        """
        Called for every image after it has been generated.
        """

        pass

A
AUTOMATIC 已提交
216 217 218 219
    def postprocess(self, p, processed, *args):
        """
        This function is called after processing ends for AlwaysVisible scripts.
        args contains all values returned by components from ui()
220 221 222 223
        """

        pass

224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240
    def before_component(self, component, **kwargs):
        """
        Called before a component is created.
        Use elem_id/label fields of kwargs to figure out which component it is.
        This can be useful to inject your own components somewhere in the middle of vanilla UI.
        You can return created components in the ui() function to add them to the list of arguments for your processing functions
        """

        pass

    def after_component(self, component, **kwargs):
        """
        Called after a component is created. Same as above.
        """

        pass

241 242 243 244
    def on_before_component(self, callback, *, elem_id):
        """
        Calls callback before a component is created. The callback function is called with a single argument of type OnComponent.

245 246
        May be called in show() or ui() - but it may be too late in latter as some components may already be created.

247 248 249
        This function is an alternative to before_component in that it also cllows to run before a component is created, but
        it doesn't require to be called for every created component - just for the one you need.
        """
A
AUTOMATIC1111 已提交
250 251
        if self.on_before_component_elem_id is None:
            self.on_before_component_elem_id = []
252 253 254 255 256 257 258

        self.on_before_component_elem_id.append((elem_id, callback))

    def on_after_component(self, callback, *, elem_id):
        """
        Calls callback after a component is created. The callback function is called with a single argument of type OnComponent.
        """
A
AUTOMATIC1111 已提交
259 260
        if self.on_after_component_elem_id is None:
            self.on_after_component_elem_id = []
261 262 263

        self.on_after_component_elem_id.append((elem_id, callback))

A
AUTOMATIC 已提交
264
    def describe(self):
265
        """unused"""
A
AUTOMATIC 已提交
266 267
        return ""

268 269 270 271
    def elem_id(self, item_id):
        """helper function to generate id for a HTML element, constructs final id out of script name, tab and user-supplied item_id"""

        need_tabname = self.show(True) == self.show(False)
S
S-Del 已提交
272
        tabkind = 'img2img' if self.is_img2img else 'txt2img'
273
        tabname = f"{tabkind}_" if need_tabname else ""
274 275 276 277
        title = re.sub(r'[^a-z_0-9]', '', re.sub(r'\s', '_', self.title().lower()))

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

A
AUTOMATIC1111 已提交
278
    def before_hr(self, p, *args):
279 280 281 282
        """
        This function is called before hires fix start.
        """
        pass
283

284

A
AUTOMATIC1111 已提交
285 286
class ScriptBuiltinUI(Script):
    setup_for_ui_only = True
287 288 289 290 291

    def elem_id(self, item_id):
        """helper function to generate id for a HTML element, constructs final id out of tab and user-supplied item_id"""

        need_tabname = self.show(True) == self.show(False)
A
AUTOMATIC1111 已提交
292
        tabname = ('img2img' if self.is_img2img else 'txt2img') + "_" if need_tabname else ""
293 294 295 296

        return f'{tabname}{item_id}'


297 298 299 300 301 302 303 304 305 306 307 308
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


ScriptFile = namedtuple("ScriptFile", ["basedir", "filename", "path"])
309 310 311

scripts_data = []
postprocessing_scripts_data = []
312
ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir", "module"])
313

314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345
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

346

347
def list_scripts(scriptdirname, extension, *, include_extensions=True):
348
    scripts = {}
349

350 351
    loaded_extensions = {ext.canonical_name: ext for ext in extensions.active()}
    loaded_extensions_scripts = {ext.canonical_name: [] for ext in extensions.active()}
352

353
    # build script dependency map
354 355 356
    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)):
W
bug fix  
wfjsw 已提交
357 358 359
            if not os.path.isfile(os.path.join(root_script_basedir, filename)):
                continue

360 361 362 363 364
            if os.path.splitext(filename)[1].lower() != extension:
                continue

            script_file = ScriptFile(paths.script_path, filename, os.path.join(root_script_basedir, filename))
            scripts[filename] = ScriptWithDependencies(filename, script_file, [], [], [])
365

366 367
    if include_extensions:
        for ext in extensions.active():
368 369
            extension_scripts_list = ext.list_files(scriptdirname, extension)
            for extension_script in extension_scripts_list:
W
bug fix  
wfjsw 已提交
370 371 372
                if not os.path.isfile(extension_script.path):
                    continue

373
                script_canonical_name = ("builtin/" if ext.is_builtin else "") + ext.canonical_name + "/" + extension_script.filename
374 375
                relative_path = scriptdirname + "/" + extension_script.filename

376 377 378 379 380 381 382
                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),
                )
383

384 385
                scripts[script_canonical_name] = script
                loaded_extensions_scripts[ext.canonical_name].append(script)
386

387
    for script_canonical_name, script in scripts.items():
388 389
        # load before requires inverse dependency
        # in this case, append the script name into the load_after list of the specified script
390
        for load_before in script.load_before:
W
wfjsw 已提交
391
            # if this requires an individual script to be loaded before
392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408
            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)
409

410
    dependencies = {}
411

412 413 414 415
    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)
416

417
        dependencies[script_canonical_name] = script.load_after
418

419 420
    ordered_scripts = topological_sort(dependencies)
    scripts_list = [scripts[script_canonical_name].file for script_canonical_name in ordered_scripts]
421

422
    return scripts_list
423

424

A
AUTOMATIC 已提交
425 426 427
def list_files_with_name(filename):
    res = []

428
    dirs = [paths.script_path] + [ext.path for ext in extensions.active()]
A
AUTOMATIC 已提交
429 430 431 432 433 434

    for dirpath in dirs:
        if not os.path.isdir(dirpath):
            continue

        path = os.path.join(dirpath, filename)
T
Tong Zeng 已提交
435
        if os.path.isfile(path):
A
AUTOMATIC 已提交
436 437 438 439 440
            res.append(path)

    return res


441 442 443
def load_scripts():
    global current_basedir
    scripts_data.clear()
444
    postprocessing_scripts_data.clear()
445 446
    script_callbacks.clear_callbacks()

447
    scripts_list = list_scripts("scripts", ".py") + list_scripts("modules/processing_scripts", ".py", include_extensions=False)
448 449

    syspath = sys.path
450

451
    def register_scripts_from_module(module):
A
AUTOMATIC 已提交
452
        for script_class in module.__dict__.values():
H
huchenlei 已提交
453
            if not inspect.isclass(script_class):
454 455 456 457 458 459 460
                continue

            if issubclass(script_class, Script):
                scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))
            elif issubclass(script_class, scripts_postprocessing.ScriptPostprocessing):
                postprocessing_scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))

461 462 463
    # here the scripts_list is already ordered
    # processing_script is not considered though
    for scriptfile in scripts_list:
464
        try:
465 466 467 468
            if scriptfile.basedir != paths.script_path:
                sys.path = [scriptfile.basedir] + sys.path
            current_basedir = scriptfile.basedir

469 470
            script_module = script_loading.load_module(scriptfile.path)
            register_scripts_from_module(script_module)
471 472

        except Exception:
473
            errors.report(f"Error loading script: {scriptfile.filename}", exc_info=True)
474

475 476 477
        finally:
            sys.path = syspath
            current_basedir = paths.script_path
A
AUTOMATIC 已提交
478
            timer.startup_timer.record(scriptfile.filename)
479

480 481 482 483 484 485
    global scripts_txt2img, scripts_img2img, scripts_postproc

    scripts_txt2img = ScriptRunner()
    scripts_img2img = ScriptRunner()
    scripts_postproc = scripts_postprocessing.ScriptPostprocessingRunner()

486 487 488

def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
    try:
489
        return func(*args, **kwargs)
490
    except Exception:
491
        errors.report(f"Error calling: {filename}/{funcname}", exc_info=True)
492 493 494 495

    return default


A
AUTOMATIC 已提交
496 497 498
class ScriptRunner:
    def __init__(self):
        self.scripts = []
499 500
        self.selectable_scripts = []
        self.alwayson_scripts = []
ふぁ 已提交
501
        self.titles = []
502
        self.title_map = {}
503
        self.infotext_fields = []
504
        self.paste_field_names = []
505
        self.inputs = [None]
A
AUTOMATIC 已提交
506

507 508 509 510 511 512
        self.on_before_component_elem_id = {}
        """dict of callbacks to be called before an element is created; key=elem_id, value=list of callbacks"""

        self.on_after_component_elem_id = {}
        """dict of callbacks to be called after an element is created; key=elem_id, value=list of callbacks"""

513
    def initialize_scripts(self, is_img2img):
514 515
        from modules import scripts_auto_postprocessing

516 517 518 519
        self.scripts.clear()
        self.alwayson_scripts.clear()
        self.selectable_scripts.clear()

520 521
        auto_processing_scripts = scripts_auto_postprocessing.create_auto_preprocessing_script_data()

A
AUTOMATIC 已提交
522 523 524
        for script_data in auto_processing_scripts + scripts_data:
            script = script_data.script_class()
            script.filename = script_data.path
525 526
            script.is_txt2img = not is_img2img
            script.is_img2img = is_img2img
527
            script.tabname = "img2img" if is_img2img else "txt2img"
A
AUTOMATIC 已提交
528

529
            visibility = script.show(script.is_img2img)
A
AUTOMATIC 已提交
530

531 532 533 534
            if visibility == AlwaysVisible:
                self.scripts.append(script)
                self.alwayson_scripts.append(script)
                script.alwayson = True
A
AUTOMATIC 已提交
535

536 537 538
            elif visibility:
                self.scripts.append(script)
                self.selectable_scripts.append(script)
A
AUTOMATIC 已提交
539

540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561
        self.apply_on_before_component_callbacks()

    def apply_on_before_component_callbacks(self):
        for script in self.scripts:
            on_before = script.on_before_component_elem_id or []
            on_after = script.on_after_component_elem_id or []

            for elem_id, callback in on_before:
                if elem_id not in self.on_before_component_elem_id:
                    self.on_before_component_elem_id[elem_id] = []

                self.on_before_component_elem_id[elem_id].append((callback, script))

            for elem_id, callback in on_after:
                if elem_id not in self.on_after_component_elem_id:
                    self.on_after_component_elem_id[elem_id] = []

                self.on_after_component_elem_id[elem_id].append((callback, script))

            on_before.clear()
            on_after.clear()

562
    def create_script_ui(self, script):
A
AUTOMATIC 已提交
563

564 565
        script.args_from = len(self.inputs)
        script.args_to = len(self.inputs)
566

567 568 569 570 571 572 573 574
        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

575
        controls = wrap_call(script.ui, script.filename, "ui", script.is_img2img)
A
AUTOMATIC 已提交
576

577 578
        if controls is None:
            return
A
AUTOMATIC 已提交
579

580
        script.name = wrap_call(script.title, script.filename, "title", default=script.filename).lower()
A
AUTOMATIC 已提交
581

582
        api_args = []
A
AUTOMATIC 已提交
583

584 585
        for control in controls:
            control.custom_script_source = os.path.basename(script.filename)
A
AUTOMATIC 已提交
586

587
            arg_info = api_models.ScriptArg(label=control.label or "")
588

589 590 591 592
            for field in ("value", "minimum", "maximum", "step"):
                v = getattr(control, field, None)
                if v is not None:
                    setattr(arg_info, field, v)
A
AUTOMATIC1111 已提交
593

594 595 596
            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]
A
AUTOMATIC 已提交
597

598
            api_args.append(arg_info)
A
AUTOMATIC 已提交
599

600 601 602 603 604 605
        script.api_info = api_models.ScriptInfo(
            name=script.name,
            is_img2img=script.is_img2img,
            is_alwayson=script.alwayson,
            args=api_args,
        )
A
AUTOMATIC 已提交
606

607 608
        if script.infotext_fields is not None:
            self.infotext_fields += script.infotext_fields
A
AUTOMATIC 已提交
609

610 611
        if script.paste_field_names is not None:
            self.paste_field_names += script.paste_field_names
612

613 614
        self.inputs += controls
        script.args_to = len(self.inputs)
A
AUTOMATIC 已提交
615

616 617 618
    def setup_ui_for_section(self, section, scriptlist=None):
        if scriptlist is None:
            scriptlist = self.alwayson_scripts
619

620 621 622
        for script in scriptlist:
            if script.alwayson and script.section != section:
                continue
A
AUTOMATIC 已提交
623

624 625 626
            if script.create_group:
                with gr.Group(visible=script.alwayson) as group:
                    self.create_script_ui(script)
627

628 629 630
                script.group = group
            else:
                self.create_script_ui(script)
631

632 633
    def prepare_ui(self):
        self.inputs = [None]
634

635
    def setup_ui(self):
636 637
        all_titles = [wrap_call(script.title, script.filename, "title") or script.filename for script in self.scripts]
        self.title_map = {title.lower(): script for title, script in zip(all_titles, self.scripts)}
638
        self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.selectable_scripts]
639

640 641 642 643 644 645
        self.setup_ui_for_section(None)

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

        self.setup_ui_for_section(None, self.selectable_scripts)
646

A
AUTOMATIC 已提交
647
        def select_script(script_index):
648
            selected_script = self.selectable_scripts[script_index - 1] if script_index>0 else None
A
AUTOMATIC 已提交
649

650
            return [gr.update(visible=selected_script == s) for s in self.selectable_scripts]
A
AUTOMATIC 已提交
651

ふぁ 已提交
652
        def init_field(title):
653 654
            """called when an initial value is set from ui-config.json to show script's UI components"""

ふぁ 已提交
655
            if title == 'None':
ふぁ 已提交
656
                return
657

ふぁ 已提交
658
            script_index = self.titles.index(title)
659
            self.selectable_scripts[script_index].group.visible = True
ふぁ 已提交
660 661

        dropdown.init_field = init_field
662

A
AUTOMATIC 已提交
663 664 665
        dropdown.change(
            fn=select_script,
            inputs=[dropdown],
666
            outputs=[script.group for script in self.selectable_scripts]
A
AUTOMATIC 已提交
667
        )
A
AUTOMATIC 已提交
668

E
EllangoK 已提交
669
        self.script_load_ctr = 0
670

E
EllangoK 已提交
671 672 673 674 675 676 677 678 679 680
        def onload_script_visibility(params):
            title = params.get('Script', None)
            if title:
                title_index = self.titles.index(title)
                visibility = title_index == self.script_load_ctr
                self.script_load_ctr = (self.script_load_ctr + 1) % len(self.titles)
                return gr.update(visible=visibility)
            else:
                return gr.update(visible=False)

681 682
        self.infotext_fields.append((dropdown, lambda x: gr.update(value=x.get('Script', 'None'))))
        self.infotext_fields.extend([(script.group, onload_script_visibility) for script in self.selectable_scripts])
E
EllangoK 已提交
683

684
        self.apply_on_before_component_callbacks()
685

686
        return self.inputs
A
AUTOMATIC 已提交
687

688
    def run(self, p, *args):
A
AUTOMATIC 已提交
689
        script_index = args[0]
A
AUTOMATIC 已提交
690

A
AUTOMATIC 已提交
691 692
        if script_index == 0:
            return None
A
AUTOMATIC 已提交
693

694
        script = self.selectable_scripts[script_index-1]
A
AUTOMATIC 已提交
695

A
AUTOMATIC 已提交
696 697
        if script is None:
            return None
A
AUTOMATIC 已提交
698

A
AUTOMATIC 已提交
699 700
        script_args = args[script.args_from:script.args_to]
        processed = script.run(p, *script_args)
A
AUTOMATIC 已提交
701

702 703
        shared.total_tqdm.clear()

A
AUTOMATIC 已提交
704
        return processed
A
AUTOMATIC 已提交
705

706 707 708 709 710 711 712 713
    def before_process(self, p):
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.before_process(p, *script_args)
            except Exception:
                errors.report(f"Error running before_process: {script.filename}", exc_info=True)

A
AUTOMATIC 已提交
714
    def process(self, p):
715 716 717 718 719
        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:
720
                errors.report(f"Error running process: {script.filename}", exc_info=True)
A
AUTOMATIC 已提交
721

722 723 724 725 726 727
    def before_process_batch(self, p, **kwargs):
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.before_process_batch(p, *script_args, **kwargs)
            except Exception:
728
                errors.report(f"Error running before_process_batch: {script.filename}", exc_info=True)
729

730 731 732 733 734 735 736 737
    def after_extra_networks_activate(self, p, **kwargs):
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.after_extra_networks_activate(p, *script_args, **kwargs)
            except Exception:
                errors.report(f"Error running after_extra_networks_activate: {script.filename}", exc_info=True)

738
    def process_batch(self, p, **kwargs):
A
Artem Zagidulin 已提交
739 740 741
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
742
                script.process_batch(p, *script_args, **kwargs)
A
Artem Zagidulin 已提交
743
            except Exception:
744
                errors.report(f"Error running process_batch: {script.filename}", exc_info=True)
A
Artem Zagidulin 已提交
745

A
AUTOMATIC 已提交
746 747 748 749 750 751
    def postprocess(self, p, processed):
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.postprocess(p, processed, *script_args)
            except Exception:
752
                errors.report(f"Error running postprocess: {script.filename}", exc_info=True)
753

754 755 756 757 758 759
    def postprocess_batch(self, p, images, **kwargs):
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.postprocess_batch(p, *script_args, images=images, **kwargs)
            except Exception:
760
                errors.report(f"Error running postprocess_batch: {script.filename}", exc_info=True)
761

L
ljleb 已提交
762 763 764 765 766 767 768 769
    def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, **kwargs):
        for script in self.alwayson_scripts:
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.postprocess_batch_list(p, pp, *script_args, **kwargs)
            except Exception:
                errors.report(f"Error running postprocess_batch_list: {script.filename}", exc_info=True)

770 771 772 773 774 775
    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:
776
                errors.report(f"Error running postprocess_image: {script.filename}", exc_info=True)
777

778
    def before_component(self, component, **kwargs):
A
AUTOMATIC1111 已提交
779 780 781 782 783
        for callback, script in self.on_before_component_elem_id.get(kwargs.get("elem_id"), []):
            try:
                callback(OnComponent(component=component))
            except Exception:
                errors.report(f"Error running on_before_component: {script.filename}", exc_info=True)
784

785 786 787 788
        for script in self.scripts:
            try:
                script.before_component(component, **kwargs)
            except Exception:
789
                errors.report(f"Error running before_component: {script.filename}", exc_info=True)
790 791

    def after_component(self, component, **kwargs):
A
AUTOMATIC1111 已提交
792 793 794 795 796
        for callback, script in self.on_after_component_elem_id.get(component.elem_id, []):
            try:
                callback(OnComponent(component=component))
            except Exception:
                errors.report(f"Error running on_after_component: {script.filename}", exc_info=True)
797

798 799 800 801
        for script in self.scripts:
            try:
                script.after_component(component, **kwargs)
            except Exception:
802
                errors.report(f"Error running after_component: {script.filename}", exc_info=True)
803

804 805 806
    def script(self, title):
        return self.title_map.get(title.lower())

A
AUTOMATIC 已提交
807
    def reload_sources(self, cache):
D
DepFA 已提交
808
        for si, script in list(enumerate(self.scripts)):
809 810 811 812 813 814 815 816 817
            args_from = script.args_from
            args_to = script.args_to
            filename = script.filename

            module = cache.get(filename, None)
            if module is None:
                module = script_loading.load_module(script.filename)
                cache[filename] = module

A
AUTOMATIC 已提交
818
            for script_class in module.__dict__.values():
819 820 821 822 823
                if type(script_class) == type and issubclass(script_class, Script):
                    self.scripts[si] = script_class()
                    self.scripts[si].filename = filename
                    self.scripts[si].args_from = args_from
                    self.scripts[si].args_to = args_to
A
AUTOMATIC 已提交
824

825 826 827 828 829 830 831 832
    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)

A
AUTOMATIC1111 已提交
833
    def setup_scrips(self, p, *, is_ui=True):
A
AUTOMATIC1111 已提交
834
        for script in self.alwayson_scripts:
A
AUTOMATIC1111 已提交
835 836 837
            if not is_ui and script.setup_for_ui_only:
                continue

A
AUTOMATIC1111 已提交
838 839 840 841 842 843
            try:
                script_args = p.script_args[script.args_from:script.args_to]
                script.setup(p, *script_args)
            except Exception:
                errors.report(f"Error running setup: {script.filename}", exc_info=True)

844

845 846 847
scripts_txt2img: ScriptRunner = None
scripts_img2img: ScriptRunner = None
scripts_postproc: scripts_postprocessing.ScriptPostprocessingRunner = None
848
scripts_current: ScriptRunner = None
D
DepFA 已提交
849

850

D
DepFA 已提交
851
def reload_script_body_only():
A
AUTOMATIC 已提交
852 853 854
    cache = {}
    scripts_txt2img.reload_sources(cache)
    scripts_img2img.reload_sources(cache)
D
DepFA 已提交
855

D
DepFA 已提交
856

857
reload_scripts = load_scripts  # compatibility alias