# LoRA-Safe TorchCompile Node for ComfyUI `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). 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 strict leaf submodules by default. Strict Flux names targeted by default: - `double_blocks.{i}.img_attn.qkv` - `double_blocks.{i}.img_attn.proj` - `double_blocks.{i}.txt_attn.qkv` - `double_blocks.{i}.txt_attn.proj` - `single_blocks.{i}.linear1` - `single_blocks.{i}.linear2` Optional widening controls: - `allow_flux_inferred_targets` (`false` by default) - `false`: use only strict names above. - `true`: allow inferred Flux leaf targets (including `img_mlp/txt_mlp` pattern matches) for compatibility across variants. - `fallback_to_full_model_if_no_targets` (`false` by default) - `false`: if transformer-only target discovery finds nothing, skip compile instead of silently widening scope. - `true`: if discovery finds nothing, fall back to compiling `diffusion_model`. 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, behavior is controlled by `fallback_to_full_model_if_no_targets` (default `false`, meaning skip compile). --- ## 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 (default strict mode): targets only `qkv/proj` attention leaves plus `linear1/linear2` leaves. - `allow_flux_inferred_targets` (`BOOLEAN`, default `false`) - `false` (recommended): strict Flux target set only. - `true`: allow inferred Flux leaf targets (`img_mlp/txt_mlp`-style matches) when preferred names differ. - `fallback_to_full_model_if_no_targets` (`BOOLEAN`, default `false`) - `false` (recommended): if transformer-only discovery finds nothing, skip compile rather than widening scope. - `true`: if transformer-only discovery finds nothing, compile `diffusion_model`. --- ## 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); otherwise start with `false` and compare. - `allow_flux_inferred_targets`: `false` (recommended for debugging and stable targeting). - `fallback_to_full_model_if_no_targets`: `false` (recommended to avoid silent whole-model escalation). 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.