From 284bbcd67b22d52376284452699d7516e17825cb Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 25 Mar 2023 09:25:13 -0400 Subject: [PATCH] update modules --- modules/sd_models.py | 47 +++--- modules/shared.py | 41 ++++-- scripts/postprocessing_upscale.py | 28 ++-- webui.py | 231 +++++++++++------------------- 4 files changed, 146 insertions(+), 201 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index 86218c08a..23665c26e 100644 --- a/modules/sd_models.py +++ b/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 \ No newline at end of file + return sd_model diff --git a/modules/shared.py b/modules/shared.py index 73ce77d43..327d5be90 100644 --- a/modules/shared.py +++ b/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) diff --git a/scripts/postprocessing_upscale.py b/scripts/postprocessing_upscale.py index 11eab31a5..069dfd1ac 100644 --- a/scripts/postprocessing_upscale.py +++ b/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]) diff --git a/webui.py b/webui.py index 30f3e4a1f..cd817d985 100644 --- a/webui.py +++ b/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__":