nodes/xmarre-TorchCompileModel_LoRASafe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

docs: clarify Flux MLP leaf-target behavior

Browse Source
This commit is contained in:
xmarre
2026-03-06 04:37:33 +01:00
parent 18a6182d3f
commit 397edb5df7
2 changed files with 324 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+186 -8
View File
@@ -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.
torch_compile_lora_safe.py
+138 -13
View File
@@ -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"]