refactor model compile and add stable-fast

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2023-11-11
+## Update for 2023-11-13
 
 - **Diffusers**  
   - **LCM** support for any *SD 1.5* or *SD-XL* model!  
@@ -13,6 +13,10 @@
     - set cfg to 1 or 2  
   - Add `cli/lcm-convert.py` script to convert any SD 1.5 or SD-XL model to LCM model  
     by baking in LORA and uploading to Huggingface, thanks @Disty0  
+  - Support for [Stable Fast](https://github.com/chengzeyi/stable-fast) model compile on *Windows/Linux/WSL2* with *CUDA*  
+    See [Wiki:Benchmark](https://github.com/vladmandic/automatic/wiki/Benchmark) for details and comparisment  
+    of different backends, precision modes, advanced settings and compile modes  
+    hint: **100+ it/s** on RTX4090 with no special tweaks  
   - Add additional pipeline types for manual model loads when loading from `safetensors`  
   - Updated logic for calculating **steps** when using base/hires/refiner workflows  
   - Safe model offloading for non-standard models  
@@ -25,6 +29,7 @@
 - **General**:  
   - Reworked parser when pasting previously generated images/prompts  
     includes all `txt2img`, `img2img` and `override` params  
+  - Reworked **model compile**
   - Add refiner options to XYZ Grid  
   - Support custom upscalers in subfolders  
   - Support `--ckpt none` to skip loading a model  
@@ -35,8 +40,8 @@
   - Fix img2img init image save  
   - More uniform models paths  
   - Safe scripts callback execution  
-  - Improve extension compatibility  
-  - Improve BF16 support  
+  - Improved extension compatibility  
+  - Improved BF16 support  
 
 ## Update for 2023-11-06
 

cli/run-benchmark.py

@@ -28,8 +28,7 @@ options = Map({
 
 
 # batch = [1, 1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128]
-# batch = [1, 1, 2, 4, 8, 12, 16]
-batch = [4, 4]
+batch = [1, 1, 2, 4, 8, 12, 16]
 oom = 0
 
 

modules/sd_hijack.py

@@ -178,7 +178,7 @@ class StableDiffusionModelHijack:
             except Exception as err:
                 shared.log.warning(f"IPEX Optimize not supported: {err}")
 
-        if opts.cuda_compile and opts.cuda_compile_backend != 'none' and shared.backend == shared.Backend.ORIGINAL:
+        if (opts.cuda_compile or opts.cuda_compile_vae or opts.cuda_compile_upscaler) and shared.opts.cuda_compile_backend != 'none' and shared.backend == shared.Backend.ORIGINAL:
             try:
                 import logging
                 shared.log.info(f"Compiling pipeline={m.model.__class__.__name__} mode={opts.cuda_compile_backend}")
@@ -199,6 +199,8 @@ class StableDiffusionModelHijack:
                     hidet.torch.dynamo_config.use_tensor_core(True)
                     hidet.torch.dynamo_config.search_space(2)
                 m.model = torch.compile(m.model, mode=opts.cuda_compile_mode, backend=opts.cuda_compile_backend, fullgraph=opts.cuda_compile_fullgraph, dynamic=False)
+                from installer import setup_logging
+                setup_logging()
                 shared.log.info("Model complilation done.")
             except Exception as err:
                 shared.log.warning(f"Model compile not supported: {err}")

modules/sd_models.py

@@ -20,8 +20,7 @@ import tomesd
 from transformers import logging as transformers_logging
 import ldm.modules.midas as midas
 from ldm.util import instantiate_from_config
-from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config
-from modules.sd_hijack_inpainting import do_inpainting_hijack
+from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config, sd_models_compile, sd_hijack_inpainting
 from modules.timer import Timer
 from modules.memstats import memory_stats
 from modules.paths import models_path, script_path
@@ -119,18 +118,6 @@ class CheckpointInfo:
         return self.shorthash
 
 
-#Used by OpenVINO, can be used with TensorRT or Olive
-class CompiledModelState:
-    def __init__(self):
-        self.first_pass = True
-        self.height = 512
-        self.width = 512
-        self.batch_size = 1
-        self.partition_id = 0
-        self.cn_model = []
-        self.lora_model = []
-
-
 class NoWatermark:
     def apply_watermark(self, img):
         return img
@@ -594,7 +581,7 @@ model_data = ModelData()
 
 def change_backend():
     shared.log.info(f'Backend changed: {shared.backend}')
-    shared.log.warning('Server restart required to apply all changes')
+    shared.log.warning('Full server restart required to apply all changes')
     if shared.backend == shared.Backend.ORIGINAL:
         change_from = shared.Backend.DIFFUSERS
     else:
@@ -673,55 +660,6 @@ def detect_pipeline(f: str, op: str = 'model'):
     return pipeline, guess
 
 
-def compile_diffusers(sd_model):
-    try:
-        if shared.opts.ipex_optimize:
-            import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
-            sd_model.unet.training = False
-            sd_model.unet = ipex.optimize(sd_model.unet, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
-            if hasattr(sd_model, 'vae'):
-                sd_model.vae.training = False
-                sd_model.vae = ipex.optimize(sd_model.vae, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
-            if hasattr(sd_model, 'movq'):
-                sd_model.movq.training = False
-                sd_model.movq = ipex.optimize(sd_model.movq, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
-            shared.log.info("Applied IPEX Optimize.")
-    except Exception as err:
-        shared.log.warning(f"IPEX Optimize not supported: {err}")
-
-    try:
-        if (shared.opts.cuda_compile or shared.opts.cuda_compile_vae or shared.opts.cuda_compile_upscaler) and shared.opts.cuda_compile_backend != 'none':
-            shared.log.info(f"Compiling pipeline={sd_model.__class__.__name__} shape={8 * sd_model.unet.config.sample_size} mode={shared.opts.cuda_compile_backend}")
-            import torch._dynamo # pylint: disable=unused-import,redefined-outer-name
-            if shared.opts.cuda_compile_backend == "openvino_fx":
-                torch._dynamo.reset() # pylint: disable=protected-access
-                from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import
-                openvino_clear_caches()
-                torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access
-                if shared.compiled_model_state is None:
-                    shared.compiled_model_state = CompiledModelState()
-                shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False
-            log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access
-            if hasattr(torch, '_logging'):
-                torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access
-            torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access
-            torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access
-            if shared.opts.cuda_compile:
-                sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init
-            if shared.opts.cuda_compile_vae:
-                if hasattr(sd_model, 'vae'):
-                    sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init
-                if hasattr(sd_model, 'movq'):
-                    sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init
-            from installer import setup_logging
-            setup_logging()
-            if shared.opts.cuda_compile_precompile:
-                sd_model("dummy prompt")
-            shared.log.info("Complilation done.")
-    except Exception as err:
-        shared.log.warning(f"Model compile not supported: {err}")
-
-
 def set_diffuser_options(sd_model, vae, op: str):
     if sd_model is None:
         shared.log.warning(f'{op} is not loaded')
@@ -778,10 +716,8 @@ def set_diffuser_options(sd_model, vae, op: str):
             sd_model.vae = vae
         if shared.opts.diffusers_vae_upcast != 'default':
             if shared.opts.diffusers_vae_upcast == 'true':
-                # sd_model.vae.config["force_upcast"] = True
                 sd_model.vae.config.force_upcast = True
             else:
-                # sd_model.vae.config["force_upcast"] = False
                 sd_model.vae.config.force_upcast = False
             if shared.opts.no_half_vae:
                 devices.dtype_vae = torch.float32
@@ -789,6 +725,18 @@ def set_diffuser_options(sd_model, vae, op: str):
         shared.log.debug(f'Setting {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}')
     if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'):
         sd_model.enable_xformers_memory_efficient_attention()
+
+    if shared.opts.diffusers_eval:
+        if hasattr(sd_model, "unet"):
+            sd_model.unet.requires_grad_(False)
+            sd_model.unet.eval()
+        if hasattr(sd_model, "vae"):
+            sd_model.vae.requires_grad_(False)
+            sd_model.vae.eval()
+        if hasattr(sd_model, "text_encoder"):
+            sd_model.text_encoder.requires_grad_(False)
+            sd_model.text_encoder.eval()
+
     if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'):
         shared.log.debug(f'Setting {op}: enable channels last')
         sd_model.unet.to(memory_format=torch.channels_last)
@@ -976,7 +924,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             elif not getattr(sd_model, 'has_accelerate', False):
                 sd_model.to(devices.device)
 
-            compile_diffusers(sd_model)
+            sd_models_compile.compile_diffusers(sd_model)
 
         if sd_model is None:
             shared.log.error('Diffuser model not loaded')
@@ -1106,7 +1054,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
             current_checkpoint_info = model_data.sd_refiner.sd_checkpoint_info
             unload_model_weights(op=op)
 
-    do_inpainting_hijack()
+    sd_hijack_inpainting.do_inpainting_hijack()
     devices.set_cuda_params()
     if already_loaded_state_dict is not None:
         state_dict = already_loaded_state_dict

modules/sd_models_compile.py

@@ -0,0 +1,130 @@
+import time
+import logging
+import torch
+from modules import shared, devices
+from installer import setup_logging
+
+
+#Used by OpenVINO, can be used with TensorRT or Olive
+class CompiledModelState:
+    def __init__(self):
+        self.first_pass = True
+        self.height = 512
+        self.width = 512
+        self.batch_size = 1
+        self.partition_id = 0
+        self.cn_model = []
+        self.lora_model = []
+
+
+def optimize_ipex(sd_model):
+    try:
+        t0 = time.time()
+        import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
+        sd_model.unet.training = False
+        sd_model.unet = ipex.optimize(sd_model.unet, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
+        if hasattr(sd_model, 'vae'):
+            sd_model.vae.training = False
+            sd_model.vae = ipex.optimize(sd_model.vae, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
+        if hasattr(sd_model, 'movq'):
+            sd_model.movq.training = False
+            sd_model.movq = ipex.optimize(sd_model.movq, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
+        t1 = time.time()
+        shared.log.info(f"Model compile: mode=IPEX-optimize time={t1-t0:.2f}")
+    except Exception as e:
+        shared.log.warning(f"Model compile: task=IPEX-optimize error: {e}")
+
+
+def optimize_openvino():
+    try:
+        from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import
+        openvino_clear_caches()
+        torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access
+        if shared.compiled_model_state is None:
+            shared.compiled_model_state = CompiledModelState()
+        shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False
+    except Exception as e:
+        shared.log.warning(f"Model compile: task=OpenVINO: {e}")
+
+
+def compile_stablefast(sd_model):
+    try:
+        import sfast.compilers.stable_diffusion_pipeline_compiler as sf
+    except Exception as e:
+        shared.log.warning(f'Model compile using stable-fast: {e}')
+        return sd_model
+    config = sf.CompilationConfig.Default()
+    try:
+        import xformers # pylint: disable=unused-import
+        config.enable_xformers = True
+    except Exception:
+        pass
+    try:
+        import triton # pylint: disable=unused-import
+        config.enable_triton = True
+    except Exception:
+        pass
+    import warnings
+    warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
+    config.enable_cuda_graph = shared.opts.cuda_compile_fullgraph
+    config.enable_jit_freeze = shared.opts.diffusers_eval
+    try:
+        t0 = time.time()
+        sd_model = sf.compile(sd_model, config)
+        setup_logging() # compile messes with logging so reset is needed
+        if shared.opts.cuda_compile_precompile:
+            sd_model("dummy prompt")
+        t1 = time.time()
+        shared.log.info(f"Model compile: task=Stable-fast config={config.__dict__} time={t1-t0:.2f}")
+    except Exception as e:
+        shared.log.info(f"Model compile: task=Stable-fast error: {e}")
+    return sd_model
+
+
+def compile_torch(sd_model):
+    try:
+        import torch._dynamo # pylint: disable=unused-import,redefined-outer-name
+        torch._dynamo.reset() # pylint: disable=protected-access
+        if shared.opts.ipex_optimize:
+            optimize_ipex(sd_model)
+        if shared.opts.cuda_compile_backend == "openvino_fx":
+            optimize_openvino()
+        log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access
+        if hasattr(torch, '_logging'):
+            torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access
+        torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access
+        torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access
+        t0 = time.time()
+        if shared.opts.cuda_compile:
+            sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+        if shared.opts.cuda_compile_vae:
+            if hasattr(sd_model, 'vae'):
+                sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+            if hasattr(sd_model, 'movq'):
+                sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+        setup_logging() # compile messes with logging so reset is needed
+        if shared.opts.cuda_compile_precompile:
+            sd_model("dummy prompt")
+        t1 = time.time()
+        shared.log.info(f"Model compile: time={t1-t0:.2f}")
+    except Exception as e:
+        shared.log.warning(f"Model compile error: {e}")
+    return sd_model
+
+
+def compile_diffusers(sd_model):
+    if not (shared.opts.cuda_compile or shared.opts.cuda_compile_vae or shared.opts.cuda_compile_upscaler):
+        return
+    if not hasattr(sd_model, 'unet') or not hasattr(sd_model.unet, 'config'):
+        shared.log.warning('Model compile enabled but model has no Unet')
+        return
+    if shared.opts.cuda_compile_backend == 'none':
+        shared.log.warning('Model compile enabled but no backend specified')
+        return
+    size = 8*getattr(sd_model.unet.config, 'sample_size', 0)
+    shared.log.info(f"Model compile: pipeline={sd_model.__class__.__name__} shape={size} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} fullgraph={shared.opts.cuda_compile_fullgraph} unet={shared.opts.cuda_compile} vae={shared.opts.cuda_compile_vae} upscaler={shared.opts.cuda_compile_upscaler}")
+    if shared.opts.cuda_compile_backend == 'stable-fast':
+        sd_model = compile_stablefast(sd_model)
+    else:
+        sd_model = compile_torch(sd_model)
+    return sd_model

modules/shared.py

@@ -288,10 +288,10 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "cuda_compile": OptionInfo(True if cmd_opts.use_openvino else False, "Compile UNet"),
     "cuda_compile_vae": OptionInfo(True if cmd_opts.use_openvino else False, "Compile VAE"),
     "cuda_compile_upscaler": OptionInfo(True if cmd_opts.use_openvino else False, "Compile upscaler"),
-    "cuda_compile_backend": OptionInfo("openvino_fx" if cmd_opts.use_openvino else "none", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx']}),
+    "cuda_compile_backend": OptionInfo("openvino_fx" if cmd_opts.use_openvino else "none", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx', 'stable-fast']}),
     "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune']}),
     "cuda_compile_fullgraph": OptionInfo(False, "Model compile fullgraph"),
-    "cuda_compile_precompile": OptionInfo(False, "Model compile precompile"),
+    "cuda_compile_precompile": OptionInfo(True, "Model compile precompile"),
     "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"),
     "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"),
 
@@ -348,6 +348,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_vae_load_variant": OptionInfo("default", "Diffusers VAE loading variant", gr.Radio, {"choices": ['default', 'fp32', 'fp16']}),
     "custom_diffusers_pipeline": OptionInfo('', 'Load custom Diffusers pipeline'),
     "diffusers_lora_loader": OptionInfo("diffusers" if cmd_opts.use_openvino else "sequential apply", "Diffusers LoRA loading variant", gr.Radio, {"choices": ['diffusers', 'sequential apply', 'merge and apply']}),
+    "diffusers_eval": OptionInfo(True, "Force model eval"),
     "diffusers_force_zeros": OptionInfo(True, "Force zeros for prompts when empty"),
     "diffusers_aesthetics_score": OptionInfo(False, "Require aesthetics score"),
     "diffusers_force_inpaint": OptionInfo(False, 'Diffusers force inpaint pipeline'),