diff --git a/cli/model-jit.py b/cli/model-jit.py new file mode 100755 index 000000000..25acf96b4 --- /dev/null +++ b/cli/model-jit.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python +import os +import time +import functools +import argparse +import logging +import warnings +from dataclasses import dataclass + +logging.getLogger("DeepSpeed").disabled = True +warnings.filterwarnings(action="ignore", category=FutureWarning) +warnings.filterwarnings(action="ignore", category=DeprecationWarning) + +import torch +import diffusers + +n_warmup = 5 +n_traces = 10 +n_runs = 100 +args = {} +pipe = None +log = logging.getLogger("sd") + + +def setup_logging(): + from rich.theme import Theme + from rich.logging import RichHandler + from rich.console import Console + from rich.traceback import install + log.setLevel(logging.DEBUG) + console = Console(log_time=True, log_time_format='%H:%M:%S-%f', theme=Theme({ "traceback.border": "black", "traceback.border.syntax_error": "black", "inspect.value.border": "black" })) + logging.basicConfig(level=logging.ERROR, format='%(asctime)s | %(name)s | %(levelname)s | %(module)s | %(message)s', handlers=[logging.NullHandler()]) # redirect default logger to null + rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=False, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=logging.DEBUG, console=console) + rh.setLevel(logging.DEBUG) + log.addHandler(rh) + logging.getLogger("diffusers").setLevel(logging.ERROR) + logging.getLogger("torch").setLevel(logging.ERROR) + warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning) + install(console=console, extra_lines=1, max_frames=10, width=console.width, word_wrap=False, indent_guides=False, suppress=[]) + + +def generate_inputs(): + if args.type == 'sd15': + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 + encoder_hidden_states = torch.randn(2, 77, 768).half().cuda() + return sample, timestep, encoder_hidden_states + if args.type == 'sdxl': + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 + encoder_hidden_states = torch.randn(2, 77, 768).half().cuda() + text_embeds = torch.randn(1, 77, 2048).half().cuda() + return sample, timestep, encoder_hidden_states, text_embeds + + +def load_model(): + log.info(f'versions: torch={torch.__version__} diffusers={diffusers.__version__}') + diffusers_load_config = { + "low_cpu_mem_usage": True, + "torch_dtype": torch.float16, + "safety_checker": None, + "requires_safety_checker": False, + "load_safety_checker": False, + "load_connected_pipeline": True, + "use_safetensors": True, + } + pipeline = diffusers.StableDiffusionPipeline if args.type == 'sd15' else diffusers.StableDiffusionXLPipeline + global pipe # pylint: disable=global-statement + t0 = time.time() + pipe = pipeline.from_single_file(args.model, **diffusers_load_config).to('cuda') + size = os.path.getsize(args.model) + log.info(f'load: model={args.model} type={args.type} time={time.time() - t0:.3f}s size={size / 1024 / 1024:.3f}mb') + + +def load_trace(fn: str): + + @dataclass + class UNet2DConditionOutput: + sample: torch.FloatTensor + + class TracedUNet(torch.nn.Module): + def __init__(self): + super().__init__() + self.in_channels = pipe.unet.in_channels + self.device = pipe.unet.device + + def forward(self, latent_model_input, t, encoder_hidden_states): + sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0] + return UNet2DConditionOutput(sample=sample) + + t0 = time.time() + unet_traced = torch.jit.load(fn) + pipe.unet = TracedUNet() + size = os.path.getsize(fn) + log.info(f'load: optimized={fn} time={time.time() - t0:.3f}s size={size / 1024 / 1024:.3f}mb') + + +def trace_model(): + log.info(f'tracing model: {args.model}') + torch.set_grad_enabled(False) + unet = pipe.unet + unet.eval() + # unet.to(memory_format=torch.channels_last) # use channels_last memory format + unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default + + # warmup + t0 = time.time() + for _ in range(n_warmup): + with torch.inference_mode(): + inputs = generate_inputs() + _output = unet(*inputs) + log.info(f'warmup: time={time.time() - t0:.3f}s passes={n_warmup}') + + # trace + t0 = time.time() + unet_traced = torch.jit.trace(unet, inputs, check_trace=True) + unet_traced.eval() + log.info(f'trace: time={time.time() - t0:.3f}s') + + # optimize graph + t0 = time.time() + for _ in range(n_traces): + with torch.inference_mode(): + inputs = generate_inputs() + _output = unet_traced(*inputs) + log.info(f'optimize: time={time.time() - t0:.3f}s passes={n_traces}') + + # save the model + if args.save: + t0 = time.time() + basename, _ext = os.path.splitext(args.model) + fn = f"{basename}.pt" + unet_traced.save(fn) + size = os.path.getsize(fn) + log.info(f'save: optimized={fn} time={time.time() - t0:.3f}s size={size / 1024 / 1024:.3f}mb') + return fn + + pipe.unet = unet_traced + return None + + +def benchmark_model(msg: str): + with torch.inference_mode(): + inputs = generate_inputs() + torch.cuda.synchronize() + for n in range(n_runs): + if n > n_runs / 10: + t0 = time.time() + _output = pipe.unet(*inputs) + torch.cuda.synchronize() + t1 = time.time() + log.info(f"benchmark unet: {t1 - t0:.3f}s passes={n_runs} type={msg}") + return t1 - t0 + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description = 'SD.Next') + parser.add_argument('--model', type=str, default='', required=True, help='model path') + parser.add_argument('--type', type=str, default='sd15', choices=['sd15', 'sdxl'], required=False, help='model type, default: %(default)s') + parser.add_argument('--benchmark', default = False, action='store_true', help = "run benchmarks, default: %(default)s") + parser.add_argument('--trace', default = True, action='store_true', help = "run jit tracing, default: %(default)s") + parser.add_argument('--save', default = False, action='store_true', help = "save optimized unet, default: %(default)s") + args = parser.parse_args() + setup_logging() + log.info('sdnext model jit tracing') + if not os.path.isfile(args.model): + log.error(f"invalid model path: {args.model}") + exit(1) + load_model() + if args.benchmark: + time0 = benchmark_model('original') + unet_saved = trace_model() + if unet_saved is not None: + load_trace(unet_saved) + if args.benchmark: + time1 = benchmark_model('traced') + log.info(f'benchmark speedup: {100 * (time0 - time1) / time0:.3f}%') diff --git a/cli/train.py b/cli/train.py index 5e3c7f99e..84fb3d592 100755 --- a/cli/train.py +++ b/cli/train.py @@ -70,7 +70,6 @@ def parse_args(): global args # pylint: disable=global-statement parser = argparse.ArgumentParser(description = 'SD.Next Train') - group_server = parser.add_argument_group('Server') group_server.add_argument('--server', type=str, default='http://127.0.0.1:7860', required=False, help='server url, default: %(default)s') group_server.add_argument('--user', type=str, default=None, required=False, help='server url, default: %(default)s') diff --git a/javascript/settings.js b/javascript/settings.js index 8eeb50b4a..d9e96dcf8 100644 --- a/javascript/settings.js +++ b/javascript/settings.js @@ -100,11 +100,8 @@ onAfterUiUpdate(async () => { gradioApp().querySelectorAll('#tab_settings .tabitem').forEach((section) => { section.querySelectorAll('.dirtyable').forEach((setting) => { const visible = setting.innerText.toLowerCase().includes(e.target.value.toLowerCase()) || setting.id.toLowerCase().includes(e.target.value.toLowerCase()); - if (!visible) { - setting.style.display = 'none'; - } else { - setting.style.removeProperty('display'); - } + if (!visible) setting.style.display = 'none'; + else setting.style.removeProperty('display'); }); }); }, 50); diff --git a/modules/lora b/modules/lora index 59c9a8e7a..4c6f3125f 160000 --- a/modules/lora +++ b/modules/lora @@ -1 +1 @@ -Subproject commit 59c9a8e7aec9dfc910d4ee2a11cebc01786b5eda +Subproject commit 4c6f3125fce31f86aa949026030e923f3d7d826b diff --git a/modules/processing.py b/modules/processing.py index 800de9fb2..a2c11040d 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -932,7 +932,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): if self.hr_upscaler is not None: self.extra_generation_params["Hires upscaler"] = self.hr_upscaler self.extra_generation_params["Secondary sampler"] = self.latent_sampler - shared.log.debug(f'Init hires: upscaler={self.hr_upscaler} sampler={self.latent_sampler} width={self.hr_upscale_to_x} height={self.hr_upscale_to_y}') + shared.log.debug(f'Init hires: upscaler={self.hr_upscaler} sampler={self.latent_sampler} resize={self.hr_resize_x}x{self.hr_resize_y} upscale={self.hr_upscale_to_x}x{self.hr_upscale_to_y}') def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): diff --git a/modules/sd_models.py b/modules/sd_models.py index 80a8c3339..bb15fd35d 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -145,7 +145,10 @@ def checkpoint_tiles(use_short=False): # pylint: disable=unused-argument def list_models(): checkpoints_list.clear() checkpoint_aliases.clear() - ext_filter=[".safetensors"] if shared.opts.sd_disable_ckpt else [".ckpt", ".safetensors"] + if shared.opts.sd_disable_ckpt or shared.backend == shared.Backend.DIFFUSERS: + ext_filter = [".safetensors"] + else: + ext_filter = [".ckpt", ".safetensors"] model_list = modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"]) if shared.backend == shared.Backend.DIFFUSERS: model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir) @@ -1125,11 +1128,12 @@ def unload_model_weights(op='model'): def apply_token_merging(sd_model, token_merging_ratio=0): current_token_merging_ratio = getattr(sd_model, 'applied_token_merged_ratio', 0) - if token_merging_ratio is None or current_token_merging_ratio == token_merging_ratio: + if token_merging_ratio is None or current_token_merging_ratio is None or current_token_merging_ratio == token_merging_ratio: return if current_token_merging_ratio > 0: tomesd.remove_patch(sd_model) if token_merging_ratio > 0: + shared.log.debug(f'Applying token merging: ratio={token_merging_ratio}') tomesd.apply_patch( sd_model, ratio=token_merging_ratio, diff --git a/modules/ui.py b/modules/ui.py index 7f5159835..af6eb6819 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1018,13 +1018,11 @@ def create_ui(startup_timer = None): with gr.Tabs(elem_id="system") as system_tabs: global ui_system_tabs # pylint: disable=global-statement ui_system_tabs = system_tabs - with gr.TabItem("Settings", id="system_settings", elem_id="system_settings_tab"): - + with gr.TabItem("Settings", id="system_settings", elem_id="tab_settings"): with gr.Row(): settings_submit = gr.Button(value="Apply settings", variant='primary', elem_id="settings_submit") preview_theme = gr.Button(value="Preview theme", variant='primary', elem_id="settings_preview_theme") defaults_submit = gr.Button(value="Restore defaults", variant='primary', elem_id="defaults_submit") - with gr.Row(): _settings_search = gr.Text(label="Search", elem_id="settings_search")