Update README for cudagraph modes

README.md

@@ -64,14 +64,24 @@ Strict Flux names targeted by default:
 - `single_blocks.{i}.linear2`
 
 For cudagraph-capable Flux runs (`backend = cudagraphs`, or `backend = inductor`
-with `disable_cudagraphs = false`), the target set is narrowed to the
-cudagraph-safe qkv subset only:
+with `disable_cudagraphs = false`), the behavior splits into two safe modes:
 
-- `double_blocks.{i}.img_attn.qkv`
-- `double_blocks.{i}.txt_attn.qkv`
+- `compile_transformer_only = true`
+  - target set is narrowed to the cudagraph-safe qkv subset only:
+    - `double_blocks.{i}.img_attn.qkv`
+    - `double_blocks.{i}.txt_attn.qkv`
+- `compile_transformer_only = false`
+  - the node does **not** compile the outer `diffusion_model` wrapper path.
+  - instead, it widens to a **Flux full-leaf compile** across the model's
+    parameterized leaf modules.
 
-In that mode, inferred Flux targets and full-model fallback are intentionally
-disabled to avoid widening back into unstable replay paths.
+This avoids full-model cudagraph capture through Flux / ComfyUI
+`WrapperExecutor` paths, which are prone to graph-break / replay failures when
+runtime patches such as Sage attention are active.
+
+In cudagraph-capable Flux mode, inferred Flux targets and full-model fallback
+are intentionally disabled to avoid widening back into unstable outer-wrapper
+paths.
 
 Optional widening controls:
 
@@ -145,12 +155,15 @@ If no known transformer compile targets are found, behavior is controlled by
   - Set to `false` only when you specifically want Inductor cudagraphs.
 
 - `compile_transformer_only` (`BOOLEAN`)
-  - `false` (default): compile the entire `diffusion_model`.
+  - `false` (default): compile the entire `diffusion_model` for non-Flux
+    models.
   - `true`: compile discovered transformer targets only.
   - Flux behavior (default strict mode): targets only `qkv/proj` attention
     leaves plus `linear1/linear2` leaves.
-  - For cudagraph-capable Flux runs, discovery is narrowed further to
-    `img_attn.qkv` and `txt_attn.qkv` only.
+  - For cudagraph-capable Flux runs:
+    - `true` -> narrowed further to `img_attn.qkv` and `txt_attn.qkv` only.
+    - `false` -> widens to a Flux **full-leaf compile** instead of compiling
+      the outer `diffusion_model` wrapper path.
 
 - `allow_flux_inferred_targets` (`BOOLEAN`, default `false`)
   - `false` (recommended): strict Flux target set only.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name        = "TorchCompileModel_LoRASafe"
 description = "LoRA‐Safe TorchCompile node for ComfyUI"
-version     = "0.1.6"
+version     = "0.1.7"
 license     = { file = "LICENSE" }
 dependencies = []
 

torch_compile_lora_safe.py

@@ -198,6 +198,15 @@ class TorchCompileModel_LoRASafe:
             return False
         return any(True for _ in module.parameters(recurse=False))
 
+    @classmethod
+    def _discover_flux_all_leaf_keys(cls, diffusion_model: nn.Module):
+        keys = []
+        for name, module in diffusion_model.named_modules():
+            if not name or not cls._is_leaf_compile_module(module):
+                continue
+            keys.append(f"diffusion_model.{name}")
+        return list(dict.fromkeys(keys))
+
     @staticmethod
     def _is_flux_model(diffusion_model: nn.Module) -> bool:
         return hasattr(diffusion_model, "double_blocks") and hasattr(
@@ -411,8 +420,9 @@ class TorchCompileModel_LoRASafe:
             backend, mode, fullgraph, dynamic, disable_cudagraphs
         )
 
+        is_flux_model = self._is_flux_model(dm)
+
         if compile_transformer_only:
-            is_flux_model = self._is_flux_model(dm)
             compile_keys = self._discover_compile_keys(
                 dm,
                 allow_flux_inferred_targets=(
@@ -441,12 +451,31 @@ class TorchCompileModel_LoRASafe:
                             "skipping compile instead of falling back to whole model."
                         )
                     return (m,)
+        elif cudagraph_capable and is_flux_model:
+            compile_keys = self._discover_flux_all_leaf_keys(dm)
+            if not compile_keys:
+                _LOG.warning(
+                    "TorchCompileModel_LoRASafe: Flux full-leaf cudagraph compile requested, "
+                    "but no leaf targets were found; skipping compile instead of compiling the "
+                    "whole WrapperExecutor-driven diffusion_model path."
+                )
+                return (m,)
+            _LOG.warning(
+                "TorchCompileModel_LoRASafe: Flux whole-model cudagraph compile requested; "
+                "using wide leaf-module compile instead to avoid WrapperExecutor / "
+                "torch.compiler.disable replay failures."
+            )
         else:
             compile_keys = ["diffusion_model"]
 
-        if cudagraph_capable and compile_transformer_only and is_flux_model:
-            _LOG.warning("Flux cudagraph targets: %s", compile_keys[:16])
-            _LOG.warning("Flux cudagraph target count: %d", len(compile_keys))
+        if cudagraph_capable and is_flux_model:
+            label = (
+                "Flux cudagraph transformer targets"
+                if compile_transformer_only
+                else "Flux cudagraph wide-leaf targets"
+            )
+            _LOG.warning("%s: %s", label, compile_keys[:16])
+            _LOG.warning("%s count: %d", label, len(compile_keys))
 
         if _set_torch_compile_wrapper is not None:
             try: