From 7187fdd5321863c2ed66d0d7f187206b263f2126 Mon Sep 17 00:00:00 2001 From: Haoming Date: Mon, 11 May 2026 14:49:38 +0800 Subject: [PATCH] compile --- .../sd_forge_compile/scripts/compile.py | 94 +++++++++---------- 1 file changed, 42 insertions(+), 52 deletions(-) diff --git a/extensions-builtin/sd_forge_compile/scripts/compile.py b/extensions-builtin/sd_forge_compile/scripts/compile.py index 902859b7..9056e3f1 100644 --- a/extensions-builtin/sd_forge_compile/scripts/compile.py +++ b/extensions-builtin/sd_forge_compile/scripts/compile.py @@ -1,4 +1,7 @@ +# https://github.com/Comfy-Org/ComfyUI/blob/master/comfy_extras/nodes_torch_compile.py + import logging +from functools import wraps from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -19,12 +22,14 @@ except ImportError: else: TRITON_AVAILABLE = True +_COMPILE_CONFIG_KEY = "_torch_compile_config" +_ORIG_APPLY_KEY = "_orig_apply_model" + logger = logging.getLogger("compile") setup_logger(logger) def skip_torch_compile_dict(guard_entries): - # https://github.com/Comfy-Org/ComfyUI/blob/master/comfy_extras/nodes_torch_compile.py#L5 return [("transformer_options" not in entry.name) for entry in guard_entries] @@ -62,66 +67,38 @@ class TorchCompileForForge(scripts.Script): _indynamic = "Require recompilation if Resolution / Batch Size is changed" _no_malloc = "Does not work with --cuda-malloc" - gr.Markdown( - rf""" + gr.Markdown(rf""" **torch.compile** speeds up the inference by compiling the model ahead of time - **guard_filter_fn:** Compile the Fastest ; {_indynamic} - **dynamic:** {_dynamic} ; Slower to Compile - **max-autotune:** Best Runtime Speed ; {_indynamic} ; {_no_malloc} - **max-autotune-no-cudagraphs:** {_dynamic} ; Faster than **dynamic** ; Even Slower to Compile - **reduce-overhead:** Similar to **max-autotune** ; {_indynamic} ; {_no_malloc} - """ - ) + """) return [preset] - @staticmethod - def restore(kmodel: "KModel"): - model = get_attr(kmodel, "_model_backup") - set_attr_raw(kmodel, "diffusion_model", model) - del kmodel._compile_config - del kmodel._compiled_backup - del kmodel._model_backup - - def before_process_batch(self, p, *args, **kwargs): - # temporarily restores the original model so LoRA can apply - # (otherwise "keys mismatched") - - kmodel: "KModel" = p.sd_model.forge_objects.unet.model - if not hasattr(kmodel, "_compile_config"): - return - - c_model = get_attr(kmodel, "diffusion_model") - set_attr_raw(kmodel, "_compiled_backup", c_model) - model = get_attr(kmodel, "_model_backup") - set_attr_raw(kmodel, "diffusion_model", model) - def process_batch(self, p, preset: str, **kwargs): kmodel: "KModel" = p.sd_model.forge_objects.unet.model - compiled: tuple[str, str] = getattr(kmodel, "_compile_config", None) - enable: bool = (compiled is not None) if preset == "Automatic" else (preset != "Disable") + prev_config: tuple[str] = getattr(kmodel, _COMPILE_CONFIG_KEY, None) + enable: bool = (prev_config is not None) if preset == "Automatic" else (preset != "Disable") if not enable: - if compiled is not None: - self.restore(kmodel) + self._remove_compile_wrapper(kmodel) return if preset in ("max-autotune", "reduce-overhead") and cmd_args.cuda_malloc: logger.error(f"{preset} does not support --cuda-malloc\nModel is not compiled...") return - _config: tuple[str, str] = (preset, p.sd_model.current_lora_hash) + _config: tuple[str] = (preset,) + if _config == prev_config: + return - if compiled is not None: - if preset in (compiled[0], "Automatic") and _config[1] == compiled[1]: - _model = get_attr(kmodel, "_compiled_backup") - set_attr_raw(kmodel, "diffusion_model", _model) - del kmodel._compiled_backup - return + setattr(kmodel, _COMPILE_CONFIG_KEY, _config) - self.restore(kmodel) - - kmodel._compile_config = _config + if prev_config is not None: + self._remove_compile_wrapper(kmodel) match preset: case "guard_filter_fn": @@ -135,17 +112,30 @@ class TorchCompileForForge(scripts.Script): case "reduce-overhead": config = dict(backend="inductor", mode="reduce-overhead", dynamic=False, fullgraph=False) - kmodel: "KModel" = p.sd_model.forge_objects.unet.detach() - model = get_attr(kmodel, "diffusion_model") - set_attr_raw(kmodel, "_model_backup", model) - - # patch LoRA ahead-of-time - p.sd_model.forge_objects.unet.refresh_loras() - - set_attr_raw( - kmodel, - "diffusion_model", - torch.compile(model, **config), - ) + self._wrap_apply_model(kmodel, config) logger.info(f"Model Compiled ({preset})") + + @staticmethod + def _wrap_apply_model(kmodel: "KModel", compile_config: dict): + setattr(kmodel, _ORIG_APPLY_KEY, kmodel.apply_model) + + @wraps(kmodel._orig_apply_model) + def apply_model_with_compile(*args, **kwargs): + orig_model = get_attr(kmodel, "diffusion_model") + compiled = torch.compile(orig_model, **compile_config) + set_attr_raw(kmodel, "diffusion_model", compiled) + try: + return kmodel._orig_apply_model(*args, **kwargs) + finally: + set_attr_raw(kmodel, "diffusion_model", orig_model) + + kmodel.apply_model = apply_model_with_compile + + @staticmethod + def _remove_compile_wrapper(kmodel: "KModel"): + if (orig := getattr(kmodel, _ORIG_APPLY_KEY, None)) is not None: + kmodel.apply_model = orig + delattr(kmodel, _ORIG_APPLY_KEY) + delattr(kmodel, _COMPILE_CONFIG_KEY) + logger.info("Model Decompiled")