aboutsummaryrefslogtreecommitdiff
path: root/modules/scripts.py
blob: be348a704818ef0a97ce58a0882b33940794fe60 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import os
import sys
import traceback

import modules.ui as ui
import gradio as gr

from modules.processing import StableDiffusionProcessing

class Script:
    filename = None
    args_from = None
    args_to = None

    def title(self):
        raise NotImplementedError()

    def ui(self, is_img2img):
        pass

    def run(self, *args):
        raise NotImplementedError()

    def describe(self):
        return ""


scripts = []


def load_scripts(basedir):
    for filename in os.listdir(basedir):
        path = os.path.join(basedir, filename)

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

        with open(path, "r", encoding="utf8") as file:
            text = file.read()

        from types import ModuleType
        compiled = compile(text, path, 'exec')
        module = ModuleType(filename)
        exec(compiled, module.__dict__)

        for key, script_class in module.__dict__.items():
            if type(script_class) == type and issubclass(script_class, Script):
                obj = script_class()
                obj.filename = path

                scripts.append(obj)


def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
    try:
        res = func(*args, **kwargs)
        return res
    except Exception:
        print(f"Error calling: {filename}/{funcname}", file=sys.stderr)
        print(traceback.format_exc(), file=sys.stderr)

    return default


def setup_ui(is_img2img):
    titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in scripts]

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

    inputs = [dropdown]

    for script in scripts:
        script.args_from = len(inputs)
        controls = script.ui(is_img2img)

        for control in controls:
            control.visible = False

        inputs += controls
        script.args_to = len(inputs)

    def select_script(index):
        if index > 0:
            script = scripts[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))]

    dropdown.change(
        fn=select_script,
        inputs=[dropdown],
        outputs=inputs
    )

    return inputs


def run(p: StableDiffusionProcessing, *args):
    script_index = args[0] - 1

    if script_index < 0 or script_index >= len(scripts):
        return None

    script = scripts[script_index]

    script_args = args[script.args_from:script.args_to]
    processed = script.run(p, *script_args)

    return processed