test model jit tracing

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}%')

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')

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);

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):
 

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,

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")