TorchCompileModel_LoRASafe

torch.compile is one of the best free speedups in ComfyUI - the Inductor backend trims a real chunk off your iteration time on most GPUs, which is why it shows up in so many optimized video and Flux workflows alongside TEA-Cache and Sage-Attention. But there's a catch that's quietly wasted hours of people's time: ComfyUI's stock TorchCompileModel node compiles the model before your LoRA stack, TEA-Cache hooks, Sage-Attention patches and the rest get a chance to inject their weights. Those patches end up outside the compiled graph. Your image still renders. Your LoRA just… doesn't load. No error, no warning - you get the base model with a compile speedup and a result that doesn't match the workflow you downloaded.

TorchCompileModel_LoRASafe is a drop-in replacement for that stock node, written precisely because of that failure mode (the author, marres on Reddit / xmarre on GitHub, hit it himself on Wan 2.1). The fix is about when compilation happens: instead of compiling at patch time, it compiles lazily on the first real forward pass - by which point every LoRA, TEA-Cache hook and Sage-Attention patch is already sitting in the model. Speedup and patches both. The release thread's money quote came from a user running Chroma with Sage-Attention who found this was the only node that reproduced the same output to the pixel with LoRAs enabled, where every other compile approach drifted worse on each generation until they force-unloaded models.

How it works

The node clones the model with disable_dynamic=True (matching stock behavior) and grabs the diffusion_model object. By default it compiles the whole thing - but Flux is where the design gets interesting. Flux block forwards are aggressively dynamic (transformer_options get mutated, patches swapped in and out), so whole-block compile tends to graph-break or replay-fail, especially when Sage-Attention is active. So for Flux it compiles leaf submodules instead: double_blocks.{i}.img_attn.qkv / .proj, txt_attn.qkv / .proj, and single_blocks.{i}.linear1 / linear2. On cudagraph-capable runs it narrows even further to the qkv leaves only, because those are the ones that replay safely - that's the FLUX.2 KLEIN 9B compatibility work in the README.

Under the hood, each target gets wrapped in a lazy module that calls torch.compile on first forward under a lock. On newer ComfyUI builds it prefers comfy_api.torch_helpers.set_torch_compile_wrapper, and it feeds ComfyUI's skip_torch_compile_dict in as a guard_filter_fn so those dynamic dicts don't trigger constant recompiles.

The inputs that matter

• model - your already-patched model. This node goes last in the patcher chain, right before the KSampler.
• backend - start with inductor. cudagraphs and nvfuser need CUDA (nvfuser availability varies by PyTorch build).
• disable_cudagraphs (default true) - passes {'triton.cudagraphs': False} for the Inductor backend. Keep it on unless you hit a cudagraph-specific stability problem; it's the most common fix for cudaMallocAsync crashes.
• compile_transformer_only (default false) - true compiles discovered transformer targets instead of the whole diffusion model. Start true for Flux/Klein, false elsewhere, and compare.
• mode - default first; reduce-overhead / max-autotune squeeze more out at the cost of a much longer warmup.
• fullgraph and dynamic - both default false; leave them alone until you're chasing the last few percent.
• The two Flux knobs, allow_flux_inferred_targets and fallback_to_full_model_if_no_targets, default to false, and the defaults are the right ones: strict target names, and skip compiling rather than silently widening scope. If no targets are found with the fallback off, the node just… doesn't compile. That's a feature - a no-op beats a compile that breaks your output.

The single output is a MODEL you wire straight into the KSampler. That's the whole API.

Installing it

It's one file with zero dependencies (the pyproject lists none), so install is trivial:

cd ComfyUI/custom_nodes
git clone https://github.com/xmarre/TorchCompileModel_LoRASafe

then restart ComfyUI and find TorchCompileModel_LoRASafe under model / optimisation 🛠️ - or just search "LoRA-Safe TorchCompile" in ComfyUI Manager. One note: the README's "copy the lora_safe_compile folder" step is stale from the pastebin days - the repo root is the custom node, so a plain clone works (an early __init__.py typo briefly broke installs, fixed since).

What goes wrong

• The first run is painfully slow. That's compile warmup; judge it on the second generation, not the first.
• UserDefinedObjectVariable has no attribute 'proxy' - a torch._dynamo bug on pre-2.5 PyTorch. Update PyTorch and it goes away.
• Crashes or instability - fall back to inductor, keep disable_cudagraphs=true, and if it's still unhappy, flip compile_transformer_only to true to shrink the compile surface.

The honest caveat: the stock node and KJNodes' PatchModelPatcherOrder were later patched to address the same ordering problem, so the gap has narrowed. But the community verdict holds up - this node stays more deterministic across generations, and if you're chasing reproducible output with LoRAs plus a compile backend, that consistency is the real reason to reach for it.