Add support for LTX 2.5 - #15499
Conversation
For now pure pytorch is slow, comfy-kitchen version already verified to be 2-3x faster
Model-side support for checkpoints trained with generated keyframes, without the nodes that drive it. model_detection sets use_keyframes_abs_pos_embedding by probing the state dict for a keyframes_abs_pos_embedding tensor, so the flag follows the checkpoint rather than needing to be configured. LTXBaseModel allocates the parameter only when that flag is set and leaves it None otherwise, keeping checkpoints without it unchanged. LTXVModel applies the embedding after patchify_proj to the tokens that encode a single standalone pixel frame, and guards the keyframe_idxs token count against the latent's tokens-per-frame: a mismatch means the appended frames were recorded at a different spatial resolution and their positions would land on the wrong tokens, so it raises instead of corrupting the sample silently. LTXV and LTXAV pass a generated_keyframes conditioning entry through to the model as a CONDConstant. Cherry-picked from 7454d436 by akvochko, dropping the comfy_extras/nodes_lt.py node changes so this carries only the loader and model-side support. Co-authored-by: Adam Oster <>
📝 WalkthroughWalkthroughThe PR adds an LTX 2.4 diffusion VAE decoder and integrates its detection, configuration, and tiled decoding. LTX video and audio models gain feed-forward bias controls, keyframe positional embeddings, prompt AdaLN gating, and STG self-attention options. New nodes provide modality guidance, dual CFG, spatio-temporal guidance, and duration prediction. Gemma 4 text encoder routing, tokenizer support, projection configuration, and LTX 2.4 prompt formatting are also added. 🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (2 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 15
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy_extras/nodes_lt.py`:
- Around line 1161-1169: In the duration prediction block, remove the
torch.no_grad wrapper and the manual head.to(device) call. Keep duration_head
placement exclusively through load_models_gpu, while preserving the existing
context device/dtype conversion and token preprocessing before calculating
seconds.
In `@comfy_extras/nodes_textgen.py`:
- Around line 259-263: Update the final cleanup regex in the text-generation
sanitization flow to remove the exact Gemma 4 control markers <|think|> and
<turn|>, alongside the existing marker patterns. Keep the surrounding reasoning
removal and trimming behavior unchanged.
In `@comfy/ldm/lightricks/duration_head.py`:
- Around line 75-80: Update the frame-bound handling around the duration
calculation in the affected function to validate whether the interval from
min_frames through max_frames contains any causal-grid value of the form 8k + 1
(using time_scale). Raise a clear error before clamping when no compatible frame
count exists, and preserve snapping and clamping behavior for valid intervals.
In `@comfy/ldm/lightricks/embeddings_connector.py`:
- Line 128: Update the model construction around both Embeddings1DConnector
instances to forward the model configuration’s connector_ff_bias value into each
connector’s ff_bias parameter instead of relying on the default; if the
configuration does not expose this setting, remove the unused connector_ff_bias
parameter consistently.
In `@comfy/ldm/lightricks/model.py`:
- Around line 788-791: In the initialization branch for
keyframes_abs_pos_embedding, replace the zero-filled tensor allocation with an
uninitialized torch.empty allocation while preserving the same shape, dtype,
device, Parameter wrapper, and None behavior when
use_keyframes_abs_pos_embedding is disabled.
- Around line 465-474: Update the attention logic around the self_attn and
stg_skip_self_attn check so self-attention computes q with self.to_q(x) only in
the non-STG path; preserve the STG bypass as out = v and avoid running the
discarded query projection on that branch.
In `@comfy/ldm/lightricks/vae/audio_vae.py`:
- Line 188: Update the audio latent count calculation in the relevant method to
use math.ceil instead of round when converting frames_number and frame_rate into
latent units. Preserve the existing float conversion and multiplication by
self.latents_per_second so fractional counts always round up.
In `@comfy/ldm/lightricks/vae/na_diffusion_decoder.py`:
- Around line 512-514: Update the decode method’s fixed-seed noise handling to
expose a selectable seed through vae_options, using seed 0 as the default to
preserve reproducibility. Read the configured seed when creating the
torch.Generator and remove the unresolved TODO.
- Around line 19-29: Remove the einops.rearrange import and replace every
rearrange call in patchify, unpatchify, and LinearPixelShuffleUpsample.forward
with equivalent native tensor operations using reshape/view, permute, and
unflatten. Preserve the existing dimension ordering, including applying the
inverse permutation in unpatchify and the equivalent view/permute transformation
in LinearPixelShuffleUpsample.forward.
- Around line 35-50: Remove the local rms_norm helper and the RMSNorm
implementation in this module, then use the shared comfy.rmsnorm.rms_norm
through the normalization modules while preserving q_norm.weight and
k_norm.weight access in NeighborhoodAttention3D.forward. Eliminate the
torch.nn.functional capability probe and float32 fallback, and retain
checkpoint-populated parameters using torch.empty where applicable.
- Around line 363-382: Update the tiled temporal decode flow in
VAE.decode_tiled_3d/CausalDiffusionVAE.decode so tile position controls
forward_pre_diffusion flags: use (drop_leading_frame=True, pad_trailing=False)
for the origin tile, (False, False) for middle tiles, and (False, True) for the
final tile. Ensure non-tiled decoding retains its existing behavior.
In `@comfy/sd.py`:
- Around line 587-602: Guard the `sd["decoder.conv_in.weight"]` access in the
LTX 2.4 decoder branch by checking that key exists before deriving
`self.latent_channels`; otherwise fail clearly rather than raising an unhandled
`KeyError`. Keep the existing `decoder.conv_in_x_t.weight` detector and branch
ordering unchanged.
- Around line 1246-1257: Update the tiled fallback budget calculation in the
surrounding decode method to use free memory after soft_empty_cache(), rather
than 80% of get_total_memory(self.device). Preserve the existing tile sizing
loops and ensure the dynamic sizing remains compatible with all 3D VAE paths
lacking handles_tiling, including Wan, Mochi, and CogVideoX.
- Around line 1262-1281: Update _tile_bounded_shape to handle three-dimensional
latent shapes before indexing spatial dimensions: clamp s[2] using tile_x, and
change the existing non-five-dimensional branch to apply only to
four-dimensional shapes. Preserve the current five-dimensional handling and
avoid accessing s[3] for 3D audio shapes.
In `@comfy/text_encoders/lt.py`:
- Around line 144-147: Add an else branch to the text projection selection in
the constructor, after the single_linear and dual_linear cases, that raises a
clear ValueError for any unsupported text_projection_type. Ensure unsupported
values cannot leave text_embedding_projection uninitialized and fall through to
unprojected output.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 2aa2a003-f5fb-4fb9-8d97-eaf767a8b9bb
📒 Files selected for processing (15)
comfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/duration_head.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/text_encoders/gemma4.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_model_patch.pycomfy_extras/nodes_textgen.py
📜 Review details
⏰ Context from checks skipped due to timeout. (7)
- GitHub Check: test (ubuntu-latest)
- GitHub Check: test (windows-latest)
- GitHub Check: test (macos-latest)
- GitHub Check: test (macos-latest)
- GitHub Check: test (ubuntu-latest)
- GitHub Check: Run Pylint
- GitHub Check: test (windows-2022)
🧰 Additional context used
📓 Path-based instructions (7)
**/*
📄 CodeRabbit inference engine (AGENTS.md)
**/*: Keep changes small, direct, and limited to the narrowest necessary code path and smallest number of files.
Prefer practical fixes, minimal dependencies, and existing repository patterns; remove obsolete, dead, unreachable, or unused code.
Preserve existing APIs, node names, model-loading behavior, file layout, and workflow compatibility unless replacement is explicitly intended.
Core ComfyUI must not add outbound internet requests, telemetry, tracking, reporting, remote configuration, or background network activity. User-authorized model downloads are limited to the requested artifact and must exclude telemetry and unrelated metadata.
Files:
comfy/model_detection.pycomfy_extras/nodes_model_patch.pycomfy/ldm/lightricks/duration_head.pycomfy_extras/nodes_lt_audio.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy_extras/nodes_textgen.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Keep state and capability flags on the object that owns the behavior. Prefer explicit parent-owned attributes over probing child objects withgetattr; use child checks only when the child owns the delegated behavior.
Preserve shared method signatures, argument order, return shapes, side effects, and error behavior unless every affected caller and interface is intentionally updated.
Do not add unused compatibility parameters, flags, attributes, constructor options, or model-specific options to shared helpers; keep one-off behavior at the integration boundary.
Normalize third-party return conventions at integration boundaries so core code receives the expected type and shape; avoid undocumented caller-side unwrapping.
Do not addtorch.no_grad,torch.inference_mode, or inference-mode wrappers. Do not add model freeze/unfreeze toggles; only disable globally enabled inference mode when a training path requires gradients.
Remove inference-only training behavior such as dropout while preserving checkpoint and state-dict compatibility; usenn.Identitywhen deleting a module would alter keys or ordering.
Keep imports at module scope except established optional-backend probes or imports required to avoid cycles; avoid unnecessarytry/exceptblocks and use specific exceptions with useful fallbacks.
Do not add workarounds for unsupported library versions, especially PyTorch exception-and-float-cast retries, unless a comment names the exact versions still requiring them.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently degrading output.
Match local style, keep comments sparse and useful, and remove comments that merely restate obvious code.
Treat dtype, device placement, VRAM use, and offloading as correctness concerns across CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low-VRAM environments.
Prefer existing ComfyUI and Comfy Kitchen operations, quantization helpers, cast/offload helpe...
Files:
comfy/model_detection.pycomfy_extras/nodes_model_patch.pycomfy/ldm/lightricks/duration_head.pycomfy_extras/nodes_lt_audio.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy_extras/nodes_textgen.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
**/*.{py,json}
📄 CodeRabbit inference engine (AGENTS.md)
Treat legacy combo,
io.Combo, andio.DynamicCombovalues affecting filesystem access as untrusted; revalidate them at load/save boundaries withfolder_paths, containment checks, or fixed allowlists.
Files:
comfy/model_detection.pycomfy_extras/nodes_model_patch.pycomfy/ldm/lightricks/duration_head.pycomfy_extras/nodes_lt_audio.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy_extras/nodes_textgen.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
**/*.{py,md,txt,json}
📄 CodeRabbit inference engine (AGENTS.md)
Keep warning and info messages short and actionable, remove noisy or misleading logging, and make documentation edits concise, factual, and tied to changed behavior.
Files:
comfy/model_detection.pycomfy_extras/nodes_model_patch.pycomfy/ldm/lightricks/duration_head.pycomfy_extras/nodes_lt_audio.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy_extras/nodes_textgen.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
**
⚙️ CodeRabbit configuration file
**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing awith:block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.
Files:
comfy/model_detection.pycomfy_extras/nodes_model_patch.pycomfy/ldm/lightricks/duration_head.pycomfy_extras/nodes_lt_audio.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy_extras/nodes_textgen.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
comfy/**
⚙️ CodeRabbit configuration file
comfy/**: Core ML/diffusion engine. Focus on:
- Backward compatibility (breaking changes affect all custom nodes)
- Memory management and GPU resource handling
- Performance implications in hot paths
- Thread safety for concurrent execution
Files:
comfy/model_detection.pycomfy/ldm/lightricks/duration_head.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
comfy_extras/**
⚙️ CodeRabbit configuration file
comfy_extras/**: Community-contributed extra nodes. Focus on:
- Consistency with node patterns (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY)
- No breaking changes to existing node interfaces
Files:
comfy_extras/nodes_model_patch.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_textgen.pycomfy_extras/nodes_lt.py
🧠 Learnings (9)
📚 Learning: 2026-02-21T14:01:41.482Z
Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.
Applied to files:
comfy/model_detection.pycomfy_extras/nodes_model_patch.pycomfy/ldm/lightricks/duration_head.pycomfy_extras/nodes_lt_audio.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy_extras/nodes_textgen.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy_extras/nodes_lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
📚 Learning: 2026-05-13T12:31:45.069Z
Learnt from: rattus128
Repo: Comfy-Org/ComfyUI PR: 13802
File: comfy/pinned_memory.py:19-30
Timestamp: 2026-05-13T12:31:45.069Z
Learning: When reviewing code that uses comfy/pinned_memory.py’s `HostBuffer.extend(size=..., reallocate=...)`: by default (`reallocate` is not True / False), `extend(size=...)` is a *relative increment* that grows the buffer by `size` bytes—so slicing like `[offset:offset+size]` after `hostbuf.extend(size=size)` is correct and the argument should not be rewritten to `offset + size`. Only in the single-segment reallocation mode (`reallocate=True`, e.g., as used by `resize_pin_buffer()` in `comfy/model_management.py`) should `size` be treated as an *absolute target* and the call/arguments should be checked accordingly.
Applied to files:
comfy/model_detection.pycomfy/ldm/lightricks/duration_head.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
📚 Learning: 2026-08-06T22:18:59.719Z
Learnt from: kijai
Repo: Comfy-Org/ComfyUI PR: 15362
File: comfy/ldm/wan/model_animate2.py:186-223
Timestamp: 2026-08-06T22:18:59.719Z
Learning: When reviewing ComfyUI quantization code, treat `comfy.quant_ops.TensorWiseINT8Layout` and `comfy.quant_ops.TensorCoreConvRotW4A4Layout` as re-exports from `comfy_kitchen`. Validate their behavior against the re-exported `comfy_kitchen` implementations rather than assuming they are local fallback classes.
Applied to files:
comfy/model_detection.pycomfy/ldm/lightricks/duration_head.pycomfy/text_encoders/gemma4.pycomfy/model_base.pycomfy/ldm/lightricks/vae/audio_vae.pycomfy/sd.pycomfy/ldm/lightricks/embeddings_connector.pycomfy/text_encoders/lt.pycomfy/ldm/lightricks/av_model.pycomfy/ldm/lightricks/model.pycomfy/ldm/lightricks/vae/na_diffusion_decoder.py
📚 Learning: 2026-03-04T14:05:31.426Z
Learnt from: jtydhr88
Repo: Comfy-Org/ComfyUI PR: 12757
File: comfy_extras/nodes_custom_sampler.py:1069-1089
Timestamp: 2026-03-04T14:05:31.426Z
Learning: In the ComfyUI sampling pipeline, treat percent_to_sigma(0.0) as a sentinel value (999999999.9) that means starting from pure noise. This is consistent with BasicScheduler via calculate_sigmas. The SamplingPercentToSigma node’s return_actual_sigma flag differentiates this sentinel from sigma_max. Reviewers should not flag CurveToSigmas or similar nodes that rely on percent_to_sigma as bugs; downstream samplers are expected to handle the sentinel correctly. When reviewing related sampling-related code, assume this sentinel semantics unless there is explicit handling for a real sigma_max.
Applied to files:
comfy_extras/nodes_model_patch.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_textgen.pycomfy_extras/nodes_lt.py
📚 Learning: 2026-04-04T13:29:15.653Z
Learnt from: kijai
Repo: Comfy-Org/ComfyUI PR: 13258
File: comfy_extras/nodes_frame_interpolation.py:151-189
Timestamp: 2026-04-04T13:29:15.653Z
Learning: In this ComfyUI codebase, node `execute()` inference is already run under a global `torch.inference_mode()` context established in the execution engine (e.g., `execution.py` around line ~732). During review, avoid recommending changes that wrap node inference loops in `torch.inference_mode()`—it is already applied, so such suggestions are likely redundant.
Applied to files:
comfy_extras/nodes_model_patch.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_textgen.pycomfy_extras/nodes_lt.py
📚 Learning: 2026-05-09T18:40:40.199Z
Learnt from: kijai
Repo: Comfy-Org/ComfyUI PR: 13813
File: comfy_extras/nodes_wandancer.py:868-872
Timestamp: 2026-05-09T18:40:40.199Z
Learning: When building video/temporal decoding nodes that call ComfyUI’s VAE.decode (comfy/sd.py), leverage VAE.decode’s existing VRAM-aware chunking along dim 0. Reshape or transpose the latent so the temporal dimension T is folded into dim 0 (e.g., transform a latent of shape [B, T, C, H, W] into [B*T, C, H, W] before calling vae.decode). This lets VAE.decode do chunked decoding without needing an explicit per-frame loop inside the node itself.
Applied to files:
comfy_extras/nodes_model_patch.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_textgen.pycomfy_extras/nodes_lt.py
📚 Learning: 2026-05-20T00:10:14.673Z
Learnt from: Pauan
Repo: Comfy-Org/ComfyUI PR: 13997
File: comfy_extras/nodes_string.py:12-25
Timestamp: 2026-05-20T00:10:14.673Z
Learning: In the ComfyUI `comfy_extras/` codebase, some nodes intentionally ship with a default input string that references parameters that may not yet be connected. If the default would raise a `KeyError` (e.g., examples like `MathExpression` default `a + b`, or `StringFormat` default `{a}` with `min=0` and autogrow inputs), treat it as an intentional “hint default” UX pattern, not a bug. During review, do not flag this behavior or recommend changing `min` to `1` or altering the default to an empty string solely to avoid the `KeyError`.
Applied to files:
comfy_extras/nodes_model_patch.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_textgen.pycomfy_extras/nodes_lt.py
📚 Learning: 2026-07-26T18:37:44.213Z
Learnt from: jtydhr88
Repo: Comfy-Org/ComfyUI PR: 15090
File: comfy_extras/nodes_video.py:230-235
Timestamp: 2026-07-26T18:37:44.213Z
Learning: In ComfyUI node implementations under `comfy_extras`, do not add regular widget/prompt inputs to `fingerprint_inputs` if they are already included in the node cache signature via `comfy_execution/caching.py:get_immediate_node_signature` (it records every non-link prompt input as `(key, inputs[key])`). Reserve `fingerprint_inputs` only for out-of-band state that can change without changing the prompt inputs (e.g., the selected source file’s modification time). For example, inputs like `LoadVideo.edit` should not be redundantly added to `fingerprint_inputs`; use it only for things not represented in prompt inputs.
Applied to files:
comfy_extras/nodes_model_patch.pycomfy_extras/nodes_lt_audio.pycomfy_extras/nodes_textgen.pycomfy_extras/nodes_lt.py
📚 Learning: 2026-04-23T13:22:31.631Z
Learnt from: kijai
Repo: Comfy-Org/ComfyUI PR: 13531
File: comfy_extras/nodes_lt.py:715-722
Timestamp: 2026-04-23T13:22:31.631Z
Learning: In the ComfyUI LTXV audio pipeline, when preparing waveforms for `AudioVAE.encode()`, resample inputs to the AudioVAE public interface sample rate: `vae_sample_rate = getattr(audio_vae, "audio_sample_rate", 44100)`. Do not resample to `first_stage_model.sample_rate` (often ~16000 Hz), since that is the VAE’s internal mel-spectrogram rate and is handled inside the VAE. Ensure the resampling/`VAEEncodeAudio.execute()` path uses `vae_sample_rate` to match `AudioVAE.encode()` expectations.
Applied to files:
comfy_extras/nodes_lt.py
🔇 Additional comments (24)
comfy/text_encoders/gemma4.py (1)
1446-1449: LGTM!Also applies to: 1481-1493
comfy/text_encoders/lt.py (2)
111-143: LGTM!Also applies to: 181-286
84-90: 🎯 Functional CorrectnessUse the explicit
min_lengthargument
clip.tokenize(..., min_length=1)overrides the tokenizer instance value of1024. The new default affects only calls that omitmin_length.> Likely an incorrect or invalid review comment.comfy_extras/nodes_lt_audio.py (1)
176-176: LGTM!comfy_extras/nodes_textgen.py (2)
229-232: 🎯 Functional Correctness | ⚡ Quick winUse the configured encoder key for prompt routing.
This check probes
clip.tokenizer.clip_nameand silently selects the Gemma 3 format when that child attribute is absent or renamed. Route on the explicittext_encoder_keyconfigured byLTXAVTEModel, and reject unsupported keys clearly.As per coding guidelines, “Keep state and capability flags on the object that owns the behavior. Prefer explicit parent-owned attributes over probing child objects with
getattr.”[ suggest_essential_refactor]
Source: Coding guidelines
1-1: LGTM!Also applies to: 156-211
comfy/ldm/lightricks/vae/na_diffusion_decoder.py (3)
53-56: LGTM!Also applies to: 59-62, 65-105, 108-124, 169-193, 195-208, 211-227, 229-250, 253-278, 280-296, 298-362, 384-394, 421-459
337-361: 🚀 Performance & ScalabilityDo not change the decoder linear layers. The
nn.Lineardeclarations anddisable_offload = Truesetting predate this PR; this PR only changes the neighborhood-attention implementation.> Likely an incorrect or invalid review comment.
126-166: 🗄️ Data Integrity & IntegrationNo change needed.
na3d(..., None, 1.0)uses an explicit score scale, andrms_rope_accepts the supplied argument order and mutatesqandkin place.> Likely an incorrect or invalid review comment.comfy/sd.py (1)
14-14: LGTM!Also applies to: 1752-1766, 1934-1957
comfy/ldm/lightricks/duration_head.py (1)
1-68: LGTM!comfy_extras/nodes_model_patch.py (1)
13-13: LGTM!Also applies to: 300-303
comfy/ldm/lightricks/av_model.py (3)
99-100: LGTM!Also applies to: 183-186, 418-427, 460-461, 486-486, 617-618
988-988: LGTM!Also applies to: 1012-1012
937-945: 🎯 Functional CorrectnessConfirm the intended STG modality —
stg_skip_self_attnreaches bothself.attn1andself.audio_attn1, so selected blocks perturb video and audio self-attention. Use modality-specific flags if STG is video-only.comfy/ldm/lightricks/model.py (4)
306-321: LGTM!Also applies to: 511-511, 527-527, 1090-1090
726-728: LGTM!Also applies to: 758-760, 798-798
1120-1128: LGTM!Also applies to: 1171-1184, 1217-1227
1186-1215: 🩺 Stability & AvailabilityNo length check is needed for
slots.orig_shapeandgrid_maskuse the same latent, including appended guides; context-window resizing also returns matching video and guide mask lengths. No in-repository caller producesgenerated_keyframeswith a mismatchedkeyframe_idxsshape.> Likely an incorrect or invalid review comment.comfy/ldm/lightricks/embeddings_connector.py (1)
53-53: LGTM!Also applies to: 78-78
comfy/model_detection.py (1)
400-400: LGTM!comfy_extras/nodes_lt.py (3)
5-12: LGTM!Also applies to: 940-991, 994-1050, 1097-1120, 1194-1197
1123-1153: LGTM!
1066-1090: 🎯 Functional CorrectnessKeep the last-dimension split.
pack_latentsproduces(batch, 1, total_elements)and places video elements before audio elements. Therefore,out[..., v:]selects the audio elements correctly.> Likely an incorrect or invalid review comment.
| comfy.model_management.load_models_gpu([model, duration_head]) | ||
| device = model.load_device | ||
| head = head.to(device) | ||
| with torch.no_grad(): | ||
| context = context.to(device=device, dtype=model.model.get_dtype_inference()) | ||
| processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False)) | ||
| video_tokens = processed[..., :dm.cross_attention_dim].float() | ||
| audio_tokens = processed[..., dm.cross_attention_dim:].float() | ||
| seconds = float(head(video_tokens, audio_tokens)[0]) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Remove the torch.no_grad wrapper and the manual device move.
Two problems exist in this block:
- Line 1164 adds
torch.no_grad(). The repository forbids inference-mode wrappers, and nodeexecutealready runs under a global inference mode. - Line 1163 moves the head with
head.to(device). Line 1161 already loadsduration_headthroughcomfy.model_management.load_models_gpu. The manual move duplicates the placement and bypasses the patcher.
🛠️ Proposed fix
comfy.model_management.load_models_gpu([model, duration_head])
device = model.load_device
- head = head.to(device)
- with torch.no_grad():
- context = context.to(device=device, dtype=model.model.get_dtype_inference())
- processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False))
- video_tokens = processed[..., :dm.cross_attention_dim].float()
- audio_tokens = processed[..., dm.cross_attention_dim:].float()
- seconds = float(head(video_tokens, audio_tokens)[0])
+ context = context.to(device=device, dtype=model.model.get_dtype_inference())
+ processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False))
+ video_tokens = processed[..., :dm.cross_attention_dim].float()
+ audio_tokens = processed[..., dm.cross_attention_dim:].float()
+ seconds = float(head(video_tokens, audio_tokens)[0])As per coding guidelines: "Do not add torch.no_grad, torch.inference_mode, or inference-mode wrappers." Based on learnings: node execute() inference already runs under a global torch.inference_mode() context established in the execution engine.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| comfy.model_management.load_models_gpu([model, duration_head]) | |
| device = model.load_device | |
| head = head.to(device) | |
| with torch.no_grad(): | |
| context = context.to(device=device, dtype=model.model.get_dtype_inference()) | |
| processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False)) | |
| video_tokens = processed[..., :dm.cross_attention_dim].float() | |
| audio_tokens = processed[..., dm.cross_attention_dim:].float() | |
| seconds = float(head(video_tokens, audio_tokens)[0]) | |
| comfy.model_management.load_models_gpu([model, duration_head]) | |
| device = model.load_device | |
| context = context.to(device=device, dtype=model.model.get_dtype_inference()) | |
| processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False)) | |
| video_tokens = processed[..., :dm.cross_attention_dim].float() | |
| audio_tokens = processed[..., dm.cross_attention_dim:].float() | |
| seconds = float(head(video_tokens, audio_tokens)[0]) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy_extras/nodes_lt.py` around lines 1161 - 1169, In the duration
prediction block, remove the torch.no_grad wrapper and the manual
head.to(device) call. Keep duration_head placement exclusively through
load_models_gpu, while preserving the existing context device/dtype conversion
and token preprocessing before calculating seconds.
Sources: Coding guidelines, Learnings
| text = out.args[0] | ||
| text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL) | ||
| if "</think>" in text: # unclosed/truncated reasoning: keep what follows the last close | ||
| text = text.rsplit("</think>", 1)[-1] | ||
| text = re.sub(r"</?think>|<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?", "", text).strip() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Remove the Gemma 4 control markers that this prompt format uses.
The sanitizer matches <think> and opening <|turn> markers. It does not match <|think|> or <turn|>, which the Gemma 4 prompt creates on Lines 240-246. If generation emits either control marker, the node returns it as prompt text. Extend the cleanup patterns for the exact Gemma 4 marker forms.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy_extras/nodes_textgen.py` around lines 259 - 263, Update the final
cleanup regex in the text-generation sanitization flow to remove the exact Gemma
4 control markers <|think|> and <turn|>, alongside the existing marker patterns.
Keep the surrounding reasoning removal and trimming behavior unchanged.
| min_frames = max(1, round(min_seconds * frame_rate)) | ||
| max_frames = round(max_seconds * frame_rate) | ||
| raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames)) | ||
| frames = (raw_frames - 1) // time_scale * time_scale + 1 | ||
| if frames < min_frames: | ||
| frames = min(-(-(min_frames - 1) // time_scale) * time_scale + 1, max_frames) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject bounds that contain no causal-grid frame count.
Lines 75-80 can return a frame count that is not 8k + 1. For min_seconds=1, max_seconds=1, frame_rate=24, and time_scale=8, line 80 returns 24.
Validate the snapped bounds before clamping. Raise a clear error when the interval contains no VAE-compatible frame count.
Proposed fix
def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds, time_scale=8):
- min_frames = max(1, round(min_seconds * frame_rate))
- max_frames = round(max_seconds * frame_rate)
- raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames))
- frames = (raw_frames - 1) // time_scale * time_scale + 1
- if frames < min_frames:
- frames = min(-(-(min_frames - 1) // time_scale) * time_scale + 1, max_frames)
- return frames
+ min_frames = max(1, round(min_seconds * frame_rate))
+ max_frames = round(max_seconds * frame_rate)
+ min_grid_frames = -(-(min_frames - 1) // time_scale) * time_scale + 1
+ max_grid_frames = (max_frames - 1) // time_scale * time_scale + 1
+ if min_grid_frames > max_grid_frames:
+ raise ValueError("Duration bounds contain no VAE-compatible frame count")
+ raw_frames = max(min_grid_frames, min(round(seconds * frame_rate), max_grid_frames))
+ return (raw_frames - 1) // time_scale * time_scale + 1As per path instructions, AGENTS.md says that latent layout is a correctness concern.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| min_frames = max(1, round(min_seconds * frame_rate)) | |
| max_frames = round(max_seconds * frame_rate) | |
| raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames)) | |
| frames = (raw_frames - 1) // time_scale * time_scale + 1 | |
| if frames < min_frames: | |
| frames = min(-(-(min_frames - 1) // time_scale) * time_scale + 1, max_frames) | |
| def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds, time_scale=8): | |
| min_frames = max(1, round(min_seconds * frame_rate)) | |
| max_frames = round(max_seconds * frame_rate) | |
| min_grid_frames = -(-(min_frames - 1) // time_scale) * time_scale + 1 | |
| max_grid_frames = (max_frames - 1) // time_scale * time_scale + 1 | |
| if min_grid_frames > max_grid_frames: | |
| raise ValueError("Duration bounds contain no VAE-compatible frame count") | |
| raw_frames = max(min_grid_frames, min(round(seconds * frame_rate), max_grid_frames)) | |
| return (raw_frames - 1) // time_scale * time_scale + 1 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/ldm/lightricks/duration_head.py` around lines 75 - 80, Update the
frame-bound handling around the duration calculation in the affected function to
validate whether the interval from min_frames through max_frames contains any
causal-grid value of the form 8k + 1 (using time_scale). Raise a clear error
before clamping when no compatible frame count exists, and preserve snapping and
clamping behavior for valid intervals.
Source: Path instructions
| causal_temporal_positioning=False, | ||
| num_learnable_registers: Optional[int] = 128, | ||
| apply_gated_attention=False, | ||
| connector_ff_bias=True, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find every construction of Embeddings1DConnector and any use of connector_ff_bias.
rg -n -C6 'Embeddings1DConnector\(' --type=py
rg -n 'connector_ff_bias' --type=pyRepository: Comfy-Org/ComfyUI
Length of output: 155
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- tracked candidate files ---'
git ls-files | grep -E '(^|/)(embeddings_connector|av_model|lt)\.py$' || true
printf '%s\n' '--- identifier search ---'
rg -n -C4 --hidden --glob '*.py' \
'Embeddings1DConnector|connector_ff_bias|connector_ff' . || true
printf '%s\n' '--- repository status and recent diff summary ---'
git status --short
git diff --statRepository: Comfy-Org/ComfyUI
Length of output: 6488
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- connector definition ---'
cat -n comfy/ldm/lightricks/embeddings_connector.py | sed -n '113,165p'
printf '%s\n' '--- av_model caller setup ---'
cat -n comfy/ldm/lightricks/av_model.py | sed -n '515,580p'
printf '%s\n' '--- lt.py compatibility callers ---'
cat -n comfy/text_encoders/lt.py | sed -n '140,180p'
printf '%s\n' '--- configuration references ---'
rg -n -C3 --hidden \
'connector_(ff_bias|attention_head_dim|num_attention_heads|num_layers|split_rope)|audio_connector_|video_connector_' \
--glob '*.py' --glob '*.json' --glob '*.yaml' --glob '*.yml' . || true
printf '%s\n' '--- AST caller/signature probe ---'
python3 - <<'PY'
import ast
from pathlib import Path
targets = [
Path("comfy/ldm/lightricks/av_model.py"),
Path("comfy/text_encoders/lt.py"),
]
for path in targets:
tree = ast.parse(path.read_text())
print(path)
for node in ast.walk(tree):
if isinstance(node, ast.Call):
fn = node.func
name = fn.id if isinstance(fn, ast.Name) else None
if name == "Embeddings1DConnector":
keywords = [kw.arg for kw in node.keywords if kw.arg is not None]
print(" call line", node.lineno, "keywords:", keywords)
connector = ast.parse(Path("comfy/ldm/lightricks/embeddings_connector.py").read_text())
for node in connector.body:
if isinstance(node, ast.ClassDef) and node.name == "Embeddings1DConnector":
init = next(
item for item in node.body
if isinstance(item, ast.FunctionDef) and item.name == "__init__"
)
print("Embeddings1DConnector.__init__ parameters:",
[arg.arg for arg in init.args.args])
PYRepository: Comfy-Org/ComfyUI
Length of output: 12538
Forward connector_ff_bias from the model configuration.
Embeddings1DConnector consumes this value as ff_bias, but all repository callers use the default. Pass the configuration value to both connectors, or remove the parameter if configuration does not support it.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/ldm/lightricks/embeddings_connector.py` at line 128, Update the model
construction around both Embeddings1DConnector instances to forward the model
configuration’s connector_ff_bias value into each connector’s ff_bias parameter
instead of relying on the default; if the configuration does not expose this
setting, remove the unused connector_ff_bias parameter consistently.
Source: Coding guidelines
| self_attn = context is None | ||
| q = self.to_q(x) | ||
| context = x if context is None else context | ||
| k = self.to_k(context) | ||
| v = self.to_v(context) | ||
|
|
||
| q = self.q_norm(q) | ||
| k = self.k_norm(k) | ||
|
|
||
| # These norms span all heads, so the per-head RMS+RoPE kernel is not equivalent. | ||
| if pe is not None: | ||
| if k_pe is None and q.shape == k.shape: | ||
| q, k = apply_rotary_emb_qk(q, k, pe) | ||
| else: | ||
| q = apply_rotary_emb(q, pe) | ||
| k = apply_rotary_emb(k, pe if k_pe is None else k_pe) | ||
|
|
||
| if mask is None: | ||
| out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options) | ||
| elif isinstance(mask, GuideAttentionMask): | ||
| out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options) | ||
| # Spatio-Temporal Guidance (STG) perturbation: for the flagged self-attention | ||
| # layers, the attention degrades to a passthrough of the value projection (out = V). | ||
| if self_attn and transformer_options.get("stg_skip_self_attn", False): | ||
| out = v |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Skip the query projection when STG bypasses the attention.
q = self.to_q(x) runs before the STG check. On the STG branch the result is discarded. The query projection is a full linear over all tokens, and it runs for every selected block on every STG pass.
Move the to_q call into the non-STG branch.
♻️ Proposed refactor
self_attn = context is None
- q = self.to_q(x)
context = x if context is None else context
- k = self.to_k(context)
v = self.to_v(context)
# Spatio-Temporal Guidance (STG) perturbation: for the flagged self-attention
# layers, the attention degrades to a passthrough of the value projection (out = V).
if self_attn and transformer_options.get("stg_skip_self_attn", False):
out = v
else:
- q = self.q_norm(q)
- k = self.k_norm(k)
+ q = self.q_norm(self.to_q(x))
+ k = self.k_norm(self.to_k(context))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self_attn = context is None | |
| q = self.to_q(x) | |
| context = x if context is None else context | |
| k = self.to_k(context) | |
| v = self.to_v(context) | |
| q = self.q_norm(q) | |
| k = self.k_norm(k) | |
| # These norms span all heads, so the per-head RMS+RoPE kernel is not equivalent. | |
| if pe is not None: | |
| if k_pe is None and q.shape == k.shape: | |
| q, k = apply_rotary_emb_qk(q, k, pe) | |
| else: | |
| q = apply_rotary_emb(q, pe) | |
| k = apply_rotary_emb(k, pe if k_pe is None else k_pe) | |
| if mask is None: | |
| out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options) | |
| elif isinstance(mask, GuideAttentionMask): | |
| out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options) | |
| # Spatio-Temporal Guidance (STG) perturbation: for the flagged self-attention | |
| # layers, the attention degrades to a passthrough of the value projection (out = V). | |
| if self_attn and transformer_options.get("stg_skip_self_attn", False): | |
| out = v | |
| self_attn = context is None | |
| context = x if context is None else context | |
| v = self.to_v(context) | |
| # Spatio-Temporal Guidance (STG) perturbation: for the flagged self-attention | |
| # layers, the attention degrades to a passthrough of the value projection (out = V). | |
| if self_attn and transformer_options.get("stg_skip_self_attn", False): | |
| out = v | |
| else: | |
| q = self.q_norm(self.to_q(x)) | |
| k = self.k_norm(self.to_k(context)) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/ldm/lightricks/model.py` around lines 465 - 474, Update the attention
logic around the self_attn and stg_skip_self_attn check so self-attention
computes q with self.to_q(x) only in the non-STG path; preserve the STG bypass
as out = v and avoid running the discarded query projection on that branch.
| # Fixed-seed noise so decodes are reproducible TODO: expose? | ||
| generator = torch.Generator(device=x.device) | ||
| generator.manual_seed(0) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Address the fixed-seed TODO.
decode hardcodes seed 0. Every decode of the same latent yields identical noise, which is good for reproducibility, but the TODO signals that the seed should be selectable. Do you want me to open an issue to track exposing the seed through vae_options?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/ldm/lightricks/vae/na_diffusion_decoder.py` around lines 512 - 514,
Update the decode method’s fixed-seed noise handling to expose a selectable seed
through vae_options, using seed 0 as the default to preserve reproducibility.
Read the configured seed when creating the torch.Generator and remove the
unresolved TODO.
| elif "decoder.conv_in_x_t.weight" in sd: # lightricks LTX 2.4 diffusion VAE decoder | ||
| vae_config = None | ||
| if metadata is not None and "config" in metadata: | ||
| vae_config = json.loads(metadata["config"]).get("vae", None) | ||
| self.first_stage_model = comfy.ldm.lightricks.vae.na_diffusion_decoder.CausalDiffusionVAE(config=vae_config) | ||
| self.latent_channels = sd["decoder.conv_in.weight"].shape[1] | ||
| self.latent_dim = 3 | ||
| self.disable_offload = True | ||
| self.crop_input = False # generic crop would narrow the frame axis by the 32x spatial ratio | ||
| self.memory_used_decode = lambda shape, dtype: (1700 * shape[2] * shape[3] * shape[4] * (8 * 8 * 8)) * model_management.dtype_size(dtype) | ||
| self.memory_used_encode = lambda shape, dtype: (80 * max(shape[2], 7) * shape[3] * shape[4]) * model_management.dtype_size(dtype) | ||
| self.upscale_ratio = (lambda a: max(0, a * 8 - 7), 32, 32) | ||
| self.upscale_index_formula = (8, 32, 32) | ||
| self.downscale_ratio = (lambda a: max(0, math.floor((a + 7) / 8)), 32, 32) | ||
| self.downscale_index_formula = (8, 32, 32) | ||
| self.working_dtypes = [torch.bfloat16, torch.float32] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard the decoder.conv_in.weight lookup.
The branch condition tests decoder.conv_in_x_t.weight, but line 592 dereferences a different key, decoder.conv_in.weight. A checkpoint that carries the diffusion decoder without that exact key raises KeyError during VAE construction instead of failing with a clear message.
LTX_24_VAE_CONFIG already declares in_channels: 128, and the constructed NADiffusionDecoder uses the config value, so the state-dict read can disagree with the built model.
🛡️ Proposed fix
- self.latent_channels = sd["decoder.conv_in.weight"].shape[1]
+ self.latent_channels = self.first_stage_model.decoder.conv_in.in_featuresAs per coding guidelines: "Model detectors must inspect linear weight shapes using only the first dimension, guard every dereferenced state-dict key, and order specific signatures before broad fallbacks." The branch ordering before the generic decoder.conv_in.weight case at line 603 is correct.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/sd.py` around lines 587 - 602, Guard the `sd["decoder.conv_in.weight"]`
access in the LTX 2.4 decoder branch by checking that key exists before deriving
`self.latent_channels`; otherwise fail clearly rather than raising an unhandled
`KeyError`. Keep the existing `decoder.conv_in_x_t.weight` detector and branch
ordering unchanged.
Source: Coding guidelines
| # Reserve as much as an untiled decode could use (capped by what the device can provide), then size the tiles to fill that reservation: | ||
| # shrink the temporal tile until one tile fits, then grow the spatial tile while it still fits. | ||
| budget = min(memory_used, int(model_management.get_total_memory(self.device) * 0.8)) | ||
| model_management.load_models_gpu([self.patcher], memory_required=budget, force_full_load=self.disable_offload) | ||
| tile_t = samples_in.shape[2] | ||
| est = lambda tt, txy: self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, txy, txy, tt), self.vae_dtype) | ||
| while tile_t > 2 and est(tile_t, tile) > budget: | ||
| tile_t = -(-tile_t // 2) | ||
| while tile * 2 <= max(samples_in.shape[3], samples_in.shape[4]) and est(tile_t, tile * 2) <= budget: | ||
| tile *= 2 | ||
| overlap = tile // 4 | ||
| pixel_samples = self.decode_tiled_3d(samples_in, tile_t=tile_t, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Size the tiled fallback budget from free memory, not total memory.
This branch runs only after an untiled decode raised OOM. Line 1248 derives the budget from model_management.get_total_memory(self.device) * 0.8. Total memory includes memory already held by other loaded models and by the caller's tensors, so the budget can exceed what the device can actually supply. The tile loops then grow tile until it fills that budget, and the fallback can OOM again.
Use the free memory reported after soft_empty_cache() instead.
🛡️ Proposed fix
- budget = min(memory_used, int(model_management.get_total_memory(self.device) * 0.8))
+ budget = min(memory_used, int(model_management.get_free_memory(self.device) * 0.8))This change also affects every 3D VAE that lacks handles_tiling, not only the new LTX 2.4 path. The previous code used a fixed tile = 256 // spacial_compression_decode(). Please confirm the new dynamic sizing was validated on the existing Wan, Mochi, and CogVideoX decode paths.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Reserve as much as an untiled decode could use (capped by what the device can provide), then size the tiles to fill that reservation: | |
| # shrink the temporal tile until one tile fits, then grow the spatial tile while it still fits. | |
| budget = min(memory_used, int(model_management.get_total_memory(self.device) * 0.8)) | |
| model_management.load_models_gpu([self.patcher], memory_required=budget, force_full_load=self.disable_offload) | |
| tile_t = samples_in.shape[2] | |
| est = lambda tt, txy: self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, txy, txy, tt), self.vae_dtype) | |
| while tile_t > 2 and est(tile_t, tile) > budget: | |
| tile_t = -(-tile_t // 2) | |
| while tile * 2 <= max(samples_in.shape[3], samples_in.shape[4]) and est(tile_t, tile * 2) <= budget: | |
| tile *= 2 | |
| overlap = tile // 4 | |
| pixel_samples = self.decode_tiled_3d(samples_in, tile_t=tile_t, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) | |
| # Reserve as much as an untiled decode could use (capped by what the device can provide), then size the tiles to fill that reservation: | |
| # shrink the temporal tile until one tile fits, then grow the spatial tile while it still fits. | |
| budget = min(memory_used, int(model_management.get_free_memory(self.device) * 0.8)) | |
| model_management.load_models_gpu([self.patcher], memory_required=budget, force_full_load=self.disable_offload) | |
| tile_t = samples_in.shape[2] | |
| est = lambda tt, txy: self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, txy, txy, tt), self.vae_dtype) | |
| while tile_t > 2 and est(tile_t, tile) > budget: | |
| tile_t = -(-tile_t // 2) | |
| while tile * 2 <= max(samples_in.shape[3], samples_in.shape[4]) and est(tile_t, tile * 2) <= budget: | |
| tile *= 2 | |
| overlap = tile // 4 | |
| pixel_samples = self.decode_tiled_3d(samples_in, tile_t=tile_t, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/sd.py` around lines 1246 - 1257, Update the tiled fallback budget
calculation in the surrounding decode method to use free memory after
soft_empty_cache(), rather than 80% of get_total_memory(self.device). Preserve
the existing tile sizing loops and ensure the dynamic sizing remains compatible
with all 3D VAE paths lacking handles_tiling, including Wan, Mochi, and
CogVideoX.
| def _tile_bounded_shape(self, shape, tile_x, tile_y, tile_t): | ||
| """Clamp a latent shape to one tile for memory estimates: peak memory of a tiled decode is per-tile. Only caller-provided tile dims are clamped.""" | ||
| s = list(shape) | ||
| if len(s) == 5: | ||
| if tile_t is not None: | ||
| s[2] = min(s[2], tile_t) | ||
| if tile_y is not None: | ||
| s[3] = min(s[3], tile_y) | ||
| if tile_x is not None: | ||
| s[4] = min(s[4], tile_x) | ||
| else: | ||
| if tile_y is not None: | ||
| s[2] = min(s[2], tile_y) | ||
| if tile_x is not None: | ||
| s[3] = min(s[3], tile_x) | ||
| return tuple(s) | ||
|
|
||
| def decode_tiled(self, samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None): | ||
| self.throw_exception_if_invalid() | ||
| memory_used = self.memory_used_decode(samples.shape, self.vae_dtype) #TODO: calculate mem required for tile | ||
| memory_used = self.memory_used_decode(self._tile_bounded_shape(samples.shape, tile_x, tile_y, tile_t), self.vae_dtype) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm which nodes call VAE.decode_tiled with tile_x on 1D-latent VAEs.
rg -n -C5 '\.decode_tiled\s*\(' --type=py -g '!comfy/sd.py'Repository: Comfy-Org/ComfyUI
Length of output: 155
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- changed lines in comfy/sd.py ---'
git diff --unified=20 -- comfy/sd.py | sed -n '1,220p'
printf '%s\n' '--- decode_tiled definitions and call sites ---'
rg -n -C8 'decode_tiled|extra_1d_channel|decode_1d' --glob '*.py' .
printf '%s\n' '--- relevant method body ---'
sed -n '1215,1325p' comfy/sd.pyRepository: Comfy-Org/ComfyUI
Length of output: 42650
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from pathlib import Path
p = Path("comfy/sd.py")
text = p.read_text()
start = text.index(" def _tile_bounded_shape")
end = text.index("\n def ", start + 5)
print(text[start:end])
PYRepository: Comfy-Org/ComfyUI
Length of output: 844
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
def bounded(shape, tile_x, tile_y, tile_t):
s = list(shape)
if len(s) == 5:
if tile_t is not None:
s[2] = min(s[2], tile_t)
if tile_y is not None:
s[3] = min(s[3], tile_y)
if tile_x is not None:
s[4] = min(s[4], tile_x)
else:
if tile_y is not None:
s[2] = min(s[2], tile_y)
if tile_x is not None:
s[3] = min(s[3], tile_x)
return tuple(s)
cases = [
((1, 2, 2_000), 256, 256, None), # 1D audio latent
((1, 256, 1, 2_000), 256, 256, None), # extra_1d_channel layout
((1, 4, 2, 64, 64), 32, 32, None), # video latent
]
for shape, tile_x, tile_y, tile_t in cases:
try:
print(shape, "->", bounded(shape, tile_x, tile_y, tile_t))
except Exception as e:
print(shape, "->", type(e).__name__, str(e))
PY
printf '%s\n' '--- audio VAE shape-producing code ---'
rg -n -C12 'extra_1d_channel|latent_dim = 1|vae_decode_audio|encode_tiled_1d|decode_tiled_1d' comfy comfy_extras nodes.py tests-unit --glob '*.py' | sed -n '1,320p'Repository: Comfy-Org/ComfyUI
Length of output: 25835
Guard 3D audio latent shapes before indexing. decode_tiled() calls _tile_bounded_shape() before its 1D dispatch. With a shape such as (1, 2, 2000) and tile_x set, the helper raises IndexError at s[3]. Handle 3D shapes by clamping s[2] with tile_x, and restrict the existing branch to four-dimensional shapes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/sd.py` around lines 1262 - 1281, Update _tile_bounded_shape to handle
three-dimensional latent shapes before indexing spatial dimensions: clamp s[2]
using tile_x, and change the existing non-five-dimensional branch to apply only
to four-dimensional shapes. Preserve the current five-dimensional handling and
avoid accessing s[3] for 3D audio shapes.
| if self.text_projection_type == "single_linear": | ||
| self.text_embedding_projection = operations.Linear(3840 * 49, 3840, bias=False, dtype=dtype, device=device) | ||
| self.text_embedding_projection = operations.Linear(projection_in_dim, video_projection_dim, bias=video_projection_bias, dtype=dtype, device=device) | ||
| elif self.text_projection_type == "dual_linear": | ||
| self.text_embedding_projection = DualLinearProjection(3840 * 49, 4096, 2048, dtype=dtype, device=device, operations=operations) | ||
| self.text_embedding_projection = DualLinearProjection(projection_in_dim, video_projection_dim, audio_projection_dim, video_bias=video_projection_bias, audio_bias=audio_projection_bias, dtype=dtype, device=device, operations=operations) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject unsupported projection types.
If text_projection_type is not single_linear or dual_linear, the constructor does not create text_embedding_projection. encode_token_weights then returns unprojected encoder output instead of failing with a clear error. Add an else branch that raises ValueError.
As per coding guidelines, “Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently degrading output.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/text_encoders/lt.py` around lines 144 - 147, Add an else branch to the
text projection selection in the constructor, after the single_linear and
dual_linear cases, that raises a clear ValueError for any unsupported
text_projection_type. Ensure unsupported values cannot leave
text_embedding_projection uninitialized and fall through to unprojected output.
Source: Coding guidelines
|
To avoid confusion, replacing |
5 participants
No description provided.