compile

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