Merge branch 'master' into fix/save-3d-advanced-output-item

.github/workflows/cla.yml

@@ -22,7 +22,7 @@ jobs:
       - name: Build author-only allowlist
         id: allowlist
         if: >
-          github.event_name == 'pull_request_target' ||
+          (github.event_name == 'pull_request_target' && github.event.action != 'closed') ||
           (github.event_name == 'issue_comment' && github.event.issue.pull_request && (
             github.event.comment.body == 'recheck' ||
             github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
@@ -48,10 +48,10 @@ jobs:
           fi
 
       - name: CLA Assistant
-        # Run on PR events, on "recheck" comment, or when someone posts the signing phrase.
+        # Run on open/update PR events, on "recheck", or when someone posts the signing phrase.
         # IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
         if: >
-          github.event_name == 'pull_request_target' ||
+          (github.event_name == 'pull_request_target' && github.event.action != 'closed') ||
           (github.event_name == 'issue_comment' && github.event.issue.pull_request && (
             github.event.comment.body == 'recheck' ||
             github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
@@ -62,6 +62,8 @@ jobs:
           # PAT required to write to the centralized signatures repo.
           PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
         with:
+          lock-pullrequest-aftermerge: false
+
           # Where the CLA document lives (shown to contributors)
           path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
 
@@ -94,3 +96,25 @@ jobs:
 
           custom-allsigned-prcomment: |
             ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
+
+      - name: Lock PR containing a CLA signature or bot approval
+        if: github.event_name == 'pull_request_target' && github.event.action == 'closed' && github.event.pull_request.merged == true
+        uses: actions/github-script@v7
+        with:
+          retries: 3
+          script: |
+            const pr = context.payload.pull_request;
+            const issue = { ...context.repo, issue_number: pr.number };
+            const comments = await github.paginate(github.rest.issues.listComments, {
+              ...issue,
+              per_page: 100,
+            });
+            const signed = comments.some(comment =>
+              (comment.user?.id === pr.user.id &&
+                comment.body?.trim().toLowerCase() === 'i have read and agree to the contributor license agreement') ||
+              (comment.user?.login === 'github-actions[bot]' &&
+                comment.body?.startsWith('✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.'))
+            );
+            if (signed) {
+              await github.rest.issues.lock(issue);
+            }

comfy/controlnet.py

@@ -501,7 +501,7 @@ def controlnet_config(sd, model_options={}):
 
     operations = model_options.get("custom_operations", None)
     if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)
+        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, load_device=load_device, disable_fast_fp8=True)
 
     offload_device = comfy.model_management.unet_offload_device()
     return model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device
@@ -585,7 +585,7 @@ def load_controlnet_sd35(sd, model_options={}):
 
     operations = model_options.get("custom_operations", None)
     if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)
+        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, load_device=load_device, disable_fast_fp8=True)
 
     control_model = comfy.cldm.dit_embedder.ControlNetEmbedder(img_size=None,
                                                                patch_size=2,
@@ -683,7 +683,7 @@ def load_controlnet_qwen_fun(sd, model_options={}):
 
     operations = model_options.get("custom_operations", None)
     if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)
+        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, load_device=load_device, disable_fast_fp8=True)
 
     in_features = sd["control_img_in.weight"].shape[1]
     inner_dim = sd["control_img_in.weight"].shape[0]
@@ -838,7 +838,7 @@ def load_controlnet_state_dict(state_dict, model=None, model_options={}):
     manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
     operations = model_options.get("custom_operations", None)
     if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype)
+        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, load_device=load_device, disable_fast_fp8=True)
 
     controlnet_config["operations"] = operations
     controlnet_config["dtype"] = unet_dtype

comfy/latent_formats.py

@@ -9,6 +9,7 @@ class LatentFormat:
     latent_rgb_factors_bias = None
     latent_rgb_factors_reshape = None
     taesd_decoder_name = None
+    compile_preview = False
     spacial_downscale_ratio = 8
     temporal_downscale_ratio = 1
 
@@ -625,6 +626,7 @@ class MiniMaxH3Video(LatentFormat):
     temporal_downscale_ratio = 4
     scale_factor = 1.0
     taesd_decoder_name = "taeh3"
+    compile_preview = True
 
     latent_rgb_factors = [
         [-0.018555,  0.024344, -0.017536],

comfy/ldm/cosmos/predict2.py

@@ -881,13 +881,6 @@ class MiniTrainDIT(nn.Module):
         t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder[1](self.t_embedder[0](timesteps_B_T).to(x_B_T_H_W_D.dtype))
         t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)
 
-        # for logging purpose
-        affline_scale_log_info = {}
-        affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach()
-        self.affline_scale_log_info = affline_scale_log_info
-        self.affline_emb = t_embedding_B_T_D
-        self.crossattn_emb = crossattn_emb
-
         if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
             assert (
                 x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape

comfy/ldm/lightricks/av_model.py

@@ -938,7 +938,7 @@ class LTXAVModel(LTXVModel):
         stg_self_attn_blocks = transformer_options.get("stg_self_attn_blocks", ())
 
         # Process transformer blocks
-        comfy.model_prefetch.malloc_graph_begin(self, vx.device)
+        comfy.model_prefetch.malloc_graph_begin(vx.device)
         for i, block in enumerate(self.transformer_blocks):
             comfy.model_prefetch.prefetch_queue_pop(
                 prefetch_queue, vx.device, block, malloc_scope="block"

comfy/ldm/minimax/model.py

@@ -573,7 +573,7 @@ class MiniMaxH3Model(nn.Module):
         compile_allocations = comfy.model_prefetch.malloc_graph_enabled(x[0].device)
         if compile_allocations:
             out = [torch.empty_like(x[0]), torch.empty_like(x[1])]
-            comfy.model_prefetch.malloc_graph_begin(self, x[0].device)
+            comfy.model_prefetch.malloc_graph_begin(x[0].device)
         graph_out = comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
@@ -588,6 +588,12 @@ class MiniMaxH3Model(nn.Module):
         else:
             out = graph_out
 
+        # Masked rows predict at mask * sigma; scale their velocity to match the outer x0 conversion.
+        if denoise_mask is not None:
+            out[0] = out[0] * denoise_mask
+        if audio_denoise_mask is not None:
+            out[1] = out[1] * audio_denoise_mask
+
         if scale != 1.0:
             # d/d(sigma_v) of the carried variable
             out[1] = ((1.0 - scale) * (audio_src * carry)

comfy/ldm/minimax_music/ar.py

@@ -297,7 +297,7 @@ class MiniMaxMusic3AR(nn.Module):
                         break
 
             if frame_index:
-                comfy.model_prefetch.malloc_graph_begin(self, device)
+                comfy.model_prefetch.malloc_graph_begin(device)
             c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask)
             if pending_code is None:
                 pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device)

comfy/ldm/trellis2/model.py

@@ -768,10 +768,7 @@ def _coords_to_proj_world(coords: torch.Tensor, resolution: int, mesh_scale: tor
         norm = coords[:, 1:].to(torch.float32) / (resolution - 1) * 2.0 - 1.0
     R = _PROJ_GRID_ROTATION.to(device=coords.device, dtype=torch.float32)
     rotated = norm @ R.T
-    if mesh_scale.ndim == 0:
-        scale_per_voxel = mesh_scale.expand(coords.shape[0])
-    else:
-        scale_per_voxel = mesh_scale.to(coords.device)[batch_ids]
+    scale_per_voxel = mesh_scale.to(coords.device).reshape(-1)[batch_ids % mesh_scale.numel()]
     world = rotated / scale_per_voxel.unsqueeze(-1) / 2.0
     return world, batch_ids
 
@@ -800,20 +797,22 @@ def _back_project_to_tokens(
 ) -> torch.Tensor:
     if coords_world.dim() == 2:
         assert batch_ids is not None
-        B = transform_matrix.shape[0]
+        n_cond = transform_matrix.shape[0]
         out = torch.zeros((coords_world.shape[0], feature_map.shape[1]),
                           device=feature_map.device, dtype=feature_map.dtype)
-        for b in range(B):
+        # multi-seed: latent sample b uses conditioning b % n_cond
+        for b in range(int(batch_ids.max().item()) + 1):
             mask = batch_ids == b
             if not mask.any():
                 continue
+            c = b % n_cond
             p = coords_world[mask].unsqueeze(0)
             uv, _, _ = _project_points_to_image(
-                p, transform_matrix[b:b+1], camera_angle_x[b:b+1], image_resolution)
+                p, transform_matrix[c:c+1], camera_angle_x[c:c+1], image_resolution)
             uv_ndc = (uv + 0.5) / image_resolution * 2.0 - 1.0
             # padding_mode='border' is load-bearing: masking out-of-frame voxels confuses
             # the SS DiT (~half the voxels go to zero, producing low poly + rotation drift).
-            sampled = _sample_features(feature_map[b:b+1], uv_ndc)
+            sampled = _sample_features(feature_map[c:c+1], uv_ndc)
             sampled = sampled.squeeze(0).transpose(0, 1)
             out[mask] = sampled
         return out
@@ -846,25 +845,14 @@ def compute_stage_proj_feats(
     batch_size: Optional[int] = None,
     device=None,
 ) -> torch.Tensor:
-    """Back-project a Pixal3D stage's feature maps onto its target voxel/grid coords.
-
-    For sparse (shape / texture) stages: pass ``coords`` (with ``coord_resolution``).
-    Returns ``[N_voxels, C]`` per-voxel features with channel count =
-    LR channels + optional HR channels.
-
-    For the dense SS stage: pass ``dense_grid_resolution`` (16) + ``batch_size``.
-    Returns ``[B, R^3, C]`` features for the dense grid.
-
-    """
+    """Back-project a stage's feature maps onto sparse coords [N, C] or the dense SS grid [B, R^3, C]; views at stride num_views are averaged."""
     if device is None:
         device = coords.device if coords is not None else proj_pack["mesh_scale"].device
     mesh_scale = proj_pack["mesh_scale"].to(device)
     T = proj_pack["transform_matrix"].to(device)
     cam_angle = proj_pack["camera_angle_x"].to(device)
+    num_views = int(proj_pack.get("num_views", 1))
     feat_map_lr, feat_map_hr, image_resolution = _select_stage_entry(proj_pack, stage)
-    feat_map_lr = feat_map_lr.to(device)
-    if feat_map_hr is not None:
-        feat_map_hr = feat_map_hr.to(device)
 
     if coords is not None:
         if coord_resolution is None:
@@ -877,13 +865,18 @@ def compute_stage_proj_feats(
                                               device=device, dtype=torch.float32)
         batch_ids = None
 
-    proj_lr = _back_project_to_tokens(coords_world, feat_map_lr, T, cam_angle,
+    out = None
+    for v in range(num_views):
+        sel = slice(v, None, num_views)
+        proj = _back_project_to_tokens(coords_world, feat_map_lr[sel].to(device), T[sel], cam_angle[sel],
                                        image_resolution=image_resolution, batch_ids=batch_ids)
-    if feat_map_hr is not None:
-        proj_hr = _back_project_to_tokens(coords_world, feat_map_hr, T, cam_angle,
-                                           image_resolution=image_resolution, batch_ids=batch_ids)
-        return torch.cat([proj_lr, proj_hr], dim=-1)
-    return proj_lr
+        if feat_map_hr is not None:
+            proj_hr = _back_project_to_tokens(coords_world, feat_map_hr[sel].to(device), T[sel], cam_angle[sel],
+                                              image_resolution=image_resolution, batch_ids=batch_ids)
+            proj = torch.cat([proj, proj_hr], dim=-1)
+        # average views in fp32 like upstream
+        out = proj.float() if out is None else out.add_(proj)
+    return out.div_(num_views).to(proj.dtype)
 
 
 def _shape_proj_cond(global_cond: torch.Tensor, image_attn_mode: str,
@@ -917,6 +910,11 @@ def _shape_proj_cond(global_cond: torch.Tensor, image_attn_mode: str,
             f"sub-model expects {proj_in_channels}.{hint}"
         )
 
+    # multi-seed: latent sample i uses conditioning i % B
+    if batch_ids is None and logical_batch is not None and proj_feats.shape[0] != logical_batch:
+        reps = -(-logical_batch // proj_feats.shape[0])
+        proj_feats = proj_feats.repeat((reps,) + (1,) * (proj_feats.ndim - 1))[:logical_batch]
+
     # CFG-duplicate proj_feats to match the model's eval batch.
     if eval_batch is not None and logical_batch is not None and eval_batch > logical_batch:
         repeats = eval_batch // logical_batch
@@ -1125,7 +1123,7 @@ class Trellis2(nn.Module):
 
         else: # structure
             struct_attn = self.image_attn_mode_structure
-            logical_batch_ss = proj_feats.shape[0] if proj_feats is not None else x.shape[0]
+            logical_batch_ss = x.shape[0] // len(cond_or_uncond) if cond_or_uncond else x.shape[0]
             struct_cond = context
             if struct_attn != "global":
                 struct_cond = _shape_proj_cond(context, struct_attn, proj_feats,

comfy/model_management.py

@@ -499,19 +499,31 @@ try:
 
             can_use_flash_attention() evaluates runtime eligibility for the given
             parameters; on a ROCm build that includes checking the gpu arch against the
-            kernel images AOTriton was compiled for. Querying it avoids assuming where
-            those images live inside the torch install. The probe tensor is shaped and
+            arches AOTriton was built for. Querying it avoids assuming where the kernel
+            images live inside the torch install. The probe tensor is shaped and
             typed to pass the unrelated SDPA checks, so False means no hardware support
             rather than a rejected shape.
+
+            It answers True on a supported arch whose kernel image was never shipped,
+            and that only fails at launch, without raising. So run one attention
+            through the flash backend and force the pending error check.
             """
             try:
+                device = get_torch_device()
                 if not torch.backends.cuda.is_flash_attention_available():  # not built with flash attention
                     return False
-                q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device())
+                q = torch.zeros((1, 1, 8, 64), dtype=torch.float16, device=device)
                 params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False)
-                return torch.backends.cuda.can_use_flash_attention(params, False)
-            except (AttributeError, RuntimeError, TypeError) as e:
-                logging.warning("Could not query aotriton support: {}".format(e))
+                if not torch.backends.cuda.can_use_flash_attention(params, False):
+                    return False
+                from torch.nn.attention import SDPBackend, sdpa_kernel
+                with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+                    torch.nn.functional.scaled_dot_product_attention(q, q, q)
+                torch.cuda.synchronize()
+                torch.zeros(1, device=device).add_(1).item()  # raises if the launch above failed
+                return True
+            except Exception as e:
+                logging.warning("Could not run flash attention, disabling it: {}".format(e))
                 return False
 
         logging.info("AMD arch: {}".format(arch))
@@ -1373,7 +1385,6 @@ LARGEST_CASTED_WEIGHT = (None, 0)
 STREAM_AIMDO_CAST_BUFFERS = {}
 LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
 CROSS_STEP_STATE = weakref.WeakSet()
-MALLOC_GRAPH_MODULES = weakref.WeakSet()
 
 DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
 
@@ -1459,9 +1470,6 @@ def reset_cast_buffers():
 
     STREAM_CAST_BUFFERS.clear()
     STREAM_AIMDO_CAST_BUFFERS.clear()
-    for module in MALLOC_GRAPH_MODULES:
-        del module._comfy_malloc_graph
-    MALLOC_GRAPH_MODULES.clear()
     soft_empty_cache()
 
 def get_offload_stream(device):
@@ -2022,6 +2030,24 @@ def supports_fp64(device=None):
 
     return True
 
+def supports_int8_compute(device=None):
+    # The eager comfy_kitchen backend implements int8 weight-only quantized
+    # matmul via torch._int_mm, which PyTorch does not implement for MPS.
+    # https://github.com/pytorch/pytorch/issues/141287
+    if (device is not None and is_device_mps(device)) or mps_mode():
+        return False
+
+    if is_intel_xpu():
+        return False
+
+    if is_directml_enabled():
+        return False
+
+    if is_ixuca():
+        return False
+
+    return True
+
 def extended_fp16_support():
     # TODO: check why some models work with fp16 on newer torch versions but not on older
     if torch_version_numeric < (2, 7):

comfy/model_prefetch.py

@@ -1,9 +1,9 @@
-import contextlib
 import logging
 import threading
 import warnings
 import weakref
 
+import comfy_kitchen as ck
 import torch
 
 import comfy_aimdo.malloc_graph
@@ -14,11 +14,11 @@ import comfy.model_management
 import comfy.ops
 
 PREFETCH_QUEUES = []
-GRAPH_MODULES = weakref.WeakSet()
 GRAPH_WARMED_MODULES = weakref.WeakSet()
 GRAPH_CAPTURE_STREAMS = {}
-ACTIVE_MALLOC_GRAPHS = {}
+MALLOC_GRAPHS = {}
 MALLOC_GRAPH_BREAKS = 0
+MALLOC_GRAPH_ROGUES = 0
 MALLOC_GRAPH_USED = False
 
 def _malloc_graph_break():
@@ -29,40 +29,62 @@ def _malloc_graph_break():
 def malloc_graph_enabled(device):
     return not args.disable_comfy_compiler and comfy.memory_management.aimdo_enabled and comfy.model_management.is_device_cuda(device)
 
-@contextlib.contextmanager
-def pause_malloc_graph(sync=False):
-    graph = ACTIVE_MALLOC_GRAPHS.get(threading.get_ident())
-    if graph is not None:
-        graph.pause(sync=sync)
-    try:
-        yield
-    finally:
-        if graph is not None:
-            graph.resume(sync=sync)
+class _PauseMallocGraph:
+    def __init__(self, sync=False):
+        self.sync = sync
 
-def malloc_graph_begin(module, device):
+    def __enter__(self):
+        graph = MALLOC_GRAPHS.get(threading.get_ident())
+        if graph is not None and graph._comfy_active:
+            graph.pause(sync=self.sync)
+
+    def __exit__(self, *args):
+        graph = MALLOC_GRAPHS.get(threading.get_ident())
+        if graph is not None and graph._comfy_active:
+            graph.resume(sync=self.sync)
+
+def pause_malloc_graph(sync=False):
+    return _PauseMallocGraph(sync)
+
+def malloc_graph_begin(device):
     global MALLOC_GRAPH_USED
     if not malloc_graph_enabled(device):
         return
-    graph = getattr(module, "_comfy_malloc_graph", None)
+    thread_id = threading.get_ident()
+    graph = MALLOC_GRAPHS.get(thread_id)
     if graph is None:
         graph = comfy_aimdo.malloc_graph.record(
             comfy.model_management.current_stream(device), args.assert_graph_breaks
         )
-        module._comfy_malloc_graph = graph
-        comfy.model_management.MALLOC_GRAPH_MODULES.add(module)
+        graph._comfy_cuda_graph_modules = weakref.WeakSet()
+        MALLOC_GRAPHS[thread_id] = graph
     else:
         graph.push()
-    ACTIVE_MALLOC_GRAPHS[threading.get_ident()] = graph
+    if hasattr(ck, "set_allocation_context"):
+        ck.set_allocation_context(pause_malloc_graph())
+    graph._comfy_active = True
     MALLOC_GRAPH_USED = True
 
 def malloc_graph_end():
     thread_id = threading.get_ident()
-    graph = ACTIVE_MALLOC_GRAPHS.get(thread_id)
-    if graph is not None:
+    graph = MALLOC_GRAPHS.get(thread_id)
+    if graph is not None and graph._comfy_active:
         if graph.pop():
             _malloc_graph_break()
-        ACTIVE_MALLOC_GRAPHS.pop(thread_id)
+        graph._comfy_active = False
+
+def cleanup_malloc_graph():
+    global MALLOC_GRAPH_ROGUES
+
+    graph = MALLOC_GRAPHS.pop(threading.get_ident(), None)
+    if graph is not None:
+        if graph._comfy_active:
+            graph.abort()
+            graph._comfy_active = False
+        for module in graph._comfy_cuda_graph_modules:
+            _drop_graph(module)
+        MALLOC_GRAPH_ROGUES += graph.rogue_count
+        del graph
 
 def cleanup_prefetched_modules(module, comfy_modules):
     for s in comfy_modules:
@@ -95,11 +117,10 @@ def _drop_graph(module):
 def cleanup_prefetch_queues():
     global PREFETCH_QUEUES
     global MALLOC_GRAPH_BREAKS
+    global MALLOC_GRAPH_ROGUES
     global MALLOC_GRAPH_USED
 
-    graph = ACTIVE_MALLOC_GRAPHS.pop(threading.get_ident(), None)
-    if graph is not None:
-        graph.abort()
+    cleanup_malloc_graph()
     for queue in PREFETCH_QUEUES:
         for entry in queue:
             if entry is None or not isinstance(entry, tuple):
@@ -109,17 +130,17 @@ def cleanup_prefetch_queues():
             if comfy_modules is not None:
                 cleanup_prefetched_modules(prefetched_module, comfy_modules)
     PREFETCH_QUEUES = []
-    for module in GRAPH_MODULES:
-        _drop_graph(module)
-    GRAPH_MODULES.clear()
     GRAPH_WARMED_MODULES.clear()
     if MALLOC_GRAPH_USED:
-        logging.info("Comfy model compiler graph breaks: %d", MALLOC_GRAPH_BREAKS)
+        logging.info("Comfy model compiler graph breaks: %d, rogues: %d", MALLOC_GRAPH_BREAKS, MALLOC_GRAPH_ROGUES)
     MALLOC_GRAPH_BREAKS = 0
+    MALLOC_GRAPH_ROGUES = 0
     MALLOC_GRAPH_USED = False
 
 def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None, malloc_scope=None):
-    malloc_graph = ACTIVE_MALLOC_GRAPHS.get(threading.get_ident())
+    malloc_graph = MALLOC_GRAPHS.get(threading.get_ident())
+    if malloc_graph is not None and not malloc_graph._comfy_active:
+        malloc_graph = None
     enable_graph = enable_graph and malloc_graph is not None and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) and getattr(module, "_v_block", None) is not None
     if queue is None:
         if malloc_graph is not None and malloc_scope is not None:
@@ -223,7 +244,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
                 comfy.model_management.current_stream(device).wait_stream(capture_stream)
                 graph.replay()
                 module._comfy_graph = {"graph": graph, "signature": signature}
-                GRAPH_MODULES.add(module)
+                malloc_graph._comfy_cuda_graph_modules.add(module)
                 return
         if capture_stream is None:
             core()

comfy/multigpu.py

@@ -11,6 +11,7 @@ if TYPE_CHECKING:
 import comfy.utils
 import comfy.patcher_extension
 import comfy.model_management
+import comfy.model_prefetch
 
 
 class MultiGPUThreadPool:
@@ -46,18 +47,21 @@ class MultiGPUThreadPool:
                     return
                 result_q.put((None, e))
             return
-        while True:
-            item = work_q.get()
-            if item is None:
-                break
-            fn, args, kwargs = item
-            try:
-                result = fn(*args, **kwargs)
-                result_q.put((result, None))
-            except comfy.model_management.InterruptProcessingException as e:
-                result_q.put((None, e))
-            except Exception as e:
-                result_q.put((None, e))
+        try:
+            while True:
+                item = work_q.get()
+                if item is None:
+                    break
+                fn, args, kwargs = item
+                try:
+                    result = fn(*args, **kwargs)
+                    result_q.put((result, None))
+                except comfy.model_management.InterruptProcessingException as e:
+                    result_q.put((None, e))
+                except Exception as e:
+                    result_q.put((None, e))
+        finally:
+            comfy.model_prefetch.cleanup_malloc_graph()
 
     def submit(self, device: torch.device, fn, *args, **kwargs):
         self._work_queues[device].put((fn, args, kwargs))

comfy/ops.py

@@ -1287,6 +1287,16 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
     return sd
 
 
+class MixedPrecisionOp(CastWeightBiasOp):
+    quant_format = None
+
+    def can_use_quantized_matmul(self, disabled_formats):
+        return (self.quant_format in QUANT_ALGOS
+                and not self._full_precision_mm_config
+                and self.quant_format not in self._disabled_formats
+                and self.quant_format not in disabled_formats)
+
+
 def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_precision_mm=False, disabled=[]):
     class MixedPrecisionOps(manual_cast):
         _quant_config = quant_config
@@ -1294,7 +1304,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
         _full_precision_mm = full_precision_mm
         _disabled = disabled
 
-        class Linear(torch.nn.Module, CastWeightBiasOp):
+        class Linear(torch.nn.Module, MixedPrecisionOp):
             _disabled_formats = disabled
 
             def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None):
@@ -1342,6 +1352,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                         compute_dtype=compute_dtype,
                         want_requant=want_requant,
                     ) as (weight, bias):
+                        if self._full_precision_mm and isinstance(weight, QuantizedTensor):
+                            weight = weight.dequantize()
                         return self._forward(input, weight, bias)
 
                 with CastBiasWeightContext(
@@ -1444,7 +1456,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
             def _apply(self, fn, recurse=True):  # This is to get torch.compile + moving weights to another device working
                 return _quantized_apply(self, fn, recurse)
 
-        class MoEExperts(torch.nn.Module, CastWeightBiasOp):
+        class MoEExperts(torch.nn.Module, MixedPrecisionOp):
             """Container for E quantized expert weights, indexed via expert_weight(i).
 
             The bank lives on self.weight as a single 3D tensor — either a
@@ -1648,24 +1660,45 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
 
     return MixedPrecisionOps
 
-def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, model_config=None):
-    fp8_compute = comfy.model_management.supports_fp8_compute(load_device) # TODO: if we support more ops this needs to be more granular
-    nvfp4_compute = comfy.model_management.supports_nvfp4_compute(load_device)
-    mxfp8_compute = comfy.model_management.supports_mxfp8_compute(load_device)
+def get_disabled_quant_formats(device=None):
+    disabled = set()
+    if not comfy.model_management.supports_nvfp4_compute(device):
+        disabled.add("nvfp4")
+    if not comfy.model_management.supports_mxfp8_compute(device):
+        disabled.add("mxfp8")
+    if not comfy.model_management.supports_fp8_compute(device):
+        disabled.add("float8_e4m3fn")
+        disabled.add("float8_e5m2")
+    if not comfy.model_management.supports_int8_compute(device):
+        disabled.add("int8_tensorwise")
+        disabled.add("convrot_w4a4")
+        disabled.add("asym_w4a8_int8")
+    return disabled
 
+
+@contextlib.contextmanager
+def use_quantized_matmul(model, device):
+    disabled = get_disabled_quant_formats(device)
+    previous = []
+    try:
+        for module in model.modules():
+            if isinstance(module, MixedPrecisionOp) and module.can_use_quantized_matmul(disabled):
+                previous.append((module, module._full_precision_mm))
+                module._full_precision_mm = False
+        yield
+    finally:
+        for module, full_precision_mm in previous:
+            module._full_precision_mm = full_precision_mm
+
+
+def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, model_config=None):
     if model_config and hasattr(model_config, 'quant_config') and model_config.quant_config:
         logging.info("Using mixed precision operations")
-        disabled = set()
-        if not nvfp4_compute:
-            disabled.add("nvfp4")
-        if not mxfp8_compute:
-            disabled.add("mxfp8")
-        if not fp8_compute:
-            disabled.add("float8_e4m3fn")
-            disabled.add("float8_e5m2")
+        disabled = get_disabled_quant_formats(load_device)
         logging.info("Native ops: {} {}".format(", ".join(QUANT_ALGOS.keys() - disabled), ", emulated ops: {}".format(", ".join(disabled)) if len(disabled) > 0 else ""))
         return mixed_precision_ops(model_config.quant_config, compute_dtype, disabled=disabled)
 
+    fp8_compute = comfy.model_management.supports_fp8_compute(load_device)
     if (
         fp8_compute and
         (fp8_optimizations or PerformanceFeature.Fp8MatrixMultiplication in args.fast) and

comfy/sd.py

@@ -35,6 +35,7 @@ import os
 
 import comfy.utils
 import comfy.ops
+import comfy.model_prefetch
 
 from . import clip_vision
 from . import gligen
@@ -474,7 +475,7 @@ class CLIP:
         self.cond_stage_model.set_clip_options({"layer": None})
         self.cond_stage_model.set_clip_options({"execution_device": device})
 
-        with model_management.cuda_device_context(device):
+        with model_management.cuda_device_context(device), comfy.ops.use_quantized_matmul(self.cond_stage_model, device):
             return self.cond_stage_model.generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed, presence_penalty=presence_penalty)
 
     def decode(self, token_ids, skip_special_tokens=True):
@@ -1227,7 +1228,8 @@ class VAE:
         with model_management.cuda_device_context(self.device):
             try:
                 memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
-                model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
+                with comfy.model_prefetch.pause_malloc_graph():
+                    model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
                 free_memory = self.patcher.get_free_memory(self.device)
                 batch_number = int(free_memory / memory_used)
                 batch_number = max(1, batch_number)
@@ -1235,7 +1237,8 @@ class VAE:
                 # Pre-allocate output for VAEs that support direct buffer writes
                 preallocated = False
                 if getattr(self.first_stage_model, 'comfy_has_chunked_io', False):
-                    pixel_samples = torch.empty(self.first_stage_model.decode_output_shape(samples_in.shape), device=self.output_device, dtype=self.vae_output_dtype())
+                    with comfy.model_prefetch.pause_malloc_graph():
+                        pixel_samples = torch.empty(self.first_stage_model.decode_output_shape(samples_in.shape), device=self.output_device, dtype=self.vae_output_dtype())
                     preallocated = True
 
                 for x in range(0, samples_in.shape[0], batch_number):
@@ -1245,7 +1248,8 @@ class VAE:
                     else:
                         out = self.first_stage_model.decode(samples, **vae_options).to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True)
                         if pixel_samples is None:
-                            pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype())
+                            with comfy.model_prefetch.pause_malloc_graph():
+                                pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype())
                         pixel_samples[x:x+batch_number].copy_(out)
                         del out
                     self.process_output(pixel_samples[x:x+batch_number])

comfy/text_encoders/ace15.py

@@ -162,7 +162,8 @@ def generate_audio_codes(model, positive, negative, min_tokens=1, max_tokens=102
     else:
         ids = [positive]
 
-    return sample_manual_loop_no_classes(model, ids, cfg_scale=cfg_scale, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, seed=seed, min_tokens=min_tokens, max_new_tokens=max_tokens)
+    with comfy.ops.use_quantized_matmul(model, model.execution_device):
+        return sample_manual_loop_no_classes(model, ids, cfg_scale=cfg_scale, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, seed=seed, min_tokens=min_tokens, max_new_tokens=max_tokens)
 
 
 class ACE15Tokenizer(sd1_clip.SD1Tokenizer):

comfy/text_encoders/gemma4.py

@@ -528,17 +528,18 @@ class Gemma4Transformer(nn.Module):
                         and comfy.model_management.is_device_cuda(x.device))
         decode_bias = None
         decode_masks = None
-        if decode:
+        if fixed_kv:
+            # prefill must advance the device-side write position of the global caches too
             prepared = set()
             for kv in past_key_values:
                 if isinstance(kv, FixedKV) and id(kv.position) not in prepared:
                     kv.prepare(seq_len)
                     prepared.add(id(kv.position))
-            if mask is not None:
-                decode_masks = {}
-                for kv in past_key_values:
-                    if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
-                        decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
+        if decode and mask is not None:
+            decode_masks = {}
+            for kv in past_key_values:
+                if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
+                    decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
         if compiled_decode:
             capacities = tuple(sorted({kv.key.shape[2] for kv in past_key_values if isinstance(kv, FixedKV)}))
             valid = past_len + 1

comfy/text_encoders/llama.py

@@ -1036,7 +1036,7 @@ class BaseGenerate:
         for step in tqdm(range(max_length), desc="Generating tokens"):
             if step > 0:
                 if compile_allocations:
-                    comfy.model_prefetch.malloc_graph_begin(self, device)
+                    comfy.model_prefetch.malloc_graph_begin(device)
                 embeds = self.model.embed_tokens(decode_tokens).to(execution_dtype)
                 current_input_ids = decode_tokens if initial_input_ids is not None else None
                 position_ids = torch.tensor([[next_pos]], device=device) if next_pos is not None else None

comfy_api_nodes/nodes_openai.py

@@ -402,13 +402,18 @@ class OpenAIGPTImage1(IO.ComfyNode):
         return IO.NodeOutput(await validate_and_cast_response(response))
 
 
-def _gpt_image_shared_inputs():
+GPT_IMAGE_QUALITIES = ("low", "medium", "high")
+GPT_IMAGE_25_QUALITIES = ("low", "medium", "high", "xhigh", "max")
+GPT_IMAGE_MODELS = ("gpt-image-2.5-flare", "gpt-image-2.5-sunburst", "gpt-image-2", "gpt-image-1.5", "gpt-image-1")
+
+
+def _gpt_image_shared_inputs(qualities: tuple[str, ...] = GPT_IMAGE_QUALITIES):
     """Inputs shared by all GPT Image models (quality + reference images + mask)."""
     return [
         IO.Combo.Input(
             "quality",
             default="low",
-            options=["low", "medium", "high"],
+            options=list(qualities),
             tooltip="Image quality, affects cost and generation time.",
         ),
         IO.Autogrow.Input(
@@ -448,13 +453,58 @@ def _gpt_image_legacy_model_inputs():
     ]
 
 
+def _gpt_image_2_model_inputs(backgrounds: tuple[str, ...], qualities: tuple[str, ...]):
+    return [
+        IO.Combo.Input(
+            "size",
+            default="auto",
+            options=[
+                "auto",
+                "1024x1024",
+                "1024x1536",
+                "1536x1024",
+                "2048x2048",
+                "2048x1152",
+                "1152x2048",
+                "3840x2160",
+                "2160x3840",
+                "Custom",
+            ],
+            tooltip="Image size. Select 'Custom' to use the custom width and height.",
+        ),
+        IO.Int.Input(
+            "custom_width",
+            default=1024,
+            min=480,
+            max=3840,
+            step=16,
+            tooltip="Used only when `size` is 'Custom'. Must be a multiple of 16.",
+        ),
+        IO.Int.Input(
+            "custom_height",
+            default=1024,
+            min=480,
+            max=3840,
+            step=16,
+            tooltip="Used only when `size` is 'Custom'. Must be a multiple of 16.",
+        ),
+        IO.Combo.Input(
+            "background",
+            default="auto",
+            options=list(backgrounds),
+            tooltip="Return image with or without background.",
+        ),
+        *_gpt_image_shared_inputs(qualities),
+    ]
+
+
 class OpenAIGPTImageNodeV2(IO.ComfyNode):
 
     @classmethod
     def define_schema(cls):
         return IO.Schema(
             node_id="OpenAIGPTImageNodeV2",
-            display_name="OpenAI GPT Image 2",
+            display_name="OpenAI GPT Image 2.5",
             category="partner/image/OpenAI",
             description="Generates images via OpenAI's GPT Image endpoint.",
             inputs=[
@@ -467,50 +517,17 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "gpt-image-2.5-flare",
+                            _gpt_image_2_model_inputs(("auto", "opaque", "transparent"), GPT_IMAGE_25_QUALITIES),
+                        ),
+                        IO.DynamicCombo.Option(
+                            "gpt-image-2.5-sunburst",
+                            _gpt_image_2_model_inputs(("auto", "opaque", "transparent"), GPT_IMAGE_25_QUALITIES),
+                        ),
                         IO.DynamicCombo.Option(
                             "gpt-image-2",
-                            [
-                                IO.Combo.Input(
-                                    "size",
-                                    default="auto",
-                                    options=[
-                                        "auto",
-                                        "1024x1024",
-                                        "1024x1536",
-                                        "1536x1024",
-                                        "2048x2048",
-                                        "2048x1152",
-                                        "1152x2048",
-                                        "3840x2160",
-                                        "2160x3840",
-                                        "Custom",
-                                    ],
-                                    tooltip="Image size. Select 'Custom' to use the custom width and height.",
-                                ),
-                                IO.Int.Input(
-                                    "custom_width",
-                                    default=1024,
-                                    min=1024,
-                                    max=3840,
-                                    step=16,
-                                    tooltip="Used only when `size` is 'Custom'. Must be a multiple of 16.",
-                                ),
-                                IO.Int.Input(
-                                    "custom_height",
-                                    default=1024,
-                                    min=1024,
-                                    max=3840,
-                                    step=16,
-                                    tooltip="Used only when `size` is 'Custom'. Must be a multiple of 16.",
-                                ),
-                                IO.Combo.Input(
-                                    "background",
-                                    default="auto",
-                                    options=["auto", "opaque"],
-                                    tooltip="Return image with or without background.",
-                                ),
-                                *_gpt_image_shared_inputs(),
-                            ],
+                            _gpt_image_2_model_inputs(("auto", "opaque"), GPT_IMAGE_QUALITIES),
                         ),
                         IO.DynamicCombo.Option("gpt-image-1.5", _gpt_image_legacy_model_inputs()),
                         IO.DynamicCombo.Option("gpt-image-1", _gpt_image_legacy_model_inputs()),
@@ -544,7 +561,7 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "model.quality", "n"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "model.quality", "model.size", "n"], input_groups=["model.images"]),
                 expr="""
                 (
                   $ranges := {
@@ -559,22 +576,66 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                       "high":   [0.133, 0.22]
                     },
                     "gpt-image-2": {
-                      "low":    [0.0058, 0.0228],
-                      "medium": [0.0492, 0.2016],
-                      "high":   [0.198, 0.804]
+                      "low":    [0.0019, 0.0237],
+                      "medium": [0.0186, 0.2135],
+                      "high":   [0.0744, 0.8539]
+                    },
+                    "gpt-image-2.5-flare": {
+                      "low":    [0.0023, 0.0283],
+                      "medium": [0.0056, 0.0636],
+                      "high":   [0.0222, 0.2544],
+                      "xhigh":  [0.0388, 0.4523],
+                      "max":    [0.0887, 1.0175]
+                    },
+                    "gpt-image-2.5-sunburst": {
+                      "low":    [0.0023, 0.0283],
+                      "medium": [0.0056, 0.0636],
+                      "high":   [0.0222, 0.2544],
+                      "xhigh":  [0.0388, 0.4523],
+                      "max":    [0.0887, 1.0175]
                     }
                   };
-                  $range := $lookup($lookup($ranges, widgets.model), $lookup(widgets, "model.quality"));
+                  $presets := {
+                    "gpt-image-2": {
+                      "low": {"1024x1024": 0.0071, "1024x1536": 0.0057, "1536x1024": 0.0057, "2048x2048": 0.0143, "2048x1152": 0.0057, "1152x2048": 0.0057, "3840x2160": 0.0134, "2160x3840": 0.0134},
+                      "medium": {"1024x1024": 0.0632, "1024x1536": 0.0494, "1536x1024": 0.0494, "2048x2048": 0.1284, "2048x1152": 0.0509, "1152x2048": 0.0509, "3840x2160": 0.1201, "2160x3840": 0.1201},
+                      "high": {"1024x1024": 0.2529, "1024x1536": 0.1976, "1536x1024": 0.1976, "2048x2048": 0.5138, "2048x1152": 0.2034, "1152x2048": 0.2034, "3840x2160": 0.4803, "2160x3840": 0.4803}
+                    },
+                    "gpt-image-2.5": {
+                      "low": {"1024x1024": 0.0084, "1024x1536": 0.0068, "1536x1024": 0.0068, "2048x2048": 0.0170, "2048x1152": 0.0067, "1152x2048": 0.0067, "3840x2160": 0.0159, "2160x3840": 0.0159},
+                      "medium": {"1024x1024": 0.0188, "1024x1536": 0.0147, "1536x1024": 0.0147, "2048x2048": 0.0383, "2048x1152": 0.0157, "1152x2048": 0.0157, "3840x2160": 0.0371, "2160x3840": 0.0371},
+                      "high": {"1024x1024": 0.0753, "1024x1536": 0.0589, "1536x1024": 0.0589, "2048x2048": 0.1531, "2048x1152": 0.0606, "1152x2048": 0.0606, "3840x2160": 0.1431, "2160x3840": 0.1431},
+                      "xhigh": {"1024x1024": 0.1339, "1024x1536": 0.1055, "1536x1024": 0.1055, "2048x2048": 0.2721, "2048x1152": 0.1077, "1152x2048": 0.1077, "3840x2160": 0.2544, "2160x3840": 0.2544},
+                      "max": {"1024x1024": 0.3013, "1024x1536": 0.2354, "1536x1024": 0.2354, "2048x2048": 0.6123, "2048x1152": 0.2424, "1152x2048": 0.2424, "3840x2160": 0.5724, "2160x3840": 0.5724}
+                    }
+                  };
+                  $perImage := {
+                    "gpt-image-1": [0.0019, 0.0019],
+                    "gpt-image-1.5": [0.0016, 0.0016],
+                    "gpt-image-2": [0.0098, 0.0147],
+                    "gpt-image-2.5-flare": [0.0117, 0.0176],
+                    "gpt-image-2.5-sunburst": [0.0117, 0.0176]
+                  };
+                  $model := widgets.model;
+                  $family := ($model = "gpt-image-2.5-flare" or $model = "gpt-image-2.5-sunburst") ? "gpt-image-2.5" : $model;
+                  $qualityRaw := $lookup(widgets, "model.quality");
+                  $quality := ($qualityRaw != null) ? $qualityRaw : "";
+                  $sizeRaw := $lookup(widgets, "model.size");
+                  $size := ($sizeRaw != null) ? $sizeRaw : "";
+                  $range := $lookup($lookup($ranges, $model), $quality);
+                  $preset := $lookup($lookup($lookup($presets, $family), $quality), $size);
+                  $out := ($preset != null) ? [$preset, $preset] : $range;
+                  $image := $lookup($perImage, $model);
+                  $refsRaw := $lookup(inputGroups, "model.images");
+                  $refs := ($refsRaw != null) ? $refsRaw : 0;
                   $nRaw := widgets.n;
                   $n := ($nRaw != null and $nRaw != 0) ? $nRaw : 1;
-                  ($n = 1)
-                    ? {"type":"range_usd","min_usd": $range[0], "max_usd": $range[1], "format": {"approximate": true}}
-                    : {
-                        "type":"range_usd",
-                        "min_usd": $range[0] * $n,
-                        "max_usd": $range[1] * $n,
-                        "format": { "suffix": "/Run", "approximate": true }
-                      }
+                  $min := ($out[0] + $refs * $image[0]) * $n;
+                  $max := ($out[1] + $refs * $image[1]) * $n;
+                  $format := ($n = 1) ? {"approximate": true} : {"suffix": "/Run", "approximate": true};
+                  ($min = $max)
+                    ? {"type": "usd", "usd": $min, "format": $format}
+                    : {"type": "range_usd", "min_usd": $min, "max_usd": $max, "format": $format}
                 )
                 """,
             ),
@@ -626,7 +687,7 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                 )
             size = f"{custom_width}x{custom_height}"
 
-        if model_id not in ("gpt-image-1", "gpt-image-1.5", "gpt-image-2"):
+        if model_id not in GPT_IMAGE_MODELS:
             raise ValueError(f"Unknown model: {model_id}")
 
         if image_tensors:

comfy_extras/nodes_load_3d.py

@@ -156,16 +156,16 @@ class Preview3DAdvanced(IO.ComfyNode):
                 ),
                 IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
                 IO.Load3D.Input("viewport_state"),
-                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
-                IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
-                IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
+                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True, tooltip="Viewport camera information: position, look-at target, zoom, and type."),
+                IO.Int.Input("width", default=1024, min=1, max=4096, step=1, tooltip="Render width of the viewport in pixels."),
+                IO.Int.Input("height", default=1024, min=1, max=4096, step=1, tooltip="Render height of the viewport in pixels."),
             ],
             outputs=[
-                IO.File3DAny.Output(display_name="model_3d"),
-                IO.Load3DModelInfo.Output(display_name="model_3d_info"),
-                IO.Load3DCamera.Output(display_name="camera_info"),
-                IO.Int.Output(display_name="width"),
-                IO.Int.Output(display_name="height"),
+                IO.File3DAny.Output(display_name="model_3d", tooltip="3D model file (glb/obj/stl/etc.) from an upstream 3D node."),
+                IO.Load3DModelInfo.Output(display_name="model_3d_info", tooltip="Placement of each model in the scene: position, rotation, and scale (Y-up world space)."),
+                IO.Load3DCamera.Output(display_name="camera_info", tooltip="Viewport camera information: position, look-at target, zoom, and type."),
+                IO.Int.Output(display_name="width", tooltip="Render width of the viewport in pixels."),
+                IO.Int.Output(display_name="height", tooltip="Render height of the viewport in pixels."),
             ],
         )
 
@@ -355,15 +355,15 @@ class Load3DAdvanced(IO.ComfyNode):
             inputs=[
                 IO.Combo.Input("model_file", options=["none"] + sorted(files), upload=IO.UploadType.model),
                 IO.Load3D.Input("viewport_state"),
-                IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
-                IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
+                IO.Int.Input("width", default=1024, min=1, max=4096, step=1, tooltip="Render width of the viewport in pixels."),
+                IO.Int.Input("height", default=1024, min=1, max=4096, step=1, tooltip="Render height of the viewport in pixels."),
             ],
             outputs=[
-                IO.File3DAny.Output(display_name="model_3d"),
-                IO.Load3DModelInfo.Output(display_name="model_3d_info"),
-                IO.Load3DCamera.Output(display_name="camera_info"),
-                IO.Int.Output(display_name="width"),
-                IO.Int.Output(display_name="height"),
+                IO.File3DAny.Output(display_name="model_3d", tooltip="Loaded 3D model file (glb/obj/stl/etc.)."),
+                IO.Load3DModelInfo.Output(display_name="model_3d_info", tooltip="Placement of each model in the scene: position, rotation, and scale (Y-up world space)."),
+                IO.Load3DCamera.Output(display_name="camera_info", tooltip="Viewport camera information: position, look-at target, zoom, and type."),
+                IO.Int.Output(display_name="width", tooltip="Render width of the viewport in pixels."),
+                IO.Int.Output(display_name="height", tooltip="Render height of the viewport in pixels."),
             ],
         )
 

comfy_extras/nodes_lt.py

@@ -428,6 +428,43 @@ class LTXVAddGuide(io.ComfyNode):
 
         return latent_image, noise_mask
 
+    @classmethod
+    def attach_guide_latent(cls, positive, negative, latent_image, noise_mask, guide_latent, frame_idx, strength,
+                            scale_factors, latent_downscale_factor=1, causal_fix=None, attention_mask=None):
+        """Dilate a guide latent onto the canvas, pin it, and record its attention entry.
+
+        Keeps the three values that have to agree in one place: the pre-dilation
+        shape recorded on the attention entry, the post-dilation token count, and
+        the downscale factor handed to append_keyframe. context_windows re-derives
+        the factor from the first two, so they must not drift apart.
+        """
+        guide_latent_shape = list(guide_latent.shape[2:])  # pre-dilation [F, H, W] for spatial-mask downsampling
+        guide_mask = None
+        if latent_downscale_factor > 1:
+            guide_latent, guide_mask = cls.dilate_latent(guide_latent, latent_downscale_factor)
+
+        positive, negative, latent_image, noise_mask = cls.append_keyframe(
+            positive,
+            negative,
+            frame_idx,
+            latent_image,
+            noise_mask,
+            guide_latent,
+            strength,
+            scale_factors,
+            guide_mask=guide_mask,
+            latent_downscale_factor=latent_downscale_factor,
+            causal_fix=causal_fix,
+        )
+
+        # Track this guide for per-reference attention control.
+        pre_filter_count = guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
+        positive, negative = _append_guide_attention_entry(
+            positive, negative, pre_filter_count, guide_latent_shape, strength=strength,
+            attention_mask=attention_mask,
+        )
+        return positive, negative, latent_image, noise_mask
+
     @classmethod
     def execute(cls, positive, negative, vae, latent, image, frame_idx, strength, attention_mask=None, iclora_parameters=None) -> io.NodeOutput:
         scale_factors = vae.downscale_index_formula
@@ -462,32 +499,20 @@ class LTXVAddGuide(io.ComfyNode):
             t = t[:, :, 1:, :, :]
             image = image[1:]
 
-        guide_latent_shape = list(t.shape[2:])  # pre-dilation [F, H, W] for spatial-mask downsampling
-        guide_mask = None
-        if latent_downscale_factor > 1:
-            t, guide_mask = cls.dilate_latent(t, latent_downscale_factor)
-
         frame_idx, latent_idx = cls.get_latent_index(positive, latent_length, len(image), frame_idx, scale_factors, latent_shape=latent_image.shape)
         assert latent_idx + t.shape[2] <= latent_length, "Conditioning frames exceed the length of the latent sequence."
 
-        positive, negative, latent_image, noise_mask = cls.append_keyframe(
+        positive, negative, latent_image, noise_mask = cls.attach_guide_latent(
             positive,
             negative,
-            frame_idx,
             latent_image,
             noise_mask,
             t,
+            frame_idx,
             strength,
             scale_factors,
-            guide_mask=guide_mask,
             latent_downscale_factor=latent_downscale_factor,
             causal_fix=causal_fix,
-        )
-
-        # Track this guide for per-reference attention control.
-        pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4]
-        positive, negative = _append_guide_attention_entry(
-            positive, negative, pre_filter_count, guide_latent_shape, strength=strength,
             attention_mask=attention_mask,
         )
 
@@ -496,6 +521,130 @@ class LTXVAddGuide(io.ComfyNode):
     generate = execute  # TODO: remove
 
 
+class LTXVAddLatentGuide(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVAddLatentGuide",
+            display_name="LTXV Add Latent Guide",
+            category="model/conditioning/ltxv",
+            description="Pins an already-encoded latent as a guide, for when the guide comes out of an "
+                        "earlier stage rather than an image. Same effect as LTXV Add Guide without the "
+                        "VAE decode/encode round trip. A guide that is spatially smaller than the target "
+                        "(an IC-LoRA or detailing reference) is dilated onto a sparse grid, and its RoPE "
+                        "end positions are expanded by the same ratio so it covers the target canvas "
+                        "instead of addressing only the top-left corner of it.",
+            search_aliases=["latent guide", "add latent guide", "guide latent"],
+            inputs=[
+                io.Conditioning.Input("positive"),
+                io.Conditioning.Input("negative"),
+                io.Vae.Input("vae"),
+                io.Latent.Input("latent", tooltip="Target video latent the guide is pinned onto."),
+                io.Latent.Input(
+                    "guiding_latent",
+                    tooltip="Guide latent. Its spatial size must divide the target's by the same whole "
+                            "number on both axes; equal size pins it as-is, half size is treated as an "
+                            "x2 IC-LoRA reference.",
+                ),
+                io.Int.Input(
+                    "latent_idx",
+                    default=0,
+                    min=-9999,
+                    max=9999,
+                    tooltip="Latent frame index to start the guide at, counted in latent frames rather "
+                            "than pixel frames. Negative values place the guide on frames before the "
+                            "start of the latent, not counted back from its end.",
+                ),
+                io.Float.Input(
+                    "strength",
+                    default=1.0,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip="Capped at 1.0. A dilated guide marks its padding positions with a "
+                            "negative denoise mask so the model drops them; above 1.0 the kept "
+                            "positions would go negative too and the whole guide would be dropped. "
+                            "Amplify beyond 1.0 with attention_mask instead.",
+                ),
+                io.Mask.Input(
+                    "attention_mask",
+                    optional=True,
+                    tooltip="Optional pixel-space spatial mask. Controls per-region "
+                            "conditioning influence via self-attention, multiplied by strength.",
+                ),
+            ],
+            outputs=[
+                io.Conditioning.Output(display_name="positive"),
+                io.Conditioning.Output(display_name="negative"),
+                io.Latent.Output(display_name="latent"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, positive, negative, vae, latent, guiding_latent, latent_idx, strength, attention_mask=None) -> io.NodeOutput:
+        scale_factors = vae.downscale_index_formula
+        latent_image = latent["samples"]
+        noise_mask = get_noise_mask(latent)
+        guide_latent = guiding_latent["samples"]
+
+        for name, samples in (("latent", latent_image), ("guiding_latent", guide_latent)):
+            if samples.ndim != 5:
+                raise ValueError(
+                    f"{name} must be a 5D video latent (batch, channels, frames, height, width), "
+                    f"got shape {list(samples.shape)}."
+                )
+
+        guide_frames = guide_latent.shape[2]
+        latent_frames = latent_image.shape[2]
+        if latent_idx + guide_frames > latent_frames:
+            raise ValueError(
+                f"Guide of {guide_frames} latent frame(s) at latent_idx {latent_idx} runs past the "
+                f"end of the {latent_frames}-frame latent. Negative values are allowed and place the "
+                f"guide before the start of the latent."
+            )
+
+        if latent_image.shape[3] % guide_latent.shape[3] != 0 or latent_image.shape[4] % guide_latent.shape[4] != 0:
+            raise ValueError(
+                f"Guiding latent spatial size {guide_latent.shape[3]}x{guide_latent.shape[4]} must divide "
+                f"the target size {latent_image.shape[3]}x{latent_image.shape[4]} by a whole number."
+            )
+
+        height_scale = latent_image.shape[3] // guide_latent.shape[3]
+        width_scale = latent_image.shape[4] // guide_latent.shape[4]
+        # dilate_latent and append_keyframe take a single IC-LoRA downscale factor for both
+        # axes, so a non-square ratio would mis-place RoPE on one of them.
+        if height_scale != width_scale:
+            raise ValueError(
+                f"Guiding latent spatial ratio must be square, got height x{height_scale} and "
+                f"width x{width_scale} ({guide_latent.shape[3]}x{guide_latent.shape[4]} -> "
+                f"{latent_image.shape[3]}x{latent_image.shape[4]})."
+            )
+
+        time_scale_factor = scale_factors[0]
+        if latent_idx <= 0:
+            frame_idx = latent_idx * time_scale_factor
+        else:
+            frame_idx = 1 + (latent_idx - 1) * time_scale_factor
+
+        positive, negative, latent_image, noise_mask = LTXVAddGuide.attach_guide_latent(
+            positive,
+            negative,
+            latent_image,
+            noise_mask,
+            guide_latent,
+            frame_idx,
+            strength,
+            scale_factors,
+            latent_downscale_factor=width_scale,
+            attention_mask=attention_mask,
+        )
+
+        out = latent.copy()
+        out["samples"] = latent_image
+        out["noise_mask"] = noise_mask
+        return io.NodeOutput(positive, negative, out)
+
+
 class LTXVCropGuides(io.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -1186,6 +1335,7 @@ class LtxvExtension(ComfyExtension):
             LTXVScheduler,
             GetICLoRAParameters,
             LTXVAddGuide,
+            LTXVAddLatentGuide,
             LTXVPreprocess,
             LTXVCropGuides,
             LTXVConcatAVLatent,

comfy_extras/nodes_model_advanced.py

@@ -7,6 +7,7 @@ import comfy.ldm.modules.attention
 import nodes
 import torch
 import node_helpers
+from comfy_api.latest import io
 
 
 class LCM(comfy.model_sampling.EPS):
@@ -366,26 +367,34 @@ class ModelComputeDtype:
         return (m, )
 
 
-class ModelAttentionBackend:
+class ModelAttentionBackend(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
         backends = ["pytorch attention"]
         if comfy.ldm.modules.attention.COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE:
             backends.append("comfy kitchen attention")
-        return {"required": {"model": ("MODEL",),
-                             "attention": (backends,),
-                             }}
+        return io.Schema(
+            node_id="ModelAttentionBackend",
+            display_name="Model Attention Backend",
+            category="model/patch",
+            is_experimental=True,
+            description="Selects the dense attention implementation for the model. When used with Block Sparse Attention, this backend is used whenever sparse attention is inactive or unsupported.",
+            inputs=[
+                io.Model.Input("model", tooltip="The model to patch."),
+                io.Combo.Input("attention", display_name="backend", options=backends, default="pytorch attention",
+                               tooltip="The dense attention backend. Comfy Kitchen attention uses quantized INT8 attention and is available only on Nvidia and AMD GPUs."),
+            ],
+            outputs=[
+                io.Model.Output(display_name="model", tooltip="The model with the selected attention backend."),
+            ],
+        )
 
     @classmethod
-    def VALIDATE_INPUTS(s, attention):
+    def validate_inputs(cls, attention):
         return True
 
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "model/patch"
-
-    def patch(self, model, attention):
+    @classmethod
+    def execute(cls, model, attention):
         attention_name = {
             "comfy kitchen attention": "comfy_kitchen_int8",
             "pytorch attention": "pytorch",
@@ -396,7 +405,7 @@ class ModelAttentionBackend:
             attention_function = comfy.ldm.modules.attention.get_attention_function("pytorch")
         m = model.clone()
         m.set_model_optimized_attention(attention_function)
-        return (m, )
+        return io.NodeOutput(m)
 
 
 NODE_CLASS_MAPPINGS = {

comfy_extras/nodes_model_patch.py

@@ -283,7 +283,7 @@ class ModelPatchLoader:
                 )
                 manual_cast_dtype = comfy.model_management.unet_manual_cast(
                     dtype, load_device, supported_dtypes=[torch.bfloat16, torch.float32])
-                operations = comfy.ops.pick_operations(dtype, manual_cast_dtype)
+                operations = comfy.ops.pick_operations(dtype, manual_cast_dtype, load_device=load_device)
 
             num_blocks = 0
             while "control_blocks.{}.after_proj.weight".format(num_blocks) in sd:

comfy_extras/nodes_save_3d.py

@@ -930,17 +930,17 @@ class Save3DAdvanced(IO.ComfyNode):
                 ),
                 IO.String.Input("filename_prefix", default="3d/ComfyUI"),
                 IO.Load3D.Input("viewport_state"),
-                IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
-                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
-                IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
-                IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
+                IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True, tooltip="Placement of each model in the scene: position, rotation, and scale (Y-up world space)."),
+                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True, tooltip="Viewport camera information: position, look-at target, zoom, and type."),
+                IO.Int.Input("width", default=1024, min=1, max=4096, step=1, tooltip="Render width of the viewport in pixels."),
+                IO.Int.Input("height", default=1024, min=1, max=4096, step=1, tooltip="Render height of the viewport in pixels."),
             ],
             outputs=[
-                IO.File3DAny.Output(display_name="model_3d"),
-                IO.Load3DModelInfo.Output(display_name="model_3d_info"),
-                IO.Load3DCamera.Output(display_name="camera_info"),
-                IO.Int.Output(display_name="width"),
-                IO.Int.Output(display_name="height"),
+                IO.File3DAny.Output(display_name="model_3d", tooltip="3D model file (glb/obj/stl/etc.) from an upstream 3D node."),
+                IO.Load3DModelInfo.Output(display_name="model_3d_info", tooltip="Placement of each model in the scene: position, rotation, and scale (Y-up world space)."),
+                IO.Load3DCamera.Output(display_name="camera_info", tooltip="Viewport camera information: position, look-at target, zoom, and type."),
+                IO.Int.Output(display_name="width", tooltip="Render width of the viewport in pixels."),
+                IO.Int.Output(display_name="height", tooltip="Render height of the viewport in pixels."),
             ],
         )
 

comfy_extras/nodes_sparse_attention.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import logging
 import re
+import weakref
 
 import comfy_kitchen as ck
 import torch
@@ -142,11 +143,11 @@ class SparseAttnPatch:
 
     def vsa_rope_freqs(self, rope_freqs, plan):
         hit = self.vsa_rope
-        if hit is not None and hit[0] is rope_freqs and hit[1] is plan:
+        if hit is not None and hit[0]() is rope_freqs and hit[1] is plan:
             return hit[2]
         padded = rope_freqs.new_zeros((1, plan["n"]) + tuple(rope_freqs.shape[2:]))
         padded[0, plan["inv"]] = rope_freqs[0]
-        self.vsa_rope = (rope_freqs, plan, padded)
+        self.vsa_rope = (weakref.ref(rope_freqs), plan, padded)
         return padded
 
 
@@ -257,6 +258,7 @@ def h3_sparse_attention(attn, x, rope_freqs, transformer_options, patch: SparseA
         if patch.vsa:
             plan = patch.vsa_plan(transformer_options["minimax_h3_layout"], x.device)
             n = plan["n"]
+            freqs = patch.vsa_rope_freqs(rope_freqs, plan)
 
         key = (block_index, n, tuple(transformer_options.get("uuids", ())))   # statistics per conditioning branch
         pooled = patch.pooled.get(key)
@@ -268,7 +270,6 @@ def h3_sparse_attention(attn, x, rope_freqs, transformer_options, patch: SparseA
             )
 
     if patch.vsa:
-        freqs = patch.vsa_rope_freqs(rope_freqs, plan)
         sink = sink_q = (0, plan["n_prefix"])
         extra = {"tail": False, "block_len": plan["block_len"]}
         gate = attn.to_gate_compress
@@ -354,40 +355,42 @@ class BlockSparseAttention(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="BlockSparseAttention",
-            display_name="Block Sparse Attention",
-            category="advanced/model",
+            display_name="Model Sparse Attention",
+            category="model/patch",
             is_experimental=True,
-            description="Block-sparse attention through comfy_kitchen: each query block attends a selected subset of key blocks exactly, reducing attention compute. "
-                        "The relative speed gain grows with sequence length since short sequences are usually faster dense. "
-                        "Outside the active schedule, dense_blocks and under min_tokens, the model uses the active dense model attention backend.",
+            search_aliases=["Block Sparse Attention"],
+            description="Applies block-sparse attention to eligible model attention layers, reducing compute for long sequences. "
+                        "The speed gain grows with sequence length since short sequences are usually faster dense. "
+                        "Outside the start/end_percent, dense_blocks and under min_tokens, the model uses the dense model attention backend. "
+                        "Use the node Model Attention Backend to select that fallback.",
             inputs=[
-                io.Model.Input("model"),
-                io.DynamicCombo.Input("selection", options=[
-                    io.DynamicCombo.Option("Sol-Attn (adaptive tau)", [
+                io.Model.Input("model", tooltip="The model to patch."),
+                io.DynamicCombo.Input("selection", display_name="method", options=[
+                    io.DynamicCombo.Option("sol-attn", [
                         io.Float.Input("tau", default=1.3, min=0.0, max=4.0, step=0.05,
                                        tooltip="Threshold in score-distribution sigmas. Higher is sparser: "
                                                "1.0 keeps ~16% of key blocks exact, 1.5 ~7%, 2.0 ~2.7%."),
                     ]),
-                    io.DynamicCombo.Option("top-k (SLA)", [
+                    io.DynamicCombo.Option("sla", [
                         io.Float.Input("keep_percent", default=10.0, min=0.5, max=95.0, step=0.5,
                                        tooltip="Percent of key blocks each query block keeps exactly (sinks and "
                                                "the diagonal ride on top). The selection SLA-style LoRAs are "
                                                "distilled against; without such a LoRA higher is closer to dense."),
                     ]),
-                    io.DynamicCombo.Option("VSA (FastVideo)", [
+                    io.DynamicCombo.Option("vsa", [
                         io.Float.Input("keep_percent", default=10.0, min=0.5, max=95.0, step=0.5,
                                        tooltip="Percent of video cubes each query cube keeps; FastH3-VSA "
                                                "checkpoints are trained at 10. Uses the model's to_gate_compress "
                                                "layers for the coarse branch when present."),
                     ]),
-                ], tooltip="How exact key blocks are chosen. "
-                           "Sol-Attn: per head/block adaptive threshold. "
-                           "top-k (SLA): fixed keep_percent everywhere, recommended only with trained weights. "
-                           "VSA (FastVideo): FastH3-VSA's cube tiling and coarse branch, requires weights trained for it."),
+                ], tooltip="Method used to choose key blocks for full token-level attention. "
+                           "sol-attn: Sparsifying Online Attention uses a training-free adaptive threshold for each attention head and query block. "
+                           "sla: Sparse-Linear Attention keeps a fixed percentage of the highest-scoring key blocks; use only with model weights trained for this pattern. "
+                           "vsa: Video Sparse Attention (FastVideo) uses 3D video-cube tiling and a learned coarse attention branch; requires FastH3 model weights."),
                 io.Float.Input("start_percent", default=0.2, min=0.0, max=1.0, step=0.01,
-                               tooltip="Dense before this point of the schedule."),
+                               tooltip="Percentage point when sparse attention begins. Before this point, attention stays dense."),
                 io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.01,
-                               tooltip="Dense after this point of the schedule."),
+                               tooltip="Percentage point when sparse attention ends. After this point, attention returns to dense."),
                 io.String.Input("dense_blocks", default="", advanced=True,
                                 tooltip="Transformer blocks that always run dense, e.g. '0, 1, 47-49'."),
                 io.Int.Input("min_tokens", default=12288, min=0, max=1 << 20, step=512, advanced=True,
@@ -401,22 +404,30 @@ class BlockSparseAttention(io.ComfyNode):
                                tooltip="MiniMax-H3 only. exact_kv: every query attends the packed text/audio/"
                                        "reference rows exactly (~3% cost). exact_kv_and_rows: additionally runs "
                                        "the target-audio query rows dense (keeps generated audio intact)."),
-                io.Boolean.Input("verbose", default=False, advanced=True),
+                io.Boolean.Input("verbose", default=False, advanced=True,
+                                 tooltip="Logs whether each attention shape used sparse attention or why it stayed dense."),
             ],
-            outputs=[io.Model.Output()],
+            outputs=[io.Model.Output(display_name="model", tooltip="The model with block-sparse attention applied.")],
         )
 
     @classmethod
     def execute(cls, model, selection, start_percent, end_percent, dense_blocks="", min_tokens=12288,
                 extra_tokens=0, sink_conditioning="exact_kv_and_rows", verbose=False) -> io.NodeOutput:
         mode = selection["selection"]
-        return io.NodeOutput(apply_block_sparse_attention(
-            model, tau=selection.get("tau", 1.3),
-            topk_ratio=0.0 if mode == "Sol-Attn (adaptive tau)" else selection["keep_percent"] / 100.0,
-            vsa=mode == "VSA (FastVideo)",
-            start_percent=start_percent, end_percent=end_percent, min_tokens=min_tokens,
-            dense_blocks=parse_block_list(dense_blocks), sink_conditioning=sink_conditioning,
-            extra_tokens=extra_tokens, verbose=verbose))
+        patched_model = apply_block_sparse_attention(
+            model,
+            tau=selection.get("tau", 1.3),
+            topk_ratio=0.0 if mode == "sol-attn" else selection["keep_percent"] / 100.0,
+            vsa=mode == "vsa",
+            start_percent=start_percent,
+            end_percent=end_percent,
+            min_tokens=min_tokens,
+            dense_blocks=parse_block_list(dense_blocks),
+            sink_conditioning=sink_conditioning,
+            extra_tokens=extra_tokens,
+            verbose=verbose,
+        )
+        return io.NodeOutput(patched_model)
 
 
 class BlockSparseAttentionExtension(ComfyExtension):

comfy_extras/nodes_trellis2.py

@@ -7,6 +7,7 @@ from comfy_extras.nodes_mesh_postprocess import pack_variable_mesh_batch
 import comfy.latent_formats
 import comfy.model_management
 import comfy.utils
+import logging
 import math
 import torch
 
@@ -690,6 +691,73 @@ def _dino_encode_batch(clip_vision_model, image, out_device, *, want_patches=Fal
         out["composites"] = composite_list
     return out
 
+def _naf_upsample(naf_model, lr_feat, composites, image_size, naf_target, out_device, compute_device):
+    """NAF-upsample each item's DINO patch grid to naf_target, guided by its composite."""
+    if naf_model is None:
+        return None
+    comfy.model_management.load_model_gpu(naf_model)
+    inner = naf_model.model
+    model_dtype = next(inner.parameters()).dtype
+    out = torch.empty((len(composites), lr_feat.shape[1], *naf_target), device=out_device, dtype=model_dtype)
+    for i, c in enumerate(composites):
+        img_i = comfy.utils.common_upscale(c, image_size, image_size, "lanczos", "disabled").to(compute_device, model_dtype)
+        lr_i = lr_feat[i:i + 1].to(compute_device, model_dtype)
+        inner(img_i, lr_i, naf_target, output=out[i:i + 1])
+    return out
+
+
+def _build_pixal3d_conditioning(clip_vision_model, image, transform_matrix, camera_angle_x, mesh_scale, num_views=1):
+    """Per-item inputs hold B*num_views entries with each object's views consecutive; mesh_scale holds B."""
+    naf_model = clip_vision_model.naf
+    out_device = comfy.model_management.intermediate_device()
+    compute_device = comfy.model_management.get_torch_device()
+
+    cond = _dino_encode_batch(clip_vision_model, image, out_device, want_patches=True)
+    batch_size = cond["batch_size"] // num_views
+    fm_512_dino, fm_1024_dino = cond["patches_512"], cond["patches_1024"]
+    composite_list = cond["composites"]
+
+    # NAF HR targets per stage: shape_512=512, shape_1024=512, tex_1024=1024
+    hr_shape_512  = _naf_upsample(naf_model, fm_512_dino,  composite_list, 512,  (512, 512), out_device, compute_device)
+    hr_shape_1024 = _naf_upsample(naf_model, fm_1024_dino, composite_list, 1024, (512, 512), out_device, compute_device)
+    hr_tex_1024   = _naf_upsample(naf_model, fm_1024_dino, composite_list, 1024, (1024, 1024), out_device, compute_device)
+
+    # CLS + register tokens averaged over each object's views
+    global_512 = cond["global_512"].unflatten(0, (batch_size, num_views)).mean(dim=1)
+    global_1024 = cond["global_1024"].unflatten(0, (batch_size, num_views)).mean(dim=1)
+
+    proj_pack = {
+        "stages": {
+            "ss":         {"feature_map": fm_512_dino,    "feature_map_hr": None,         "image_resolution": 512},
+            "shape_512":  {"feature_map": fm_512_dino,    "feature_map_hr": hr_shape_512, "image_resolution": 512},
+            "shape_1024": {"feature_map": fm_1024_dino,   "feature_map_hr": hr_shape_1024,"image_resolution": 1024},
+            "tex_1024":   {"feature_map": fm_1024_dino,   "feature_map_hr": hr_tex_1024,  "image_resolution": 1024},
+        },
+        "transform_matrix": transform_matrix.to(out_device),
+        "camera_angle_x": camera_angle_x.to(out_device),
+        "mesh_scale": mesh_scale.to(out_device),
+        "num_views": num_views,
+        "patch_size": 16,
+    }
+
+    # global_512 -> SS/shape_512 cross-attn; global_1024 -> shape_1024/tex_1024.
+    ss_proj_feats = compute_stage_proj_feats(
+        proj_pack, "ss", dense_grid_resolution=16, batch_size=batch_size,
+        device=compute_device,
+    )
+    base_extras = {
+        "embeds": global_1024, "proj_feat_pack": proj_pack,
+        "trellis2_proj_feats": ss_proj_feats,
+    }
+    neg_extras = {
+        "embeds": torch.zeros_like(global_1024), "proj_feat_pack": proj_pack,
+        "trellis2_proj_feats": ss_proj_feats,
+    }
+    positive = [[global_512, base_extras]]
+    negative = [[torch.zeros_like(global_512), neg_extras]]
+    return IO.NodeOutput(positive, negative)
+
+
 class Pixal3DConditioning(IO.ComfyNode):
 
     @classmethod
@@ -715,79 +783,85 @@ class Pixal3DConditioning(IO.ComfyNode):
 
     @classmethod
     def execute(cls, clip_vision_model, image, camera_angle_x) -> IO.NodeOutput:
-        naf_model = clip_vision_model.naf
-        out_device = comfy.model_management.intermediate_device()
-        compute_device = comfy.model_management.get_torch_device()
-
-        cond = _dino_encode_batch(clip_vision_model, image, out_device, want_patches=True)
-        batch_size = cond["batch_size"]
-        global_512, global_1024 = cond["global_512"], cond["global_1024"]
-        fm_512_dino, fm_1024_dino = cond["patches_512"], cond["patches_1024"]
-        composite_list = cond["composites"]
-
-        # The LR DINO grid AND the NAF HR grid are sampled separately
-        # NAF targets per stage: shape_512=512, shape_1024=512, tex_1024=1024.
-        def _naf_hr(lr_feat, composites, image_size, naf_target):
-            if naf_model is None or naf_target is None:
-                return None
-            comfy.model_management.load_model_gpu(naf_model)
-            inner = naf_model.model
-            model_dtype = next(inner.parameters()).dtype  # set at load time (see clip_vision NAF)
-            hrs = []
-            for i, c in enumerate(composites):
-                img_i = comfy.utils.common_upscale(c, image_size, image_size, "lanczos", "disabled")\
-                    .to(compute_device).to(model_dtype)
-                lr_i = lr_feat[i:i + 1].to(compute_device).to(model_dtype)
-                output = torch.empty((1, lr_i.shape[1], *naf_target), device=out_device, dtype=model_dtype)
-                hr_i = inner(img_i, lr_i, naf_target, output=output)
-                hrs.append(hr_i)
-            return torch.cat(hrs, dim=0)
-
-        hr_shape_512  = _naf_hr(fm_512_dino,  composite_list, 512,  (512, 512))
-        hr_shape_1024 = _naf_hr(fm_1024_dino, composite_list, 1024, (512, 512))
-        hr_tex_1024   = _naf_hr(fm_1024_dino, composite_list, 1024, (1024, 1024))
-
+        batch_size = image.shape[0]
         # distance_from_fov: grid_point (-1, 0, 0) projects to pixel (0, image_resolution-1).
         # FOV widget is in degrees for UX; trig + downstream projection expect radians.
         camera_angle_x = math.radians(float(camera_angle_x))
         distance = 0.5 / math.tan(camera_angle_x / 2.0)
-        cam_angle_t = torch.tensor([camera_angle_x] * batch_size, device=out_device, dtype=torch.float32)
-        dist_t = torch.tensor([distance] * batch_size, device=out_device, dtype=torch.float32)
-        scale_t = torch.ones(batch_size, device=out_device, dtype=torch.float32)
-        T = build_proj_transform_matrix(dist_t, batch_size, device=out_device, dtype=torch.float32)
+        cam_angle_t = torch.full((batch_size,), camera_angle_x)
+        dist_t = torch.full((batch_size,), distance)
+        T = build_proj_transform_matrix(dist_t, batch_size, dist_t.device)
+        return _build_pixal3d_conditioning(clip_vision_model, image, T, cam_angle_t, torch.ones(batch_size))
 
-        proj_pack = {
-            "stages": {
-                "ss":         {"feature_map": fm_512_dino,    "feature_map_hr": None,         "image_resolution": 512},
-                "shape_512":  {"feature_map": fm_512_dino,    "feature_map_hr": hr_shape_512, "image_resolution": 512},
-                "shape_1024": {"feature_map": fm_1024_dino,   "feature_map_hr": hr_shape_1024,"image_resolution": 1024},
-                "tex_1024":   {"feature_map": fm_1024_dino,   "feature_map_hr": hr_tex_1024,  "image_resolution": 1024},
-            },
-            "transform_matrix": T,
-            "camera_angle_x": cam_angle_t,
-            "mesh_scale": scale_t,
-            "distance": dist_t,
-            "patch_size": 16,
-        }
 
-        # global_512 → SS/shape_512 cross-attn; global_1024 → shape_1024/tex_1024.
-        ss_proj_feats = compute_stage_proj_feats(
-            proj_pack, "ss", dense_grid_resolution=16, batch_size=batch_size,
-            device=compute_device,
+_VIEW_AZIMUTHS = {"front": 0.0, "left": 90.0, "back": 180.0, "right": 270.0}
+_VIEW_PAD = 1.1  # unit cube spans 1/1.1 of the frame, upstream's example rig
+
+
+def _orbit_camera_to_world(azimuths_deg, elevations_deg, distance):
+    """Z-up orbit cameras looking at the origin; azimuth 0 / elevation 0 is the front view."""
+    az = torch.deg2rad(torch.tensor(azimuths_deg, dtype=torch.float32))
+    el = torch.deg2rad(torch.tensor(elevations_deg, dtype=torch.float32))
+    back = torch.stack([torch.sin(az) * torch.cos(el), -torch.cos(az) * torch.cos(el), torch.sin(el)], dim=-1)
+    right = torch.stack([torch.cos(az), torch.sin(az), torch.zeros_like(az)], dim=-1)
+    c2w = torch.eye(4).repeat(az.shape[0], 1, 1)
+    c2w[:, :3, :3] = torch.stack([right, torch.cross(back, right, dim=-1), back], dim=-1)
+    c2w[:, :3, 3] = back * distance
+    return c2w
+
+
+class Pixal3DMultiViewConditioning(IO.ComfyNode):
+    """Fixed orbit rig: front, left, back and right views 90 degrees apart, used as framed."""
+
+    @classmethod
+    def define_schema(cls):
+        views = [IO.Image.Input(name, optional=True,
+                                tooltip=f"Square view of the object's {name} side, with alpha or on a black background, "
+                                        "framed like the rig: the object spans about 1/1.1 of the frame at its widest, "
+                                        "the same scale in every view. The first connected view (front, left, back, "
+                                        "right order) is the front the mesh is posed to.")
+                 for name in _VIEW_AZIMUTHS]
+        return IO.Schema(
+            node_id="Pixal3DMultiViewConditioning",
+            display_name="Pixal3D Multi-View Conditioning",
+            category="model/conditioning/trellis2",
+            inputs=[IO.ClipVision.Input("clip_vision_model", tooltip="DINOv3 ViT-L/16 ClipVision with bundled NAF weights."),
+                    IO.Float.Input("fov", default=20.0, min=1.0, max=170.0, step=0.01, round=False,
+                                   tooltip="Horizontal FOV in degrees of the views as framed: 20 for rig renders and most "
+                                           "multi-view generators, or MoGeGeometryToFOV on one of the views for photos.")]
+                   + views,
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ],
         )
-        neg_global = torch.zeros_like(global_512)
-        neg_embeds = torch.zeros_like(global_1024)
-        base_extras = {
-            "embeds": global_1024, "proj_feat_pack": proj_pack,
-            "trellis2_proj_feats": ss_proj_feats,
-        }
-        neg_extras = {
-            "embeds": neg_embeds, "proj_feat_pack": proj_pack,
-            "trellis2_proj_feats": ss_proj_feats,
-        }
-        positive = [[global_512, base_extras]]
-        negative = [[neg_global, neg_extras]]
-        return IO.NodeOutput(positive, negative)
+
+    @classmethod
+    def execute(cls, clip_vision_model, fov, front=None, left=None, back=None, right=None) -> IO.NodeOutput:
+        views = {"front": front, "left": left, "back": back, "right": right}
+        names = [name for name in _VIEW_AZIMUTHS if views[name] is not None]
+        if not names:
+            raise ValueError("Pixal3DMultiViewConditioning needs at least one view")
+        batch_size = views[names[0]].shape[0]
+        num_views = len(names)
+        # the first connected view is the front (upstream re-bases the rig onto view 0; on an orbit that is an azimuth shift)
+        if names[0] != "front":
+            logging.warning(f"Pixal3DMultiViewConditioning: no front view, the mesh will be posed with the {names[0]} view as its front")
+        azimuths = [_VIEW_AZIMUTHS[name] - _VIEW_AZIMUTHS[names[0]] for name in names]
+        items = []
+        for b in range(batch_size):
+            for name in names:
+                view = views[name][b % views[name].shape[0]][None]
+                if view.shape[-1] == 4:
+                    view = view[..., :3] * view[..., 3:4]
+                if view.shape[1:3] != (1024, 1024):
+                    view = comfy.utils.common_upscale(view.movedim(-1, 1), 1024, 1024, "lanczos", "disabled").movedim(1, -1)
+                items.append(view)
+        fov = math.radians(fov)
+        c2w = _orbit_camera_to_world(azimuths, [0.0] * num_views, _VIEW_PAD * 0.5 / math.tan(fov / 2.0))
+        return _build_pixal3d_conditioning(clip_vision_model, torch.cat(items, dim=0), c2w.repeat(batch_size, 1, 1),
+                                           torch.full((batch_size * num_views,), fov), torch.ones(batch_size),
+                                           num_views=num_views)
 
 
 class Trellis2Extension(ComfyExtension):
@@ -796,6 +870,7 @@ class Trellis2Extension(ComfyExtension):
         return [
             Trellis2Conditioning,
             Pixal3DConditioning,
+            Pixal3DMultiViewConditioning,
             Trellis2ShapeStage,
             EmptyTrellis2LatentStructure,
             Trellis2TextureStage,

latent_preview.py

@@ -4,6 +4,7 @@ from comfy.cli_args import args, LatentPreviewMethod
 from comfy.taesd.taesd import TAESD
 from comfy.sd import VAE
 import comfy.model_management
+import comfy.model_prefetch
 import folder_paths
 import comfy.utils
 import logging
@@ -45,8 +46,16 @@ class TAESDPreviewerImpl(LatentPreviewer):
         return preview_to_image(x_sample)
 
 class TAEHVPreviewerImpl(TAESDPreviewerImpl):
+    def __init__(self, taesd, compile_preview=False):
+        super().__init__(taesd)
+        self.compile_preview = compile_preview
+
     def decode_latent_to_preview(self, x0):
-        x_sample = self.taesd.decode(x0[:1, :, :1])[0][0]
+        samples = x0[:1, :, :1]
+        if self.compile_preview and comfy.model_prefetch.malloc_graph_enabled(self.taesd.device):
+            comfy.model_prefetch.malloc_graph_begin(self.taesd.device)
+        x_sample = self.taesd.decode(samples)[0][0]
+        comfy.model_prefetch.malloc_graph_end()
         return preview_to_image(x_sample, do_scale=False)
 
 class Latent2RGBPreviewer(LatentPreviewer):
@@ -97,7 +106,7 @@ def get_previewer(device, latent_format):
                 if latent_format.taesd_decoder_name in VIDEO_TAES:
                     taesd = VAE(comfy.utils.load_torch_file(taesd_decoder_path))
                     taesd.first_stage_model.show_progress_bar = False
-                    previewer = TAEHVPreviewerImpl(taesd)
+                    previewer = TAEHVPreviewerImpl(taesd, compile_preview=latent_format.compile_preview)
                 else:
                     taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device)
                     previewer = TAESDPreviewerImpl(taesd)

main.py

@@ -196,9 +196,9 @@ def execute_prestartup_script():
         return False
 
     node_paths = folder_paths.get_folder_paths("custom_nodes")
+    node_prestartup_times = []
     for custom_node_path in node_paths:
         possible_modules = os.listdir(custom_node_path)
-        node_prestartup_times = []
 
         for possible_module in possible_modules:
             module_path = os.path.join(custom_node_path, possible_module)

nodes.py

@@ -2298,7 +2298,9 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                     NODE_CLASS_MAPPINGS[name] = node_cls
                     node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
             if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
-                NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
+                for name, display_name in module.NODE_DISPLAY_NAME_MAPPINGS.items():
+                    if name not in ignore:
+                        NODE_DISPLAY_NAME_MAPPINGS[name] = display_name
             return True
         # V3 Extension Definition
         elif hasattr(module, "comfy_entrypoint"):
@@ -2325,8 +2327,8 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                     if schema.node_id not in ignore:
                         NODE_CLASS_MAPPINGS[schema.node_id] = node_cls
                         node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
-                    if schema.display_name is not None:
-                        NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
+                        if schema.display_name is not None:
+                            NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
                 return True
             except Exception as e:
                 logging.warning(f"Error while calling comfy_entrypoint in {module_path}: {e}")
@@ -2458,6 +2460,7 @@ async def init_builtin_extra_nodes():
         "nodes_minimax_music.py",
         "nodes_minimax_h3.py",
         "nodes_lt.py",
+        "nodes_lt_keyframes.py",
         "nodes_hooks.py",
         "nodes_multigpu.py",
         "nodes_load_3d.py",

requirements.txt

@@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0
 filelock
 av>=17.0.0
 comfy-kitchen==0.2.33
-comfy-aimdo==0.5.2
+comfy-aimdo==0.5.3
 requests
 simpleeval>=1.0.0
 blake3

tests-unit/comfy_extras_test/nodes_lt_keyframes_test.py

@@ -0,0 +1,936 @@
+"""Unit tests for native LTXV generated-keyframe nodes and Freeze Latent.
+
+They cover keyframe placement, conditioning metadata, guide conversion, and freeze-mask behavior.
+"""
+
+from __future__ import annotations
+
+import sys
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+mock_nodes = MagicMock()
+mock_nodes.MAX_RESOLUTION = 16384
+mock_server = MagicMock()
+
+
+def _conditioning_get_any_value(conditioning, key, default=None):
+    for t in conditioning:
+        if key in t[1]:
+            return t[1][key]
+    return default
+
+
+def _get_noise_mask(latent):
+    noise_mask = latent.get("noise_mask", None)
+    latent_image = latent["samples"]
+    if noise_mask is None:
+        batch_size, _, latent_length, _, _ = latent_image.shape
+        noise_mask = torch.ones(
+            (batch_size, 1, latent_length, 1, 1),
+            dtype=torch.float32,
+            device=latent_image.device,
+        )
+    else:
+        noise_mask = noise_mask.clone()
+    return noise_mask
+
+
+def _get_keyframe_idxs(cond, latent_shape=None):
+    keyframe_idxs = _conditioning_get_any_value(cond, "keyframe_idxs", None)
+    if keyframe_idxs is None:
+        return None, 0
+    if latent_shape is not None and len(latent_shape) == 5:
+        tokens_per_frame = latent_shape[-2] * latent_shape[-1]
+        num_keyframes = keyframe_idxs.shape[2] // tokens_per_frame
+        return keyframe_idxs, num_keyframes
+    return keyframe_idxs, 0
+
+
+def _append_guide_attention_entry(positive, negative, pre_filter_count, latent_shape, strength=1.0, attention_mask=None):
+    import node_helpers
+
+    new_entry = {
+        "pre_filter_count": pre_filter_count,
+        "strength": strength,
+        "pixel_mask": None,
+        "latent_shape": latent_shape,
+    }
+    results = []
+    for cond in (positive, negative):
+        existing = []
+        for t in cond:
+            found = t[1].get("guide_attention_entries", None)
+            if found is not None:
+                existing = found
+                break
+        results.append(
+            node_helpers.conditioning_set_values(cond, {"guide_attention_entries": [*existing, new_entry]})
+        )
+    return results[0], results[1]
+
+
+class _StubAddGuide:
+    calls = []
+
+    @classmethod
+    def append_keyframe(
+        cls,
+        positive,
+        negative,
+        frame_idx,
+        latent_image,
+        noise_mask,
+        guiding_latent,
+        strength,
+        scale_factors,
+        **kwargs,
+    ):
+        cls.calls.append({"method": "append_keyframe", "frame_idx": int(frame_idx), "strength": strength})
+        mask = torch.full(
+            (noise_mask.shape[0], 1, guiding_latent.shape[2], noise_mask.shape[3], noise_mask.shape[4]),
+            max(0.0, 1.0 - strength),
+            dtype=noise_mask.dtype,
+            device=noise_mask.device,
+        )
+        return (
+            positive,
+            negative,
+            torch.cat([latent_image, guiding_latent], dim=2),
+            torch.cat([noise_mask, mask], dim=2),
+        )
+
+    @classmethod
+    def execute(cls, positive, negative, vae, latent, image, frame_idx, strength, **kwargs):
+        cls.calls.append({"method": "execute", "frame_idx": int(frame_idx), "strength": strength, "image": image})
+        samples = latent["samples"]
+        out = latent.copy()
+        extra = torch.zeros(
+            (samples.shape[0], samples.shape[1], 1, samples.shape[3], samples.shape[4]),
+            dtype=samples.dtype,
+            device=samples.device,
+        )
+        out["samples"] = torch.cat([samples, extra], dim=2)
+        return _NodeOutput(positive, negative, out)
+
+
+class _NodeOutput:
+    def __init__(self, *args):
+        self.args = args
+
+    def __getitem__(self, index):
+        return self.args[index]
+
+
+_nodes_lt_stub = MagicMock()
+_nodes_lt_stub.conditioning_get_any_value = _conditioning_get_any_value
+_nodes_lt_stub.get_noise_mask = _get_noise_mask
+_nodes_lt_stub.get_keyframe_idxs = _get_keyframe_idxs
+_nodes_lt_stub._append_guide_attention_entry = _append_guide_attention_entry
+_nodes_lt_stub.LTXVAddGuide = _StubAddGuide
+
+def _import_keyframes_against_stub():
+    """Import the module under test with comfy_extras.nodes_lt stubbed, then put it back.
+
+    Only the stubbed keys are restored, not the whole of sys.modules: patch.dict
+    restores the entire dict on exit, which evicts every module imported inside the
+    block and forces a later re-import. Re-importing torch internals raises on
+    duplicate TORCH_LIBRARY registration, which broke running this file alongside
+    nodes_lt_test.py.
+    """
+    stubs = {
+        "nodes": mock_nodes,
+        "server": mock_server,
+        "comfy_extras.nodes_lt": _nodes_lt_stub,
+    }
+    saved = {name: sys.modules.get(name) for name in stubs}
+    sys.modules.update(stubs)
+    try:
+        import comfy_extras.nodes_lt_keyframes as module
+
+        return module
+    finally:
+        for name, original in saved.items():
+            if original is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = original
+        # The imported module stays bound to the stubs above, so drop it from the cache
+        # rather than let a later import pick up a stub-backed copy. The reference
+        # returned to this module keeps working.
+        sys.modules.pop("comfy_extras.nodes_lt_keyframes", None)
+
+
+keyframes = _import_keyframes_against_stub()
+
+
+def _zeros(shape):
+    return torch.zeros(shape)
+
+
+def _empty_121():
+    return {"samples": _zeros((1, 2, 16, 2, 1))}
+
+
+def _empty_241():
+    return {"samples": _zeros((1, 2, 31, 2, 1))}
+
+
+def _cond(**extra):
+    return [({}, dict(extra))]
+
+
+def _vae():
+    return SimpleNamespace(downscale_index_formula=(8, 32, 32))
+
+
+def _mask(shape, occupied):
+    tensor = torch.ones(shape)
+    for frame in occupied:
+        tensor[:, :, frame] = 0.0
+    return tensor
+
+
+def _keyframe_idxs_at(starts, tokens_per_frame=1):
+    times = []
+    for start in starts:
+        times.extend([start] * tokens_per_frame)
+    n = len(times)
+    coords = torch.zeros((1, 3, n, 2))
+    for i, start in enumerate(times):
+        coords[0, 0, i, 0] = float(start)
+        coords[0, 0, i, 1] = float(start + 1)
+        coords[0, 1, i, 1] = 1.0
+        coords[0, 2, i, 1] = 1.0
+    return coords
+
+
+@contextmanager
+def _stub_get_keyframe_idxs(idxs, num_guide_frames):
+    original = keyframes.get_keyframe_idxs
+    keyframes.get_keyframe_idxs = lambda cond, shape=None: (idxs, num_guide_frames)
+    try:
+        yield
+    finally:
+        keyframes.get_keyframe_idxs = original
+
+
+@contextmanager
+def _stub_keyframe_coords():
+    original = keyframes.LTXVAddGeneratedKeyframes.keyframe_coords
+
+    def _fake(cls, latent, frame_index, scale_factors):
+        return torch.zeros((latent.shape[0], 3, latent.shape[3] * latent.shape[4], 2))
+
+    keyframes.LTXVAddGeneratedKeyframes.keyframe_coords = classmethod(_fake)
+    try:
+        yield
+    finally:
+        keyframes.LTXVAddGeneratedKeyframes.keyframe_coords = original
+
+
+class TestPlacementHelpers:
+    def test_detailing_positions_121_24(self):
+        assert keyframes.detailing_positions(121, 24) == [24, 48, 72, 96, 120]
+        assert keyframes.free_detailing_slots(121, 24, occupied=set()) == [24, 48, 72, 96, 120]
+        assert keyframes.free_detailing_slots(241, 24, occupied={0, 48, 96, 144, 192, 240}) == [
+            24, 72, 120, 168, 216
+        ]
+
+    def test_free_slots_skip_last_frame_when_occupied(self):
+        assert keyframes.free_detailing_slots(121, 24, occupied={120}) == [24, 48, 72, 96]
+
+    def test_free_slots_rejects_when_every_candidate_is_occupied(self):
+        with pytest.raises(ValueError, match="already has an image keyframe"):
+            keyframes.free_detailing_slots(121, 24, occupied={24, 48, 72, 96, 120})
+
+    def test_scale_frame_indices_temporal_x2(self):
+        assert keyframes.scale_frame_indices([24, 48, 72, 96, 120], 121, 241) == [
+            48, 96, 144, 192, 240
+        ]
+        with pytest.raises(ValueError, match="from a 1-frame"):
+            keyframes.scale_frame_indices([0], 1, 241)
+        with pytest.raises(ValueError, match="onto a 1-frame"):
+            keyframes.scale_frame_indices([24], 121, 1)
+
+    def test_scale_frame_indices_rejects_collapsed_duplicates(self):
+        with pytest.raises(ValueError, match="collapsed"):
+            keyframes.scale_frame_indices([0, 1], 121, 3)
+
+    def test_detailing_positions_keeps_last_skips_zero(self):
+        positions = keyframes.detailing_positions(121, 24.0)
+        assert positions[0] != 0
+        assert positions[-1] == 120
+
+    def test_detailing_positions_rejects_nonpositive_interval(self):
+        with pytest.raises(ValueError, match="interval_frames"):
+            keyframes.detailing_positions(121, 0)
+
+    def test_detailing_positions_rejects_one_frame_canvas(self):
+        with pytest.raises(ValueError, match="no pixel frames"):
+            keyframes.detailing_positions(1, 24)
+        with pytest.raises(ValueError, match="no pixel frames"):
+            keyframes.free_detailing_slots(1, 24, occupied=set())
+
+    def test_keyframes_from_video_stacking_shape(self):
+        samples = torch.arange(1 * 2 * 4 * 2 * 1, dtype=torch.float32).reshape(1, 2, 4, 2, 1)
+        stacked = keyframes.keyframes_from_video(samples, [8, 16, 24], temporal_scale=8)
+        assert stacked.shape == (1, 2, 3, 2, 1)
+        assert torch.equal(stacked[:, :, 0:1], samples[:, :, 1:2])
+        assert torch.equal(stacked[:, :, 1:2], samples[:, :, 2:3])
+        assert torch.equal(stacked[:, :, 2:3], samples[:, :, 3:4])
+
+    def test_keyframes_from_video_rejects_non_video_and_bad_scale(self):
+        with pytest.raises(ValueError, match="plain 5D video latent"):
+            keyframes.keyframes_from_video([0], [8], 8)
+        with pytest.raises(ValueError, match="temporal_scale"):
+            keyframes.keyframes_from_video(_zeros((1, 2, 4, 2, 1)), [8], 0)
+        with pytest.raises(ValueError, match="no frames to copy"):
+            keyframes.keyframes_from_video(_zeros((1, 2, 0, 2, 1)), [8], 8)
+
+    def test_nearest_latent_index_clamps(self):
+        assert keyframes.nearest_latent_index(0, 8, 4) == 0
+        assert keyframes.nearest_latent_index(8, 8, 4) == 1
+        assert keyframes.nearest_latent_index(999, 8, 4) == 3
+
+    def test_should_copy_nearest_video_frames(self):
+        assert keyframes.should_copy_nearest_video_frames(31, 5, False, False) is True
+        assert keyframes.should_copy_nearest_video_frames(5, 5, False, False) is False
+        assert keyframes.should_copy_nearest_video_frames(4, 5, False, False) is False
+        assert keyframes.should_copy_nearest_video_frames(31, 5, True, False) is False
+        assert keyframes.should_copy_nearest_video_frames(31, None, False, False) is False
+        assert keyframes.should_copy_nearest_video_frames(1, 5, False, True) is False
+
+    def test_parse_frame_index_list_validates_count_range_and_duplicates(self):
+        assert keyframes._parse_frame_index_list(
+            "24, 48", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them"
+        ) == [24, 48]
+        assert keyframes._parse_frame_index_list(
+            "24 48", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them"
+        ) == [24, 48]
+        with pytest.raises(ValueError, match="lists 1"):
+            keyframes._parse_frame_index_list(
+                "24", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them"
+            )
+        with pytest.raises(ValueError, match="same pixel frame"):
+            keyframes._parse_frame_index_list(
+                "24,24", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them"
+            )
+        with pytest.raises(ValueError, match="must lie between"):
+            keyframes._parse_frame_index_list(
+                "0,24", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them"
+            )
+        with pytest.raises(ValueError, match="could not parse"):
+            keyframes._parse_frame_index_list(
+                "24,abc", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them"
+            )
+
+    def test_parse_frame_index_list_allows_omitted_count(self):
+        assert keyframes._parse_frame_index_list(
+            "24,48,72", "frame_indices", None, 1, 120, "unused", "to auto-place"
+        ) == [24, 48, 72]
+
+    def test_parse_frame_index_list_rejects_empty_separator_only(self):
+        with pytest.raises(ValueError, match="is empty"):
+            keyframes._parse_frame_index_list(
+                ",", "frame_indices", None, 1, 120, "unused", "to place them from interval_frames"
+            )
+        with pytest.raises(ValueError, match="is empty"):
+            keyframes._parse_frame_index_list(
+                " , , ", "frame_indices", None, 1, 120, "unused", "to place them from interval_frames"
+            )
+
+    def test_add_parse_frame_indices_allows_last_frame(self):
+        assert keyframes.LTXVAddGeneratedKeyframes.parse_frame_indices("24,120", 121) == [24, 120]
+        with pytest.raises(ValueError, match="no pixel frames"):
+            keyframes.LTXVAddGeneratedKeyframes.parse_frame_indices("1", 1)
+
+    def test_occupied_from_nonzero_samples_without_mask(self):
+        samples = _zeros((1, 2, 16, 2, 1))
+        samples[0, 0, 0, 0, 0] = 1.0
+        taken = keyframes.occupied_pixel_frames({"samples": samples}, 8, 121)
+        assert 0 in taken
+        assert 120 not in taken
+
+    def test_occupied_prefers_noise_mask_over_nonzero_samples(self):
+        samples = _zeros((1, 2, 16, 2, 1))
+        samples[0, 0, 0, 0, 0] = 1.0
+        latent = {"samples": samples, "noise_mask": _mask((1, 1, 16, 1, 1), occupied=set())}
+        assert keyframes.occupied_pixel_frames(latent, 8, 121) == set()
+
+    def test_occupied_ignores_appended_guide_frames(self):
+        latent = {
+            "samples": _zeros((1, 2, 21, 2, 1)),
+            "noise_mask": _mask((1, 1, 21, 1, 1), occupied={0, 16, 17, 18, 19, 20}),
+        }
+        taken = keyframes.occupied_pixel_frames(latent, 8, 121, video_latent_frames=16)
+        assert taken == {0}
+
+    def test_pixel_frames_from_keyframe_idxs_uses_start_not_exclusive_end(self):
+        idxs = _keyframe_idxs_at([24])
+        assert idxs[0, 0, :, 0].tolist() == [24.0]
+        assert idxs[0, 0, :, 1].tolist() == [25.0]
+        assert keyframes.pixel_frames_from_keyframe_idxs(idxs) == {24}
+        assert keyframes.pixel_frames_from_keyframe_idxs(None) == set()
+
+    def test_pixel_frames_from_keyframe_idxs_rejects_malformed(self):
+        with pytest.raises((TypeError, AttributeError, IndexError, ValueError)):
+            keyframes.pixel_frames_from_keyframe_idxs("not-a-tensor")
+        with pytest.raises((TypeError, ValueError)):
+            keyframes._as_int_set(object())
+
+
+class TestNativeSchemas:
+    def test_generated_keyframe_nodes_use_ltxv_conditioning_category(self):
+        for cls, node_id, display_name in (
+            (
+                keyframes.LTXVAddGeneratedKeyframes,
+                "LTXVAddGeneratedKeyframes",
+                "LTXV Add Generated Keyframes",
+            ),
+            (
+                keyframes.LTXVSeparateGeneratedKeyframes,
+                "LTXVSeparateGeneratedKeyframes",
+                "LTXV Separate Generated Keyframes",
+            ),
+            (
+                keyframes.LTXVGeneratedKeyframesToGuides,
+                "LTXVGeneratedKeyframesToGuides",
+                "LTXV Generated Keyframes to Guides",
+            ),
+        ):
+            schema = cls.define_schema()
+            assert schema.node_id == node_id
+            assert schema.display_name == display_name
+            assert schema.category == "model/conditioning/ltxv"
+            assert "dfr" in schema.search_aliases
+
+    def test_freeze_latent_uses_ltxv_latent_category(self):
+        schema = keyframes.LTXVFreezeLatent.define_schema()
+        assert schema.node_id == "LTXVFreezeLatent"
+        assert schema.display_name == "LTXV Freeze Latent"
+        assert schema.category == "model/latent/ltxv"
+
+
+class TestAddGeneratedKeyframes:
+    def test_rejects_non_video_latent(self):
+        with pytest.raises(ValueError, match="plain video latent"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), {"samples": torch.zeros(1, 2, 16, 2)}
+            )
+
+    def test_execute_rejects_separator_only_frame_indices(self):
+        with pytest.raises(ValueError, match="is empty"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), _empty_121(), frame_indices=","
+            )
+
+    def test_execute_rejects_one_frame_canvas(self):
+        with pytest.raises(ValueError, match="no pixel frames"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), {"samples": _zeros((1, 2, 1, 2, 1))}
+            )
+
+    def test_rejects_rescaled_or_noncontiguous_existing_keyframes(self):
+        latent = _empty_121()
+        with pytest.raises(ValueError, match="rescaled"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(
+                    generated_keyframes={
+                        "tokens_per_frame": 99,
+                        "first_latent_frame": 16,
+                        "num_keyframes": 0,
+                    }
+                ),
+                _cond(),
+                _vae(),
+                latent,
+            )
+        with pytest.raises(ValueError, match="contiguous"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(
+                    generated_keyframes={
+                        "tokens_per_frame": 2,
+                        "first_latent_frame": 10,
+                        "num_keyframes": 3,
+                    }
+                ),
+                _cond(),
+                _vae(),
+                latent,
+            )
+
+    def test_execute_appends_zero_keyframes_on_t(self):
+        with _stub_keyframe_coords():
+            _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), _empty_121()
+            )
+        assert out["samples"].shape == (1, 2, 21, 2, 1)
+        assert out["noise_mask"].shape[2] == 21
+        assert torch.all(out["noise_mask"][:, :, 16:21] == 1.0)
+
+    def test_execute_copies_nearest_frames_from_longer_video(self):
+        video = {"samples": torch.arange(1 * 2 * 16 * 2 * 1, dtype=torch.float32).reshape(1, 2, 16, 2, 1)}
+        with _stub_keyframe_coords():
+            _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                frame_indices="24,48,72,96,120",
+                keyframes=video,
+            )
+        assert out["samples"].shape[2] == 21
+        stacked = out["samples"][:, :, 16:21]
+        source = video["samples"]
+        assert torch.equal(stacked[:, :, 0:1], source[:, :, 3:4])
+        assert torch.equal(stacked[:, :, 4:5], source[:, :, 15:16])
+
+    def test_execute_keeps_stacked_keyframes_when_t_equals_count(self):
+        stacked = {"samples": torch.arange(1 * 2 * 5 * 2 * 1, dtype=torch.float32).reshape(1, 2, 5, 2, 1)}
+        with _stub_keyframe_coords():
+            _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                frame_indices="24,48,72,96,120",
+                keyframes=stacked,
+            )
+        assert out["samples"].shape[2] == 21
+        assert torch.equal(out["samples"][:, :, 16:21], stacked["samples"])
+
+    def test_execute_reshapes_batched_single_frame_keyframes(self):
+        batched = {"samples": _zeros((5, 2, 1, 2, 1))}
+        with _stub_keyframe_coords():
+            _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                frame_indices="24,48,72,96,120",
+                keyframes=batched,
+            )
+        assert out["samples"].shape[2] == 21
+
+    def test_execute_records_density_slots_and_canvas_length(self):
+        with _stub_keyframe_coords():
+            positive, _negative, _out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), _empty_121()
+            )
+        record = positive[0][1]["generated_keyframes"]
+        assert record["frame_indices"] == [24, 48, 72, 96, 120]
+        assert record["num_pixel_frames"] == 121
+        assert record["num_keyframes"] == 5
+        assert record["first_latent_frame"] == 16
+        assert record["guide_entry_index"] == 0
+        entries = positive[0][1]["guide_attention_entries"]
+        assert len(entries) == 1
+        assert entries[0]["pre_filter_count"] == 5 * 2 * 1
+        assert entries[0]["latent_shape"] == [5, 2, 1]
+
+    def test_execute_copies_from_video_using_auto_slots(self):
+        video = {"samples": torch.arange(1 * 2 * 16 * 2 * 1, dtype=torch.float32).reshape(1, 2, 16, 2, 1)}
+        with _stub_keyframe_coords():
+            _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), _empty_121(), keyframes=video
+            )
+        stacked = out["samples"][:, :, 16:21]
+        source = video["samples"]
+        assert torch.equal(stacked[:, :, 0:1], source[:, :, 3:4])
+        assert torch.equal(stacked[:, :, 4:5], source[:, :, 15:16])
+
+    def test_execute_replaces_stacked_tokens_on_current_canvas(self):
+        stacked = {
+            "samples": torch.arange(1 * 2 * 5 * 2 * 1, dtype=torch.float32).reshape(1, 2, 5, 2, 1),
+            "generated_keyframe_indices": [24, 48, 72, 96, 120],
+            "generated_keyframe_num_frames": 121,
+        }
+        with _stub_keyframe_coords():
+            positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), _empty_121(), keyframes=stacked
+            )
+        assert torch.equal(out["samples"][:, :, 16:21], stacked["samples"])
+        assert positive[0][1]["generated_keyframes"]["frame_indices"] == [24, 48, 72, 96, 120]
+
+    def test_execute_skips_i2v_last_frame_noise_mask(self):
+        latent = {
+            "samples": _zeros((1, 2, 16, 2, 1)),
+            "noise_mask": _mask((1, 1, 16, 1, 1), occupied={15}),
+        }
+        with _stub_keyframe_coords():
+            positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), latent
+            )
+        indices = positive[0][1]["generated_keyframes"]["frame_indices"]
+        assert 120 not in indices
+        assert indices == [24, 48, 72, 96]
+        assert out["samples"].shape[2] == 20
+
+    def test_execute_replaces_stacked_tokens_on_longer_canvas(self):
+        stacked = {
+            "samples": _zeros((1, 2, 5, 2, 1)),
+            "generated_keyframe_indices": [24, 48, 72, 96, 120],
+            "generated_keyframe_num_frames": 121,
+        }
+        latent = _empty_241()
+        latent["noise_mask"] = _mask((1, 1, 31, 1, 1), occupied={0})
+        with _stub_keyframe_coords():
+            positive, _negative, _out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), latent, keyframes=stacked
+            )
+        indices = positive[0][1]["generated_keyframes"]["frame_indices"]
+        assert indices != [24, 48, 72, 96, 120]
+        assert indices == [24, 48, 72, 96, 120, 144, 168, 192, 216, 240]
+
+    def test_execute_skips_existing_guide_keyframe_idxs(self):
+        latent = {
+            "samples": _zeros((1, 2, 36, 2, 1)),
+            "noise_mask": _mask((1, 1, 36, 1, 1), occupied={0}),
+        }
+        idxs = _keyframe_idxs_at([48, 96, 144, 192, 240])
+        with _stub_keyframe_coords(), _stub_get_keyframe_idxs(idxs, 5):
+            positive, _negative, _out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), latent
+            )
+        assert positive[0][1]["generated_keyframes"]["frame_indices"] == [24, 72, 120, 168, 216]
+
+    def test_execute_rejects_occupied_manual_indices(self):
+        latent = {
+            "samples": _zeros((1, 2, 16, 2, 1)),
+            "noise_mask": _mask((1, 1, 16, 1, 1), occupied={15}),
+        }
+        with pytest.raises(ValueError, match="reuses pixel frame"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), latent, frame_indices="24,120"
+            )
+
+    def test_execute_rejects_wrong_spatial_size_keyframes(self):
+        stacked = {"samples": _zeros((1, 2, 5, 4, 4))}
+        with pytest.raises(ValueError, match="whole latent frames"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                frame_indices="24,48,72,96,120",
+                keyframes=stacked,
+            )
+
+    def test_execute_rejects_too_many_stacked_keyframes(self):
+        stacked = {
+            "samples": _zeros((1, 2, 6, 2, 1)),
+            "generated_keyframe_indices": [24, 48, 72, 96, 120, 8],
+        }
+        with pytest.raises(ValueError, match="only 5 free slot"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                frame_indices="24,48,72,96,120",
+                keyframes=stacked,
+            )
+
+    def test_execute_rejects_non_5d_keyframes(self):
+        with pytest.raises(ValueError, match="5 dimensional"):
+            keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                frame_indices="24",
+                keyframes={"samples": torch.zeros(1, 2, 1, 2)},
+            )
+
+    def test_execute_grows_existing_generated_block(self):
+        latent = _empty_121()
+        with _stub_keyframe_coords():
+            positive, negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                _cond(), _cond(), _vae(), latent, frame_indices="24,48"
+            )
+            positive, negative, out = keyframes.LTXVAddGeneratedKeyframes.execute(
+                positive, negative, _vae(), out, frame_indices="72,96"
+            )
+        record = positive[0][1]["generated_keyframes"]
+        assert record["frame_indices"] == [24, 48, 72, 96]
+        assert record["num_keyframes"] == 4
+        assert record["first_latent_frame"] == 16
+        assert out["samples"].shape[2] == 20
+        entries = positive[0][1]["guide_attention_entries"]
+        assert len(entries) == 1
+        assert entries[0]["pre_filter_count"] == 4 * 2 * 1
+        assert entries[0]["latent_shape"] == [4, 2, 1]
+
+
+class TestSeparateGeneratedKeyframes:
+    def test_requires_generated_keyframes(self):
+        with pytest.raises(ValueError, match="no generated keyframes"):
+            keyframes.LTXVSeparateGeneratedKeyframes.execute(_cond(), _cond(), _empty_121())
+
+    def test_execute_peels_keyframes_and_indices(self):
+        video = _zeros((1, 2, 16, 2, 1))
+        keys = torch.arange(1, 1 + 1 * 2 * 5 * 2 * 1, dtype=torch.float32).reshape(1, 2, 5, 2, 1)
+        samples = torch.cat([video, keys], dim=2)
+        record = {
+            "first_latent_frame": 16,
+            "num_keyframes": 5,
+            "frame_indices": [24, 48, 72, 96, 120],
+            "num_pixel_frames": 121,
+            "guide_entry_index": 0,
+            "tokens_per_frame": 2,
+        }
+        positive, negative, latent, peeled = keyframes.LTXVSeparateGeneratedKeyframes.execute(
+            _cond(
+                generated_keyframes=record,
+                guide_attention_entries=[{"keep": False}, {"keep": True}],
+            ),
+            _cond(generated_keyframes=record),
+            {"samples": samples},
+        )
+        assert latent["samples"].shape == (1, 2, 16, 2, 1)
+        assert peeled["samples"].shape == (1, 2, 5, 2, 1)
+        assert peeled["generated_keyframe_indices"] == [24, 48, 72, 96, 120]
+        assert peeled["generated_keyframe_num_frames"] == 121
+        assert torch.equal(peeled["samples"], keys)
+        assert positive[0][1]["generated_keyframes"] is None
+        assert positive[0][1]["guide_attention_entries"] == [{"keep": True}]
+        assert negative[0][1]["generated_keyframes"] is None
+
+    def test_execute_keyframes_to_batch(self):
+        samples = _zeros((1, 2, 18, 2, 1))
+        record = {
+            "first_latent_frame": 16,
+            "num_keyframes": 2,
+            "frame_indices": [24, 48],
+            "guide_entry_index": 0,
+            "tokens_per_frame": 2,
+        }
+        _p, _n, _latent, peeled = keyframes.LTXVSeparateGeneratedKeyframes.execute(
+            _cond(generated_keyframes=record),
+            _cond(generated_keyframes=record),
+            {"samples": samples},
+            keyframes_to_batch=True,
+        )
+        assert peeled["samples"].shape == (2, 2, 1, 2, 1)
+
+    def test_rejects_token_mismatch_and_short_latent(self):
+        record = {
+            "first_latent_frame": 16,
+            "num_keyframes": 5,
+            "frame_indices": [24, 48, 72, 96, 120],
+            "guide_entry_index": 0,
+            "tokens_per_frame": 99,
+        }
+        with pytest.raises(ValueError, match="rescaled"):
+            keyframes.LTXVSeparateGeneratedKeyframes.execute(
+                _cond(generated_keyframes=record),
+                _cond(generated_keyframes=record),
+                _empty_121(),
+            )
+        record = dict(record)
+        record["tokens_per_frame"] = 2
+        with pytest.raises(ValueError, match="only has"):
+            keyframes.LTXVSeparateGeneratedKeyframes.execute(
+                _cond(generated_keyframes=record),
+                _cond(generated_keyframes=record),
+                _empty_121(),
+            )
+
+    def test_strip_guide_entry(self):
+        remaining = keyframes.LTXVSeparateGeneratedKeyframes.strip_guide_entry(
+            [({}, {"guide_attention_entries": [{"a": 1}, {"b": 2}]})], 0
+        )
+        assert remaining == [{"b": 2}]
+        empty = keyframes.LTXVSeparateGeneratedKeyframes.strip_guide_entry(
+            [({}, {"guide_attention_entries": [{"a": 1}]})], 0
+        )
+        assert empty is None
+        with pytest.raises(ValueError, match="recorded guide entry"):
+            keyframes.LTXVSeparateGeneratedKeyframes.strip_guide_entry(
+                [({}, {"guide_attention_entries": [{"a": 1}]})], 5
+            )
+
+    def test_rejects_non_video_latent(self):
+        record = {
+            "first_latent_frame": 0,
+            "num_keyframes": 1,
+            "frame_indices": [24],
+            "guide_entry_index": 0,
+            "tokens_per_frame": 2,
+        }
+        with pytest.raises(ValueError, match="plain video latent"):
+            keyframes.LTXVSeparateGeneratedKeyframes.execute(
+                _cond(generated_keyframes=record),
+                _cond(generated_keyframes=record),
+                {"samples": torch.zeros(1, 2, 16, 2)},
+            )
+
+
+class TestGeneratedKeyframesToGuides:
+    def test_requires_recorded_indices(self):
+        with pytest.raises(ValueError, match="does not carry generated keyframe positions"):
+            keyframes.LTXVGeneratedKeyframesToGuides.execute(
+                _cond(), _cond(), _vae(), _empty_121(), {"samples": _zeros((1, 2, 5, 2, 1))}, 1.0
+            )
+
+    def test_rejects_unseparated_conditioning(self):
+        with pytest.raises(ValueError, match="still carries generated keyframes"):
+            keyframes.LTXVGeneratedKeyframesToGuides.execute(
+                _cond(generated_keyframes={"num_keyframes": 1}),
+                _cond(),
+                _vae(),
+                _empty_121(),
+                {"samples": _zeros((1, 2, 1, 2, 1)), "generated_keyframe_indices": [24]},
+                1.0,
+            )
+
+    def test_rejects_non_video_and_batched_canvas(self):
+        kf = {"samples": _zeros((1, 2, 1, 2, 1)), "generated_keyframe_indices": [24]}
+        with pytest.raises(ValueError, match="plain video latent"):
+            keyframes.LTXVGeneratedKeyframesToGuides.execute(
+                _cond(), _cond(), _vae(), {"samples": torch.zeros(1, 2, 16, 2)}, kf, 1.0
+            )
+        with pytest.raises(ValueError, match="batch size of 1"):
+            keyframes.LTXVGeneratedKeyframesToGuides.execute(
+                _cond(), _cond(), _vae(), {"samples": _zeros((2, 2, 16, 2, 1))}, kf, 1.0
+            )
+
+    def test_pins_same_size_keyframes_via_append(self):
+        _StubAddGuide.calls.clear()
+        kf = {
+            "samples": _zeros((1, 2, 2, 2, 1)),
+            "generated_keyframe_indices": [24, 48],
+            "generated_keyframe_num_frames": 121,
+        }
+        positive, negative, out = keyframes.LTXVGeneratedKeyframesToGuides.execute(
+            _cond(), _cond(), _vae(), _empty_121(), kf, 1.0
+        )
+        assert out["samples"].shape[2] == 18
+        assert torch.all(out["noise_mask"][:, :, 16:] == 0.0)
+        assert [call["frame_idx"] for call in _StubAddGuide.calls] == [24, 48]
+        assert all(call["method"] == "append_keyframe" for call in _StubAddGuide.calls)
+        entries = positive[0][1]["guide_attention_entries"]
+        assert len(entries) == 2
+
+    def test_scales_indices_after_temporal_x2(self):
+        _StubAddGuide.calls.clear()
+        kf = {
+            "samples": _zeros((1, 2, 2, 2, 1)),
+            "generated_keyframe_indices": [24, 120],
+            "generated_keyframe_num_frames": 121,
+        }
+        keyframes.LTXVGeneratedKeyframesToGuides.execute(
+            _cond(), _cond(), _vae(), _empty_241(), kf, 1.0
+        )
+        assert [call["frame_idx"] for call in _StubAddGuide.calls] == [48, 240]
+
+    def test_override_frame_indices(self):
+        _StubAddGuide.calls.clear()
+        kf = {
+            "samples": _zeros((1, 2, 2, 2, 1)),
+            "generated_keyframe_indices": [24, 48],
+            "generated_keyframe_num_frames": 121,
+        }
+        keyframes.LTXVGeneratedKeyframesToGuides.execute(
+            _cond(), _cond(), _vae(), _empty_121(), kf, 0.5, override_frame_indices="32,96"
+        )
+        assert [call["frame_idx"] for call in _StubAddGuide.calls] == [32, 96]
+        assert all(call["strength"] == 0.5 for call in _StubAddGuide.calls)
+
+    def test_resize_path_decodes_and_calls_add_guide(self):
+        _StubAddGuide.calls.clear()
+        vae = _vae()
+        decoded = []
+
+        def decode(samples):
+            decoded.append(tuple(samples.shape))
+            return torch.zeros((samples.shape[0], 8, 8, 3))
+
+        vae.decode = decode
+        kf = {
+            "samples": _zeros((1, 2, 2, 4, 4)),
+            "generated_keyframe_indices": [24, 48],
+            "generated_keyframe_num_frames": 121,
+        }
+        _p, _n, out = keyframes.LTXVGeneratedKeyframesToGuides.execute(
+            _cond(), _cond(), vae, _empty_121(), kf, 1.0
+        )
+        assert decoded == [(2, 2, 1, 4, 4)]
+        assert [call["method"] for call in _StubAddGuide.calls] == ["execute", "execute"]
+        assert out["samples"].shape[2] == 18
+
+    def test_rejects_count_mismatch(self):
+        kf = {
+            "samples": _zeros((1, 2, 2, 2, 1)),
+            "generated_keyframe_indices": [24],
+            "generated_keyframe_num_frames": 121,
+        }
+        with pytest.raises(ValueError, match="2 keyframes for 1 recorded"):
+            keyframes.LTXVGeneratedKeyframesToGuides.execute(
+                _cond(), _cond(), _vae(), _empty_121(), kf, 1.0
+            )
+
+    def test_override_rejects_separator_only_indices(self):
+        kf = {
+            "samples": _zeros((1, 2, 2, 2, 1)),
+            "generated_keyframe_indices": [24, 48],
+            "generated_keyframe_num_frames": 121,
+        }
+        with pytest.raises(ValueError, match="is empty"):
+            keyframes.LTXVGeneratedKeyframesToGuides.execute(
+                _cond(), _cond(), _vae(), _empty_121(), kf, 1.0, override_frame_indices=","
+            )
+
+
+class TestFreezeLatent:
+    def test_video_and_audio_masks(self):
+        video = keyframes.LTXVFreezeLatent.execute({"samples": _zeros((2, 4, 8, 3, 5))})[0]
+        assert video["noise_mask"].shape == (2, 1, 8, 1, 1)
+        assert video["noise_mask"].device.type == "cpu"
+        assert torch.all(video["noise_mask"] == 0)
+        audio = keyframes.LTXVFreezeLatent.execute({"samples": _zeros((1, 8, 16, 4))})[0]
+        assert audio["noise_mask"].shape == (1, 1, 16, 1)
+        assert torch.all(audio["noise_mask"] == 0)
+
+    def test_preserves_extra_latent_keys(self):
+        out = keyframes.LTXVFreezeLatent.execute(
+            {"samples": _zeros((1, 4, 8, 2, 2)), "downscale_ratio_spacial": 32}
+        )[0]
+        assert out["downscale_ratio_spacial"] == 32
+
+    def test_rejects_av_and_wrong_rank(self):
+        with pytest.raises(ValueError, match="plain tensor"):
+            keyframes.LTXVFreezeLatent.execute({"samples": [0.0]})
+        with pytest.raises(ValueError, match="4D audio or 5D video"):
+            keyframes.LTXVFreezeLatent.execute({"samples": _zeros((1, 2, 3))})
+
+
+class TestKeyframeCoords:
+    def test_single_pixel_span_at_requested_index(self):
+        latent = torch.zeros((1, 4, 1, 2, 2))
+        coords = keyframes.LTXVAddGeneratedKeyframes.keyframe_coords(latent, 24, (8, 32, 32))
+        assert coords.shape[0] == 1
+        assert coords.shape[1] == 3
+        assert coords.shape[-1] == 2
+        starts = coords[0, 0, :, 0]
+        ends = coords[0, 0, :, 1]
+        assert torch.all(starts == 24)
+        assert torch.all(ends == 25)
+
+
+def test_extension_registers_all_four_nodes():
+    import asyncio
+
+    ext = asyncio.run(keyframes.comfy_entrypoint())
+    names = [cls.__name__ for cls in asyncio.run(ext.get_node_list())]
+    assert names == [
+        "LTXVAddGeneratedKeyframes",
+        "LTXVSeparateGeneratedKeyframes",
+        "LTXVGeneratedKeyframesToGuides",
+        "LTXVFreezeLatent",
+    ]

tests-unit/comfy_extras_test/nodes_lt_test.py

@@ -0,0 +1,242 @@
+"""Unit tests for LTXVAddLatentGuide and the guide-attachment path it shares with LTXVAddGuide.
+
+The RoPE arithmetic runs for real here: only ``nodes`` and ``server`` are stubbed, so
+``append_keyframe``, ``dilate_latent`` and ``_append_guide_attention_entry`` are the
+real implementations and the assertions are on actual keyframe coordinates.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+# Stub nodes/server for the import only, then restore exactly those keys. patch.dict is
+# not used because it restores the whole of sys.modules on exit, which evicts everything
+# imported inside the block and forces a re-import that trips duplicate TORCH_LIBRARY
+# registration. Leaving the stubs installed is equally wrong: pytest imports every test
+# module at collection time, so a lingering MagicMock "nodes" breaks later modules that
+# use the real one. Same shape as tests-unit/comfy_extras_test/image_stitch_test.py.
+_stubs = {"nodes": MagicMock(MAX_RESOLUTION=16384), "server": MagicMock()}
+_saved = {name: sys.modules.get(name) for name in _stubs}
+sys.modules.update(_stubs)
+try:
+    import comfy_extras.nodes_lt as nodes_lt
+finally:
+    for _name, _original in _saved.items():
+        if _original is None:
+            sys.modules.pop(_name, None)
+        else:
+            sys.modules[_name] = _original
+
+LATENT_CHANNELS = 128
+SCALE_FACTORS = (8, 32, 32)
+TIME, HEIGHT, WIDTH = 0, 1, 2
+START, END = 0, 1
+
+
+def _vae():
+    return SimpleNamespace(downscale_index_formula=SCALE_FACTORS)
+
+
+def _latent(frames, height, width):
+    return {"samples": torch.zeros((1, LATENT_CHANNELS, frames, height, width))}
+
+
+def _cond():
+    return [({}, {})]
+
+
+def _add_latent_guide(guide_hw, latent_hw=(4, 4), guide_frames=1, latent_frames=3, latent_idx=0):
+    positive, negative, latent = nodes_lt.LTXVAddLatentGuide.execute(
+        _cond(),
+        _cond(),
+        _vae(),
+        _latent(latent_frames, *latent_hw),
+        _latent(guide_frames, *guide_hw),
+        latent_idx,
+        1.0,
+    )
+    metadata = positive[0][1]
+    return metadata["keyframe_idxs"], metadata["guide_attention_entries"], latent
+
+
+def _axis(keyframe_idxs, axis, bound):
+    return keyframe_idxs[0, axis, :, bound].tolist()
+
+
+def test_same_size_guide_spans_one_patch_per_token():
+    """A 1:1 guide gets no offset: each token's end is one scale factor past its start.
+
+    The inverse of the case below, so a factor derived wrongly at 1:1 is caught too.
+    """
+    keyframe_idxs, entries, _ = _add_latent_guide(guide_hw=(4, 4))
+
+    for axis in (HEIGHT, WIDTH):
+        starts = _axis(keyframe_idxs, axis, START)
+        assert _axis(keyframe_idxs, axis, END) == [s + SCALE_FACTORS[axis] for s in starts]
+
+    assert entries[0]["latent_shape"] == [1, 4, 4]
+
+
+def test_half_size_guide_expands_only_the_end_positions():
+    """An x2 guide keeps its start positions and pushes each end out by one scale factor.
+
+    That is what makes the dilated reference cover the whole canvas. Leaving the factor
+    at 1 keeps same-size coordinates while each token encodes a larger patch, so the
+    reference addresses only the top-left corner of the target.
+    """
+    same_size, _, _ = _add_latent_guide(guide_hw=(4, 4))
+    downscaled, entries, _ = _add_latent_guide(guide_hw=(2, 2))
+
+    # Dilation puts the small guide on the same sparse grid, so token count is unchanged.
+    assert downscaled.shape == same_size.shape
+
+    for axis in (HEIGHT, WIDTH):
+        assert _axis(downscaled, axis, START) == _axis(same_size, axis, START)
+        expected = [e + SCALE_FACTORS[axis] for e in _axis(same_size, axis, END)]
+        assert _axis(downscaled, axis, END) == expected
+
+    # Time is never touched by the spatial offset.
+    assert _axis(downscaled, TIME, START) == _axis(same_size, TIME, START)
+    assert _axis(downscaled, TIME, END) == _axis(same_size, TIME, END)
+
+    assert entries[0]["latent_shape"] == [1, 2, 2]
+
+
+def test_attention_entry_lets_context_windows_rederive_the_factor():
+    """The entry keeps the pre-dilation shape while the token count is post-dilation.
+
+    ``context_windows`` divides the post-dilation guide height by the entry's
+    ``latent_shape`` height to recover the downscale factor, so these two must not drift
+    apart or windowed and non-windowed sampling disagree on the guide's RoPE.
+    """
+    _, entries, latent = _add_latent_guide(guide_hw=(2, 2))
+
+    entry = entries[0]
+    assert entry["latent_shape"] == [1, 2, 2]  # pre-dilation
+    assert entry["pre_filter_count"] == 1 * 4 * 4  # post-dilation
+    assert latent["samples"].shape[3] // entry["latent_shape"][1] == 2
+
+
+@pytest.mark.parametrize(
+    "latent_idx, expected_start", [(-1, -8), (0, 0), (1, 1), (2, 9)]
+)
+def test_latent_idx_maps_onto_pixel_frames(latent_idx, expected_start):
+    """latent_idx is in latent frames, and negatives sit before the start of the latent.
+
+    The first latent frame covers a single pixel frame, so the mapping is 0 -> 0, 1 -> 1,
+    then 8 apart. Negative values are not counted back from the end.
+    """
+    keyframe_idxs, _, _ = _add_latent_guide(
+        guide_hw=(4, 4), latent_frames=8, latent_idx=latent_idx
+    )
+
+    assert set(_axis(keyframe_idxs, TIME, START)) == {expected_start}
+
+
+@pytest.mark.parametrize(
+    "kwargs, message",
+    [
+        (dict(guide_hw=(2, 4)), "square"),
+        (dict(guide_hw=(3, 3)), "whole number"),
+        (dict(guide_hw=(4, 4), latent_idx=99), "runs past the end"),
+        (dict(guide_hw=(4, 4), guide_frames=5), "runs past the end"),
+    ],
+)
+def test_unusable_guides_are_rejected(kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        _add_latent_guide(**kwargs)
+
+
+def test_non_5d_guiding_latent_is_rejected():
+    """An image-model latent would otherwise fail with a bare IndexError on shape[4]."""
+    with pytest.raises(ValueError, match="5D video latent"):
+        nodes_lt.LTXVAddLatentGuide.execute(
+            _cond(),
+            _cond(),
+            _vae(),
+            _latent(3, 4, 4),
+            {"samples": torch.zeros((1, LATENT_CHANNELS, 4, 4))},
+            0,
+            1.0,
+        )
+
+
+def test_attention_mask_reaches_the_guide_entry():
+    """Only coverage that the optional input is forwarded to the guide entry."""
+    mask = torch.full((1, 128, 128), 0.5)
+    positive, negative, _ = nodes_lt.LTXVAddLatentGuide.execute(
+        _cond(), _cond(), _vae(), _latent(3, 4, 4), _latent(1, 2, 2), 0, 1.0, attention_mask=mask
+    )
+
+    # Stored as (1, 1, F, H, W) for downstream self-attention masking.
+    assert positive[0][1]["guide_attention_entries"][0]["pixel_mask"].shape == (1, 1, 1, 128, 128)
+    # Positive and negative each get their own entry, so neither leaks into the other.
+    assert len(negative[0][1]["guide_attention_entries"]) == 1
+
+
+def test_node_is_registered_with_a_loadable_schema():
+    """Both failure modes here are invisible to every other test in this file.
+
+    A bad io.Schema keyword only surfaces when the node is registered, and a node left
+    out of the extension list simply does not exist in ComfyUI. The strength cap is
+    asserted here because raising it re-exposes the missing clamp in append_keyframe's
+    guide_mask branch, where a dilated guide is dropped above 1.0.
+    """
+    schema = nodes_lt.LTXVAddLatentGuide.define_schema()
+    inputs = {inp.id: inp for inp in schema.inputs}
+
+    assert schema.node_id == "LTXVAddLatentGuide"
+    assert inputs["strength"].max == 1.0
+    assert inputs["attention_mask"].optional is True
+
+    node_list = asyncio.run(nodes_lt.LtxvExtension().get_node_list())
+    assert nodes_lt.LTXVAddLatentGuide in node_list
+
+
+@pytest.mark.parametrize(
+    "iclora_parameters, expected_shape, expected_extra_end",
+    [
+        (None, [1, 4, 4], 0),
+        ({"reference_downscale_factor": 2}, [1, 2, 2], SCALE_FACTORS[HEIGHT]),
+    ],
+)
+def test_add_guide_image_path_still_routes_through_the_shared_helper(
+    iclora_parameters, expected_shape, expected_extra_end
+):
+    """LTXVAddGuide must be unchanged by sharing attach_guide_latent with the latent node."""
+
+    class _Vae:
+        downscale_index_formula = SCALE_FACTORS
+
+        def encode(self, pixels):
+            frames, height, width, _ = pixels.shape
+            return torch.zeros(
+                (1, LATENT_CHANNELS, (frames - 1) // SCALE_FACTORS[TIME] + 1, height // 32, width // 32)
+            )
+
+    positive, _, _ = nodes_lt.LTXVAddGuide.execute(
+        _cond(),
+        _cond(),
+        _Vae(),
+        _latent(3, 4, 4),
+        torch.zeros((1, 4 * 32, 4 * 32, 3)),
+        0,
+        1.0,
+        iclora_parameters=iclora_parameters,
+    )
+
+    metadata = positive[0][1]
+    entry = metadata["guide_attention_entries"][0]
+    assert entry["latent_shape"] == expected_shape
+    assert entry["pre_filter_count"] == 1 * 4 * 4
+
+    keyframe_idxs = metadata["keyframe_idxs"]
+    starts = _axis(keyframe_idxs, HEIGHT, START)
+    ends = _axis(keyframe_idxs, HEIGHT, END)
+    assert ends == [s + SCALE_FACTORS[HEIGHT] + expected_extra_end for s in starts]

tests-unit/comfy_quant/test_mixed_precision.py

@@ -3,6 +3,7 @@ import torch
 import sys
 import os
 import json
+from types import SimpleNamespace
 
 # Add comfy to path
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
@@ -284,6 +285,75 @@ class TestMixedPrecisionOps(unittest.TestCase):
         saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
         self.assertTrue(saved_conf["convrot"])
 
+    def test_int8_disabled_on_unsupported_device_falls_back_to_full_precision(self):
+        """On a device that can't run comfy_kitchen's fast int8 matmul (e.g. MPS,
+        which lacks aten::_int_mm), pick_operations must mark int8 formats as
+        disabled so layers dequantize instead of taking the fast quantized path."""
+        import comfy.model_management as mm
+
+        orig_supports_int8 = mm.supports_int8_compute
+        mm.supports_int8_compute = lambda device=None: False
+        try:
+            model_config = SimpleNamespace(quant_config={"layer": {"format": "int8_tensorwise"}})
+            operations = ops.pick_operations(torch.bfloat16, torch.bfloat16, model_config=model_config)
+
+            torch.manual_seed(789)
+            weight = torch.randn(16, 256, dtype=torch.bfloat16)
+            bias = torch.randn(16, dtype=torch.bfloat16)
+            q_weight = QuantizedTensor.from_float(weight, "TensorWiseINT8Layout", per_channel=True)
+            state_dict = {
+                "layer.weight": q_weight._qdata,
+                "layer.bias": bias,
+                "layer.weight_scale": q_weight._params.scale,
+            }
+            layer_quant_config = {"layer": {"format": "int8_tensorwise"}}
+            state_dict, _ = comfy.utils.convert_old_quants(
+                state_dict,
+                metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})},
+            )
+
+            model = torch.nn.Module()
+            model.layer = operations.Linear(256, 16, device="cpu", dtype=torch.bfloat16)
+            model.load_state_dict(state_dict, strict=False)
+
+            self.assertIsInstance(model.layer.weight, QuantizedTensor)
+            # The layer must be forced onto the full-precision (dequantized)
+            # path since the fast int8 path isn't usable on this device.
+            self.assertTrue(model.layer._full_precision_mm)
+
+            # The weight's orig_dtype matches the compute dtype here (both bfloat16),
+            # so cast_bias_weight's dtype-change check alone won't dequantize it. Confirm
+            # the module still hands a real Tensor (not a QuantizedTensor) to the plain
+            # linear() call, since dispatching a QuantizedTensor there would route back
+            # into the disabled fast int8 matmul instead of the full-precision fallback.
+            seen_weight_types = []
+            orig_module_forward = model.layer._forward
+            def _capturing_forward(input, weight, bias, _orig=orig_module_forward):
+                seen_weight_types.append(type(weight))
+                return _orig(input, weight, bias)
+            model.layer._forward = _capturing_forward
+
+            input_tensor = torch.randn(4, 256, dtype=torch.bfloat16)
+            output = model.layer(input_tensor)
+            self.assertEqual(output.shape, (4, 16))
+            self.assertEqual(seen_weight_types, [torch.Tensor])
+        finally:
+            mm.supports_int8_compute = orig_supports_int8
+
+    def test_supports_int8_compute_treats_mps_mode_as_unsupported_when_device_is_none(self):
+        """Call sites (like pick_operations' default) may omit load_device. On an
+        MPS machine that must still report int8 as unsupported instead of
+        silently defaulting to True, matching supports_fp64's handling of the
+        same device=None case (see Comfy-Org/ComfyUI#16136)."""
+        import comfy.model_management as mm
+
+        orig_cpu_state = mm.cpu_state
+        mm.cpu_state = mm.CPUState.MPS
+        try:
+            self.assertFalse(mm.supports_int8_compute(None))
+        finally:
+            mm.cpu_state = orig_cpu_state
+
     def test_convrot_w4a4_loads_into_params(self):
         """ConvRot W4A4 checkpoints must load as the dedicated kitchen layout."""
         if "convrot_w4a4" not in QUANT_ALGOS:

tests-unit/comfy_test/test_minimax_h3_model.py

@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+
+from comfy.ldm.minimax.model import MiniMaxH3Model, time_shift_sigma
+from comfy.model_sampling import CONST
+
+
+def make_model(video_output, audio_output):
+    model = MiniMaxH3Model.__new__(MiniMaxH3Model)
+    nn.Module.__init__(model)
+    model.sigma_shift_video = 12.0
+    model.sigma_shift_audio = 3.0
+    model._forward = lambda *args, **kwargs: [video_output.clone(), audio_output.clone()]
+    return model
+
+
+def test_forward_scales_velocity_to_mask_timestep():
+    video_output = torch.full((1, 2, 1, 2, 2), 2.0)
+    audio_output = torch.full((1, 2, 2, 3), 3.0)
+    video_mask = torch.tensor([[[[[1.0, 0.75], [0.5, 0.25]]]]])
+    audio_mask = torch.tensor([[[[1.0, 0.5, 0.25], [0.75, 0.5, 0.0]]]])
+    sigma = torch.tensor([0.5])
+    clean = torch.arange(video_output.numel(), dtype=torch.float32).reshape_as(video_output)
+    model_input = clean + sigma.reshape(1, 1, 1, 1, 1) * video_mask * video_output
+    model = make_model(video_output, audio_output)
+
+    out = model(
+        [model_input, torch.zeros_like(audio_output)],
+        sigma * 1000.0,
+        torch.empty(1, 1, 1),
+        minimax_payload={"audio_scale": 1.0},
+        denoise_mask=video_mask,
+        audio_denoise_mask=audio_mask,
+    )
+
+    torch.testing.assert_close(out[0], video_output * video_mask)
+    torch.testing.assert_close(out[1], audio_output * audio_mask)
+    denoised = CONST.calculate_denoised(None, sigma, out[0], model_input)
+    torch.testing.assert_close(denoised, clean)
+
+
+def test_forward_scales_audio_velocity_before_carry_conversion():
+    video_output = torch.ones((1, 1, 1, 1, 1))
+    audio_output = torch.full((1, 1, 2, 2), 3.0)
+    audio_src = torch.full_like(audio_output, 2.0)
+    audio_mask = torch.tensor([[[[0.75, 0.5], [0.25, 0.0]]]])
+    model = make_model(video_output, audio_output)
+    sigma_v = torch.tensor(0.5)
+    sigma_a = time_shift_sigma(sigma_v, 12.0, 3.0)
+    carry = sigma_a / sigma_v
+
+    out = model(
+        [torch.zeros_like(video_output), audio_src],
+        sigma_v.reshape(1) * 1000.0,
+        torch.empty(1, 1, 1),
+        minimax_payload={"audio_scale": 4.0},
+        audio_denoise_mask=audio_mask,
+    )
+
+    expected = -3.0 * audio_src * carry + (1.0 + 3.0 * sigma_a) * audio_output * audio_mask
+    torch.testing.assert_close(out[1], expected)

tests-unit/execution_test/preview_compiler_test.py

@@ -0,0 +1,57 @@
+from unittest.mock import MagicMock, Mock
+
+import pytest
+
+import latent_preview
+from comfy import latent_formats
+
+
+def test_minimax_h3_enables_preview_compiler():
+    assert latent_formats.MiniMaxH3Video.compile_preview
+    assert latent_formats.MiniMaxH3AV.compile_preview
+    assert not latent_formats.HunyuanVideo.compile_preview
+
+
+def test_video_preview_compiles_decode(monkeypatch):
+    taesd = Mock()
+    taesd.device = "cuda:0"
+    taesd.decode.return_value = [[Mock()]]
+    previewer = latent_preview.TAEHVPreviewerImpl(taesd, compile_preview=True)
+    x0 = MagicMock()
+    samples = Mock(shape=(1, 24, 1, 30, 52))
+    x0.__getitem__.return_value = samples
+
+    monkeypatch.setattr(latent_preview, "preview_to_image", Mock())
+    monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_enabled", Mock(return_value=True))
+    calls = []
+    taesd.decode.side_effect = lambda value: calls.append(("decode", value)) or [[Mock()]]
+    begin = Mock(side_effect=lambda device: calls.append(("begin", device)))
+    end = Mock(side_effect=lambda: calls.append(("end",)))
+    monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_begin", begin)
+    monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_end", end)
+
+    previewer.decode_latent_to_preview(x0)
+    assert calls == [
+        ("begin", "cuda:0"),
+        ("decode", samples),
+        ("end",),
+    ]
+
+
+def test_video_preview_leaves_failed_compiler_scope_for_execution_cleanup(monkeypatch):
+    taesd = Mock()
+    taesd.device = "cuda:0"
+    taesd.decode.side_effect = RuntimeError("decode failed")
+    previewer = latent_preview.TAEHVPreviewerImpl(taesd, compile_preview=True)
+    x0 = MagicMock()
+
+    monkeypatch.setattr(latent_preview, "preview_to_image", Mock())
+    monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_enabled", Mock(return_value=True))
+    monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_begin", Mock())
+    end = Mock()
+    monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_end", end)
+
+    with pytest.raises(RuntimeError, match="decode failed"):
+        previewer.decode_latent_to_preview(x0)
+
+    end.assert_not_called()

tests-unit/main_prestartup_test.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import ast
+import importlib
+import logging
+import os
+from pathlib import Path
+from types import SimpleNamespace
+import time
+
+import folder_paths
+
+
+def _load_execute_prestartup_script():
+    main_path = Path(__file__).resolve().parents[1] / "main.py"
+    module = ast.parse(main_path.read_text(), filename=str(main_path))
+    function = next(node for node in module.body if isinstance(node, ast.FunctionDef) and node.name == "execute_prestartup_script")
+    compiled = compile(ast.Module(body=[function], type_ignores=[]), filename=str(main_path), mode="exec")
+    namespace = {
+        "args": SimpleNamespace(disable_all_custom_nodes=False, whitelist_custom_nodes=[], enable_manager=False),
+        "folder_paths": folder_paths,
+        "importlib": importlib,
+        "logging": logging,
+        "os": os,
+        "time": time,
+    }
+    exec(compiled, namespace)  # noqa: S102 - trusted AST extracted from main.py itself, not external input
+    return namespace["execute_prestartup_script"]
+
+
+def _load_prestartup_script_for_paths(monkeypatch, custom_nodes_paths: list[str]):
+    monkeypatch.setattr(
+        folder_paths,
+        "get_folder_paths",
+        lambda name: list(custom_nodes_paths) if name == "custom_nodes" else [],
+    )
+    return _load_execute_prestartup_script()
+
+
+def _make_pack(root: Path, name: str) -> Path:
+    pack = root / name
+    pack.mkdir(parents=True)
+    (pack / "prestartup_script.py").write_text("VALUE = 1\n")
+    return pack
+
+
+def test_execute_prestartup_script_handles_empty_custom_nodes_paths(monkeypatch):
+    execute_prestartup_script = _load_prestartup_script_for_paths(monkeypatch, [])
+
+    execute_prestartup_script()
+
+
+def test_execute_prestartup_script_keeps_all_timing_entries(monkeypatch, tmp_path):
+    first_custom_nodes = tmp_path / "custom_nodes_1"
+    second_custom_nodes = tmp_path / "custom_nodes_2"
+    pack_one = _make_pack(first_custom_nodes, "pack_one")
+    pack_two = _make_pack(second_custom_nodes, "pack_two")
+
+    execute_prestartup_script = _load_prestartup_script_for_paths(monkeypatch, [str(first_custom_nodes), str(second_custom_nodes)])
+
+    messages: list[str] = []
+    monkeypatch.setattr(logging, "info", lambda message, *args, **kwargs: messages.append(message))
+
+    execute_prestartup_script()
+
+    joined = "\n".join(messages)
+    assert str(pack_one) in joined
+    assert str(pack_two) in joined

tests-unit/nodes_test/test_ignore_display_name.py

@@ -0,0 +1,62 @@
+import sys
+
+import pytest
+import torch
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+import nodes
+
+
+pytestmark = pytest.mark.asyncio
+
+
+@pytest.fixture(autouse=True)
+def _restore_node_mappings():
+    class_mappings = dict(nodes.NODE_CLASS_MAPPINGS)
+    display_name_mappings = dict(nodes.NODE_DISPLAY_NAME_MAPPINGS)
+    try:
+        yield
+    finally:
+        nodes.NODE_CLASS_MAPPINGS.clear()
+        nodes.NODE_CLASS_MAPPINGS.update(class_mappings)
+        nodes.NODE_DISPLAY_NAME_MAPPINGS.clear()
+        nodes.NODE_DISPLAY_NAME_MAPPINGS.update(display_name_mappings)
+        sys.modules.pop("test_v1_custom_node", None)
+        sys.modules.pop("test_v3_custom_node", None)
+
+
+async def test_load_custom_node_skips_display_names_for_ignored_nodes(tmp_path, monkeypatch):
+    v1_module = tmp_path / "test_v1_custom_node.py"
+    v1_module.write_text(
+        "NODE_CLASS_MAPPINGS = {\"LeakTest\": object}\n"
+        "NODE_DISPLAY_NAME_MAPPINGS = {\"LeakTest\": \"Leak Test\"}\n",
+    )
+
+    v3_module = tmp_path / "test_v3_custom_node.py"
+    v3_module.write_text(
+        "from comfy_api.latest import ComfyExtension\n\n"
+        "class LeakTestV3Node:\n"
+        "    @classmethod\n"
+        "    def GET_SCHEMA(cls):\n"
+        "        class Schema:\n"
+        "            node_id = \"LeakTestV3\"\n"
+        "            display_name = \"Leak Test V3\"\n\n"
+        "        return Schema()\n\n\n"
+        "class TestExtension(ComfyExtension):\n"
+        "    async def get_node_list(self):\n"
+        "        return [LeakTestV3Node]\n\n\n"
+        "async def comfy_entrypoint():\n"
+        "    return TestExtension()\n",
+    )
+
+    monkeypatch.syspath_prepend(str(tmp_path))
+
+    assert await nodes.load_custom_node(str(v1_module), ignore={"LeakTest"})
+    assert await nodes.load_custom_node(str(v3_module), ignore={"LeakTestV3"})
+
+    assert "LeakTest" not in nodes.NODE_DISPLAY_NAME_MAPPINGS
+    assert "LeakTestV3" not in nodes.NODE_DISPLAY_NAME_MAPPINGS