fluxdev/stable-diffusion-webui-forge

1from __future__ import annotations2 3import os4import time5 6from modules import timer7from modules import initialize_util8from modules import initialize9from threading import Thread10from modules_forge.initialization import initialize_forge11from modules_forge import main_thread12 13 14startup_timer = timer.startup_timer15startup_timer.record("launcher")16 17initialize_forge()18 19initialize.imports()20 21initialize.check_versions()22 23initialize.initialize()24 25 26def create_api(app):27    from modules.api.api import Api28    from modules.call_queue import queue_lock29 30    api = Api(app, queue_lock)31    return api32 33 34def api_only_worker():35    from fastapi import FastAPI36    from modules.shared_cmd_options import cmd_opts37 38    app = FastAPI()39    initialize_util.setup_middleware(app)40    api = create_api(app)41 42    from modules import script_callbacks43    script_callbacks.before_ui_callback()44    script_callbacks.app_started_callback(None, app)45 46    print(f"Startup time: {startup_timer.summary()}.")47    api.launch(48        server_name=initialize_util.gradio_server_name(),49        port=cmd_opts.port if cmd_opts.port else 7861,50        root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else ""51    )52 53 54def webui_worker():55    from modules.shared_cmd_options import cmd_opts56 57    launch_api = cmd_opts.api58 59    from modules import shared, ui_tempdir, script_callbacks, ui, progress, ui_extra_networks60 61    while 1:62        if shared.opts.clean_temp_dir_at_start:63            ui_tempdir.cleanup_tmpdr()64            startup_timer.record("cleanup temp dir")65 66        script_callbacks.before_ui_callback()67        startup_timer.record("scripts before_ui_callback")68 69        shared.demo = ui.create_ui()70        startup_timer.record("create ui")71 72        if not cmd_opts.no_gradio_queue:73            shared.demo.queue(64)74 75        gradio_auth_creds = list(initialize_util.get_gradio_auth_creds()) or None76 77        auto_launch_browser = False78        if os.getenv('SD_WEBUI_RESTARTING') != '1':79            if shared.opts.auto_launch_browser == "Remote" or cmd_opts.autolaunch:80                auto_launch_browser = True81            elif shared.opts.auto_launch_browser == "Local":82                auto_launch_browser = not cmd_opts.webui_is_non_local83 84        app, local_url, share_url = shared.demo.launch(85            share=True,86            server_name=initialize_util.gradio_server_name(),87            server_port=cmd_opts.port,88            ssl_keyfile=cmd_opts.tls_keyfile,89            ssl_certfile=cmd_opts.tls_certfile,90            ssl_verify=cmd_opts.disable_tls_verify,91            debug=cmd_opts.gradio_debug,92            auth=gradio_auth_creds,93            inbrowser=auto_launch_browser,94            prevent_thread_lock=True,95            allowed_paths=cmd_opts.gradio_allowed_path,96            app_kwargs={97                "docs_url": "/docs",98                "redoc_url": "/redoc",99            },100            root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else "",101        )102 103        startup_timer.record("gradio launch")104 105        # gradio uses a very open CORS policy via app.user_middleware, which makes it possible for106        # an attacker to trick the user into opening a malicious HTML page, which makes a request to the107        # running web ui and do whatever the attacker wants, including installing an extension and108        # running its code. We disable this here. Suggested by RyotaK.109        app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']110 111        initialize_util.setup_middleware(app)112 113        progress.setup_progress_api(app)114        ui.setup_ui_api(app)115 116        if launch_api:117            create_api(app)118 119        ui_extra_networks.add_pages_to_demo(app)120 121        startup_timer.record("add APIs")122 123        with startup_timer.subcategory("app_started_callback"):124            script_callbacks.app_started_callback(shared.demo, app)125 126        timer.startup_record = startup_timer.dump()127        print(f"Startup time: {startup_timer.summary()}.")128 129        try:130            while True:131                server_command = shared.state.wait_for_server_command(timeout=5)132                if server_command:133                    if server_command in ("stop", "restart"):134                        break135                    else:136                        print(f"Unknown server command: {server_command}")137        except KeyboardInterrupt:138            print('Caught KeyboardInterrupt, stopping...')139            server_command = "stop"140 141        if server_command == "stop":142            print("Stopping server...")143            # If we catch a keyboard interrupt, we want to stop the server and exit.144            shared.demo.close()145            break146 147        # disable auto launch webui in browser for subsequent UI Reload148        os.environ.setdefault('SD_WEBUI_RESTARTING', '1')149 150        print('Restarting UI...')151        shared.demo.close()152        time.sleep(0.5)153        startup_timer.reset()154        script_callbacks.app_reload_callback()155        startup_timer.record("app reload callback")156        script_callbacks.script_unloaded_callback()157        startup_timer.record("scripts unloaded callback")158        initialize.initialize_rest(reload_script_modules=True)159 160 161def api_only():162    Thread(target=api_only_worker, daemon=True).start()163 164 165def webui():166    Thread(target=webui_worker, daemon=True).start()167 168 169if __name__ == "__main__":170    from modules.shared_cmd_options import cmd_opts171 172    if cmd_opts.nowebui:173        api_only()174    else:175        webui()176 177    main_thread.loop()178