stable-diffusion-webui

1
from __future__ import annotations
2

3
import os
4
import time
5

6
from modules import timer
7
from modules import initialize_util
8
from modules import initialize
9

10
startup_timer = timer.startup_timer
11
startup_timer.record("launcher")
12

13
initialize.imports()
14

15
initialize.check_versions()
16

17

18
def create_api(app):
19
    from modules.api.api import Api
20
    from modules.call_queue import queue_lock
21

22
    api = Api(app, queue_lock)
23
    return api
24

25

26
def api_only():
27
    from fastapi import FastAPI
28
    from modules.shared_cmd_options import cmd_opts
29

30
    initialize.initialize()
31

32
    app = FastAPI()
33
    initialize_util.setup_middleware(app)
34
    api = create_api(app)
35

36
    from modules import script_callbacks
37
    script_callbacks.before_ui_callback()
38
    script_callbacks.app_started_callback(None, app)
39

40
    print(f"Startup time: {startup_timer.summary()}.")
41
    api.launch(
42
        server_name=initialize_util.gradio_server_name(),
43
        port=cmd_opts.port if cmd_opts.port else 7861,
44
        root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else ""
45
    )
46

47

48
def webui():
49
    from modules.shared_cmd_options import cmd_opts
50

51
    launch_api = cmd_opts.api
52
    initialize.initialize()
53

54
    from modules import shared, ui_tempdir, script_callbacks, ui, progress, ui_extra_networks
55

56
    while 1:
57
        if shared.opts.clean_temp_dir_at_start:
58
            ui_tempdir.cleanup_tmpdr()
59
            startup_timer.record("cleanup temp dir")
60

61
        script_callbacks.before_ui_callback()
62
        startup_timer.record("scripts before_ui_callback")
63

64
        shared.demo = ui.create_ui()
65
        startup_timer.record("create ui")
66

67
        if not cmd_opts.no_gradio_queue:
68
            shared.demo.queue(64)
69

70
        gradio_auth_creds = list(initialize_util.get_gradio_auth_creds()) or None
71

72
        auto_launch_browser = False
73
        if os.getenv('SD_WEBUI_RESTARTING') != '1':
74
            if shared.opts.auto_launch_browser == "Remote" or cmd_opts.autolaunch:
75
                auto_launch_browser = True
76
            elif shared.opts.auto_launch_browser == "Local":
77
                auto_launch_browser = not cmd_opts.webui_is_non_local
78

79
        app, local_url, share_url = shared.demo.launch(
80
            share=cmd_opts.share,
81
            server_name=initialize_util.gradio_server_name(),
82
            server_port=cmd_opts.port,
83
            ssl_keyfile=cmd_opts.tls_keyfile,
84
            ssl_certfile=cmd_opts.tls_certfile,
85
            ssl_verify=cmd_opts.disable_tls_verify,
86
            debug=cmd_opts.gradio_debug,
87
            auth=gradio_auth_creds,
88
            inbrowser=auto_launch_browser,
89
            prevent_thread_lock=True,
90
            allowed_paths=cmd_opts.gradio_allowed_path,
91
            app_kwargs={
92
                "docs_url": "/docs",
93
                "redoc_url": "/redoc",
94
            },
95
            root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else "",
96
        )
97

98
        startup_timer.record("gradio launch")
99

100
        # gradio uses a very open CORS policy via app.user_middleware, which makes it possible for
101
        # an attacker to trick the user into opening a malicious HTML page, which makes a request to the
102
        # running web ui and do whatever the attacker wants, including installing an extension and
103
        # running its code. We disable this here. Suggested by RyotaK.
104
        app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']
105

106
        initialize_util.setup_middleware(app)
107

108
        progress.setup_progress_api(app)
109
        ui.setup_ui_api(app)
110

111
        if launch_api:
112
            create_api(app)
113

114
        ui_extra_networks.add_pages_to_demo(app)
115

116
        startup_timer.record("add APIs")
117

118
        with startup_timer.subcategory("app_started_callback"):
119
            script_callbacks.app_started_callback(shared.demo, app)
120

121
        timer.startup_record = startup_timer.dump()
122
        print(f"Startup time: {startup_timer.summary()}.")
123

124
        try:
125
            while True:
126
                server_command = shared.state.wait_for_server_command(timeout=5)
127
                if server_command:
128
                    if server_command in ("stop", "restart"):
129
                        break
130
                    else:
131
                        print(f"Unknown server command: {server_command}")
132
        except KeyboardInterrupt:
133
            print('Caught KeyboardInterrupt, stopping...')
134
            server_command = "stop"
135

136
        if server_command == "stop":
137
            print("Stopping server...")
138
            # If we catch a keyboard interrupt, we want to stop the server and exit.
139
            shared.demo.close()
140
            break
141

142
        # disable auto launch webui in browser for subsequent UI Reload
143
        os.environ.setdefault('SD_WEBUI_RESTARTING', '1')
144

145
        print('Restarting UI...')
146
        shared.demo.close()
147
        time.sleep(0.5)
148
        startup_timer.reset()
149
        script_callbacks.app_reload_callback()
150
        startup_timer.record("app reload callback")
151
        script_callbacks.script_unloaded_callback()
152
        startup_timer.record("scripts unloaded callback")
153
        initialize.initialize_rest(reload_script_modules=True)
154

155

156
if __name__ == "__main__":
157
    from modules.shared_cmd_options import cmd_opts
158

159
    if cmd_opts.nowebui:
160
        api_only()
161
    else:
162
        webui()
163