automatic/webui.py

import io
import os
import sys
import time
import glob
import signal
import asyncio
import logging
import importlib
import contextlib
from threading import Thread
from installer import git_commit, custom_excepthook, version
from modules.logger import log
from modules import timer
import modules.errors
import modules.loader
import modules.hashes
import modules.paths
import modules.devices
import modules.migrate
from modules import shared
from modules.call_queue import queue_lock, wrap_queued_call, wrap_gradio_gpu_call # pylint: disable=unused-import
import modules.gr_tempdir
import modules.modeldata
import modules.extensions
import modules.modelloader
import modules.sd_checkpoint
import modules.sd_samplers
import modules.scripts_manager
import modules.scripts
import modules.sd_models
import modules.sd_vae
import modules.sd_unet
import modules.sd_hijack
import modules.model_te
import modules.progress
import modules.ui
import modules.txt2img
import modules.img2img
import modules.detailer
import modules.upscaler
import modules.upscaler_simple
import modules.upscaler_vae
import modules.upscaler_algo
import modules.upscaler_spandrel
import modules.extra_networks
import modules.ui_extra_networks
import modules.textual_inversion
import modules.script_callbacks
import modules.api.middleware


if not modules.loader.initialized:
    timer.startup.record("libraries")
modules.loader.initialized = True


sys.excepthook = custom_excepthook
local_url = None
state = shared.state
backend = shared.backend
if shared.cmd_opts.server_name:
    server_name = shared.cmd_opts.server_name
else:
    server_name = "0.0.0.0" if shared.cmd_opts.listen else None
fastapi_args = {
    "version": f'0.0.{git_commit}',
    "title": "SD.Next",
    "description": "SD.Next",
    "docs_url": None,
    "redoc_url": None,
    "openapi_url": "/openapi.json" if shared.cmd_opts.docs else None,  # only expose OpenAPI schema if docs are enabled
}


def initialize():
    log.debug('Initializing: modules')
    from concurrent.futures import ThreadPoolExecutor, as_completed

    from installer import register_sdnq
    register_sdnq(skip=True, devices=modules.devices, shared=shared) # monkey-patch sdnq to use sdnext devices and shared modules

    from modules.models_hf import hf_init
    hf_init()

    modules.sd_checkpoint.init_metadata()
    modules.hashes.load_cache()

    modules.sd_samplers.list_samplers()
    timer.startup.record("samplers")

    # run independent filesystem scans in parallel
    def _scan_vae():
        modules.sd_vae.refresh_vae_list()
    def _scan_unet():
        modules.sd_unet.refresh_unet_list()
    def _scan_te():
        modules.model_te.refresh_te_list()
    def _scan_models():
        modules.modelloader.cleanup_models()
        modules.sd_checkpoint.setup_model()
    def _scan_lora():
        from modules.lora import lora_load
        lora_load.list_available_networks()
    def _scan_upscalers():
        modules.modelloader.load_upscalers()

    scans = [_scan_vae, _scan_unet, _scan_te, _scan_models, _scan_lora, _scan_upscalers]
    with ThreadPoolExecutor(max_workers=len(scans), thread_name_prefix='sdnext-scan') as pool:
        futures = {pool.submit(fn): fn.__name__ for fn in scans}
        for future in as_completed(futures):
            name = futures[future]
            try:
                future.result()
            except Exception as e:
                log.error(f'Scan error: {name} {e}')
    from modules.sd_checkpoint import write_metadata
    write_metadata()

    timer.startup.record("scans")

    shared.prompt_styles.reload()
    timer.startup.record("styles")

    modules.detailer.initialize()
    timer.startup.record("detailer")

    modules.extensions.list_extensions()
    timer.startup.record("extensions")

    log.info('Load extensions')
    t_timer, t_total = modules.scripts_manager.load_scripts()
    modules.scripts.register_runners()
    timer.startup.record("extensions")
    timer.startup.records["extensions"] = t_total # scripts can reset the time
    log.debug(f'Extensions init time: {t_timer.summary()}')

    modules.ui_extra_networks.initialize()
    modules.ui_extra_networks.register_pages()
    modules.extra_networks.initialize()
    modules.extra_networks.register_default_extra_networks()
    timer.startup.record("networks")

    if shared.cmd_opts.test:
        from modules.models_hf import hf_check_cache
        hf_check_cache()
        timer.startup.record("huggingface")

    if shared.cmd_opts.test:
        from modules.storage import check_storage
        check_storage()
        timer.startup.record("storage")

    if shared.cmd_opts.tls_keyfile is not None and shared.cmd_opts.tls_certfile is not None:
        try:
            if not os.path.exists(shared.cmd_opts.tls_keyfile):
                log.error("Invalid path to TLS keyfile given")
            if not os.path.exists(shared.cmd_opts.tls_certfile):
                log.error(f"Invalid path to TLS certfile: '{shared.cmd_opts.tls_certfile}'")
        except TypeError:
            shared.cmd_opts.tls_keyfile = shared.cmd_opts.tls_certfile = None
            log.error("TLS setup invalid, running webui without TLS")
        else:
            log.info("Running with TLS")
        timer.startup.record("tls")

    # make the program just exit at ctrl+c without waiting for anything
    def sigint_handler(_sig, _frame):
        log.trace(f'State history: uptime={round(time.time() - shared.state.server_start)} jobs={shared.state.job_history} tasks={shared.state.task_history} latents={shared.state.latent_history} images={shared.state.image_history}')
        if modules.errors._profiler is not None: # pylint: disable=protected-access
            modules.errors.profile_print('SIGINT')
        log.info('Exiting')
        try:
            for f in glob.glob("*.lock"):
                os.remove(f)
        except Exception:
            pass
        sys.exit(0)

    signal.signal(signal.SIGINT, sigint_handler)


def load_model():
    modules.modeldata.model_data.locked = False
    autoload = shared.opts.sd_checkpoint_autoload or shared.cmd_opts.ckpt is not None
    log.info(f'Model: autoload={autoload} selected="{shared.opts.sd_model_checkpoint}"')
    if autoload:
        jobid = shared.state.begin('Load model')
        thread_model = Thread(target=lambda: shared.sd_model)
        thread_model.start()
        thread_refiner = Thread(target=lambda: shared.sd_refiner)
        thread_refiner.start()
        thread_model.join()
        thread_refiner.join()
        shared.state.end(jobid)
    timer.startup.record("checkpoint")
    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='model')), call=False)
    shared.opts.onchange("sd_model_refiner", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False)
    shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False)
    shared.opts.onchange("sd_unet", wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False)
    shared.opts.onchange("sd_unet_secondary", wrap_queued_call(lambda: modules.sd_unet.load_unet_secondary(shared.sd_model)), call=False)
    shared.opts.onchange("sd_text_encoder", wrap_queued_call(lambda: modules.sd_models.reload_text_encoder()), call=False)
    shared.opts.onchange("temp_dir", modules.gr_tempdir.on_tmpdir_changed)
    for opt in modules.sd_offload_state.offload_reapply_options:
        shared.opts.onchange(opt, wrap_queued_call(modules.sd_models.reapply_offload), call=False)
    timer.startup.record("onchange")


def create_api(app):
    log.debug('API initialize')
    from modules.api.api import Api
    api = Api(app, queue_lock)
    return api


def verbose_task_factory(loop, coro, context=None):
    """Custom task factory that intercepts and logs every created task."""
    # Retrieve the origin frame (where asyncio.create_task was called)
    frame = coro.cr_frame if hasattr(coro, 'cr_frame') and coro.cr_frame else None
    origin = f"{frame.f_code.co_filename}:{frame.f_lineno}" if frame else "unknown origin"

    # Get the coroutine function name
    coro_name = getattr(coro, '__qualname__', str(coro))

    print(f"[TASK CREATED] {coro_name} | Origin: {origin}")

    # Fallback to the default Task creation (handles Python 3.11+ context kwargs)
    if context is not None:
        return asyncio.Task(coro, loop=loop, name=coro_name, context=context)
    return asyncio.Task(coro, loop=loop, name=coro_name)


def async_policy():
    if sys.platform == "win32" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
        AsyncPolicy = asyncio.WindowsSelectorEventLoopPolicy
    else:
        AsyncPolicy = asyncio.DefaultEventLoopPolicy

    class AnyThreadEventLoopPolicy(AsyncPolicy):
        @staticmethod
        def handle_exception(loop, context):
            msg = context.get("exception", context.get("message"))
            log.error(f"AsyncIO: loop={loop}: {msg}")

        def new_event_loop(self) -> asyncio.AbstractEventLoop:
            """Ensure custom exception handler is attached whenever a loop is created."""
            loop = super().new_event_loop()
            if shared.cmd_opts.profile:
                loop.slow_callback_duration = 0.001
            loop.set_debug(shared.cmd_opts.profile)
            log.debug(f'AsyncIO: loop={loop}')
            loop.set_task_factory(verbose_task_factory)
            loop.set_exception_handler(self.handle_exception)
            return loop

        def get_event_loop(self) -> asyncio.AbstractEventLoop:
            """Get current thread's event loop, creating one if none exists (thread-safe)."""
            try:
                return super().get_event_loop()
            except (RuntimeError, AssertionError):
                loop = self.new_event_loop()
                self.set_event_loop(loop)
                return loop

    asyncio.set_event_loop_policy(AnyThreadEventLoopPolicy())

def get_external_ip():
    import socket
    try:
        ip_address = socket.gethostbyname(socket.gethostname())
        if ip_address.startswith('127.'):
            return None
        return ip_address
    except Exception:
        return None


def get_remote_ip():
    import requests
    try:
        response = requests.get('https://api.ipify.org?format=json', timeout=2)
        ip_address = response.json()['ip']
        return ip_address
    except Exception:
        return None


def start_common():
    log.debug('Server start sequence...')
    if shared.cmd_opts.data_dir is not None and len(shared.cmd_opts.data_dir) > 0:
        log.info(f'Base path: data="{shared.cmd_opts.data_dir}"')
    if shared.cmd_opts.models_dir is not None and len(shared.cmd_opts.models_dir) > 0 and shared.cmd_opts.models_dir != 'models':
        log.info(f'Base path: models="{shared.cmd_opts.models_dir}"')
    modules.paths.create_paths(shared.opts)
    async_policy()
    initialize()
    if shared.cmd_opts.backend == 'original':
        log.error('Legacy option: backend=original is no longer supported')
        shared.cmd_opts.backend = 'diffusers'
    try:
        from installer import diffusers_commit, transformers_commit
        if diffusers_commit != 'unknown':
            shared.opts.diffusers_version = diffusers_commit # update installed diffusers version
        if transformers_commit != 'unknown':
            shared.opts.transformers_version = transformers_commit # update installed transformers version
    except Exception:
        pass
    if shared.opts.clean_temp_dir_at_start:
        modules.gr_tempdir.cleanup_tmpdr()
        timer.startup.record("cleanup")


def mount_subpath(app):
    if shared.cmd_opts.subpath:
        shared.opts.subpath = shared.cmd_opts.subpath
    if shared.opts.subpath is None or len(shared.opts.subpath) == 0:
        return
    import gradio
    if not shared.opts.subpath.startswith('/'):
        shared.opts.subpath = f'/{shared.opts.subpath}'
    shared.cmd_opts.subpath = shared.opts.subpath # update cmd_opts to match opts
    gradio.mount_gradio_app(app, shared.demo, path=shared.opts.subpath)
    log.info(f'Mounted: subpath="{shared.opts.subpath}"')


def start_server(
    blocks,
    server_name: str | None = None, # pylint: disable=redefined-outer-name
    server_port: int | None = None,
    ssl_keyfile: str | None = None,
    ssl_certfile: str | None = None,
    ssl_keyfile_password: str | None = None, # pylint: disable=unused-argument
    app_kwargs: dict | None = None,
):
    from gradio.routes import App
    from modules.server import UvicornServer
    if ssl_keyfile is not None and ssl_certfile is None:
        raise ValueError("ssl_certfile must be provided if ssl_keyfile is provided.")
    server_name = server_name or "127.0.0.1"
    server_port = server_port or 7860
    url_host_name = "localhost" if server_name == "0.0.0.0" else server_name
    if server_name.startswith("[") and server_name.endswith("]"):
        server_name = server_name[1:-1]
    app = App.create_app(blocks, app_kwargs=app_kwargs)
    server = UvicornServer(
        app=app,
        listen=server_name == '0.0.0.0',
        host=server_name,
        port=server_port,
        keyfile=ssl_keyfile,
        certfile=ssl_certfile,
    )
    server.start()
    if ssl_keyfile is not None:
        path_to_local_server = f"https://{url_host_name}:{server_port}/"
    else:
        path_to_local_server = f"http://{url_host_name}:{server_port}/"
    return server_name, server_port, path_to_local_server, app, server


def start_ui():
    log.debug('UI start sequence')
    log.debug(f'UI image support: kanvas={version.get("kanvas", "unknown")}')
    modules.script_callbacks.before_ui_callback()
    timer.startup.record("before-ui")
    shared.demo = modules.ui.create_ui(timer.startup)
    timer.startup.record("ui")
    if shared.cmd_opts.disable_queue:
        log.info('Server queues disabled')
        shared.demo.progress_tracking = False
    else:
        shared.demo.queue(concurrency_count=64)

    gradio_auth_creds = []
    if shared.cmd_opts.auth:
        gradio_auth_creds += [x.strip() for x in shared.cmd_opts.auth.strip('"').replace('\n', '').split(',') if x.strip()]
    if shared.cmd_opts.auth_file:
        if not os.path.exists(shared.cmd_opts.auth_file):
            log.error(f"Invalid path to auth file: '{shared.cmd_opts.auth_file}'")
        else:
            with open(shared.cmd_opts.auth_file, 'r', encoding="utf8") as file:
                for line in file.readlines():
                    gradio_auth_creds += [x.strip() for x in line.split(',') if x.strip()]
    if len(gradio_auth_creds) > 0:
        log.info(f'Authentication enabled: users={len(list(gradio_auth_creds))}')
    auth_pairs = []
    for cred in gradio_auth_creds:
        if ':' not in cred:
            log.warning(f'Ignoring malformed auth entry: "{cred}"')
            continue
        user, password = cred.split(':', 1)
        if len(user) == 0 or len(password) == 0:
            log.warning(f'Ignoring malformed auth entry: "{cred}"')
            continue
        auth_pairs.append((user, password))

    global local_url # pylint: disable=global-statement
    stdout = io.StringIO()
    allowed_paths = [os.path.dirname(__file__)]
    if shared.cmd_opts.data_dir is not None and os.path.isdir(shared.cmd_opts.data_dir):
        allowed_paths.append(shared.cmd_opts.data_dir)
    if shared.cmd_opts.models_dir is not None and os.path.isdir(shared.cmd_opts.models_dir):
        allowed_paths.append(shared.cmd_opts.models_dir)
    if shared.cmd_opts.allowed_paths is not None:
        allowed_paths += [p for p in shared.cmd_opts.allowed_paths if os.path.isdir(p)]
    log.info(f'Server: name={server_name} port={shared.cmd_opts.port} paths={allowed_paths} ssl={shared.cmd_opts.tls_keyfile}:{shared.cmd_opts.tls_certfile} auth={len(auth_pairs)}')
    with contextlib.redirect_stdout(stdout):
        import gradio.networking
        gradio.networking.start_server = start_server
        app, local_url, share_url = shared.demo.launch( # app is FastAPI(Starlette) instance
            share=shared.cmd_opts.share,
            server_name=server_name,
            server_port=shared.cmd_opts.port if shared.cmd_opts.port != 7860 else None,
            ssl_keyfile=shared.cmd_opts.tls_keyfile,
            ssl_certfile=shared.cmd_opts.tls_certfile,
            ssl_verify=not shared.cmd_opts.tls_selfsign,
            debug=False,
            auth=auth_pairs if auth_pairs else None,
            prevent_thread_lock=True,
            max_threads=64,
            show_api=False,
            quiet=True,
            favicon_path='ui/assets/favicon.svg',
            allowed_paths=allowed_paths,
            app_kwargs=fastapi_args,
            _frontend=True and shared.cmd_opts.share,
        )

    uc = shared.demo.server.config
    get_name = lambda c: getattr(c, '__name__', c) # pylint: disable=unnecessary-lambda-assignment
    log.debug(f'Server config: loop={shared.demo.server.loop} http={get_name(uc.http_protocol_class)} ws={get_name(uc.ws_protocol_class)} interface={uc.interface} workers={uc.workers} backlog={uc.backlog} timeout_keep_alive={uc.timeout_keep_alive} timeout_notify={uc.timeout_notify} ws_max_size={uc.ws_max_size} ws_max_queue={uc.ws_max_queue} ws_ping_interval={uc.ws_ping_interval} ws_ping_timeout={uc.ws_ping_timeout}')

    if shared.cmd_opts.data_dir is not None:
        modules.gr_tempdir.register_tmp_file(shared.demo, os.path.join(shared.cmd_opts.data_dir, 'x'))
    log.info(f'Local URL: {local_url}')
    if shared.cmd_opts.listen:
        if not gradio_auth_creds:
            log.warning('Public URL: enabled without authentication')
        if shared.cmd_opts.insecure:
            log.warning('Public URL: enabled with insecure flag')
        proto = 'https' if shared.cmd_opts.tls_keyfile is not None else 'http'
        external_ip = get_external_ip()
        if external_ip is not None:
            log.info(f'External URL: {proto}://{external_ip}:{shared.cmd_opts.port}')
        public_ip = get_remote_ip()
        if public_ip is not None:
            log.info(f'Public URL: {proto}://{public_ip}:{shared.cmd_opts.port}')
    if shared.cmd_opts.docs:
        log.info(f'API docs: {local_url[:-1]}/docs') # pylint: disable=unsubscriptable-object
        log.info(f'API redocs: {local_url[:-1]}/redocs') # pylint: disable=unsubscriptable-object
    if share_url is not None:
        log.info(f'Share URL: {share_url}')
    if getattr(shared.cmd_opts, 'enso', False):
        log.info(f'Enso URL: {local_url[:-1]}/enso') # pylint: disable=unsubscriptable-object
    # log.debug(f'Gradio functions: registered={len(shared.demo.fns)}')
    shared.demo.server.wants_restart = False
    modules.api.middleware.setup_middleware(app, shared.cmd_opts)
    modules.api.middleware.setup_logging(debug=shared.cmd_opts.profile)

    timer.startup.record("launch")

    shared.api = create_api(app)
    shared.api.register()
    modules.progress.setup_progress_api()
    modules.ui_extra_networks.init_api()
    timer.startup.record("api")

    modules.script_callbacks.app_started_callback(shared.demo, app)
    timer.startup.record("app-started")

    time_sorted = sorted(modules.scripts_manager.time_setup.items(), key=lambda x: x[1], reverse=True)
    time_script = [f'{k}:{round(v,3)}' for (k,v) in time_sorted if v > 0.05]
    time_total = sum(modules.scripts_manager.time_setup.values())
    log.debug(f'Scripts setup: time={time_total:.3f} {time_script}')
    time_component = [f'{k}:{round(v,3)}' for (k,v) in modules.scripts_manager.time_component.items() if v > 0.005]
    if len(time_component) > 0:
        log.debug(f'Scripts components: {time_component}')
    return app


def webui(restart=False, _exit=False, profiler=None):
    if restart:
        modules.script_callbacks.app_reload_callback()
        modules.script_callbacks.script_unloaded_callback()

    start_common()
    app = start_ui()
    modules.script_callbacks.after_ui_callback()

    load_model()
    mount_subpath(app)
    shared.opts.save()

    if shared.cmd_opts.profile:
        for k, v in modules.script_callbacks.callback_map.items():
            log.debug(f'Registered callbacks: {k}={len(v)} {[c.script for c in v]}')
    debug = log.trace if os.environ.get('SD_SCRIPT_DEBUG', None) is not None else lambda *args, **kwargs: None
    debug('Trace: SCRIPTS')
    for m in modules.scripts_manager.scripts_data:
        debug(f'  {m}')
    debug('Loaded postprocessing scripts:')
    for m in modules.scripts_manager.postprocessing_scripts_data:
        debug(f'  {m}')
    modules.script_callbacks.print_timers()

    if shared.cmd_opts.profile:
        log.info(f"Launch time: {timer.launch.summary(min_time=0)}")
        log.info(f"Installer time: {timer.init.summary(min_time=0)}")
        log.info(f"Startup time: {timer.startup.summary(min_time=0)}")
        modules.errors._profiler = profiler # pylint: disable=protected-access
    else:
        timer.startup.add('launch', timer.launch.get_total())
        timer.startup.add('installer', timer.init.get_total())
        log.info(f"Startup time: {timer.startup.summary()}")
    timer.startup.reset()

    if _exit:
        return None

    if not restart:
        # override all loggers to use the same handlers as the main logger
        for log_item in [logging.getLogger(name) for name in logging.root.manager.loggerDict]: # pylint: disable=no-member
            if log_item.name.startswith('uvicorn') or log_item.name.startswith('sd'):
                continue
            log_item.handlers = log.handlers
        # autolaunch only on initial start
        if (shared.opts.autolaunch or shared.cmd_opts.autolaunch) and local_url is not None:
            shared.cmd_opts.autolaunch = False
            log.info('Launching browser')
            import webbrowser
            webbrowser.open(local_url, new=2, autoraise=True)
    else:
        for module in [module for name, module in sys.modules.items() if name.startswith("modules.ui")]:
            importlib.reload(module)

    return shared.demo.server


if __name__ == "__main__":
    webui()