update modules

modules/sd_models.py

@@ -9,6 +9,7 @@ from omegaconf import OmegaConf
 from os import mkdir
 from urllib import request
 import ldm.modules.midas as midas
+import io
 
 from ldm.util import instantiate_from_config
 
@@ -17,6 +18,9 @@ from modules.paths import models_path
 from modules.sd_hijack_inpainting import do_inpainting_hijack
 from modules.timer import Timer
 
+import rich
+from rich import print
+
 model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
 
@@ -234,16 +238,18 @@ def read_metadata_from_safetensors(filename):
         return res
 
 
-def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):
-    _, extension = os.path.splitext(checkpoint_file)
-    if extension.lower() == ".safetensors":
-        device = map_location or shared.weight_load_location or devices.get_optimal_device_name()
-        pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
+def read_state_dict(checkpoint_file):
+    if 'v1-5-pruned-emaonly.safetensors' in checkpoint_file:
+        pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu')
     else:
-        pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)
-
-    if print_global_state and "global_step" in pl_sd:
-        print(f"Global Step: {pl_sd['global_step']}")
+        _, extension = os.path.splitext(checkpoint_file)
+        with rich.progress.open(checkpoint_file, 'rb') as f:
+            if extension.lower() == ".safetensors":
+                buffer = f.read()
+                pl_sd = safetensors.torch.load(buffer)
+            else:
+                buffer = io.BytesIO(f.read())
+                pl_sd = torch.load(buffer, map_location='cpu')
 
     sd = get_state_dict_from_checkpoint(pl_sd)
     return sd
@@ -255,12 +261,12 @@ def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):
 
     if checkpoint_info in checkpoints_loaded:
         # use checkpoint cache
-        print(f"Loading weights [{sd_model_hash}] from cache")
+        print(f"Loading weights from cache")
         return checkpoints_loaded[checkpoint_info]
 
-    print(f"Loading weights [{sd_model_hash}] from {checkpoint_info.filename}")
+    print(f"Loading weights from {checkpoint_info.filename}")
     res = read_state_dict(checkpoint_info.filename)
-    timer.record("load weights from disk")
+    timer.record("load weights")
 
     return res
 
@@ -276,7 +282,7 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
 
     model.load_state_dict(state_dict, strict=False)
     del state_dict
-    timer.record("apply weights to model")
+    timer.record("apply weights")
 
     if shared.opts.sd_checkpoint_cache > 0:
         # cache newly loaded model
@@ -302,15 +308,12 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
         if depth_model:
             model.depth_model = depth_model
 
-        timer.record("apply half()")
-
     devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16
     devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16
     devices.dtype_unet = model.model.diffusion_model.dtype
     devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16
 
     model.first_stage_model.to(devices.dtype_vae)
-    timer.record("apply dtype to VAE")
 
     # clean up cache if limit is reached
     while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:
@@ -327,7 +330,7 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
     sd_vae.clear_loaded_vae()
     vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)
     sd_vae.load_vae(model, vae_file, vae_source)
-    timer.record("load VAE")
+    timer.record("load vae")
 
 
 def enable_midas_autodownload():
@@ -387,7 +390,7 @@ def repair_config(sd_config):
 sd1_clip_weight = 'cond_stage_model.transformer.text_model.embeddings.token_embedding.weight'
 sd2_clip_weight = 'cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight'
 
-def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_to_load_state_dict=None):
+def load_model(checkpoint_info=None, already_loaded_state_dict=None):
     from modules import lowvram, sd_hijack
     checkpoint_info = checkpoint_info or select_checkpoint()
 
@@ -440,7 +443,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, time_taken_
     else:
         sd_model.to(shared.device)
 
-    timer.record("move model to device")
+    timer.record("device move")
 
     sd_hijack.model_hijack.hijack(sd_model)
 
@@ -512,12 +515,10 @@ def reload_model_weights(sd_model=None, info=None):
 
         if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
             sd_model.to(devices.device)
-            timer.record("move model to device")
+            timer.record("device move")
 
     print(f"Weights loaded in {timer.summary()}.")
 
-    return sd_model
-
 def unload_model_weights(sd_model=None, info=None):
     from modules import lowvram, devices, sd_hijack
     timer = Timer()
@@ -536,4 +537,4 @@ def unload_model_weights(sd_model=None, info=None):
 
     print(f"Unloaded weights {timer.summary()}.")
 
-    return sd_model
+    return sd_model

modules/shared.py

@@ -5,7 +5,6 @@ import os
 import sys
 import time
 
-from PIL import Image
 import gradio as gr
 import tqdm
 
@@ -15,6 +14,7 @@ import modules.styles
 import modules.devices as devices
 from modules import localization, extensions, script_loading, errors, ui_components, shared_items
 from modules.paths import models_path, script_path, data_path
+from rich import print
 
 
 demo = None
@@ -37,7 +37,7 @@ parser.add_argument("--no-half-vae", action='store_true', help="do not switch th
 parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
 parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
 parser.add_argument("--embeddings-dir", type=str, default=os.path.join(data_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
-parser.add_argument("--textual-inversion-templates-dir", type=str, default=os.path.join(script_path, 'textual_inversion_templates'), help="directory with textual inversion templates")
+parser.add_argument("--textual-inversion-templates-dir", type=str, default=os.path.join(script_path, 'train/templates'), help="directory with textual inversion templates")
 parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")
 parser.add_argument("--localizations-dir", type=str, default=os.path.join(script_path, 'localizations'), help="localizations directory")
 parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
@@ -48,6 +48,7 @@ parser.add_argument("--always-batch-cond-uncond", action='store_true', help="dis
 parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")
 parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")
 parser.add_argument("--upcast-sampling", action='store_true', help="upcast sampling. No effect with --no-half. Usually produces similar results to --no-half with better performance while using less memory.")
+parser.add_argument("--profile", action='store_true', help="run profiler")
 parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")
 parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)
 parser.add_argument("--ngrok-region", type=str, help="The region in which ngrok should start.", default="us")
@@ -58,7 +59,7 @@ parser.add_argument("--esrgan-models-path", type=str, help="Path to directory wi
 parser.add_argument("--bsrgan-models-path", type=str, help="Path to directory with BSRGAN model file(s).", default=os.path.join(models_path, 'BSRGAN'))
 parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=os.path.join(models_path, 'RealESRGAN'))
 parser.add_argument("--clip-models-path", type=str, help="Path to directory with CLIP model file(s).", default=None)
-parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")
+parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers", default=True)
 parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")
 parser.add_argument("--xformers-flash-attention", action='store_true', help="enable xformers with Flash Attention to improve reproducibility (supported for SD2.x or variant only)")
 parser.add_argument("--deepdanbooru", action='store_true', help="does not do anything")
@@ -69,10 +70,10 @@ parser.add_argument("--sub-quad-kv-chunk-size", type=int, help="kv chunk size fo
 parser.add_argument("--sub-quad-chunk-threshold", type=int, help="the percentage of VRAM threshold for the sub-quadratic cross-attention layer optimization to use chunking", default=None)
 parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")
 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
-parser.add_argument("--opt-sdp-attention", action='store_true', help="enable scaled dot product cross-attention layer optimization; requires PyTorch 2.*")
+parser.add_argument("--opt-sdp-attention", action='store_true', help="enable scaled dot product cross-attention layer optimization; requires PyTorch 2.*", default=True)
 parser.add_argument("--opt-sdp-no-mem-attention", action='store_true', help="enable scaled dot product cross-attention layer optimization without memory efficient attention, makes image generation deterministic; requires PyTorch 2.*")
 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
-parser.add_argument("--disable-nan-check", action='store_true', help="do not check if produced images/latent spaces have nans; useful for running without a checkpoint in CI")
+parser.add_argument("--disable-nan-check", action='store_true', help="do not check if produced images/latent spaces have nans; useful for running without a checkpoint in CI", default=True)
 parser.add_argument("--use-cpu", nargs='+', help="use CPU as torch device for specified modules", default=[], type=str.lower)
 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
 parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
@@ -89,13 +90,13 @@ parser.add_argument("--gradio-inpaint-tool", type=str, help="does not do anythin
 parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")
 parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=os.path.join(data_path, 'styles.csv'))
 parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False)
-parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)
+parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default='dark')
 parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)
-parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)
+parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=True)
 parser.add_argument("--enable-console-prompts", action='store_true', help="print prompts to console when generating with txt2img and img2img", default=False)
 parser.add_argument('--vae-path', type=str, help='Checkpoint to use as VAE; setting this argument disables all settings related to VAE', default=None)
-parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)
-parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)")
+parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=True)
+parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)", default=True)
 parser.add_argument("--api-auth", type=str, help='Set authentication for API like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
 parser.add_argument("--api-log", action='store_true', help="use api-log=True to enable logging of all API requests")
 parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the API instead of the webui")
@@ -109,7 +110,7 @@ parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, req
 parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)
 parser.add_argument("--gradio-queue", action='store_true', help="does not do anything", default=True)
 parser.add_argument("--no-gradio-queue", action='store_true', help="Disables gradio queue; causes the webpage to use http requests instead of websockets; was the defaul in earlier versions")
-parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers")
+parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers", default=True)
 parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False)
 parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False)
 
@@ -153,7 +154,6 @@ devices.device, devices.device_interrogate, devices.device_gfpgan, devices.devic
     (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])
 
 device = devices.device
-weight_load_location = None if cmd_opts.lowram else "cpu"
 
 batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)
 parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
@@ -699,7 +699,7 @@ class TotalTQDM:
 
     def reset(self):
         self._tqdm = tqdm.tqdm(
-            desc="Total progress",
+            desc="Total",
             total=state.job_count * state.sampling_steps,
             position=1,
             file=progress_print_out
@@ -749,3 +749,20 @@ def html(filename):
             return file.read()
 
     return ""
+
+try:
+    from rich.pretty import install as pretty_install
+    from rich.traceback import install as traceback_install
+    from rich.console import Console
+    console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+    pretty_install(console=console)
+    traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, show_locals=True, max_frames=2)
+except:
+    console = None
+    import traceback
+
+def exception():
+    if console is not None:
+        console.print_exception(show_locals=True, max_frames=10, extra_lines=1, suppress=[gr], word_wrap=False, width=min([console.width, 200]))
+    else:
+        print(traceback.format_exc(), file=sys.stderr)

scripts/postprocessing_upscale.py

@@ -17,24 +17,20 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
     def ui(self):
         selected_tab = gr.State(value=0)
 
-        with gr.Column():
-            with FormRow():
-                with gr.Tabs(elem_id="extras_resize_mode"):
-                    with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
-                        upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
-
-                    with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:
-                        with FormRow():
-                            upscaling_resize_w = gr.Number(label="Width", value=512, precision=0, elem_id="extras_upscaling_resize_w")
-                            upscaling_resize_h = gr.Number(label="Height", value=512, precision=0, elem_id="extras_upscaling_resize_h")
-                            upscaling_crop = gr.Checkbox(label='Crop to fit', value=True, elem_id="extras_upscaling_crop")
-
-            with FormRow():
+        with gr.Tabs(elem_id="extras_resize_mode"):
+            with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
+                upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
                 extras_upscaler_1 = gr.Dropdown(label='Upscaler 1', elem_id="extras_upscaler_1", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
 
-            with FormRow():
-                extras_upscaler_2 = gr.Dropdown(label='Upscaler 2', elem_id="extras_upscaler_2", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
-                extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=0.0, elem_id="extras_upscaler_2_visibility")
+            with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:
+                with FormRow():
+                    upscaling_resize_w = gr.Number(label="Width", value=512, precision=0, elem_id="extras_upscaling_resize_w")
+                    upscaling_resize_h = gr.Number(label="Height", value=512, precision=0, elem_id="extras_upscaling_resize_h")
+                    upscaling_crop = gr.Checkbox(label='Crop to fit', value=True, elem_id="extras_upscaling_crop")
+
+        with FormRow():
+            extras_upscaler_2 = gr.Dropdown(label='Upscaler 2', elem_id="extras_upscaler_2", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
+            extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=0.0, elem_id="extras_upscaler_2_visibility")
 
         tab_scale_by.select(fn=lambda: 0, inputs=[], outputs=[selected_tab])
         tab_scale_to.select(fn=lambda: 1, inputs=[], outputs=[selected_tab])

webui.py

@@ -9,6 +9,7 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from packaging import version
+from rich import print
 
 import logging
 logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
@@ -23,10 +24,7 @@ warnings.filterwarnings(action="ignore", category=DeprecationWarning, module="py
 startup_timer.record("import torch")
 
 import gradio
-startup_timer.record("import gradio")
-
 import ldm.modules.encoders.modules
-startup_timer.record("import ldm")
 
 from modules import extra_networks, ui_extra_networks_checkpoints
 from modules import extra_networks_hypernet, ui_extra_networks_hypernets, ui_extra_networks_textual_inversion
@@ -37,7 +35,7 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__:
     torch.__long_version__ = torch.__version__
     torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)
 
-from modules import shared, devices, sd_samplers, upscaler, extensions, localization, ui_tempdir, ui_extra_networks
+from modules import shared, devices, sd_samplers, upscaler, extensions, ui_tempdir, ui_extra_networks
 import modules.codeformer_model as codeformer
 import modules.face_restoration
 import modules.gfpgan_model as gfpgan
@@ -58,7 +56,7 @@ from modules import modelloader
 from modules.shared import cmd_opts
 import modules.hypernetworks.hypernetwork
 
-startup_timer.record("other imports")
+startup_timer.record("import libraries")
 
 
 if cmd_opts.server_name:
@@ -67,42 +65,18 @@ else:
     server_name = "0.0.0.0" if cmd_opts.listen else None
 
 
-def check_versions():
-    if shared.cmd_opts.skip_version_check:
-        return
-
-    expected_torch_version = "1.13.1"
-
-    if version.parse(torch.__version__) < version.parse(expected_torch_version):
-        errors.print_error_explanation(f"""
-You are running torch {torch.__version__}.
-The program is tested to work with torch {expected_torch_version}.
-To reinstall the desired version, run with commandline flag --reinstall-torch.
-Beware that this will cause a lot of large files to be downloaded, as well as
-there are reports of issues with training tab on the latest version.
-
-Use --skip-version-check commandline argument to disable this check.
-        """.strip())
-
-    expected_xformers_version = "0.0.16rc425"
-    if shared.xformers_available:
-        import xformers
-
-        if version.parse(xformers.__version__) < version.parse(expected_xformers_version):
-            errors.print_error_explanation(f"""
-You are running xformers {xformers.__version__}.
-The program is tested to work with xformers {expected_xformers_version}.
-To reinstall the desired version, run with commandline flag --reinstall-xformers.
-
-Use --skip-version-check commandline argument to disable this check.
-            """.strip())
-
-
 def initialize():
-    check_versions()
+    if torch.cuda.is_available():
+        if torch.version.cuda: cuda_version = f'CUDA {torch.version.cuda} cuDNN {torch.backends.cudnn.version()}'
+        elif torch.version.hip: cuda_version = f'HIP {torch.version.hip}'
+        else: cuda_version = ''
+        print(f'Torch {getattr(torch, "__long_version__", torch.__version__)} {cuda_version}')
+        for device in [torch.cuda.device(i) for i in range(torch.cuda.device_count())]:
+            print(f'GPU {torch.cuda.get_device_name(device)} VRAM {round(torch.cuda.get_device_properties(device).total_memory / 1024 / 1024)} Arch {torch.cuda.get_device_capability(device)} Cores {torch.cuda.get_device_properties(device).multi_processor_count}')
+    else:
+        print(f'Torch {getattr(torch, "__long_version__", torch.__version__)} running on CPU')
 
     extensions.list_extensions()
-    localization.list_localizations(cmd_opts.localizations_dir)
     startup_timer.record("list extensions")
 
     if cmd_opts.ui_debug_mode:
@@ -112,7 +86,7 @@ def initialize():
 
     modelloader.cleanup_models()
     modules.sd_models.setup_model()
-    startup_timer.record("list SD models")
+    startup_timer.record("list models")
 
     codeformer.setup_model(cmd_opts.codeformer_models_path)
     startup_timer.record("setup codeformer")
@@ -135,18 +109,6 @@ def initialize():
     modules.textual_inversion.textual_inversion.list_textual_inversion_templates()
     startup_timer.record("refresh textual inversion templates")
 
-    try:
-        modules.sd_models.load_model()
-    except Exception as e:
-        errors.display(e, "loading stable diffusion model")
-        print("", file=sys.stderr)
-        print("Stable diffusion model failed to load, exiting", file=sys.stderr)
-        exit(1)
-    startup_timer.record("load SD checkpoint")
-
-    shared.opts.data["sd_model_checkpoint"] = shared.sd_model.sd_checkpoint_info.title
-
-    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights()))
     shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False)
     shared.opts.onchange("sd_vae_as_default", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False)
     shared.opts.onchange("temp_dir", ui_tempdir.on_tmpdir_changed)
@@ -179,16 +141,32 @@ def initialize():
         startup_timer.record("TLS")
 
     # make the program just exit at ctrl+c without waiting for anything
-    def sigint_handler(sig, frame):
-        print(f'Interrupted with signal {sig} in {frame}')
+    def sigint_handler(_sig, _frame):
+        print('Exiting')
         os._exit(0)
 
     signal.signal(signal.SIGINT, sigint_handler)
 
 
+def load_model():
+    shared.state.begin()
+    shared.state.job = 'load model'
+    try:
+        modules.sd_models.load_model()
+    except Exception as e:
+        errors.display(e, "loading stable diffusion model")
+        print("", file=sys.stderr)
+        print("Stable diffusion model failed to load, exiting", file=sys.stderr)
+        exit(1)
+    shared.opts.data["sd_model_checkpoint"] = shared.sd_model.sd_checkpoint_info.title
+    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights()))
+    shared.state.end()
+    startup_timer.record("load checkpoint")
+
+
 def setup_middleware(app):
     app.middleware_stack = None # reset current middleware to allow modifying user provided list
-    app.add_middleware(GZipMiddleware, minimum_size=1000)
+    app.add_middleware(GZipMiddleware, minimum_size=1024)
     if cmd_opts.cors_allow_origins and cmd_opts.cors_allow_origins_regex:
         app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_allow_origins.split(','), allow_origin_regex=cmd_opts.cors_allow_origins_regex, allow_methods=['*'], allow_credentials=True, allow_headers=['*'])
     elif cmd_opts.cors_allow_origins:
@@ -204,19 +182,9 @@ def create_api(app):
     return api
 
 
-def wait_on_server(demo=None):
-    while 1:
-        time.sleep(0.5)
-        if shared.state.need_restart:
-            shared.state.need_restart = False
-            time.sleep(0.5)
-            demo.close()
-            time.sleep(0.5)
-            break
-
-
 def api_only():
     initialize()
+    load_model()
 
     app = FastAPI()
     setup_middleware(app)
@@ -232,108 +200,71 @@ def webui():
     launch_api = cmd_opts.api
     initialize()
 
-    while 1:
-        if shared.opts.clean_temp_dir_at_start:
-            ui_tempdir.cleanup_tmpdr()
-            startup_timer.record("cleanup temp dir")
+    if shared.opts.clean_temp_dir_at_start:
+        ui_tempdir.cleanup_tmpdr()
+        startup_timer.record("cleanup temp dir")
 
-        modules.script_callbacks.before_ui_callback()
-        startup_timer.record("scripts before_ui_callback")
+    modules.script_callbacks.before_ui_callback()
+    startup_timer.record("scripts before_ui_callback")
 
-        shared.demo = modules.ui.create_ui()
-        startup_timer.record("create ui")
+    shared.demo = modules.ui.create_ui()
+    startup_timer.record("create ui")
 
-        if not cmd_opts.no_gradio_queue:
-            shared.demo.queue(64)
+    if not cmd_opts.no_gradio_queue:
+        shared.demo.queue(16)
 
-        gradio_auth_creds = []
-        if cmd_opts.gradio_auth:
-            gradio_auth_creds += [x.strip() for x in cmd_opts.gradio_auth.strip('"').replace('\n', '').split(',') if x.strip()]
-        if cmd_opts.gradio_auth_path:
-            with open(cmd_opts.gradio_auth_path, 'r', encoding="utf8") as file:
-                for line in file.readlines():
-                    gradio_auth_creds += [x.strip() for x in line.split(',') if x.strip()]
+    gradio_auth_creds = []
+    if cmd_opts.gradio_auth:
+        gradio_auth_creds += [x.strip() for x in cmd_opts.gradio_auth.strip('"').replace('\n', '').split(',') if x.strip()]
+    if cmd_opts.gradio_auth_path:
+        with open(cmd_opts.gradio_auth_path, 'r', encoding="utf8") as file:
+            for line in file.readlines():
+                gradio_auth_creds += [x.strip() for x in line.split(',') if x.strip()]
 
-        app, local_url, share_url = shared.demo.launch(
-            share=cmd_opts.share,
-            server_name=server_name,
-            server_port=cmd_opts.port,
-            ssl_keyfile=cmd_opts.tls_keyfile,
-            ssl_certfile=cmd_opts.tls_certfile,
-            debug=cmd_opts.gradio_debug,
-            auth=[tuple(cred.split(':')) for cred in gradio_auth_creds] if gradio_auth_creds else None,
-            inbrowser=cmd_opts.autolaunch,
-            prevent_thread_lock=True
-        )
-        for dep in shared.demo.dependencies:
-            dep['show_progress'] = False  # disable gradio css animation on component update
+    app, _local_url, _share_url = shared.demo.launch(
+        share=cmd_opts.share,
+        server_name=server_name,
+        server_port=cmd_opts.port,
+        ssl_keyfile=cmd_opts.tls_keyfile,
+        ssl_certfile=cmd_opts.tls_certfile,
+        debug=cmd_opts.gradio_debug,
+        auth=[tuple(cred.split(':')) for cred in gradio_auth_creds] if gradio_auth_creds else None,
+        inbrowser=cmd_opts.autolaunch,
+        prevent_thread_lock=True,
+        favicon_path='automatic.ico',
+    )
+    for dep in shared.demo.dependencies:
+        dep['show_progress'] = False  # disable gradio css animation on component update
 
-        # after initial launch, disable --autolaunch for subsequent restarts
-        cmd_opts.autolaunch = False
+    # app is instance of FastAPI server
+    # shared.demo.server is instance of gradio class which inherits from uvicorn.Server
+    # shared.demo.config is instance of uvicorn.Config
+    # shared.demo.app is instance of ASGIApp
 
-        startup_timer.record("gradio launch")
+    cmd_opts.autolaunch = False
 
-        # gradio uses a very open CORS policy via app.user_middleware, which makes it possible for
-        # an attacker to trick the user into opening a malicious HTML page, which makes a request to the
-        # running web ui and do whatever the attacker wants, including installing an extension and
-        # running its code. We disable this here. Suggested by RyotaK.
-        app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']
+    startup_timer.record("gradio launch")
 
-        setup_middleware(app)
+    app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']
 
-        modules.progress.setup_progress_api(app)
+    setup_middleware(app)
 
-        if launch_api:
-            create_api(app)
+    modules.progress.setup_progress_api(app)
 
-        ui_extra_networks.add_pages_to_demo(app)
+    if launch_api:
+        create_api(app)
 
-        modules.script_callbacks.app_started_callback(shared.demo, app)
-        startup_timer.record("scripts app_started_callback")
+    ui_extra_networks.add_pages_to_demo(app)
 
-        print(f"Startup time: {startup_timer.summary()}.")
+    modules.script_callbacks.app_started_callback(shared.demo, app)
+    startup_timer.record("scripts app_started_callback")
 
-        wait_on_server(shared.demo)
-        print('Restarting UI...')
+    load_model()
 
-        startup_timer.reset()
+    print(f"Startup time: {startup_timer.summary()}.")
 
-        sd_samplers.set_samplers()
-
-        modules.script_callbacks.script_unloaded_callback()
-        extensions.list_extensions()
-        startup_timer.record("list extensions")
-
-        localization.list_localizations(cmd_opts.localizations_dir)
-
-        modelloader.forbid_loaded_nonbuiltin_upscalers()
-        modules.scripts.reload_scripts()
-        startup_timer.record("load scripts")
-
-        modules.script_callbacks.model_loaded_callback(shared.sd_model)
-        startup_timer.record("model loaded callback")
-
-        modelloader.load_upscalers()
-        startup_timer.record("load upscalers")
-
-        for module in [module for name, module in sys.modules.items() if name.startswith("modules.ui")]:
-            importlib.reload(module)
-        startup_timer.record("reload script modules")
-
-        modules.sd_models.list_models()
-        startup_timer.record("list SD models")
-
-        shared.reload_hypernetworks()
-        startup_timer.record("reload hypernetworks")
-
-        ui_extra_networks.intialize()
-        ui_extra_networks.register_page(ui_extra_networks_textual_inversion.ExtraNetworksPageTextualInversion())
-        ui_extra_networks.register_page(ui_extra_networks_hypernets.ExtraNetworksPageHypernetworks())
-        ui_extra_networks.register_page(ui_extra_networks_checkpoints.ExtraNetworksPageCheckpoints())
-
-        extra_networks.initialize()
-        extra_networks.register_extra_network(extra_networks_hypernet.ExtraNetworkHypernet())
-        startup_timer.record("initialize extra networks")
+    while True:
+        time.sleep(0.1)
 
 
 if __name__ == "__main__":