diff --git a/README.md b/README.md index 5e7ca9e..f94f89f 100644 --- a/README.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 77d747a..8cd52a2 100644 --- a/pyproject.toml +++ b/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 = [] diff --git a/torch_compile_lora_safe.py b/torch_compile_lora_safe.py index b5e997e..77a3894 100644 --- a/torch_compile_lora_safe.py +++ b/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: