diff --git a/README.md b/README.md index 1205487..ecb7c85 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,189 @@ # LoRA-Safe TorchCompile Node for ComfyUI -Drop-in replacement for the stock `TorchCompileModel` node. -Compiles after all LoRA / TEA-Cache / Sage-Attention patches, so LoRAs stay -active while you still get the compile speed-up. +`TorchCompileModel_LoRASafe` is a drop-in replacement for ComfyUI's stock +`TorchCompileModel` node, designed for workflows that rely on model patching +(e.g., LoRA stacks, TEA-Cache, Sage-Attention, and related runtime patches). -## Install -1. Clone / download this repo. -2. Copy **lora_safe_compile** into `ComfyUI/custom_nodes/`. -3. Restart ComfyUI and search for **TorchCompileModel_LoRASafe** under - *model / optimisation 🛠️*. +The main difference from the stock node is **when** compilation is applied: +this node compiles after patching so your patches stay active instead of being +silently bypassed. + +--- + +## Why this exists + +In many optimized ComfyUI workflows, the diffusion model is patched before +sampling starts. If compile wrapping happens at the wrong stage, you can end up +compiling a pre-patch object and losing expected patch behavior. + +This node is built to avoid that failure mode: + +- Patch first +- Compile lazily on first real execution +- Keep patched behavior intact while still benefiting from `torch.compile` + +--- + +## Feature summary + +- **LoRA-safe compile flow** for patched diffusion models +- **Lazy compilation** (first forward pass triggers compile) +- **Whole-model compile** by default +- **Transformer-only compile mode** for architectures where compiling only block + lists is more stable or faster to warm up +- **ComfyUI-aligned stability behavior**: attempts `model.clone(disable_dynamic=True)` + and (when available) uses ComfyUI's `guard_filter_fn` to avoid unnecessary + guards on dynamic `transformer_options` dictionaries +- Backend support: + - `inductor` + - `cudagraphs` (CUDA required; generally best on newer GPUs) + - `nvfuser` (CUDA required) + +--- + +## FLUX support (including FLUX.2 KLEIN 9B) + +This node includes Flux-specific compile target discovery intended to improve +compatibility with **FLUX.2 KLEIN 9B** and other Flux-style workflows. + +For Flux-style models in ComfyUI, block forward passes are highly dynamic +(`transformer_options` mutation, dynamic patch replacement, and patch hooks +around attention). Because of that, this node **does not** compile whole +`double_blocks.{i}` / `single_blocks.{i}` forwards. + +Instead, when `compile_transformer_only = true`, Flux targets leaf tensor-heavy +submodules. It first tries preferred names and then falls back to inferred leaf +modules from each block when naming differs across ComfyUI builds. + +Preferred Flux names attempted first: + +- `double_blocks.{i}.img_attn.qkv` +- `double_blocks.{i}.img_attn.proj` +- `double_blocks.{i}.txt_attn.qkv` +- `double_blocks.{i}.txt_attn.proj` +- `double_blocks.{i}.img_mlp` +- `double_blocks.{i}.txt_mlp` +- `single_blocks.{i}.linear1` +- `single_blocks.{i}.linear2` + +Note: MLP containers (`img_mlp`/`txt_mlp`) are often not leaf modules in Flux +builds, so the node may compile inferred leaf submodules inside those MLP paths +instead of compiling the container module directly. + +Fallback inference (when preferred names are missing, or when preferred names +resolve to non-leaf containers) looks for leaf modules matching attention +`qkv/proj`, `img_mlp/txt_mlp`, and `linear1/linear2` patterns inside Flux +blocks. + +For non-Flux transformer-style models, it compiles known block containers such +as: + +- `layers` +- `transformer_blocks` +- `blocks` +- `visual_transformer_blocks` +- `text_transformer_blocks` + +If no known transformer compile targets are found, the node safely falls back +to compiling `diffusion_model` as a whole. + +--- + +## Installation + +1. Clone or download this repository. +2. Copy the `lora_safe_compile` folder into your ComfyUI `custom_nodes/` directory. +3. Restart ComfyUI. +4. Search for **TorchCompileModel_LoRASafe** in: + **model / optimisation 🛠️** + +--- + +## Node inputs (detailed) + +- `model` (`MODEL`) + - The model object to patch and compile. + +- `backend` (`inductor`, `cudagraphs`, `nvfuser`) + - `inductor`: default recommendation for most users. + - `cudagraphs`: CUDA-only; can reduce overhead in stable-shape loops. + - `nvfuser`: CUDA-only; may vary by PyTorch/CUDA environment. + +- `mode` (`default`, `reduce-overhead`, `max-autotune`) + - Passed through to `torch.compile` mode behavior. + +- `fullgraph` (`BOOLEAN`) + - Requests full-graph capture. Can improve performance in some setups, but + may reduce tolerance for graph breaks. + +- `dynamic` (`BOOLEAN`) + - Enables dynamic-shape aware compilation behavior. + - Note: the node still attempts `model.clone(disable_dynamic=True)` first to + match ComfyUI stock compile behavior when supported by the model object. + +- `disable_cudagraphs` (`BOOLEAN`, default `true`) + - For `backend = inductor`, adds + `torch.compile(..., options={'triton.cudagraphs': False})`. + - Not applied to the `cudagraphs` backend. + - Recommended if you hit `cudaMallocAsync` / cudagraph instability in + Inductor-based runs. + +- `compile_transformer_only` (`BOOLEAN`) + - `false` (default): compile the entire `diffusion_model`. + - `true`: compile discovered transformer targets only. + - Flux behavior: compiles leaf heavy submodules (qkv/proj/mlp/linear) instead + of full block forwards to avoid compiling across dynamic hook-heavy paths. + - Fallback behavior: if no recognized transformer compile targets are + detected, compile + `diffusion_model` anyway. + +--- + +## Recommended starting points + +If you're unsure where to begin: + +- Backend: `inductor` +- `disable_cudagraphs`: `true` (recommended default for stability) +- Mode: `default` +- `fullgraph`: `false` +- `dynamic`: `false` +- `compile_transformer_only`: `true` for FLUX-style models (including + FLUX.2 KLEIN 9B) to target leaf heavy modules; otherwise start with `false` + and compare. + +Then iterate one setting at a time based on stability and speed. + +--- + +## Performance and stability notes + +- First generation after graph changes can be slower (compile warm-up). +- Recompiles can occur when shapes/settings change substantially. +- Best backend can differ across GPUs, drivers, and PyTorch versions. +- If a backend fails, switch to `inductor` first. +- If you encounter `cudaMallocAsync`/cudagraph issues, keep + `disable_cudagraphs = true`. + +- Guard filtering: when ComfyUI's `skip_torch_compile_dict` is available, this + node passes it as `options['guard_filter_fn']` to reduce guard churn on + dynamic `transformer_options` dictionaries. +- If full-model compile is unstable on your graph, enable + `compile_transformer_only`. + +--- + +## Troubleshooting quick guide + +- **Issue:** Backend error on startup or first sample. + - Try `backend = inductor`. + - Keep `disable_cudagraphs = true`. + +- **Issue:** Very long first run. + - Expected in compile warm-up; test subsequent runs before evaluating. + +- **Issue:** Unstable behavior with full model compile. + - Set `compile_transformer_only = true`. + +- **Issue:** CUDA backend unavailable. + - Use `inductor`, or verify CUDA-compatible PyTorch/GPU environment. diff --git a/torch_compile_lora_safe.py b/torch_compile_lora_safe.py index 0cae1fa..a762011 100644 --- a/torch_compile_lora_safe.py +++ b/torch_compile_lora_safe.py @@ -1,7 +1,8 @@ import logging +import threading + import torch import torch.nn as nn -import threading _LOG = logging.getLogger(__name__) @@ -10,6 +11,11 @@ try: except Exception: _set_torch_compile_wrapper = None +try: + from comfy_extras.nodes_torch_compile import skip_torch_compile_dict as _skip_torch_compile_dict +except Exception: + _skip_torch_compile_dict = None + class _LazyCompiled(nn.Module): """ @@ -73,9 +79,7 @@ class _LazyCompiled(nn.Module): class TorchCompileModel_LoRASafe: """LoRA-safe torch.compile that also works with Flux block layouts.""" - _BLOCK_LAYER_TYPES = ( - "double_blocks", - "single_blocks", + _GENERIC_BLOCK_LAYER_TYPES = ( "layers", "transformer_blocks", "blocks", @@ -83,6 +87,17 @@ class TorchCompileModel_LoRASafe: "text_transformer_blocks", ) + _FLUX_DOUBLE_PREFERRED_SUFFIXES = ( + "img_attn.qkv", + "img_attn.proj", + "txt_attn.qkv", + "txt_attn.proj", + "img_mlp", + "txt_mlp", + ) + + _FLUX_SINGLE_PREFERRED_SUFFIXES = ("linear1", "linear2") + @classmethod def INPUT_TYPES(cls): return { @@ -92,13 +107,22 @@ class TorchCompileModel_LoRASafe: "mode": (["default", "reduce-overhead", "max-autotune"],), "fullgraph": ("BOOLEAN", {"default": False}), "dynamic": ("BOOLEAN", {"default": False}), + "disable_cudagraphs": ( + "BOOLEAN", + { + "default": True, + "tooltip": "True -> pass torch.compile options {'triton.cudagraphs': False} " + "for inductor backend. Useful to avoid cudaMallocAsync/cudagraph issues " + "in some environments.", + }, + ), "compile_transformer_only": ( "BOOLEAN", { "default": False, - "tooltip": "True -> compile known transformer block lists only. " - "Flux models use double_blocks/single_blocks; " - "SD-style models often use transformer_blocks/blocks.", + "tooltip": "True -> compile discovered transformer targets only. " + "Flux models compile leaf submodules inside double_blocks/single_blocks; " + "SD-style models often compile transformer_blocks/blocks.", }, ), } @@ -124,10 +148,88 @@ class TorchCompileModel_LoRASafe: f"(yours is {cap[0]}.{cap[1]})." ) + @staticmethod + def _is_leaf_compile_module(module: nn.Module) -> bool: + if not isinstance(module, nn.Module): + return False + if any(True for _ in module.children()): + return False + return any(True for _ in module.parameters(recurse=False)) + + @classmethod + def _try_preferred_suffixes(cls, block: nn.Module, suffixes): + targets = [] + for suffix in suffixes: + obj = block + ok = True + for part in suffix.split("."): + if not hasattr(obj, part): + ok = False + break + obj = getattr(obj, part) + if ok and cls._is_leaf_compile_module(obj): + targets.append(suffix) + return targets + + @classmethod + def _discover_flux_double_block_suffixes(cls, block: nn.Module): + preferred = cls._try_preferred_suffixes(block, cls._FLUX_DOUBLE_PREFERRED_SUFFIXES) + if preferred: + return preferred + + inferred = [] + for name, module in block.named_modules(): + if not name or not cls._is_leaf_compile_module(module): + continue + lname = name.lower() + parts = lname.split(".") + tail = parts[-1] + if ( + (("img_attn" in lname or "txt_attn" in lname) and tail in {"qkv", "proj"}) + or ("img_mlp" in lname) + or ("txt_mlp" in lname) + ): + inferred.append(name) + return list(dict.fromkeys(inferred)) + + @classmethod + def _discover_flux_single_block_suffixes(cls, block: nn.Module): + preferred = cls._try_preferred_suffixes(block, cls._FLUX_SINGLE_PREFERRED_SUFFIXES) + if preferred: + return preferred + + inferred = [] + for name, module in block.named_modules(): + if not name or not cls._is_leaf_compile_module(module): + continue + lname = name.lower() + if lname.endswith("linear1") or lname.endswith("linear2"): + inferred.append(name) + return list(dict.fromkeys(inferred)) + @classmethod def _discover_compile_keys(cls, diffusion_model): keys = [] - for layer_name in cls._BLOCK_LAYER_TYPES: + + # Flux special-case: compile only leaf tensor-heavy modules, not full block forward. + if hasattr(diffusion_model, "double_blocks") and hasattr(diffusion_model, "single_blocks"): + double_blocks = getattr(diffusion_model, "double_blocks") + single_blocks = getattr(diffusion_model, "single_blocks") + + if isinstance(double_blocks, (nn.ModuleList, list, tuple)): + for i, block in enumerate(double_blocks): + for suffix in cls._discover_flux_double_block_suffixes(block): + keys.append(f"diffusion_model.double_blocks.{i}.{suffix}") + + if isinstance(single_blocks, (nn.ModuleList, list, tuple)): + for i, block in enumerate(single_blocks): + for suffix in cls._discover_flux_single_block_suffixes(block): + keys.append(f"diffusion_model.single_blocks.{i}.{suffix}") + + if keys: + return list(dict.fromkeys(keys)) + + for layer_name in cls._GENERIC_BLOCK_LAYER_TYPES: if hasattr(diffusion_model, layer_name): blocks = getattr(diffusion_model, layer_name) if not isinstance(blocks, (nn.ModuleList, list, tuple)): @@ -137,7 +239,7 @@ class TorchCompileModel_LoRASafe: return list(dict.fromkeys(keys)) @staticmethod - def _build_compile_kwargs(backend, mode, fullgraph, dynamic): + def _build_compile_kwargs(backend, mode, fullgraph, dynamic, disable_cudagraphs): compile_kw = dict( backend=backend, fullgraph=fullgraph, @@ -145,21 +247,44 @@ class TorchCompileModel_LoRASafe: ) if mode != "default": compile_kw["mode"] = mode + + options = {} + if _skip_torch_compile_dict is not None: + options["guard_filter_fn"] = _skip_torch_compile_dict + if disable_cudagraphs and backend == "inductor": + options["triton.cudagraphs"] = False + if options: + compile_kw["options"] = options + return compile_kw - def patch(self, model, backend, mode, fullgraph, dynamic, compile_transformer_only): + def patch( + self, + model, + backend, + mode, + fullgraph, + dynamic, + disable_cudagraphs, + compile_transformer_only, + ): self._check_backend(backend) - m = model.clone() + try: + m = model.clone(disable_dynamic=True) + except TypeError: + m = model.clone() dm = m.get_model_object("diffusion_model") - compile_kw = self._build_compile_kwargs(backend, mode, fullgraph, dynamic) + compile_kw = self._build_compile_kwargs( + backend, mode, fullgraph, dynamic, disable_cudagraphs + ) if compile_transformer_only: compile_keys = self._discover_compile_keys(dm) if not compile_keys: _LOG.warning( - "TorchCompileModel_LoRASafe: no known transformer block lists found; " + "TorchCompileModel_LoRASafe: no known transformer compile targets found; " "falling back to whole diffusion model compile." ) compile_keys = ["diffusion_model"]