Merge remote-tracking branch 'origin/master' into fix/minimax-h3-attention-patches

# Conflicts:
#	comfy/ldm/minimax/model.py

app/assets/api/routes.py

@@ -18,7 +18,7 @@ from app.assets.api.schemas_in import (
     AssetValidationError,
     UploadError,
 )
-from app.assets.helpers import validate_blake3_hash
+from app.assets.helpers import normalize_tags, validate_blake3_hash
 from app.assets.api.upload import (
     delete_temp_file_if_exists,
     parse_multipart_upload,
@@ -117,6 +117,87 @@ def _build_validation_error_response(code: str, ve: ValidationError) -> web.Resp
     return _build_error_response(400, code, "Validation failed.", {"errors": errors})
 
 
+class InvalidTagFilterError(Exception):
+    """Invalid combination of tag-filter query parameters."""
+
+    def __init__(self, message: str, details: dict):
+        super().__init__(message)
+        self.details = details
+
+
+# Caps the per-tag EXISTS fan-out; deliberately covers the legacy spellings too.
+MAX_TAG_FILTER_TAGS = 100
+
+
+def _resolve_tag_filters(
+    q: schemas_in.ListAssetsQuery | schemas_in.TagsRefineQuery,
+) -> tuple[list[str], list[str], list[str]]:
+    """Resolve legacy (include/exclude) and new (all/any/none) tag-filter
+    spellings into effective (all, any, none) lists.
+
+    Combination validation applies only when the request uses at least one
+    new-name parameter (non-empty after normalisation); requests using only
+    the legacy names keep their historical behaviour, including degenerate
+    combinations like include_tags=a&exclude_tags=a.
+    """
+    # model_dump, not attribute access: deprecated fields warn on every attribute read.
+    legacy = q.model_dump(include={"include_tags", "exclude_tags"})
+    include_tags = normalize_tags(legacy["include_tags"])
+    exclude_tags = normalize_tags(legacy["exclude_tags"])
+    tags_all = normalize_tags(q.tags_all)
+    tags_any = normalize_tags(q.tags_any)
+    tags_none = normalize_tags(q.tags_none)
+
+    for param_name, values in (
+        ("include_tags", include_tags),
+        ("exclude_tags", exclude_tags),
+        ("tags_all", tags_all),
+        ("tags_any", tags_any),
+        ("tags_none", tags_none),
+    ):
+        if len(values) > MAX_TAG_FILTER_TAGS:
+            raise InvalidTagFilterError(
+                f"'{param_name}' lists {len(values)} tags; the maximum is "
+                f"{MAX_TAG_FILTER_TAGS}.",
+                {
+                    "parameter": param_name,
+                    "count": len(values),
+                    "max": MAX_TAG_FILTER_TAGS,
+                },
+            )
+
+    if not (tags_all or tags_any or tags_none):
+        return include_tags, [], exclude_tags
+
+    if include_tags and tags_all:
+        raise InvalidTagFilterError(
+            "Cannot combine 'include_tags' with 'tags_all'; use 'tags_all'.",
+            {"parameters": ["include_tags", "tags_all"]},
+        )
+    if exclude_tags and tags_none:
+        raise InvalidTagFilterError(
+            "Cannot combine 'exclude_tags' with 'tags_none'; use 'tags_none'.",
+            {"parameters": ["exclude_tags", "tags_none"]},
+        )
+
+    all_param, all_list = (
+        ("tags_all", tags_all) if tags_all else ("include_tags", include_tags)
+    )
+    none_param, none_list = (
+        ("tags_none", tags_none) if tags_none else ("exclude_tags", exclude_tags)
+    )
+
+    conflicting = sorted(set(all_list) & set(none_list))
+    if conflicting:
+        raise InvalidTagFilterError(
+            f"Query can never match: {', '.join(repr(t) for t in conflicting)} "
+            f"required by '{all_param}' but rejected by '{none_param}'.",
+            {"conflicting_tags": conflicting, "parameters": [all_param, none_param]},
+        )
+
+    return all_list, tags_any, none_list
+
+
 def _validate_sort_field(requested: str | None) -> str:
     if not requested:
         return "created_at"
@@ -217,6 +298,11 @@ async def list_assets_route(request: web.Request) -> web.Response:
     except ValidationError as ve:
         return _build_validation_error_response("INVALID_QUERY", ve)
 
+    try:
+        tags_all, tags_any, tags_none = _resolve_tag_filters(q)
+    except InvalidTagFilterError as e:
+        return _build_error_response(400, "INVALID_TAG_FILTER", str(e), e.details)
+
     sort = _validate_sort_field(q.sort)
     order_candidate = (q.order or "desc").lower()
     order = order_candidate if order_candidate in {"asc", "desc"} else "desc"
@@ -224,8 +310,9 @@ async def list_assets_route(request: web.Request) -> web.Response:
     try:
         result = list_assets_page(
             owner_id=USER_MANAGER.get_request_user_id(request),
-            include_tags=q.include_tags,
-            exclude_tags=q.exclude_tags,
+            include_tags=tags_all,
+            exclude_tags=tags_none,
+            any_tags=tags_any,
             name_contains=q.name_contains,
             metadata_filter=q.metadata_filter,
             limit=q.limit,
@@ -715,10 +802,16 @@ async def get_tags_refine(request: web.Request) -> web.Response:
     except ValidationError as ve:
         return _build_validation_error_response("INVALID_QUERY", ve)
 
+    try:
+        tags_all, tags_any, tags_none = _resolve_tag_filters(q)
+    except InvalidTagFilterError as e:
+        return _build_error_response(400, "INVALID_TAG_FILTER", str(e), e.details)
+
     tag_counts = list_tag_histogram(
         owner_id=USER_MANAGER.get_request_user_id(request),
-        include_tags=q.include_tags,
-        exclude_tags=q.exclude_tags,
+        include_tags=tags_all,
+        exclude_tags=tags_none,
+        any_tags=tags_any,
         name_contains=q.name_contains,
         metadata_filter=q.metadata_filter,
         limit=q.limit,

app/assets/api/schemas_in.py

@@ -50,8 +50,12 @@ class ParsedUpload:
 
 
 class ListAssetsQuery(BaseModel):
-    include_tags: list[str] = Field(default_factory=list)
-    exclude_tags: list[str] = Field(default_factory=list)
+    # Deprecated spellings: include_tags ≡ tags_all, exclude_tags ≡ tags_none.
+    include_tags: list[str] = Field(default_factory=list, deprecated=True)
+    exclude_tags: list[str] = Field(default_factory=list, deprecated=True)
+    tags_all: list[str] = Field(default_factory=list)
+    tags_any: list[str] = Field(default_factory=list)
+    tags_none: list[str] = Field(default_factory=list)
     name_contains: str | None = None
 
     # Accept either a JSON string (query param) or a dict
@@ -70,7 +74,10 @@ class ListAssetsQuery(BaseModel):
     )
     order: Literal["asc", "desc"] = "desc"
 
-    @field_validator("include_tags", "exclude_tags", mode="before")
+    @field_validator(
+        "include_tags", "exclude_tags", "tags_all", "tags_any", "tags_none",
+        mode="before",
+    )
     @classmethod
     def _split_csv_tags(cls, v):
         # Accept "a,b,c" or ["a","b"] (we are liberal in what we accept)
@@ -154,13 +161,20 @@ class CreateFromHashBody(BaseModel):
 
 
 class TagsRefineQuery(BaseModel):
-    include_tags: list[str] = Field(default_factory=list)
-    exclude_tags: list[str] = Field(default_factory=list)
+    # Deprecated spellings: include_tags ≡ tags_all, exclude_tags ≡ tags_none.
+    include_tags: list[str] = Field(default_factory=list, deprecated=True)
+    exclude_tags: list[str] = Field(default_factory=list, deprecated=True)
+    tags_all: list[str] = Field(default_factory=list)
+    tags_any: list[str] = Field(default_factory=list)
+    tags_none: list[str] = Field(default_factory=list)
     name_contains: str | None = None
     metadata_filter: dict[str, Any] | None = None
     limit: conint(ge=1, le=1000) = 100
 
-    @field_validator("include_tags", "exclude_tags", mode="before")
+    @field_validator(
+        "include_tags", "exclude_tags", "tags_all", "tags_any", "tags_none",
+        mode="before",
+    )
     @classmethod
     def _split_csv_tags(cls, v):
         if v is None:

app/assets/database/queries/asset_reference.py

@@ -268,6 +268,8 @@ def list_references_page(
     order: str | None = None,
     after_cursor_value: object | None = None,
     after_cursor_id: str | None = None,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> tuple[list[AssetReference], dict[str, list[str]], int]:
     """List references with pagination, filtering, and sorting.
 
@@ -293,7 +295,7 @@ def list_references_page(
         escaped, esc = escape_sql_like_string(name_contains)
         base = base.where(AssetReference.name.ilike(f"%{escaped}%", escape=esc))
 
-    base = apply_tag_filters(base, include_tags, exclude_tags)
+    base = apply_tag_filters(base, include_tags, exclude_tags, any_tags)
     base = apply_metadata_filter(base, metadata_filter)
 
     sort = (sort or "created_at").lower()
@@ -345,7 +347,7 @@ def list_references_page(
         count_stmt = count_stmt.where(
             AssetReference.name.ilike(f"%{escaped}%", escape=esc)
         )
-    count_stmt = apply_tag_filters(count_stmt, include_tags, exclude_tags)
+    count_stmt = apply_tag_filters(count_stmt, include_tags, exclude_tags, any_tags)
     count_stmt = apply_metadata_filter(count_stmt, metadata_filter)
 
     total = int(session.execute(count_stmt).scalar_one() or 0)

app/assets/database/queries/common.py

@@ -60,10 +60,13 @@ def apply_tag_filters(
     stmt: sa.sql.Select,
     include_tags: Sequence[str] | None = None,
     exclude_tags: Sequence[str] | None = None,
+    any_tags: Sequence[str] | None = None,
 ) -> sa.sql.Select:
-    """include_tags: every tag must be present; exclude_tags: none may be present."""
+    """include_tags: every tag must be present; any_tags: at least one must be
+    present; exclude_tags: none may be present."""
     include_tags = normalize_tags(include_tags)
     exclude_tags = normalize_tags(exclude_tags)
+    any_tags = normalize_tags(any_tags)
 
     if include_tags:
         for tag_name in include_tags:
@@ -74,6 +77,14 @@ def apply_tag_filters(
                 )
             )
 
+    if any_tags:
+        stmt = stmt.where(
+            exists().where(
+                (AssetReferenceTag.asset_reference_id == AssetReference.id)
+                & (AssetReferenceTag.tag_name.in_(any_tags))
+            )
+        )
+
     if exclude_tags:
         stmt = stmt.where(
             ~exists().where(

app/assets/database/queries/tags.py

@@ -340,6 +340,8 @@ def list_tag_counts_for_filtered_assets(
     name_contains: str | None = None,
     metadata_filter: dict | None = None,
     limit: int = 100,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> dict[str, int]:
     """Return tag counts for assets matching the given filters.
 
@@ -359,7 +361,7 @@ def list_tag_counts_for_filtered_assets(
         escaped, esc = escape_sql_like_string(name_contains)
         ref_sq = ref_sq.where(AssetReference.name.ilike(f"%{escaped}%", escape=esc))
 
-    ref_sq = apply_tag_filters(ref_sq, include_tags, exclude_tags)
+    ref_sq = apply_tag_filters(ref_sq, include_tags, exclude_tags, any_tags)
     ref_sq = apply_metadata_filter(ref_sq, metadata_filter)
     ref_sq = ref_sq.subquery()
 

app/assets/services/asset_management.py

@@ -279,6 +279,8 @@ def list_assets_page(
     sort: str = "created_at",
     order: str = "desc",
     after: str | None = None,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> ListAssetsResult:
     """List assets with optional cursor pagination.
 
@@ -317,6 +319,7 @@ def list_assets_page(
             owner_id=owner_id,
             include_tags=include_tags,
             exclude_tags=exclude_tags,
+            any_tags=any_tags,
             name_contains=name_contains,
             metadata_filter=metadata_filter,
             limit=fetch_limit,

app/assets/services/tagging.py

@@ -85,6 +85,8 @@ def list_tag_histogram(
     name_contains: str | None = None,
     metadata_filter: dict | None = None,
     limit: int = 100,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> dict[str, int]:
     with create_session() as session:
         return list_tag_counts_for_filtered_assets(
@@ -92,6 +94,7 @@ def list_tag_histogram(
             owner_id=owner_id,
             include_tags=include_tags,
             exclude_tags=exclude_tags,
+            any_tags=any_tags,
             name_contains=name_contains,
             metadata_filter=metadata_filter,
             limit=limit,

comfy/cli_args.py

@@ -149,6 +149,7 @@ attn_group.add_argument("--use-quad-cross-attention", action="store_true", help=
 attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
 attn_group.add_argument("--use-sage-attention", action="store_true", help="Use sage attention.")
 attn_group.add_argument("--use-flash-attention", action="store_true", help="Use FlashAttention.")
+attn_group.add_argument("--use-ck-attention", action="store_true", help="Use Comfy Kitchen attention.")
 
 parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
 

comfy/clip_model.py

@@ -314,13 +314,18 @@ class CLIPVisionModelProjection(torch.nn.Module):
         if "projection_dim" in config_dict:
             self.visual_projection = operations.Linear(config_dict["hidden_size"], config_dict["projection_dim"], bias=False)
         else:
-            self.visual_projection = lambda a: a
+            self.visual_projection = torch.nn.Identity()
 
         if "llava3" == config_dict.get("projector_type", None):
             self.multi_modal_projector = LlavaProjector(config_dict["hidden_size"], 4096, dtype, device, operations)
         else:
             self.multi_modal_projector = None
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        if "{}visual_projection.weight".format(prefix) not in state_dict:
+            self.visual_projection = torch.nn.Identity()
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     def forward(self, *args, **kwargs):
         x = self.vision_model(*args, **kwargs)
         out = self.visual_projection(x[2])

comfy/ldm/lightricks/av_model.py

@@ -96,6 +96,8 @@ class BasicAVTransformerBlock(nn.Module):
         attn_precision=None,
         apply_gated_attention=False,
         cross_attention_adaln=False,
+        ff_bias=True,
+        audio_ff_bias=True,
         dtype=None,
         device=None,
         operations=None,
@@ -178,10 +180,10 @@ class BasicAVTransformerBlock(nn.Module):
         )
 
         self.ff = FeedForward(
-            v_dim, dim_out=v_dim, glu=True, dtype=dtype, device=device, operations=operations
+            v_dim, dim_out=v_dim, glu=True, ff_bias=ff_bias, dtype=dtype, device=device, operations=operations
         )
         self.audio_ff = FeedForward(
-            a_dim, dim_out=a_dim, glu=True, dtype=dtype, device=device, operations=operations
+            a_dim, dim_out=a_dim, glu=True, ff_bias=audio_ff_bias, dtype=dtype, device=device, operations=operations
         )
 
         num_ada_params = ADALN_CROSS_ATTN_PARAMS_COUNT if cross_attention_adaln else ADALN_BASE_PARAMS_COUNT
@@ -413,12 +415,16 @@ class LTXAVModel(LTXVModel):
         apply_gated_attention=False,
         caption_proj_before_connector=False,
         cross_attention_adaln=False,
+        ff_bias=True,
+        audio_ff_bias=True,
+        use_prompt_adaln_single=True,
         dtype=None,
         device=None,
         operations=None,
         **kwargs,
     ):
         # Store audio-specific parameters
+        self.audio_ff_bias = audio_ff_bias
         self.audio_in_channels = audio_in_channels
         self.audio_cross_attention_dim = audio_cross_attention_dim
         self.audio_attention_head_dim = audio_attention_head_dim
@@ -451,6 +457,8 @@ class LTXAVModel(LTXVModel):
             timestep_scale_multiplier=timestep_scale_multiplier,
             caption_proj_before_connector=caption_proj_before_connector,
             cross_attention_adaln=cross_attention_adaln,
+            ff_bias=ff_bias,
+            use_prompt_adaln_single=use_prompt_adaln_single,
             dtype=dtype,
             device=device,
             operations=operations,
@@ -475,7 +483,7 @@ class LTXAVModel(LTXVModel):
             operations=self.operations,
         )
 
-        if self.cross_attention_adaln:
+        if self.cross_attention_adaln and self.use_prompt_adaln_single:
             self.audio_prompt_adaln_single = AdaLayerNormSingle(
                 self.audio_inner_dim,
                 embedding_coefficient=2,
@@ -606,6 +614,8 @@ class LTXAVModel(LTXVModel):
                     a_context_dim=self.audio_cross_attention_dim,
                     apply_gated_attention=self.apply_gated_attention,
                     cross_attention_adaln=self.cross_attention_adaln,
+                    ff_bias=self.ff_bias,
+                    audio_ff_bias=self.audio_ff_bias,
                     dtype=dtype,
                     device=device,
                     operations=self.operations,
@@ -924,9 +934,15 @@ class LTXAVModel(LTXVModel):
         blocks_replace = patches_replace.get("dit", {})
         prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.transformer_blocks), vx.device, transformer_options)
 
+        # Blocks whose self-attention should be perturbed to a value-passthrough (STG).
+        stg_self_attn_blocks = transformer_options.get("stg_self_attn_blocks", ())
+
         # Process transformer blocks
         for i, block in enumerate(self.transformer_blocks):
             comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, vx.device, block)
+            block_transformer_options = transformer_options
+            if i in stg_self_attn_blocks:
+                block_transformer_options = {**transformer_options, "stg_skip_self_attn": True}
             if ("double_block", i) in blocks_replace:
 
                 def block_wrap(args):
@@ -969,7 +985,7 @@ class LTXAVModel(LTXVModel):
                         "a_cross_scale_shift_timestep": av_ca_audio_scale_shift_timestep,
                         "v_cross_gate_timestep": av_ca_a2v_gate_noise_timestep,
                         "a_cross_gate_timestep": av_ca_v2a_gate_noise_timestep,
-                        "transformer_options": transformer_options,
+                        "transformer_options": block_transformer_options,
                         "self_attention_mask": self_attention_mask,
                         "v_prompt_timestep": v_prompt_timestep,
                         "a_prompt_timestep": a_prompt_timestep,
@@ -993,7 +1009,7 @@ class LTXAVModel(LTXVModel):
                     a_cross_scale_shift_timestep=av_ca_audio_scale_shift_timestep,
                     v_cross_gate_timestep=av_ca_a2v_gate_noise_timestep,
                     a_cross_gate_timestep=av_ca_v2a_gate_noise_timestep,
-                    transformer_options=transformer_options,
+                    transformer_options=block_transformer_options,
                     self_attention_mask=self_attention_mask,
                     v_prompt_timestep=v_prompt_timestep,
                     a_prompt_timestep=a_prompt_timestep,

comfy/ldm/lightricks/duration_head.py

@@ -0,0 +1,81 @@
+"""LTX 2.4 DurationHead: predicts the natural shot duration (in seconds) from
+the caption connector token outputs, without running the diffusion pipeline.
+"""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class AttentionPooler(nn.Module):
+    """Cross-attend ``num_queries`` learnable tokens against ``tokens``."""
+
+    def __init__(self, hidden_dim=256, num_queries=1, num_heads=4):
+        super().__init__()
+        self.num_queries = num_queries
+        self.query_tokens = nn.Parameter(torch.empty(num_queries, hidden_dim))
+        self.cross_attn = nn.MultiheadAttention(embed_dim=hidden_dim, num_heads=num_heads, batch_first=True)
+
+    def forward(self, tokens):
+        queries = self.query_tokens.unsqueeze(0).expand(tokens.shape[0], -1, -1)
+        pooled, _ = self.cross_attn(queries, tokens, tokens, need_weights=False)
+        return pooled
+
+
+class DurationHead(nn.Module):
+    """Predict duration in seconds from one or both connector outputs."""
+
+    def __init__(
+        self,
+        video_cross_attention_dim=4096,
+        audio_cross_attention_dim=2048,
+        pooler_hidden_dim=256,
+        num_queries=1,
+        num_pooler_heads=4,
+        mlp_hidden=256,
+    ):
+        super().__init__()
+        self.video_input_proj = nn.Linear(video_cross_attention_dim, pooler_hidden_dim)
+        self.video_modality_emb = nn.Parameter(torch.empty(pooler_hidden_dim))
+        self.audio_input_proj = nn.Linear(audio_cross_attention_dim, pooler_hidden_dim)
+        self.audio_modality_emb = nn.Parameter(torch.empty(pooler_hidden_dim))
+        self.attention_pooler = AttentionPooler(
+            hidden_dim=pooler_hidden_dim, num_queries=num_queries, num_heads=num_pooler_heads)
+        self.mlp_hidden = nn.Linear(pooler_hidden_dim * num_queries, mlp_hidden)
+        self.mlp_out = nn.Linear(mlp_hidden, 1)
+
+    def forward(self, video_tokens=None, audio_tokens=None):
+        """``video_tokens``: (B, T_v, 4096), ``audio_tokens``: (B, T_a, 2048);
+        at least one required. Returns duration in seconds, shape (B,)."""
+        token_groups = []
+        if video_tokens is not None:
+            token_groups.append(self.video_input_proj(video_tokens) + self.video_modality_emb)
+        if audio_tokens is not None:
+            token_groups.append(self.audio_input_proj(audio_tokens) + self.audio_modality_emb)
+        if not token_groups:
+            raise ValueError("DurationHead requires at least one of video_tokens / audio_tokens")
+        pooled = self.attention_pooler(torch.cat(token_groups, dim=1))
+        pooled = pooled.reshape(pooled.shape[0], -1)
+        hidden = F.gelu(self.mlp_hidden(pooled), approximate="tanh")
+        return self.mlp_out(hidden).squeeze(-1).exp()
+
+
+def normalize_state_dict(sd):
+    for prefix in ("model.diffusion_model.duration_head.", "duration_head."):
+        stripped = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)}
+        if stripped:
+            return stripped
+    return sd
+
+
+def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds, time_scale=8):
+    """Convert seconds to a frame count clamped to ``[min_seconds, max_seconds]``
+    and snapped (floor) to the VAE's ``8k + 1`` causal temporal grid; snapping
+    that undershoots the minimum bumps up to the next grid point instead."""
+    min_frames = max(1, round(min_seconds * frame_rate))
+    max_frames = round(max_seconds * frame_rate)
+    raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames))
+    frames = (raw_frames - 1) // time_scale * time_scale + 1
+    if frames < min_frames:
+        frames = min(-(-(min_frames - 1) // time_scale) * time_scale + 1, max_frames)
+    return frames

comfy/ldm/lightricks/embeddings_connector.py

@@ -50,6 +50,7 @@ class BasicTransformerBlock1D(nn.Module):
         context_dim=None,
         attn_precision=None,
         apply_gated_attention=False,
+        ff_bias=True,
         dtype=None,
         device=None,
         operations=None,
@@ -74,6 +75,7 @@ class BasicTransformerBlock1D(nn.Module):
             dim,
             dim_out=dim,
             glu=True,
+            ff_bias=ff_bias,
             dtype=dtype,
             device=device,
             operations=operations,
@@ -123,6 +125,7 @@ class Embeddings1DConnector(nn.Module):
         causal_temporal_positioning=False,
         num_learnable_registers: Optional[int] = 128,
         apply_gated_attention=False,
+        connector_ff_bias=True,
         dtype=None,
         device=None,
         operations=None,
@@ -148,6 +151,7 @@ class Embeddings1DConnector(nn.Module):
                     attention_head_dim,
                     context_dim=cross_attention_dim,
                     apply_gated_attention=apply_gated_attention,
+                    ff_bias=connector_ff_bias,
                     dtype=dtype,
                     device=device,
                     operations=operations,

comfy/ldm/lightricks/model.py

@@ -303,22 +303,22 @@ class NormSingleLinearTextProjection(nn.Module):
 
 
 class GELU_approx(nn.Module):
-    def __init__(self, dim_in, dim_out, dtype=None, device=None, operations=None):
+    def __init__(self, dim_in, dim_out, bias=True, dtype=None, device=None, operations=None):
         super().__init__()
-        self.proj = operations.Linear(dim_in, dim_out, dtype=dtype, device=device)
+        self.proj = operations.Linear(dim_in, dim_out, bias=bias, dtype=dtype, device=device)
 
     def forward(self, x):
         return torch.nn.functional.gelu(self.proj(x), approximate="tanh")
 
 
 class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out, mult=4, glu=False, dropout=0.0, dtype=None, device=None, operations=None):
+    def __init__(self, dim, dim_out, mult=4, glu=False, dropout=0.0, ff_bias=True, dtype=None, device=None, operations=None):
         super().__init__()
         inner_dim = int(dim * mult)
-        project_in = GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations)
+        project_in = GELU_approx(dim, inner_dim, bias=ff_bias, dtype=dtype, device=device, operations=operations)
 
         self.net = nn.Sequential(
-            project_in, nn.Dropout(dropout), operations.Linear(inner_dim, dim_out, dtype=dtype, device=device)
+            project_in, nn.Dropout(dropout), operations.Linear(inner_dim, dim_out, bias=ff_bias, dtype=dtype, device=device)
         )
 
     def forward(self, x):
@@ -462,28 +462,34 @@ class CrossAttention(nn.Module):
         )
 
     def forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
+        self_attn = context is None
         q = self.to_q(x)
         context = x if context is None else context
         k = self.to_k(context)
         v = self.to_v(context)
 
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        # These norms span all heads, so the per-head RMS+RoPE kernel is not equivalent.
-        if pe is not None:
-            if k_pe is None and q.shape == k.shape:
-                q, k = apply_rotary_emb_qk(q, k, pe)
-            else:
-                q = apply_rotary_emb(q, pe)
-                k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
-
-        if mask is None:
-            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
-        elif isinstance(mask, GuideAttentionMask):
-            out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+        # Spatio-Temporal Guidance (STG) perturbation: for the flagged self-attention
+        # layers, the attention degrades to a passthrough of the value projection (out = V).
+        if self_attn and transformer_options.get("stg_skip_self_attn", False):
+            out = v
         else:
-            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+            # These norms span all heads, so the per-head RMS+RoPE kernel is not equivalent.
+            if pe is not None:
+                if k_pe is None and q.shape == k.shape:
+                    q, k = apply_rotary_emb_qk(q, k, pe)
+                else:
+                    q = apply_rotary_emb(q, pe)
+                    k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
+
+            if mask is None:
+                out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            elif isinstance(mask, GuideAttentionMask):
+                out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            else:
+                out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
 
         # Apply per-head gating if enabled
         if self.to_gate_logits is not None:
@@ -502,7 +508,7 @@ ADALN_CROSS_ATTN_PARAMS_COUNT = 9
 
 class BasicTransformerBlock(nn.Module):
     def __init__(
-        self, dim, n_heads, d_head, context_dim=None, attn_precision=None, cross_attention_adaln=False, dtype=None, device=None, operations=None
+        self, dim, n_heads, d_head, context_dim=None, attn_precision=None, cross_attention_adaln=False, ff_bias=True, dtype=None, device=None, operations=None
     ):
         super().__init__()
 
@@ -518,7 +524,7 @@ class BasicTransformerBlock(nn.Module):
             device=device,
             operations=operations,
         )
-        self.ff = FeedForward(dim, dim_out=dim, glu=True, dtype=dtype, device=device, operations=operations)
+        self.ff = FeedForward(dim, dim_out=dim, glu=True, ff_bias=ff_bias, dtype=dtype, device=device, operations=operations)
 
         self.attn2 = CrossAttention(
             query_dim=dim,
@@ -717,6 +723,9 @@ class LTXBaseModel(torch.nn.Module, ABC):
         caption_proj_before_connector=False,
         cross_attention_adaln=False,
         caption_projection_first_linear=True,
+        ff_bias=True,
+        use_prompt_adaln_single=True,
+        use_keyframes_abs_pos_embedding=False,
         dtype=None,
         device=None,
         operations=None,
@@ -746,6 +755,9 @@ class LTXBaseModel(torch.nn.Module, ABC):
         self.caption_proj_before_connector = caption_proj_before_connector
         self.cross_attention_adaln = cross_attention_adaln
         self.caption_projection_first_linear = caption_projection_first_linear
+        self.ff_bias = ff_bias
+        self.use_prompt_adaln_single = use_prompt_adaln_single
+        self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding
 
         # Common dimensions
         self.inner_dim = num_attention_heads * attention_head_dim
@@ -773,12 +785,17 @@ class LTXBaseModel(torch.nn.Module, ABC):
             self.in_channels, self.inner_dim, bias=True, dtype=dtype, device=device
         )
 
+        if self.use_keyframes_abs_pos_embedding:
+            self.keyframes_abs_pos_embedding = nn.Parameter(torch.zeros(1, self.inner_dim, dtype=dtype, device=device))
+        else:
+            self.keyframes_abs_pos_embedding = None
+
         embedding_coefficient = ADALN_CROSS_ATTN_PARAMS_COUNT if self.cross_attention_adaln else ADALN_BASE_PARAMS_COUNT
         self.adaln_single = AdaLayerNormSingle(
             self.inner_dim, embedding_coefficient=embedding_coefficient, use_additional_conditions=False, dtype=dtype, device=device, operations=self.operations
         )
 
-        if self.cross_attention_adaln:
+        if self.cross_attention_adaln and self.use_prompt_adaln_single:
             self.prompt_adaln_single = AdaLayerNormSingle(
                 self.inner_dim, embedding_coefficient=2, use_additional_conditions=False, dtype=dtype, device=device, operations=self.operations
             )
@@ -1070,6 +1087,7 @@ class LTXVModel(LTXBaseModel):
                     self.attention_head_dim,
                     context_dim=self.cross_attention_dim,
                     cross_attention_adaln=self.cross_attention_adaln,
+                    ff_bias=self.ff_bias,
                     dtype=dtype,
                     device=device,
                     operations=self.operations,
@@ -1099,6 +1117,15 @@ class LTXVModel(LTXBaseModel):
 
         grid_mask = None
         if keyframe_idxs is not None and keyframe_idxs.shape[2] > 0:
+            tokens_per_frame = self.tokens_per_latent_frame(additional_args["orig_shape"])
+            if keyframe_idxs.shape[2] % tokens_per_frame != 0:
+                raise ValueError(
+                    f"keyframe_idxs holds {keyframe_idxs.shape[2]} tokens, which is not a whole number of "
+                    f"{tokens_per_frame}-token latent frames. The appended frames were recorded against a "
+                    "different spatial resolution than the latent being sampled, so their positions would land "
+                    "on the wrong tokens. Crop the guides and separate the generated keyframes before "
+                    "upscaling the latent."
+                )
             additional_args.update({ "orig_patchified_shape": list(x.shape)})
             denoise_mask = self.patchifier.patchify(denoise_mask)[0]
             grid_mask = ~torch.any(denoise_mask < 0, dim=-1)[0]
@@ -1141,8 +1168,64 @@ class LTXVModel(LTXBaseModel):
             additional_args["num_guide_tokens"] = keyframe_idxs.shape[2]
 
         x = self.patchify_proj(x)
+        x = self.apply_keyframes_abs_pos_embedding(
+            x,
+            pixel_coords,
+            orig_shape=additional_args["orig_shape"],
+            grid_mask=grid_mask,
+            num_guide_tokens=additional_args.get("num_guide_tokens", 0),
+            generated_keyframes=kwargs.get("generated_keyframes", None),
+        )
         return x, pixel_coords, additional_args
 
+    def tokens_per_latent_frame(self, orig_shape):
+        """Token count of a single latent frame at the given latent shape."""
+        patch_size = self.patchifier.patch_size
+        return (orig_shape[3] // patch_size[1]) * (orig_shape[4] // patch_size[2])
+
+    def keyframes_abs_pos_mask(self, pixel_coords, orig_shape, grid_mask, num_guide_tokens, generated_keyframes):
+        """Per-token mask selecting the latents that encode a single standalone pixel frame.
+
+        Returns a (batch, tokens) boolean mask over the already grid-filtered token sequence.
+        """
+        temporal_start = pixel_coords[:, 0]
+        if temporal_start.ndim == 3:  # (batch, tokens, [start, end])
+            temporal_start = temporal_start[..., 0]
+        mask = temporal_start == 0
+        if num_guide_tokens > 0:
+            mask[:, -num_guide_tokens:] = False
+
+        if generated_keyframes is not None:
+            # The temporal patch size is always 1, so one latent frame is one row of tokens.
+            tokens_per_frame = self.tokens_per_latent_frame(orig_shape)
+            if generated_keyframes["tokens_per_frame"] != tokens_per_frame:
+                raise ValueError(
+                    f"The generated keyframes were recorded at {generated_keyframes['tokens_per_frame']} tokens "
+                    f"per latent frame but this latent has {tokens_per_frame}. Separate the generated keyframes "
+                    "before upscaling the latent."
+                )
+            first_token = generated_keyframes["first_latent_frame"] * tokens_per_frame
+            num_slot_tokens = generated_keyframes["num_keyframes"] * tokens_per_frame
+            slots = torch.zeros(orig_shape[2] * tokens_per_frame, dtype=torch.bool, device=mask.device)
+            slots[first_token:first_token + num_slot_tokens] = True
+            if grid_mask is not None:
+                slots = slots[grid_mask]
+            mask = mask | slots
+
+        return mask
+
+    def apply_keyframes_abs_pos_embedding(self, x, pixel_coords, orig_shape, grid_mask, num_guide_tokens, generated_keyframes):
+        """Add the learned keyframe marker to the single-pixel-frame tokens.
+
+        A no-op for every checkpoint built without the parameter.
+        """
+        if self.keyframes_abs_pos_embedding is None:
+            return x
+
+        mask = self.keyframes_abs_pos_mask(pixel_coords, orig_shape, grid_mask, num_guide_tokens, generated_keyframes)
+        embedding = self.keyframes_abs_pos_embedding.to(device=x.device, dtype=x.dtype)
+        return x + mask.unsqueeze(-1).to(x.dtype) * embedding
+
     def _build_guide_self_attention_mask(self, x, transformer_options, merged_args):
         """Build self-attention mask for per-guide attention attenuation.
 

comfy/ldm/lightricks/vae/audio_vae.py

@@ -1,6 +1,5 @@
 import json
 from dataclasses import dataclass
-import math
 import torch
 import torchaudio
 
@@ -186,7 +185,7 @@ class AudioVAE(torch.nn.Module):
         )
 
     def num_of_latents_from_frames(self, frames_number: int, frame_rate: float) -> int:
-        return math.ceil((float(frames_number) / frame_rate) * self.latents_per_second)
+        return round((float(frames_number) / frame_rate) * self.latents_per_second)
 
     def run_vocoder(self, mel_spec: torch.Tensor) -> torch.Tensor:
         audio_channels = self.autoencoder.decoder.out_ch

comfy/ldm/lightricks/vae/na_diffusion_decoder.py

@@ -0,0 +1,515 @@
+"""LTX 2.4 diffusion video VAE decoder (NADiffusionDecoder).
+
+Port of the reference ``DiffusionVideoDecoder`` without the NATTEN dependency:
+``natten.na3d`` is replaced by ``comfy_kitchen.na3d``, which reproduces
+NATTEN's semantics (window of exactly ``kernel_size`` per query, shifted
+inward at grid boundaries, dilation 1) and dispatches cuda/triton/eager per
+device and dtype (the eager backend covers CPU and fp32).
+
+Stages 1-4 deterministically upsample the latent into a context volume via
+NA transformer blocks + linear pixel-shuffle upsamples. Stage 5 runs
+``DiffusionNABlock``s that denoise patchified noised pixels ``x_t`` guided by
+that context through AdaLN-Zero scale/shift. The 2.4 checkpoint is single-step
+``x0``: one forward pass yields the pixels directly, no Euler loop.
+
+State dict keys match the shipped checkpoints directly (fused ``attn.qkv``,
+``t_embedder.mlp.{0,2}``, ``shared_adaln.proj``); no rename pass is needed.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+from comfy.ldm.lightricks.model import get_timestep_embedding
+from .causal_video_autoencoder import Encoder, processor
+
+import comfy_kitchen
+
+# Token chunk for the SwiGLU MLP (bounds the [chunk, hidden] workspace).
+MLP_TOKEN_CHUNK = 65536
+
+
+def rms_norm(x, weight, eps=1e-6):
+    if hasattr(F, "rms_norm"):
+        return F.rms_norm(x, (x.shape[-1],), weight=weight.to(x.dtype), eps=eps)
+    x_f = x.float()
+    x_f = x_f * torch.rsqrt(x_f.pow(2).mean(-1, keepdim=True) + eps)
+    return (x_f * weight.float()).to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return rms_norm(x, self.weight, self.eps)
+
+
+def patchify(x, patch_size_hw, patch_size_t=1):
+    if patch_size_hw == 1 and patch_size_t == 1:
+        return x
+    return rearrange(x, "b c (f p) (h q) (w r) -> b (c p r q) f h w", p=patch_size_t, q=patch_size_hw, r=patch_size_hw)
+
+
+def unpatchify(x, patch_size_hw, patch_size_t=1):
+    if patch_size_hw == 1 and patch_size_t == 1:
+        return x
+    return rearrange(x, "b (c p r q) f h w -> b c (f p) (h q) (w r)", p=patch_size_t, q=patch_size_hw, r=patch_size_hw)
+
+
+# --- Absolute per-axis RoPE (matches ltx-core rope.py numerics) ---
+
+def default_rope_dim_split(head_dim):
+    d_t = (head_dim // 4) // 2 * 2
+    d_hw = (head_dim - d_t) // 2
+    if d_hw % 2 != 0:
+        d_t -= 2
+        d_hw = (head_dim - d_t) // 2
+    return (d_t, d_hw, d_hw)
+
+
+def rope_inv_freqs(dim, base=10000.0, device=None):
+    exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
+    return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(torch.float32)
+
+
+def _rope_tables(lengths, inv_freqs, device):
+    """Precompute per-axis fp32 cos/sin tables for global 0-based positions."""
+    tables = []
+    for length, inv in zip(lengths, inv_freqs):
+        pos = torch.arange(length, dtype=torch.float32, device=device)
+        ang = pos[:, None] * inv[None, :]
+        tables.append((ang.cos(), ang.sin()))
+    return tables
+
+
+def _rope_matrices_slice(tables, t0, t1, h, w):
+    """Per-token rotation matrices ``(1, ts*h*w, 1, hd/2, 2, 2)`` fp32 for
+    ``comfy_kitchen.rms_rope_`` (interleaved-pair convention), covering global
+    frames ``[t0, t1)`` of the axis-factorized tables."""
+    parts = []
+    for (c, s), sl in zip(tables, (slice(t0, t1), slice(None), slice(None))):
+        c, s = c[sl], s[sl]
+        parts.append(torch.stack([c, -s, s, c], dim=-1).reshape(c.shape[0], 1, 1, c.shape[1], 2, 2))
+    ts = t1 - t0
+    freqs = torch.cat([
+        parts[0].expand(ts, h, w, -1, 2, 2),
+        parts[1].transpose(0, 1).expand(ts, h, w, -1, 2, 2),
+        parts[2].movedim(0, 2).expand(ts, h, w, -1, 2, 2),
+    ], dim=3)
+    return freqs.reshape(1, ts * h * w, 1, -1, 2, 2)
+
+
+class NeighborhoodAttention3D(nn.Module):
+    """QKV (fused, matching checkpoint keys) + q/k RMSNorm + abs RoPE + NA."""
+
+    def __init__(self, dim, kernel_size, head_dim=64, rope_base=10000.0):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = dim // head_dim
+        self.head_dim = head_dim
+        self.kernel_size = tuple(kernel_size)
+        self.scale = head_dim ** -0.5
+        self.rope_split = default_rope_dim_split(head_dim)
+        self.rope_base = rope_base
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim, bias=True)
+        self.q_norm = RMSNorm(head_dim, eps=1e-6)
+        self.k_norm = RMSNorm(head_dim, eps=1e-6)
+
+    def forward(self, x, pre=None, add_to=None):
+        """``pre`` (per-token norm/modulate) is applied slice-wise so the full
+        pre-attention tensor is never materialized; ``add_to`` streams the
+        output projection into it in place (residual add) and returns it.
+        Both bound peak memory without changing results."""
+        batch, t, h, w, _ = x.shape
+        inv_freqs = tuple(rope_inv_freqs(d, self.rope_base, device=x.device) for d in self.rope_split)
+        tables = _rope_tables((t, h, w), inv_freqs, x.device)
+        shape = (batch, t, h, w, self.num_heads, self.head_dim)
+        q = torch.empty(shape, dtype=x.dtype, device=x.device)
+        k = torch.empty(shape, dtype=x.dtype, device=x.device)
+        v = torch.empty(shape, dtype=x.dtype, device=x.device)
+        q_weight = (self.q_norm.weight.detach() * self.scale).to(x.dtype)  # scale commutes with the rotation
+        k_weight = self.k_norm.weight.detach().to(x.dtype)
+        chunk = max(1, (2 ** 25) // max(h * w * self.dim, 1))
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            sl = x[:, t0:t1] if pre is None else pre(x[:, t0:t1])
+            qc, kc, vc = self.qkv(sl).chunk(3, dim=-1)
+            cshape = (batch, t1 - t0, h, w, self.num_heads, self.head_dim)
+            q[:, t0:t1] = qc.reshape(cshape)
+            k[:, t0:t1] = kc.reshape(cshape)
+            v[:, t0:t1] = vc.reshape(cshape)
+            freqs = _rope_matrices_slice(tables, t0, t1, h, w)
+            nt = (t1 - t0) * h * w
+            for b in range(batch):
+                comfy_kitchen.rms_rope_(
+                    q[b, t0:t1].view(1, nt, self.num_heads, self.head_dim),
+                    k[b, t0:t1].view(1, nt, self.num_heads, self.head_dim),
+                    freqs, q_weight, k_weight)
+        out = comfy_kitchen.na3d(q, k, v, list(self.kernel_size), None, 1.0)
+        del q, k, v
+        out = out.reshape(batch, t, h, w, self.dim)
+        res = add_to if add_to is not None else torch.empty_like(out)
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            if add_to is not None:
+                res[:, t0:t1] += self.proj(out[:, t0:t1])
+            else:
+                res[:, t0:t1] = self.proj(out[:, t0:t1])
+        return res
+
+
+class SwiGLU(nn.Module):
+    """``w_down(silu(w_gate(x)) * w_up(x))``, chunked over tokens to bound the
+    ``[chunk, hidden]`` workspace."""
+
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.w_up = nn.Linear(dim, hidden_dim, bias=False)
+        self.w_gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.w_down = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x, pre=None, add_to=None):
+        """``pre``/``add_to`` as in ``NeighborhoodAttention3D.forward``."""
+        _, t, h, w, _ = x.shape
+        chunk = max(1, MLP_TOKEN_CHUNK // max(h * w, 1))
+        out = add_to if add_to is not None else torch.empty_like(x)
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            sl = x[:, t0:t1] if pre is None else pre(x[:, t0:t1])
+            y = self.w_down(F.silu(self.w_gate(sl)) * self.w_up(sl))
+            if add_to is not None:
+                out[:, t0:t1] += y
+            else:
+                out[:, t0:t1] = y
+        return out
+
+
+class NABlock(nn.Module):
+    """Pre-norm transformer block: NA -> SwiGLU MLP with residual adds."""
+
+    def __init__(self, dim, kernel_size, head_dim=64, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, eps=1e-6)
+        self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim)
+        self.norm2 = RMSNorm(dim, eps=1e-6)
+        hidden = (int(dim * mlp_ratio) + 15) // 16 * 16
+        self.mlp = SwiGLU(dim, hidden)
+
+    def forward(self, x):
+        x = self.attn(x, pre=self.norm1, add_to=x)
+        return self.mlp(x, pre=self.norm2, add_to=x)
+
+
+def modulate(x, scale, shift):
+    return x * (1.0 + scale) + shift
+
+
+class AdaLNZero(nn.Module):
+    """``t_emb`` -> 7 (scale/shift/gate) chunks; gate slots unused (folded at export)."""
+
+    NUM_CHUNKS = 7
+
+    def __init__(self, dim, t_emb_dim):
+        super().__init__()
+        self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True)
+
+    def forward(self, t_emb):
+        h = self.proj(F.silu(t_emb))
+        return tuple(c[:, None, None, None, :] for c in h.chunk(self.NUM_CHUNKS, dim=-1))
+
+
+class DiffusionNABlock(nn.Module):
+    """NA + SwiGLU with shared AdaLN-Zero scale/shift (ungated residuals)."""
+
+    def __init__(self, dim, kernel_size, context_channels, head_dim=64, mlp_ratio=4.0):
+        super().__init__()
+        self.context_proj = nn.Linear(context_channels, dim, bias=True)
+        self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim))
+        self.norm1 = RMSNorm(dim, eps=1e-6)
+        self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim)
+        self.norm2 = RMSNorm(dim, eps=1e-6)
+        hidden = (int(dim * mlp_ratio) + 15) // 16 * 16
+        self.mlp = SwiGLU(dim, hidden)
+
+    def forward(self, x, latent_context, modulation):
+        scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [
+            modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS)
+        ]
+        chunk = max(1, MLP_TOKEN_CHUNK // max(x.shape[2] * x.shape[3], 1))
+        for t0 in range(0, x.shape[1], chunk):
+            x[:, t0:t0 + chunk] += self.context_proj(latent_context[:, t0:t0 + chunk])
+        x = self.attn(x, pre=lambda s: modulate(self.norm1(s), scale_msa, shift_msa), add_to=x)
+        return self.mlp(x, pre=lambda s: modulate(self.norm2(s), scale_mlp, shift_mlp), add_to=x)
+
+
+class LinearPixelShuffleUpsample(nn.Module):
+    """Linear channel-expand, then channels-last pixel shuffle."""
+
+    def __init__(self, in_channels, stride, out_channels_reduction_factor=1):
+        super().__init__()
+        self.stride = tuple(stride)
+        proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
+        self.out_channels = proj_out_channels // math.prod(stride)
+        self.proj = nn.Linear(in_channels, proj_out_channels, bias=True)
+
+    def forward(self, x, drop_leading_frame=True):
+        batch, t, h, w, _ = x.shape
+        p1, p2, p3 = self.stride
+        out = torch.empty((batch, t * p1, h * p2, w * p3, self.out_channels), dtype=x.dtype, device=x.device)
+        chunk = max(1, MLP_TOKEN_CHUNK // max(h * w, 1))
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            out[:, t0 * p1:t1 * p1] = rearrange(
+                self.proj(x[:, t0:t1]), "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c",
+                p1=p1, p2=p2, p3=p3,
+            )
+        if p1 == 2 and drop_leading_frame:
+            # The causal temporal pixel-shuffle duplicates the leading frame.
+            out = out[:, 1:]
+        return out
+
+
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal(256) -> MLP. ``mlp.{0,2}`` naming matches the checkpoint."""
+
+    def __init__(self, t_emb_dim=384, freq_dim=256):
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(freq_dim, t_emb_dim, bias=True),
+            nn.SiLU(),
+            nn.Linear(t_emb_dim, t_emb_dim, bias=True),
+        )
+
+    def forward(self, timestep, dtype):
+        emb = get_timestep_embedding(timestep.flatten(), self.freq_dim, flip_sin_to_cos=True,
+                                     downscale_freq_shift=0, scale=1)
+        return self.mlp(emb.to(dtype))
+
+
+class NADiffusionDecoder(nn.Module):
+    """Stages 1-4 (deterministic NA upsample) + stage-5 diffusion blocks.
+
+    Input latent must already be un-normalized (the wrapper applies
+    ``per_channel_statistics.un_normalize``, same as the conv VAE path).
+    """
+
+    def __init__(
+        self,
+        in_channels=128,
+        out_channels=3,
+        patch_size=4,
+        head_dim=64,
+        stage_channels=(2048, 1024, 512, 512, 256),
+        stage_depths=(4, 6, 4, 2, 8),
+        stage_kernels=((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5), (11, 11, 11)),
+        upsamples=(((1, 2, 2), 2), ((2, 1, 1), 2), ((2, 2, 2), 1), ((2, 2, 2), 2)),
+        stage5_kernel=(11, 11, 11),
+        t_emb_dim=384,
+        default_num_inference_steps=1,
+        timestep_scale_multiplier=1000.0,
+        model_output_type="x0",
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.out_channels = out_channels
+        self.timestep_scale_multiplier = timestep_scale_multiplier
+        self.model_output_type = model_output_type
+        self.register_buffer(
+            "default_inference_timesteps",
+            torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps),
+            persistent=False,
+        )
+        self.temporal_upscale = math.prod(s[0] for s, _ in upsamples)
+        self.spatial_upscale = math.prod(s[1] for s, _ in upsamples) * patch_size
+        # NATTEN-style last-frame border mitigation: replicate the last latent
+        # frame through stages 1-4, crop the appendix off the context after.
+        self.trailing_pad_latent_frames = (stage_kernels[0][0] // 2) * 2
+
+        self.conv_in = nn.Linear(in_channels, stage_channels[0], bias=True)
+
+        self.det_stages = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        for stage_i in range(len(stage_channels) - 1):
+            c = stage_channels[stage_i]
+            self.det_stages.append(nn.ModuleList(
+                [NABlock(c, stage_kernels[stage_i], head_dim=head_dim) for _ in range(stage_depths[stage_i])]
+            ))
+            stride, reduction = upsamples[stage_i]
+            self.upsamples.append(LinearPixelShuffleUpsample(c, stride, out_channels_reduction_factor=reduction))
+
+        self.t_embedder = TimestepEmbedder(t_emb_dim=t_emb_dim)
+
+        c5 = stage_channels[-1]
+        self.context_channels = c5
+        noised_pixel_channels = out_channels * (patch_size ** 2)
+        self.conv_in_x_t = nn.Linear(noised_pixel_channels, c5, bias=True)
+        self.shared_adaln = AdaLNZero(c5, t_emb_dim)
+        self.diff_blocks = nn.ModuleList([
+            DiffusionNABlock(c5, stage5_kernel, context_channels=c5, head_dim=head_dim)
+            for _ in range(stage_depths[-1])
+        ])
+        self.norm_out = RMSNorm(c5, eps=1e-6)
+        self.conv_out = nn.Linear(c5, noised_pixel_channels, bias=True)
+
+    def forward_pre_diffusion(self, z, drop_leading_frame=True, pad_trailing=True):
+        """Stages 1-4: latent -> stage-5 context, channels-last.
+
+        ``drop_leading_frame`` must be True only when ``z`` contains the
+        latent's true temporal origin (t=0); tiled callers decoding a later
+        temporal chunk pass False (the duplicate leading frame belongs solely
+        to the origin chunk). ``pad_trailing`` only for chunks containing the
+        latent's last frame."""
+        n = self.trailing_pad_latent_frames if pad_trailing else 0
+        if n > 0:
+            z = torch.cat([z, z[:, :, -1:].expand(-1, -1, n, -1, -1)], dim=2)
+        x = z.permute(0, 2, 3, 4, 1)
+        x = self.conv_in(x)
+        for stage_i, blocks in enumerate(self.det_stages):
+            for block in blocks:
+                x = block(x)
+            x = self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame)
+        if n > 0:
+            x = x[:, :-(n * self.temporal_upscale)]
+        return x
+
+    def forward_diff_step(self, context, x_t, t):
+        x = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1)
+        x = self.conv_in_x_t(x.permute(0, 2, 3, 4, 1))
+        t_emb = self.t_embedder(self.timestep_scale_multiplier * t, dtype=x.dtype)
+        modulation = self.shared_adaln(t_emb)
+        for block in self.diff_blocks:
+            x = block(x, context, modulation)
+        x = self.norm_out(x)
+        x = self.conv_out(x)
+        x = x.permute(0, 4, 1, 2, 3)
+        return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1)
+
+    def forward(self, z, generator=None, drop_leading_frame=True, pad_trailing=True):
+        context = self.forward_pre_diffusion(z, drop_leading_frame=drop_leading_frame, pad_trailing=pad_trailing)
+        batch, t5, h5, w5, _ = context.shape
+        pixel_shape = (batch, self.out_channels, t5, h5 * self.patch_size, w5 * self.patch_size)
+        x_t = torch.randn(pixel_shape, dtype=z.dtype, device=z.device, generator=generator)
+
+        timesteps = self.default_inference_timesteps.to(z.device)
+        num_steps = timesteps.shape[0]
+        for i in range(num_steps):
+            t_now = timesteps[i].expand(batch)
+            model_out = self.forward_diff_step(context, x_t, t_now)
+            if self.model_output_type == "x0":
+                x0 = model_out
+                if i == num_steps - 1:
+                    return x0
+                velocity = (x_t.float() - x0.float()) / timesteps[i]
+            else:  # "v"
+                velocity = model_out.float()
+                if i == num_steps - 1:
+                    return (x_t.float() - timesteps[i] * velocity).to(z.dtype)
+            t_next = timesteps[i + 1] if i + 1 < num_steps else torch.zeros_like(timesteps[i])
+            x_t = (x_t.float() - (timesteps[i] - t_next) * velocity).to(z.dtype)
+        return x_t
+
+
+LTX_24_VAE_CONFIG = {
+    "_class_name": "CausalDiffusionVAE",
+    "dims": 3,
+    "model_output_type": "x0",
+    "encoder": {
+        "dims": 3,
+        "in_channels": 3,
+        "out_channels": 128,
+        "blocks": [
+            ["res_x", {"num_layers": 4}],
+            ["compress_space_res", {"multiplier": 2}],
+            ["res_x", {"num_layers": 6}],
+            ["compress_time_res", {"multiplier": 2}],
+            ["res_x", {"num_layers": 4}],
+            ["compress_all_res", {"multiplier": 2}],
+            ["res_x", {"num_layers": 2}],
+            ["compress_all_res", {"multiplier": 1}],
+            ["res_x", {"num_layers": 2}],
+        ],
+        "patch_size": 4,
+        "latent_log_var": "constant",
+        "norm_layer": "pixel_norm",
+        "base_channels": 128,
+        "spatial_padding_mode": "zeros",
+    },
+    "decoder": {
+        "in_channels": 128,
+        "out_channels": 3,
+        "patch_size": 4,
+        "head_dim": 64,
+        "stage_channels": [2048, 1024, 512, 512, 256],
+        "stage_depths": [4, 6, 4, 2, 8],
+        "stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]],
+        "upsamples": [[[1, 2, 2], 2], [[2, 1, 1], 2], [[2, 2, 2], 1], [[2, 2, 2], 2]],
+        "stage5_kernel": [11, 11, 11],
+        "timestep_scale_multiplier": 1000.0,
+        "default_num_inference_steps": 1,
+    },
+}
+
+
+class CausalDiffusionVAE(nn.Module):
+    """LTX 2.4 video VAE: conv encoder (shared with the 2.0 arch) + NA
+    diffusion decoder. Interface mirrors ``causal_video_autoencoder.VideoVAE``.
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+        if config is None:
+            config = LTX_24_VAE_CONFIG
+        self.config = config
+        enc = config.get("encoder", LTX_24_VAE_CONFIG["encoder"])
+        dec = config.get("decoder", LTX_24_VAE_CONFIG["decoder"])
+        dec_defaults = LTX_24_VAE_CONFIG["decoder"]
+
+        self.encoder = Encoder(
+            dims=enc.get("dims", 3),
+            in_channels=enc.get("in_channels", 3),
+            out_channels=enc.get("out_channels", 128),
+            blocks=enc.get("blocks", LTX_24_VAE_CONFIG["encoder"]["blocks"]),
+            patch_size=enc.get("patch_size", 4),
+            latent_log_var=enc.get("latent_log_var", "constant"),
+            norm_layer=enc.get("norm_layer", "pixel_norm"),
+            spatial_padding_mode=enc.get("spatial_padding_mode", "zeros"),
+            base_channels=enc.get("base_channels", 128),
+        )
+
+        self.decoder = NADiffusionDecoder(
+            in_channels=dec.get("in_channels", 128),
+            out_channels=dec.get("out_channels", 3),
+            patch_size=dec.get("patch_size", 4),
+            head_dim=dec.get("head_dim", 64),
+            stage_channels=tuple(dec.get("stage_channels", dec_defaults["stage_channels"])),
+            stage_depths=tuple(dec.get("stage_depths", dec_defaults["stage_depths"])),
+            stage_kernels=tuple(tuple(k) for k in dec.get("stage_kernels", dec_defaults["stage_kernels"])),
+            upsamples=tuple((tuple(s), r) for s, r in dec.get("upsamples", dec_defaults["upsamples"])),
+            stage5_kernel=tuple(dec.get("stage5_kernel", dec_defaults["stage5_kernel"])),
+            t_emb_dim=dec.get("t_emb_dim", 384),
+            default_num_inference_steps=dec.get("default_num_inference_steps", 1),
+            timestep_scale_multiplier=dec.get("timestep_scale_multiplier", 1000.0),
+            model_output_type=config.get("model_output_type", "x0"),
+        )
+
+        self.per_channel_statistics = processor()
+
+    def encode(self, x, device=None):
+        x = x[:, :, :max(1, 1 + ((x.shape[2] - 1) // 8) * 8), :, :]
+        means, logvar = torch.chunk(self.encoder(x, device=device), 2, dim=1)
+        return self.per_channel_statistics.normalize(means)
+
+    def decode(self, x):
+        # Fixed-seed noise so decodes are reproducible TODO: expose?
+        generator = torch.Generator(device=x.device)
+        generator.manual_seed(0)
+        return self.decoder(self.per_channel_statistics.un_normalize(x), generator=generator)

comfy/ldm/minimax/model.py

@@ -25,7 +25,7 @@ import comfy.model_prefetch
 import comfy.ops
 import comfy.patcher_extension
 import comfy.quant_ops
-from comfy.ldm.modules.attention import optimized_attention
+from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention
 
 FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
 FRAME_RESCALE = 5.0 / 3.0
@@ -176,6 +176,7 @@ class Attention(nn.Module):
             extra_options["n_heads"] = self.heads
             extra_options["dim_head"] = self.head_dim
 
+        v = v.clone()
         if "attn1_patch" in patches:
             q = q.reshape(1, s, -1)
             k = k.reshape(1, s, -1)
@@ -193,6 +194,9 @@ class Attention(nn.Module):
             q = q.transpose(0, 1).unsqueeze(0)
             k = k.transpose(0, 1).unsqueeze(0)
             v = v.transpose(0, 1).unsqueeze(0)
+        q = AttentionTensorContainer(q)
+        k = AttentionTensorContainer(k)
+        v = AttentionTensorContainer(v)
         out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)
 
         if "attn1_output_patch" in patches:

comfy/ldm/minimax/vae.py

@@ -6,6 +6,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import comfy.model_management
 import comfy.ops
 import comfy.quant_ops
 import comfy.rmsnorm
@@ -321,6 +322,8 @@ class ViT3DDecoder(nn.Module):
 # Full VAE
 
 class MiniMaxH3VideoVAE(nn.Module):
+    comfy_has_chunked_io = True
+
     def __init__(
         self,
         in_channels=3,
@@ -389,6 +392,23 @@ class MiniMaxH3VideoVAE(nn.Module):
     def _decode_pixels(self, z):
         return self.decoder(self.post_quant_conv(z))
 
+    def _normalize_pixels(self, x):
+        return x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
+
+    def _finalize_pixels(self, part):
+        # raw decoder output -> float32 pixels in [0, 1] (the VAE wrapper's process_output is identity)
+        part = part * self.pixel_std.to(device=part.device, dtype=torch.float32)
+        return part.add_(self.pixel_mean.to(device=part.device, dtype=torch.float32)).clamp_(0.0, 1.0)
+
+    def decode_output_shape(self, input_shape):
+        b, c, t, h, w = input_shape
+        if t == 1:
+            frames = 1
+        else:
+            pad_tokens, num_chunks = self._decode_temporal_chunks(t)
+            frames = self._decode_temporal_frame_plan(t + pad_tokens, num_chunks, pad_tokens)
+        return (b, self.decoder.out_channels, frames, h * self.vae_ratio, w * self.vae_ratio)
+
     def _adaptive_encode(self, x):
         if self.tiling:
             return self.tiled_encode(x)
@@ -521,18 +541,15 @@ class MiniMaxH3VideoVAE(nn.Module):
 
     # temporal chunking
 
-    def encode_temporal(self, x):
-        if x.shape[2] % self.clip_length != 0:
-            pad_size = (-x.shape[2]) % self.clip_length
-            pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
-            x = torch.cat([x, pad_frames], dim=2)
-
-        num_chunks = x.shape[2] // self.clip_length
-
+    def encode_temporal(self, x, device):
+        # chunked input io: x may live on the CPU, clips move to the device as they encode
         z_list = []
-        for i in range(num_chunks):
-            clip_x = x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :]
-            z_list.append(self._adaptive_encode(clip_x))
+        for i in range(math.ceil(x.shape[2] / self.clip_length)):
+            clip_x = x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :].to(device)
+            if clip_x.shape[2] < self.clip_length:
+                pad_frames = clip_x[:, :, -1:].repeat(1, 1, self.clip_length - clip_x.shape[2], 1, 1)
+                clip_x = torch.cat([clip_x, pad_frames], dim=2)
+            z_list.append(self._adaptive_encode(self._normalize_pixels(clip_x)))
 
         z = torch.cat(z_list, dim=2)
         if self.token_drop > 0:
@@ -577,43 +594,42 @@ class MiniMaxH3VideoVAE(nn.Module):
         total_frames += final_overlap_frames
         return total_frames - self._decode_temporal_pad_frames(z_len, pad_tokens)
 
-    def decode_temporal(self, z):
-        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
-        split_count = int(self.token_drop > 0) + 1
-
-        pseudo_total_tokens = z.shape[2] + self.token_drop
-
-        pad_tokens = 0
-        remainder = pseudo_total_tokens % self.tokens_chunk_size
-        if remainder != 0:
-            pad_tokens = self.tokens_chunk_size - remainder
-            pseudo_total_tokens += pad_tokens
+    def _decode_temporal_chunks(self, z_len):
+        pseudo_total_tokens = z_len + self.token_drop
+        pad_tokens = (-pseudo_total_tokens) % self.tokens_chunk_size
+        pseudo_total_tokens += pad_tokens
 
         num_chunks = pseudo_total_tokens // self.tokens_chunk_size - int(self.token_drop > 0)
         if num_chunks < 1:
             # too few tokens for one chunk (e.g. T_lat == 2): pad one extra chunk
             pad_tokens += self.tokens_chunk_size
             num_chunks += 1
+        return pad_tokens, num_chunks
 
+    def decode_temporal(self, z, output_buffer=None):
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+
+        if output_buffer is None:
+            # finalized chunks stream out of VRAM so the full video never sits on the GPU
+            output_buffer = torch.empty(self.decode_output_shape(z.shape), dtype=torch.float32,
+                                        device=comfy.model_management.intermediate_device())
+
+        pad_tokens, num_chunks = self._decode_temporal_chunks(z.shape[2])
         if pad_tokens > 0:
             pad_z = z[:, :, -1:, :, :].repeat(1, 1, pad_tokens, 1, 1)
             z = torch.cat([z, pad_z], dim=2)
 
-        output_frames = self._decode_temporal_frame_plan(z.shape[2], num_chunks, pad_tokens)
-
-        dec = None
+        dec = output_buffer
         dec_overlap = None
         write_pos = 0
 
         def write_part(part):
-            nonlocal dec, write_pos
+            nonlocal write_pos
             part_frames = part.shape[2]
             if part_frames <= 0:
                 return
-            if dec is None:
-                out_shape = list(part.shape)
-                out_shape[2] = output_frames
-                dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
+            part = self._finalize_pixels(part)
             copy_frames = min(part_frames, max(0, dec.shape[2] - write_pos))
             if copy_frames > 0:
                 dec[:, :, write_pos:write_pos + copy_frames, :, :].copy_(
@@ -653,18 +669,18 @@ class MiniMaxH3VideoVAE(nn.Module):
         return dec
 
 
-    def encode(self, x):
+    def encode(self, x, device=None):
         # x: [B, 3, T, H, W] in [-1, 1] -> normalized latents [B, 24, T_lat, H/16, W/16]
         if x.ndim == 4:
             x = x.unsqueeze(2)
-
-        x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
+        if device is None:
+            device = x.device
 
         if x.shape[2] == 1:
-            moments = self._adaptive_encode(x)
+            moments = self._adaptive_encode(self._normalize_pixels(x.to(device)))
             moments = moments[:, :, -1:, :, :]
         else:
-            moments = self.encode_temporal(x)
+            moments = self.encode_temporal(x, device)
 
         mean = torch.chunk(moments.float(), 2, dim=1)[0]
 
@@ -679,18 +695,16 @@ class MiniMaxH3VideoVAE(nn.Module):
     def decode_tiled(self, z, **kwargs):
         return self.decode(z)
 
-    def decode(self, z):
-        # z: [B, 24, T_lat, H_lat, W_lat] normalized latents -> pixels [B, 3, T, H, W] in [-1, 1]
+    def decode(self, z, output_buffer=None):
+        # z: [B, 24, T_lat, H_lat, W_lat] normalized latents -> float32 pixels [B, 3, T, H, W] in [0, 1]
         latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(z)
         latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(z)
         z = z * latents_std + latents_mean
 
         if z.shape[2] == 1:
-            dec = self._adaptive_decode(z)
-            dec = dec[:, :, -1:, :, :]
-        else:
-            dec = self.decode_temporal(z)
-
-        dec = dec.float()
-        dec.mul_(self.pixel_std.to(dec)).add_(self.pixel_mean.to(dec)).clamp_(0.0, 1.0).mul_(2.0).sub_(1.0)
-        return dec
+            dec = self._finalize_pixels(self._adaptive_decode(z)[:, :, -1:, :, :])
+            if output_buffer is None:
+                return dec
+            output_buffer.copy_(dec)
+            return output_buffer
+        return self.decode_temporal(z, output_buffer)

comfy/ldm/modules/attention.py

@@ -10,6 +10,8 @@ from typing import Optional, Any, Callable, Union
 import logging
 import functools
 
+import comfy_kitchen
+
 from .diffusionmodules.util import AlphaBlender, timestep_embedding
 from .sub_quadratic_attention import efficient_dot_product_attention
 
@@ -49,6 +51,8 @@ except ImportError:
         logging.error(f"\n\nTo use the `--use-flash-attention` feature, the `flash-attn` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install flash-attn")
         exit(-1)
 
+COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE = comfy_kitchen.int8_attention_is_available()
+
 REGISTERED_ATTENTION_FUNCTIONS = {}
 def register_attention_function(name: str, func: Callable):
     # avoid replacing existing functions
@@ -145,9 +149,34 @@ def Normalize(in_channels, dtype=None, device=None):
     return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
 
 
+class AttentionTensorContainer:
+    """Single-owner tensor input consumed by an optimized attention backend."""
+
+    __slots__ = ("tensor",)
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor: torch.Tensor | None = tensor
+
+    def peek(self) -> torch.Tensor:
+        if self.tensor is None:
+            raise RuntimeError("attention tensor container has already been consumed")
+        return self.tensor
+
+    def take(self) -> torch.Tensor:
+        tensor = self.peek()
+        self.tensor = None
+        return tensor
+
+
 def wrap_attn(func):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
+        containers = None
+        if len(args) >= 3 and isinstance(args[0], AttentionTensorContainer):
+            if not isinstance(args[1], AttentionTensorContainer) or not isinstance(args[2], AttentionTensorContainer):
+                raise TypeError("q, k, and v must all be attention tensor containers")
+            containers = args[:3]
+
         remove_attn_wrapper_key = False
         try:
             if "_inside_attn_wrapper" not in kwargs:
@@ -156,11 +185,22 @@ def wrap_attn(func):
                 kwargs["_inside_attn_wrapper"] = True
                 if transformer_options is not None:
                     if "optimized_attention_override" in transformer_options:
-                        return transformer_options["optimized_attention_override"](func, *args, **kwargs)
+                        optimized_attention_override = transformer_options["optimized_attention_override"]
+                        if containers is not None:
+                            if hasattr(optimized_attention_override, "container_function"):
+                                return optimized_attention_override.container_function(*args, **kwargs)
+                            args = tuple(container.take() for container in containers) + args[3:]
+                        return optimized_attention_override(func, *args, **kwargs)
+
+            if containers is not None:
+                if wrapper.container_function is not None:
+                    return wrapper.container_function(*args, **kwargs)
+                args = tuple(container.take() for container in containers) + args[3:]
             return func(*args, **kwargs)
         finally:
             if remove_attn_wrapper_key:
                 del kwargs["_inside_attn_wrapper"]
+    wrapper.container_function = None
     return wrapper
 
 @wrap_attn
@@ -545,6 +585,63 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
             ).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head)
     return out
 
+def _comfy_kitchen_int8_inputs(q, k, v, heads, mask, skip_reshape, enable_gqa):
+    dim_head = q.shape[-1] if skip_reshape else q.shape[-1] // heads
+    b = q.shape[0]
+    if not skip_reshape:
+        q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa, expand_kv=False)
+        q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
+
+    if mask is not None:
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+
+    return q, k, v, mask, b, dim_head
+
+
+@wrap_attn
+def attention_comfy_kitchen_int8(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    q, k, v, mask, b, dim_head = _comfy_kitchen_int8_inputs(
+        q, k, v, heads, mask, skip_reshape, kwargs.get("enable_gqa", False)
+    )
+    out = comfy_kitchen.int8_attention(
+        q,
+        k,
+        v,
+        scale=kwargs.get("scale", None),
+        attn_mask=mask,
+    )
+    if not skip_output_reshape:
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+    return out
+
+
+def _attention_comfy_kitchen_int8_containers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    q = q.take()
+    k = k.take()
+    v = v.take()
+    q, k, v, mask, b, dim_head = _comfy_kitchen_int8_inputs(
+        q, k, v, heads, mask, skip_reshape, kwargs.get("enable_gqa", False)
+    )
+    quantized = comfy_kitchen.prequantize_int8_attention(
+        q,
+        k,
+        v,
+        scale=kwargs.get("scale", None),
+        attn_mask=mask,
+    )
+    del q, k, v
+    out = comfy_kitchen.int8_attention_from_prequantized(quantized)
+    if not skip_output_reshape:
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+    return out
+
+
+attention_comfy_kitchen_int8.container_function = _attention_comfy_kitchen_int8_containers
+
+
 @wrap_attn
 def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
     if kwargs.get("low_precision_attention", True) is False or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK):
@@ -775,10 +872,20 @@ else:
         logging.info("Using sub quadratic optimization for attention, if you have memory or speed issues try using: --use-split-cross-attention")
         optimized_attention = attention_sub_quad
 
+if model_management.comfy_kitchen_attention_enabled():
+    if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE:
+        logging.info("Using Comfy Kitchen attention")
+        optimized_attention = attention_comfy_kitchen_int8
+    else:
+        logging.error("Comfy Kitchen attention is unavailable. Install a Comfy Kitchen build with attention support to use --use-ck-attention.")
+        exit(-1)
+
 optimized_attention_masked = optimized_attention
 
 
 # register core-supported attention functions
+if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE:
+    register_attention_function("comfy_kitchen_int8", attention_comfy_kitchen_int8)
 if SAGE_ATTENTION_IS_AVAILABLE:
     register_attention_function("sage", attention_sage)
 if SAGE_ATTENTION3_IS_AVAILABLE:

comfy/model_base.py

@@ -1153,6 +1153,10 @@ class LTXV(BaseModel):
         if guide_attention_entries is not None:
             out['guide_attention_entries'] = comfy.conds.CONDConstant(guide_attention_entries)
 
+        generated_keyframes = kwargs.get("generated_keyframes", None)
+        if generated_keyframes is not None:
+            out['generated_keyframes'] = comfy.conds.CONDConstant(generated_keyframes)
+
         return out
 
     def process_timestep(self, timestep, x, denoise_mask=None, **kwargs):
@@ -1213,6 +1217,10 @@ class LTXAV(BaseModel):
         if ref_audio is not None:
             out['ref_audio'] = comfy.conds.CONDConstant(ref_audio)
 
+        generated_keyframes = kwargs.get("generated_keyframes", None)
+        if generated_keyframes is not None:
+            out['generated_keyframes'] = comfy.conds.CONDConstant(generated_keyframes)
+
         return out
 
     def process_timestep(self, timestep, x, denoise_mask=None, audio_denoise_mask=None, **kwargs):

comfy/model_detection.py

@@ -397,6 +397,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["cross_attention_dim"] = shape[1]
         if metadata is not None and "config" in metadata:
             dit_config.update(json.loads(metadata["config"]).get("transformer", {}))
+        dit_config["use_keyframes_abs_pos_embedding"] = '{}keyframes_abs_pos_embedding'.format(key_prefix) in state_dict_keys
         return dit_config
 
     if '{}genre_embedder.weight'.format(key_prefix) in state_dict_keys: #ACE-Step model

comfy/model_management.py

@@ -1658,6 +1658,9 @@ def unpin_memory(tensor):
 def sage_attention_enabled():
     return args.use_sage_attention
 
+def comfy_kitchen_attention_enabled():
+    return args.use_ck_attention
+
 def flash_attention_enabled():
     return args.use_flash_attention
 

comfy/model_patcher.py

@@ -685,6 +685,14 @@ class ModelPatcher:
     def set_model_attn2_output_patch(self, patch):
         self.set_model_patch(patch, "attn2_output_patch")
 
+    def set_model_optimized_attention(self, optimized_attention):
+        def optimized_attention_override(_, *args, **kwargs):
+            return optimized_attention(*args, **kwargs)
+
+        if hasattr(optimized_attention, "container_function") and optimized_attention.container_function is not None:
+            optimized_attention_override.container_function = optimized_attention.container_function
+        self.model_options["transformer_options"]["optimized_attention_override"] = optimized_attention_override
+
     def set_model_input_block_patch(self, patch):
         self.set_model_patch(patch, "input_block_patch")
 

comfy/quant_ops.py

@@ -40,7 +40,7 @@ try:
         cuda_version = tuple(map(int, str(torch.version.cuda).split('.')))
         if cuda_version < (13,):
             ck.registry.disable("cuda")
-            logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
+            logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.\nWARNING WARNING WARNING\nIf you are on nvidia 20 series and above it is required that you update your pytorch to cu130 or higher.\n")
 
     # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated
     # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7 AND a

comfy/sample.py

@@ -37,6 +37,11 @@ def prepare_noise(latent_image, seed, noise_inds=None):
 
     return noises
 
+def prepare_empty_noise(latent_image):
+    if latent_image.is_nested:
+        return comfy.nested_tensor.NestedTensor([torch.zeros_like(t, device="cpu") for t in latent_image.unbind()])
+    return torch.zeros_like(latent_image, device="cpu")
+
 def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None, downscale_ratio_temporal=None):
     if latent_image.is_nested:
         return latent_image

comfy/sd.py

@@ -11,6 +11,7 @@ from .ldm.cascade.stage_c_coder import StageC_coder
 from .ldm.audio.autoencoder import AudioOobleckVAE
 import comfy.ldm.genmo.vae.model
 import comfy.ldm.lightricks.vae.causal_video_autoencoder
+import comfy.ldm.lightricks.vae.na_diffusion_decoder
 import comfy.ldm.lightricks.vae.audio_vae
 import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
@@ -583,6 +584,22 @@ class VAE:
                 self.working_dtypes = [torch.bfloat16, torch.float32]
                 self.memory_used_encode = lambda shape, dtype: (400 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
                 self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * 16 * 16) * model_management.dtype_size(dtype)
+            elif "decoder.conv_in_x_t.weight" in sd:  # lightricks LTX 2.4 diffusion VAE decoder
+                vae_config = None
+                if metadata is not None and "config" in metadata:
+                    vae_config = json.loads(metadata["config"]).get("vae", None)
+                self.first_stage_model = comfy.ldm.lightricks.vae.na_diffusion_decoder.CausalDiffusionVAE(config=vae_config)
+                self.latent_channels = sd["decoder.conv_in.weight"].shape[1]
+                self.latent_dim = 3
+                self.disable_offload = True
+                self.crop_input = False  # generic crop would narrow the frame axis by the 32x spatial ratio
+                self.memory_used_decode = lambda shape, dtype: (1700 * shape[2] * shape[3] * shape[4] * (8 * 8 * 8)) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (80 * max(shape[2], 7) * shape[3] * shape[4]) * model_management.dtype_size(dtype)
+                self.upscale_ratio = (lambda a: max(0, a * 8 - 7), 32, 32)
+                self.upscale_index_formula = (8, 32, 32)
+                self.downscale_ratio = (lambda a: max(0, math.floor((a + 7) / 8)), 32, 32)
+                self.downscale_index_formula = (8, 32, 32)
+                self.working_dtypes = [torch.bfloat16, torch.float32]
             elif "decoder.conv_in.weight" in sd:
                 if sd['decoder.conv_in.weight'].shape[1] == 64:
                     ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
@@ -955,13 +972,21 @@ class VAE:
                 self.working_dtypes = [torch.float16, torch.float32]
                 # the model tiles internally (256px spatial, 17-frame temporal chunks)
                 self.handles_tiling = True
+                # decode finalizes straight to [0, 1] while streaming chunks out
+                self.process_output = lambda image: image
+                # one decoded temporal chunk (with overlap) is all that ever sits in VRAM
+                chunk_frames = (self.first_stage_model.tokens_chunk_size + self.first_stage_model.token_overlap) * self.first_stage_model.vae_ratio_t
+
                 def estimate_encode_memory(frames, height, width, dtype):
                     fixed = 110_000_000 if frames == 1 else 1_300_000_000
                     elements_per_pixel = 7 if frames == 1 else 9.5
+                    # only one clip of the input video is ever resident on the GPU
+                    frames = min(frames, self.first_stage_model.clip_length)
                     return (elements_per_pixel * frames * height * width + fixed) * model_management.dtype_size(dtype) * 1.03
 
                 def estimate_decode_memory(frames, height, width, dtype):
                     fixed = 110_000_000 if frames <= 22 else 270_000_000
+                    frames = min(frames, chunk_frames + 2)
                     return (9.5 * frames * height * width + fixed) * model_management.dtype_size(dtype) * 1.03
 
                 self.memory_used_encode = lambda shape, dtype: estimate_encode_memory(shape[2], shape[3], shape[4], dtype)
@@ -1198,6 +1223,7 @@ class VAE:
                 do_tile = True
 
             if do_tile:
+                pixel_samples = None
                 comfy.model_management.soft_empty_cache()
                 dims = samples_in.ndim - 2
                 if dims == 1 or self.extra_1d_channel is not None:
@@ -1213,16 +1239,48 @@ class VAE:
                     tile = 256 // self.spacial_compression_decode()
                     overlap = tile // 4
                     if self.handles_tiling:
+                        memory_used = self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, tile, tile, None), self.vae_dtype)
+                        model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
                         pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
                     else:
-                        pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+                        # Reserve as much as an untiled decode could use (capped by what the device can provide), then size the tiles to fill that reservation:
+                        # shrink the temporal tile until one tile fits, then grow the spatial tile while it still fits.
+                        budget = min(memory_used, int(model_management.get_total_memory(self.device) * 0.8))
+                        model_management.load_models_gpu([self.patcher], memory_required=budget, force_full_load=self.disable_offload)
+                        tile_t = samples_in.shape[2]
+                        est = lambda tt, txy: self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, txy, txy, tt), self.vae_dtype)
+                        while tile_t > 2 and est(tile_t, tile) > budget:
+                            tile_t = -(-tile_t // 2)
+                        while tile * 2 <= max(samples_in.shape[3], samples_in.shape[4]) and est(tile_t, tile * 2) <= budget:
+                            tile *= 2
+                        overlap = tile // 4
+                        pixel_samples = self.decode_tiled_3d(samples_in, tile_t=tile_t, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
 
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
         return pixel_samples
 
+    def _tile_bounded_shape(self, shape, tile_x, tile_y, tile_t):
+        """Clamp a latent shape to one tile for memory estimates: peak memory of a tiled decode is per-tile. Only caller-provided tile dims are clamped."""
+        s = list(shape)
+        if len(s) == 5:
+            if tile_t is not None:
+                s[2] = min(s[2], tile_t)
+            if tile_y is not None:
+                s[3] = min(s[3], tile_y)
+            if tile_x is not None:
+                s[4] = min(s[4], tile_x)
+        elif len(s) == 4 and self.extra_1d_channel is None:
+            if tile_y is not None:
+                s[2] = min(s[2], tile_y)
+            if tile_x is not None:
+                s[3] = min(s[3], tile_x)
+        elif tile_x is not None:
+            s[-1] = min(s[-1], tile_x)
+        return tuple(s)
+
     def decode_tiled(self, samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
         self.throw_exception_if_invalid()
-        memory_used = self.memory_used_decode(samples.shape, self.vae_dtype) #TODO: calculate mem required for tile
+        memory_used = self.memory_used_decode(self._tile_bounded_shape(samples.shape, tile_x, tile_y, tile_t), self.vae_dtype)
         model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
         dims = samples.ndim - 2
         args = {}
@@ -1693,12 +1751,21 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.tokenizer = comfy.text_encoders.sa3.SAT5GemmaTokenizer
             tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
         elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B, TEModel.GEMMA_4_12B):
-            variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
-                       TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
-                       TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B,
-                       TEModel.GEMMA_4_12B: comfy.text_encoders.gemma4.Gemma4_12B}[te_model]
-            clip_target.clip = comfy.text_encoders.gemma4.gemma4_te(**llama_detect(clip_data), model_class=variant)
-            clip_target.tokenizer = variant.tokenizer
+            if te_model == TEModel.GEMMA_4_12B and "text_embedding_projection.video_aggregate_embed.weight" in clip_data[0]:
+                clip_target.clip = comfy.text_encoders.lt.ltxav_te(
+                    **llama_detect(clip_data),
+                    **comfy.text_encoders.lt.sd_detect(clip_data),
+                    text_encoder_model=comfy.text_encoders.gemma4.gemma4_text_encoder_model(comfy.text_encoders.gemma4.Gemma4_12B),
+                    text_encoder_key="gemma4",
+                )
+                clip_target.tokenizer = comfy.text_encoders.lt.ltxav_gemma4_tokenizer(comfy.text_encoders.gemma4.Gemma4_12B.tokenizer)
+            else:
+                variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
+                           TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
+                           TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B,
+                           TEModel.GEMMA_4_12B: comfy.text_encoders.gemma4.Gemma4_12B}[te_model]
+                clip_target.clip = comfy.text_encoders.gemma4.gemma4_te(**llama_detect(clip_data), model_class=variant)
+                clip_target.tokenizer = variant.tokenizer
             tokenizer_data["tokenizer_json"] = clip_data[0].get("tokenizer_json", None)
         elif te_model == TEModel.GEMMA_2_2B:
             if clip_type == CLIPType.PIXELDIT:
@@ -1866,9 +1933,30 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.clip = comfy.text_encoders.kandinsky5.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage
         elif clip_type == CLIPType.LTXV:
-            clip_target.clip = comfy.text_encoders.lt.ltxav_te(**llama_detect(clip_data), **comfy.text_encoders.lt.sd_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.lt.LTXAVGemmaTokenizer
-            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
+            te_models = [detect_te_model(sd) for sd in clip_data]
+            gemma4_models = {
+                TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
+                TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
+                TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B,
+                TEModel.GEMMA_4_12B: comfy.text_encoders.gemma4.Gemma4_12B,
+            }
+            gemma4_type = next((model for model in te_models if model in gemma4_models), None)
+            if gemma4_type is None:
+                clip_target.clip = comfy.text_encoders.lt.ltxav_te(**llama_detect(clip_data), **comfy.text_encoders.lt.sd_detect(clip_data))
+                clip_target.tokenizer = comfy.text_encoders.lt.LTXAVGemmaTokenizer
+                gemma_sd = clip_data[te_models.index(TEModel.GEMMA_3_12B)] if TEModel.GEMMA_3_12B in te_models else clip_data[0]
+                tokenizer_data["spiece_model"] = gemma_sd.get("spiece_model", None)
+            else:
+                variant = gemma4_models[gemma4_type]
+                clip_target.clip = comfy.text_encoders.lt.ltxav_te(
+                    **llama_detect(clip_data),
+                    **comfy.text_encoders.lt.sd_detect(clip_data),
+                    text_encoder_model=comfy.text_encoders.gemma4.gemma4_text_encoder_model(variant),
+                    text_encoder_key="gemma4",
+                )
+                clip_target.tokenizer = comfy.text_encoders.lt.ltxav_gemma4_tokenizer(variant.tokenizer)
+                gemma_sd = clip_data[te_models.index(gemma4_type)]
+                tokenizer_data["tokenizer_json"] = gemma_sd.get("tokenizer_json", None)
         elif clip_type == CLIPType.NEWBIE:
             clip_target.clip = comfy.text_encoders.newbie.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.newbie.NewBieTokenizer

comfy/text_encoders/bpe_tokenizer.py

@@ -0,0 +1,333 @@
+"""
+Pure-Python byte-level BPE tokenizer.
+Supports loading from HuggingFace tokenizer.json (LLaMA-style)
+and from Mistral tekken JSON blobs.
+No dependency on the `transformers`, `tokenizers`, or `regex` packages.
+"""
+import base64
+import json
+import os
+import re
+import unicodedata
+
+
+# This is also the default pattern used by the previous MistralConverter path.
+_LLAMA_PATTERN = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+_CONTRACTIONS = ("'re", "'ve", "'ll", "'s", "'t", "'m", "'d")
+
+
+def _is_letter(c):
+    return unicodedata.category(c)[0] == "L"
+
+
+def _is_number(c):
+    return unicodedata.category(c)[0] == "N"
+
+
+def _is_whitespace(c):
+    return c in " \t\n\r\v\f\x85\u2028\u2029" or unicodedata.category(c) == "Zs"
+
+
+def _split_llama(text):
+    pieces = []
+    i = 0
+    while i < len(text):
+        contraction = None
+        if text[i] == "'":
+            for suffix in _CONTRACTIONS:
+                if text[i:i + len(suffix)].casefold() == suffix:
+                    contraction = text[i:i + len(suffix)]
+                    break
+        if contraction is not None:
+            pieces.append(contraction)
+            i += len(contraction)
+            continue
+
+        j = i
+        if text[j] not in "\r\n" and not _is_letter(text[j]) and not _is_number(text[j]):
+            j += 1
+        if j < len(text) and _is_letter(text[j]):
+            j += 1
+            while j < len(text) and _is_letter(text[j]):
+                j += 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        if _is_number(text[i]):
+            j = i + 1
+            while j < len(text) and j - i < 3 and _is_number(text[j]):
+                j += 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        j = i
+        if text[j] == " ":
+            j += 1
+        punct_start = j
+        while j < len(text) and not _is_whitespace(text[j]) and not _is_letter(text[j]) and not _is_number(text[j]):
+            j += 1
+        if j > punct_start:
+            while j < len(text) and text[j] in "\r\n":
+                j += 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        if _is_whitespace(text[i]):
+            j = i + 1
+            while j < len(text) and _is_whitespace(text[j]):
+                j += 1
+            last_newline = max(text.rfind("\r", i, j), text.rfind("\n", i, j))
+            if last_newline >= i:
+                j = last_newline + 1
+            elif j < len(text) and j - i > 1:
+                j -= 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        pieces.append(text[i])
+        i += 1
+    return pieces
+
+
+def _make_split_pattern(pattern_str):
+    if pattern_str != _LLAMA_PATTERN:
+        raise ValueError(f"Unsupported tokenizer split pattern: {pattern_str}")
+    return _split_llama
+
+
+def _bytes_to_unicode():
+    bs = (list(range(ord("!"), ord("~") + 1))
+          + list(range(ord("¡"), ord("¬") + 1))
+          + list(range(ord("®"), ord("ÿ") + 1)))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+class BPETokenizer:
+    """Byte-level BPE tokenizer with optional BOS prepending."""
+
+    def __init__(self, vocab, merges_by_pair, special_token_ids, pattern_str,
+                 byte_encoder, byte_decoder, bos_id=None):
+        self._vocab = vocab                          # str -> int
+        self._inv_vocab = {v: k for k, v in vocab.items()}
+        self._merges = merges_by_pair                # (str, str) -> priority int
+        self._special_token_ids = special_token_ids  # str -> int
+        self._special_ids = set(special_token_ids.values())
+        self._byte_encoder = byte_encoder
+        self._byte_decoder = byte_decoder
+        self._bos_id = bos_id
+
+        self._split = _make_split_pattern(pattern_str)
+        sorted_specials = sorted(special_token_ids.keys(), key=len, reverse=True)
+        if sorted_specials:
+            self._special_split = re.compile(
+                '(' + '|'.join(re.escape(s) for s in sorted_specials) + ')'
+            )
+        else:
+            self._special_split = None
+
+    def _bpe_encode_piece(self, chars):
+        if len(chars) <= 1:
+            return chars
+        while True:
+            min_rank = float('inf')
+            best_pair = None
+            for i in range(len(chars) - 1):
+                r = self._merges.get((chars[i], chars[i + 1]), float('inf'))
+                if r < min_rank:
+                    min_rank = r
+                    best_pair = (chars[i], chars[i + 1])
+            if best_pair is None:
+                break
+            merged = best_pair[0] + best_pair[1]
+            new_chars = []
+            i = 0
+            while i < len(chars):
+                if i < len(chars) - 1 and chars[i] == best_pair[0] and chars[i + 1] == best_pair[1]:
+                    new_chars.append(merged)
+                    i += 2
+                else:
+                    new_chars.append(chars[i])
+                    i += 1
+            chars = new_chars
+            if len(chars) == 1:
+                break
+        return chars
+
+    def _encode_raw(self, text):
+        ids = []
+        parts = self._special_split.split(text) if self._special_split else [text]
+        for part in parts:
+            if not part:
+                continue
+            if part in self._special_token_ids:
+                ids.append(self._special_token_ids[part])
+            else:
+                for piece in self._split(part):
+                    byte_chars = [self._byte_encoder[b] for b in piece.encode('utf-8')]
+                    for tok in self._bpe_encode_piece(byte_chars):
+                        ids.append(self._vocab[tok])
+        return ids
+
+    def __call__(self, text):
+        ids = self._encode_raw(text)
+        if self._bos_id is not None:
+            ids = [self._bos_id] + ids
+        return {"input_ids": ids}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, token_ids, skip_special_tokens=True):
+        buf = bytearray()
+        for tid in token_ids:
+            s = self._inv_vocab.get(tid, '')
+            if tid in self._special_ids:
+                if not skip_special_tokens:
+                    buf.extend(s.encode('utf-8'))
+            else:
+                for c in s:
+                    buf.append(self._byte_decoder[c])
+        return buf.decode('utf-8', errors='replace')
+
+
+def _extract_pattern(pretok):
+    if pretok.get('type') == 'Sequence':
+        for sub in pretok.get('pretokenizers', []):
+            if sub.get('type') == 'Split':
+                pat = sub.get('pattern', {})
+                if 'Regex' in pat:
+                    return pat['Regex']
+    elif pretok.get('type') == 'Split':
+        pat = pretok.get('pattern', {})
+        if 'Regex' in pat:
+            return pat['Regex']
+    return None
+
+
+def _extract_bos_id(post_processor, special_token_ids):
+    if post_processor.get('type') == 'TemplateProcessing':
+        single = post_processor.get('single', [])
+        if single and 'SpecialToken' in single[0]:
+            bos_str = single[0]['SpecialToken']['id']
+            return special_token_ids.get(bos_str)
+    return None
+
+
+def from_tokenizer_json(path):
+    """Load a BPETokenizer from a directory containing tokenizer.json."""
+    tok_file = os.path.join(path, 'tokenizer.json')
+    with open(tok_file, encoding='utf-8') as f:
+        data = json.load(f)
+
+    vocab = dict(data['model']['vocab'])  # str -> int
+
+    merges_by_pair = {}
+    for i, merge_str in enumerate(data['model'].get('merges', [])):
+        a, b = merge_str.split(' ', 1)
+        if (a, b) not in merges_by_pair:
+            merges_by_pair[(a, b)] = i
+
+    special_token_ids = {}
+    for tok in data.get('added_tokens', []):
+        special_token_ids[tok['content']] = tok['id']
+        vocab[tok['content']] = tok['id']  # include in vocab for inv_vocab decode
+
+    pattern = _extract_pattern(data.get('pre_tokenizer', {}))
+    if pattern is None:
+        raise ValueError(f"Could not extract regex pattern from {tok_file}")
+
+    bos_id = _extract_bos_id(data.get('post_processor', {}), special_token_ids)
+
+    byte_encoder = _bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+
+    return BPETokenizer(vocab, merges_by_pair, special_token_ids, pattern,
+                        byte_encoder, byte_decoder, bos_id=bos_id)
+
+
+def from_tekken_json(data):
+    """Build a BPETokenizer from a Mistral tekken JSON blob (bytes or str)."""
+    mistral_vocab = json.loads(data)
+    config = mistral_vocab["config"]
+
+    byte_encoder = _bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+
+    def tbts(b):
+        return "".join(byte_encoder[ord(c)] for c in b.decode("latin-1"))
+
+    special_token_offset = config["default_num_special_tokens"]
+    max_vocab = config["default_vocab_size"] - special_token_offset
+
+    raw_vocab = {}
+    for w in mistral_vocab["vocab"]:
+        r = w["rank"]
+        if r >= max_vocab:
+            continue
+        raw_vocab[base64.b64decode(w["token_bytes"])] = r + special_token_offset
+
+    special_tokens_dict = {}
+    for w in mistral_vocab["special_tokens"]:
+        if "token_bytes" in w:
+            special_tokens_dict[base64.b64decode(w["token_bytes"])] = w["rank"]
+        else:
+            special_tokens_dict[w["token_str"]] = w["rank"]
+
+    all_special = list(special_tokens_dict.keys())
+    combined = dict(special_tokens_dict)
+    combined.update(raw_vocab)
+
+    bpe_vocab = {}
+    merge_triples = []
+    for token, rank in combined.items():
+        if token not in all_special:
+            bpe_vocab[tbts(token)] = rank
+            if len(token) == 1:
+                continue
+            local = []
+            for i in range(1, len(token)):
+                pl, pr = token[:i], token[i:]
+                if pl in combined and pr in combined and (pl + pr) in combined:
+                    local.append((pl, pr, rank))
+            local.sort(key=lambda x: (combined[x[0]], combined[x[1]]))
+            merge_triples.extend(local)
+        else:
+            tok_str = token.decode("utf-8", errors="replace") if isinstance(token, bytes) else token
+            bpe_vocab[tok_str] = rank
+
+    merge_triples.sort(key=lambda v: v[2])
+
+    merges_by_pair = {}
+    for i, (pl, pr, _) in enumerate(merge_triples):
+        pair = (tbts(pl), tbts(pr))
+        if pair not in merges_by_pair:
+            merges_by_pair[pair] = i
+
+    special_str_ids = {}
+    for tok in all_special:
+        tok_str = tok.decode("utf-8", errors="replace") if isinstance(tok, bytes) else tok
+        if tok_str in bpe_vocab:
+            special_str_ids[tok_str] = bpe_vocab[tok_str]
+
+    return BPETokenizer(bpe_vocab, merges_by_pair, special_str_ids, _LLAMA_PATTERN,
+                        byte_encoder, byte_decoder, bos_id=None)
+
+
+class LlamaTokenizerFast:
+    """Drop-in replacement for transformers.LlamaTokenizerFast (read-only use)."""
+
+    @staticmethod
+    def from_pretrained(path, **kwargs):
+        return from_tokenizer_json(path)

comfy/text_encoders/flux.py

@@ -3,11 +3,10 @@ import comfy.text_encoders.t5
 import comfy.text_encoders.sd3_clip
 import comfy.text_encoders.llama
 import comfy.model_management
-from transformers import T5TokenizerFast, LlamaTokenizerFast, Qwen2Tokenizer
+from transformers import T5TokenizerFast, Qwen2Tokenizer
+from .bpe_tokenizer import from_tekken_json
 import torch
 import os
-import json
-import base64
 
 class T5XXLTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None, tokenizer_data={}):
@@ -75,45 +74,13 @@ def flux_clip(dtype_t5=None, t5_quantization_metadata=None):
 def load_mistral_tokenizer(data):
     if torch.is_tensor(data):
         data = data.numpy().tobytes()
+    return {"tokenizer_object": from_tekken_json(data)}
 
-    try:
-        from transformers.integrations.mistral import MistralConverter
-    except ModuleNotFoundError:
-        from transformers.models.pixtral.convert_pixtral_weights_to_hf import MistralConverter
-
-    mistral_vocab = json.loads(data)
-
-    special_tokens = {}
-    vocab = {}
-
-    max_vocab = mistral_vocab["config"]["default_vocab_size"]
-    max_vocab -= len(mistral_vocab["special_tokens"])
-
-    for w in mistral_vocab["vocab"]:
-        r = w["rank"]
-        if r >= max_vocab:
-            continue
-
-        vocab[base64.b64decode(w["token_bytes"])] = r
-
-    for w in mistral_vocab["special_tokens"]:
-        if "token_bytes" in w:
-            special_tokens[base64.b64decode(w["token_bytes"])] = w["rank"]
-        else:
-            special_tokens[w["token_str"]] = w["rank"]
-
-    all_special = []
-    for v in special_tokens:
-        all_special.append(v)
-
-    special_tokens.update(vocab)
-    vocab = special_tokens
-    return {"tokenizer_object": MistralConverter(vocab=vocab, additional_special_tokens=all_special).converted(), "legacy": False}
 
 class MistralTokenizerClass:
     @staticmethod
-    def from_pretrained(path, **kwargs):
-        return LlamaTokenizerFast(**kwargs)
+    def from_pretrained(path, tokenizer_object=None, **kwargs):
+        return tokenizer_object
 
 class Mistral3Tokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None, embedding_size=5120, embedding_key='mistral3_24b', tokenizer_data={}):

comfy/text_encoders/gemma4.py

@@ -1443,14 +1443,14 @@ class Gemma4UnifiedTokenizer(Gemma4Tokenizer):
 class Gemma4Model(sd1_clip.SDClipModel):
     model_class = None
     def __init__(self, device="cpu", layer="all", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
+        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
+        if llama_quantization_metadata is not None:
+            model_options = model_options.copy()
+            model_options["quantization_metadata"] = llama_quantization_metadata
         self.dtypes = set()
         self.dtypes.add(dtype)
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=self.model_class, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-    def process_tokens(self, tokens, device):
-        embeds, _, _, _ = super().process_tokens(tokens, device)
-        return embeds
-
     def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0):
         if isinstance(tokens, dict):
             tokens = next(iter(tokens.values()))
@@ -1474,8 +1474,19 @@ class Gemma4Model(sd1_clip.SDClipModel):
         return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, initial_tokens=initial_token_ids[0], presence_penalty=presence_penalty, initial_input_ids=input_ids, embeds_info=embeds_info)
 
 
+def gemma4_clip_model(model_class):
+    return type('Gemma4Model_', (Gemma4Model,), {'model_class': model_class})
+
+
+def gemma4_text_encoder_model(model_class):
+    return type('Gemma4TextEncoderModel_', (Gemma4Model,), {
+        'model_class': model_class,
+        'process_tokens': sd1_clip.SDClipModel.process_tokens,
+    })
+
+
 def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=None):
-    clip_model = type('Gemma4Model_', (Gemma4Model,), {'model_class': model_class})
+    clip_model = gemma4_clip_model(model_class)
     class Gemma4TEModel_(sd1_clip.SD1ClipModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
             if llama_quantization_metadata is not None:

comfy/text_encoders/hunyuan_video.py

@@ -2,7 +2,7 @@ from comfy import sd1_clip
 import comfy.model_management
 import comfy.text_encoders.llama
 from .hunyuan_image import HunyuanImageTokenizer
-from transformers import LlamaTokenizerFast
+from .bpe_tokenizer import LlamaTokenizerFast
 import torch
 import os
 import numbers

comfy/text_encoders/lt.py

@@ -81,6 +81,17 @@ class LTXAVGemmaTokenizer(sd1_clip.SD1Tokenizer):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma3_12b", tokenizer=Gemma3_12BTokenizer)
 
 
+def ltxav_gemma4_tokenizer(tokenizer):
+    class LTXAVGemma4Tokenizer(tokenizer):
+        def __init__(self, embedding_directory=None, tokenizer_data={}):
+            super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
+            gemma_tokenizer = getattr(self, self.clip)
+            if gemma_tokenizer.min_length == 1:
+                gemma_tokenizer.min_length = 1024
+
+    return LTXAVGemma4Tokenizer
+
+
 class Gemma3_12BModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="all", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
         llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
@@ -97,10 +108,10 @@ class Gemma3_12BModel(sd1_clip.SDClipModel):
         return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, stop_tokens=[106], presence_penalty=presence_penalty)  # 106 is <end_of_turn>
 
 class DualLinearProjection(torch.nn.Module):
-    def __init__(self, in_dim, out_dim_video, out_dim_audio, dtype=None, device=None, operations=None):
+    def __init__(self, in_dim, out_dim_video, out_dim_audio, video_bias=True, audio_bias=True, dtype=None, device=None, operations=None):
         super().__init__()
-        self.audio_aggregate_embed = operations.Linear(in_dim, out_dim_audio, bias=True, dtype=dtype, device=device)
-        self.video_aggregate_embed = operations.Linear(in_dim, out_dim_video, bias=True, dtype=dtype, device=device)
+        self.audio_aggregate_embed = operations.Linear(in_dim, out_dim_audio, bias=audio_bias, dtype=dtype, device=device)
+        self.video_aggregate_embed = operations.Linear(in_dim, out_dim_video, bias=video_bias, dtype=dtype, device=device)
 
     def forward(self, x):
         source_dim = x.shape[-1]
@@ -112,22 +123,28 @@ class DualLinearProjection(torch.nn.Module):
         return torch.cat((video, audio), dim=-1)
 
 class LTXAVTEModel(torch.nn.Module):
-    def __init__(self, dtype_llama=None, device="cpu", dtype=None, text_projection_type="single_linear", model_options={}):
+    def __init__(self, dtype_llama=None, device="cpu", dtype=None, text_projection_type="single_linear", text_encoder_model=Gemma3_12BModel, text_encoder_key="gemma3_12b", video_projection_dim=3840, audio_projection_dim=2048, video_projection_bias=None, audio_projection_bias=True, model_options={}):
         super().__init__()
         self.dtypes = set()
         self.dtypes.add(dtype)
         self.compat_mode = False
         self.text_projection_type = text_projection_type
+        self.text_encoder_key = text_encoder_key
+        self.execution_device = None
 
-        self.gemma3_12b = Gemma3_12BModel(device=device, dtype=dtype_llama, model_options=model_options, layer="all", layer_idx=None)
+        self.gemma3_12b = text_encoder_model(device=device, dtype=dtype_llama, model_options=model_options, layer="all", layer_idx=None)
         self.dtypes.add(dtype_llama)
 
         operations = self.gemma3_12b.operations # TODO
+        text_encoder_config = self.gemma3_12b.transformer.model.config
+        projection_in_dim = text_encoder_config.hidden_size * (text_encoder_config.num_hidden_layers + 1)
+        if video_projection_bias is None:
+            video_projection_bias = self.text_projection_type == "dual_linear"
 
         if self.text_projection_type == "single_linear":
-            self.text_embedding_projection = operations.Linear(3840 * 49, 3840, bias=False, dtype=dtype, device=device)
+            self.text_embedding_projection = operations.Linear(projection_in_dim, video_projection_dim, bias=video_projection_bias, dtype=dtype, device=device)
         elif self.text_projection_type == "dual_linear":
-            self.text_embedding_projection = DualLinearProjection(3840 * 49, 4096, 2048, dtype=dtype, device=device, operations=operations)
+            self.text_embedding_projection = DualLinearProjection(projection_in_dim, video_projection_dim, audio_projection_dim, video_bias=video_projection_bias, audio_bias=audio_projection_bias, dtype=dtype, device=device, operations=operations)
 
 
     def enable_compat_mode(self):  # TODO: remove
@@ -161,7 +178,7 @@ class LTXAVTEModel(torch.nn.Module):
         self.execution_device = None
 
     def encode_token_weights(self, token_weight_pairs):
-        token_weight_pairs = token_weight_pairs["gemma3_12b"]
+        token_weight_pairs = token_weight_pairs[self.text_encoder_key]
 
         out, pooled, extra = self.gemma3_12b.encode_token_weights(token_weight_pairs)
         out = out[:, :, -torch.sum(extra["attention_mask"]).item():]
@@ -189,51 +206,54 @@ class LTXAVTEModel(torch.nn.Module):
         return out.to(device=out_device, dtype=torch.float), pooled, extra
 
     def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty):
-        return self.gemma3_12b.generate(tokens["gemma3_12b"], do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty)
+        return self.gemma3_12b.generate(tokens[self.text_encoder_key], do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty)
 
     def load_sd(self, sd):
-        if "model.layers.47.self_attn.q_norm.weight" in sd:
-            return self.gemma3_12b.load_sd(sd)
-        else:
-            sdo = comfy.utils.state_dict_prefix_replace(sd, {"text_embedding_projection.aggregate_embed.weight": "text_embedding_projection.weight", "text_embedding_projection.": "text_embedding_projection."}, filter_keys=True)
-            if len(sdo) == 0:
-                sdo = sd
+        missing_all = []
+        unexpected_all = []
 
-            missing_all = []
-            unexpected_all = []
+        if "model.layers.0.self_attn.q_norm.weight" in sd:
+            gemma_sd = {k: v for k, v in sd.items() if not k.startswith("text_embedding_projection.")}
+            missing, unexpected = self.gemma3_12b.load_sd(gemma_sd)
+            missing_all.extend(missing)
+            unexpected_all.extend(unexpected)
 
-            for prefix, component in [("text_embedding_projection.", self.text_embedding_projection)]:
-                component_sd = {k.replace(prefix, ""): v for k, v in sdo.items() if k.startswith(prefix)}
-                if component_sd:
-                    missing, unexpected = component.load_state_dict(component_sd, strict=False, assign=getattr(self, "can_assign_sd", False))
-                    missing_all.extend([f"{prefix}{k}" for k in missing])
-                    unexpected_all.extend([f"{prefix}{k}" for k in unexpected])
+        sdo = comfy.utils.state_dict_prefix_replace(sd, {"text_embedding_projection.aggregate_embed.": "text_embedding_projection.", "text_embedding_projection.": "text_embedding_projection."}, filter_keys=True)
+        if len(sdo) == 0:
+            sdo = sd
 
-            if "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.2.attn1.to_q.bias" not in sd:  # TODO: remove
-                ww = sd.get("model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.bias", None)
-                if ww is not None:
-                    if ww.shape[0] == 3840:
-                        self.enable_compat_mode()
-                        sdv = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.video_embeddings_connector.": ""}, filter_keys=True)
-                        self.video_embeddings_connector.load_state_dict(sdv, strict=False, assign=getattr(self, "can_assign_sd", False))
-                        sda = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.audio_embeddings_connector.": ""}, filter_keys=True)
-                        self.audio_embeddings_connector.load_state_dict(sda, strict=False, assign=getattr(self, "can_assign_sd", False))
+        for prefix, component in [("text_embedding_projection.", self.text_embedding_projection)]:
+            component_sd = {k.replace(prefix, ""): v for k, v in sdo.items() if k.startswith(prefix)}
+            if component_sd:
+                missing, unexpected = component.load_state_dict(component_sd, strict=False, assign=getattr(self, "can_assign_sd", False))
+                missing_all.extend([f"{prefix}{k}" for k in missing])
+                unexpected_all.extend([f"{prefix}{k}" for k in unexpected])
 
-            return (missing_all, unexpected_all)
+        if "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.2.attn1.to_q.bias" not in sd:  # TODO: remove
+            ww = sd.get("model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.bias", None)
+            if ww is not None:
+                if ww.shape[0] == 3840:
+                    self.enable_compat_mode()
+                    sdv = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.video_embeddings_connector.": ""}, filter_keys=True)
+                    self.video_embeddings_connector.load_state_dict(sdv, strict=False, assign=getattr(self, "can_assign_sd", False))
+                    sda = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.audio_embeddings_connector.": ""}, filter_keys=True)
+                    self.audio_embeddings_connector.load_state_dict(sda, strict=False, assign=getattr(self, "can_assign_sd", False))
+
+        return (missing_all, unexpected_all)
 
     def memory_estimation_function(self, token_weight_pairs, device=None):
         constant = 6.0
         if comfy.model_management.should_use_bf16(device):
             constant /= 2.0
 
-        token_weight_pairs = token_weight_pairs.get("gemma3_12b", [])
+        token_weight_pairs = token_weight_pairs.get(self.text_encoder_key, [])
         m = min([sum(1 for _ in itertools.takewhile(lambda x: x[0] == 0, sub)) for sub in token_weight_pairs])
 
         num_tokens = sum(map(lambda a: len(a), token_weight_pairs)) - m
         num_tokens = max(num_tokens, 642)
         return num_tokens * constant * 1024 * 1024
 
-def ltxav_te(dtype_llama=None, llama_quantization_metadata=None, text_projection_type="single_linear"):
+def ltxav_te(dtype_llama=None, llama_quantization_metadata=None, text_projection_type="single_linear", text_encoder_model=Gemma3_12BModel, text_encoder_key="gemma3_12b", video_projection_dim=3840, audio_projection_dim=2048, video_projection_bias=None, audio_projection_bias=True):
     class LTXAVTEModel_(LTXAVTEModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
             if llama_quantization_metadata is not None:
@@ -241,16 +261,29 @@ def ltxav_te(dtype_llama=None, llama_quantization_metadata=None, text_projection
                 model_options["llama_quantization_metadata"] = llama_quantization_metadata
             if dtype_llama is not None:
                 dtype = dtype_llama
-            super().__init__(dtype_llama=dtype_llama, device=device, dtype=dtype, text_projection_type=text_projection_type, model_options=model_options)
+            super().__init__(dtype_llama=dtype_llama, device=device, dtype=dtype, text_projection_type=text_projection_type, text_encoder_model=text_encoder_model, text_encoder_key=text_encoder_key, video_projection_dim=video_projection_dim, audio_projection_dim=audio_projection_dim, video_projection_bias=video_projection_bias, audio_projection_bias=audio_projection_bias, model_options=model_options)
     return LTXAVTEModel_
 
 
 def sd_detect(state_dict_list, prefix=""):
     for sd in state_dict_list:
-        if "{}text_embedding_projection.audio_aggregate_embed.bias".format(prefix) in sd:
-            return {"text_projection_type": "dual_linear"}
-        if "{}text_embedding_projection.weight".format(prefix) in sd or "{}text_embedding_projection.aggregate_embed.weight".format(prefix) in sd:
-            return {"text_projection_type": "single_linear"}
+        video_key = "{}text_embedding_projection.video_aggregate_embed.weight".format(prefix)
+        audio_key = "{}text_embedding_projection.audio_aggregate_embed.weight".format(prefix)
+        if video_key in sd and audio_key in sd:
+            return {
+                "text_projection_type": "dual_linear",
+                "video_projection_dim": sd[video_key].shape[0],
+                "audio_projection_dim": sd[audio_key].shape[0],
+                "video_projection_bias": "{}text_embedding_projection.video_aggregate_embed.bias".format(prefix) in sd,
+                "audio_projection_bias": "{}text_embedding_projection.audio_aggregate_embed.bias".format(prefix) in sd,
+            }
+        for key in ("{}text_embedding_projection.weight".format(prefix), "{}text_embedding_projection.aggregate_embed.weight".format(prefix)):
+            if key in sd:
+                return {
+                    "text_projection_type": "single_linear",
+                    "video_projection_dim": sd[key].shape[0],
+                    "video_projection_bias": key.removesuffix("weight") + "bias" in sd,
+                }
     return {}
 
 

comfy/text_encoders/lumina2.py

@@ -49,10 +49,6 @@ class Gemma3_4B_Vision_Model(sd1_clip.SDClipModel):
 
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B_Vision, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-    def process_tokens(self, tokens, device):
-        embeds, _, _, _ = super().process_tokens(tokens, device)
-        return embeds
-
 class LuminaModel(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, model_options={}, name="gemma2_2b", clip_model=Gemma2_2BModel):
         super().__init__(device=device, dtype=dtype, name=name, clip_model=clip_model, model_options=model_options)

comfy_api_nodes/apis/grok.py

@@ -9,6 +9,7 @@ class ImageGenerationRequest(BaseModel):
     seed: int = Field(...)
     response_format: str = Field("url")
     resolution: str = Field(...)
+    quality: str | None = Field(None)
 
 
 class InputUrlObject(BaseModel):
@@ -28,6 +29,7 @@ class ImageEditRequest(BaseModel):
     seed: int = Field(...)
     response_format: str = Field("url")
     aspect_ratio: str | None = Field(...)
+    quality: str | None = Field(None)
 
 
 class VideoGenerationRequest(BaseModel):

comfy_api_nodes/apis/qwen.py

@@ -0,0 +1,46 @@
+from pydantic import BaseModel, Field
+
+
+class QwenImageContentItem(BaseModel):
+    image: str | None = Field(None)
+    text: str | None = Field(None)
+
+
+class QwenImageMessage(BaseModel):
+    role: str = Field("user")
+    content: list[QwenImageContentItem] = Field(...)
+
+
+class QwenImageInputField(BaseModel):
+    messages: list[QwenImageMessage] = Field(...)
+
+
+class QwenImageParametersField(BaseModel):
+    size: str | None = Field(None, description="Output resolution as 'width*height'; omit for the model default.")
+    n: int = Field(1, ge=1, le=6)
+    seed: int = Field(..., ge=0, le=2147483647)
+    prompt_extend: bool = Field(True)
+    watermark: bool = Field(False)
+    negative_prompt: str | None = Field(None)
+
+
+class QwenImageGenerationRequest(BaseModel):
+    model: str = Field(...)
+    input: QwenImageInputField = Field(...)
+    parameters: QwenImageParametersField = Field(...)
+
+
+class QwenImageChoice(BaseModel):
+    finish_reason: str | None = Field(None)
+    message: QwenImageMessage | None = Field(None)
+
+
+class QwenImageOutputField(BaseModel):
+    choices: list[QwenImageChoice] = Field(default_factory=list)
+
+
+class QwenImageGenerationResponse(BaseModel):
+    output: QwenImageOutputField | None = Field(None)
+    request_id: str = Field(...)
+    code: str | None = Field(None, description="Error code for the failed request.")
+    message: str | None = Field(None, description="Details about the failed request.")

comfy_api_nodes/nodes_grok.py

@@ -36,6 +36,26 @@ _GROK_VIDEO_MODEL_API_IDS = {
     "grok-imagine-video-1.5": "grok-imagine-video-1.5",
 }
 
+_GROK_IMAGE_MODEL_API_IDS = {
+    "grok-imagine-image-2.0": "grok-imagine-image-2.0",
+}
+
+_GROK_IMAGE_QUALITY_MODELS = {"grok-imagine-image-2.0"}
+
+_GROK_IMAGE_QUALITY_OPTIONS = ["medium", "low"]
+
+_GROK_IMAGE_EDIT_MAX_IMAGES = {
+    "grok-imagine-image-2.0": 3,
+    "grok-imagine-image-pro": 1,
+    "grok-imagine-image-quality": 3,
+    "grok-imagine-image": 3,
+}
+
+_GROK_IMAGE_EDIT_ASPECT_RATIO_NEEDS_MULTIPLE = {
+    "grok-imagine-image-quality",
+    "grok-imagine-image",
+}
+
 _GROK_VOICE_OPTIONS = [
     "none",
     "ara",
@@ -132,6 +152,7 @@ class GrokImageNode(IO.ComfyNode):
                 IO.Combo.Input(
                     "model",
                     options=[
+                        "grok-imagine-image-2.0",
                         "grok-imagine-image-quality",
                         "grok-imagine-image-pro",
                         "grok-imagine-image",
@@ -181,6 +202,12 @@ class GrokImageNode(IO.ComfyNode):
                     "actual results are nondeterministic regardless of seed.",
                 ),
                 IO.Combo.Input("resolution", options=["1K", "2K"], optional=True),
+                IO.Combo.Input(
+                    "quality",
+                    options=_GROK_IMAGE_QUALITY_OPTIONS,
+                    optional=True,
+                    tooltip="Quality level, supported only by the grok-imagine-image-2.0 model.",
+                ),
             ],
             outputs=[
                 IO.Image.Output(),
@@ -192,12 +219,15 @@ class GrokImageNode(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution", "quality"]),
                 expr="""
                 (
-                  $rate := widgets.model = "grok-imagine-image-quality"
-                    ? (widgets.resolution = "1k" ? 0.05 : 0.07)
-                    : ($contains(widgets.model, "pro") ? 0.07 : 0.02);
+                  $is1k := widgets.resolution = "1k";
+                  $rate := widgets.model = "grok-imagine-image-2.0"
+                    ? (widgets.quality = "low" ? ($is1k ? 0.04 : 0.06) : ($is1k ? 0.06 : 0.08))
+                    : (widgets.model = "grok-imagine-image-quality"
+                        ? ($is1k ? 0.05 : 0.07)
+                        : ($contains(widgets.model, "pro") ? 0.07 : 0.02));
                   {"type":"usd","usd": $rate * widgets.number_of_images}
                 )
                 """,
@@ -213,18 +243,20 @@ class GrokImageNode(IO.ComfyNode):
         number_of_images: int,
         seed: int,
         resolution: str = "1K",
+        quality: str = "medium",
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/xai/v1/images/generations", method="POST"),
             data=ImageGenerationRequest(
-                model=model,
+                model=_GROK_IMAGE_MODEL_API_IDS.get(model, model),
                 prompt=prompt,
                 aspect_ratio=aspect_ratio,
                 n=number_of_images,
                 seed=seed,
                 resolution=resolution.lower(),
+                quality=quality if model in _GROK_IMAGE_QUALITY_MODELS else None,
             ),
             response_model=ImageGenerationResponse,
         )
@@ -255,7 +287,9 @@ _GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS = [
 ]
 
 
-def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: bool):
+def _grok_image_edit_model_inputs(
+    *, max_ref_images: int, with_aspect_ratio: bool, with_quality: bool = False, aspect_ratio_needs_multiple: bool = True
+):
     inputs = [
         IO.Autogrow.Input(
             "images",
@@ -281,12 +315,18 @@ def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: boo
             display_mode=IO.NumberDisplay.number,
         ),
     ]
+    if with_quality:
+        inputs.append(IO.Combo.Input("quality", options=_GROK_IMAGE_QUALITY_OPTIONS))
     if with_aspect_ratio:
         inputs.append(
             IO.Combo.Input(
                 "aspect_ratio",
                 options=_GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS,
-                tooltip="Only allowed when multiple images are connected.",
+                tooltip=(
+                    "Only allowed when multiple images are connected."
+                    if aspect_ratio_needs_multiple
+                    else "Aspect ratio of the edited image."
+                ),
             )
         )
     return inputs
@@ -451,6 +491,15 @@ class GrokImageEditNodeV2(IO.ComfyNode):
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "grok-imagine-image-2.0",
+                            _grok_image_edit_model_inputs(
+                                max_ref_images=3,
+                                with_aspect_ratio=True,
+                                with_quality=True,
+                                aspect_ratio_needs_multiple=False,
+                            ),
+                        ),
                         IO.DynamicCombo.Option(
                             "grok-imagine-image-quality",
                             _grok_image_edit_model_inputs(max_ref_images=3, with_aspect_ratio=True),
@@ -488,18 +537,23 @@ class GrokImageEditNodeV2(IO.ComfyNode):
             is_api_node=True,
             price_badge=IO.PriceBadge(
                 depends_on=IO.PriceBadgeDepends(
-                    widgets=["model", "model.resolution", "model.number_of_images"],
+                    widgets=["model", "model.resolution", "model.number_of_images", "model.quality"],
                 ),
                 expr="""
                 (
-                  $isQualityModel := widgets.model = "grok-imagine-image-quality";
+                  $is20 := widgets.model = "grok-imagine-image-2.0";
                   $isPro := $contains(widgets.model, "pro");
                   $res := $lookup(widgets, "model.resolution");
                   $n := $lookup(widgets, "model.number_of_images");
-                  $rate := $isQualityModel
-                    ? ($res = "1k" ? 0.05 : 0.07)
-                    : ($isPro ? 0.07 : 0.02);
-                  $base := $isQualityModel ? 0.01 : 0.002;
+                  $is1k := $res = "1k";
+                  $rate := $is20
+                    ? ($lookup(widgets, "model.quality") = "low"
+                        ? ($is1k ? 0.04 : 0.06)
+                        : ($is1k ? 0.06 : 0.08))
+                    : (widgets.model = "grok-imagine-image-quality"
+                        ? ($is1k ? 0.05 : 0.07)
+                        : ($isPro ? 0.07 : 0.02));
+                  $base := ($is20 or widgets.model = "grok-imagine-image-quality") ? 0.01 : 0.002;
                   $output := $rate * $n;
                   $isPro
                     ? {"type":"usd","usd": $base + $output}
@@ -525,13 +579,15 @@ class GrokImageEditNodeV2(IO.ComfyNode):
 
         image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None]
         n_images = sum(get_number_of_images(t) for t in image_tensors)
+        max_images = _GROK_IMAGE_EDIT_MAX_IMAGES.get(model_id, 3)
         if n_images < 1:
             raise ValueError("At least one image is required for editing.")
-        if model_id == "grok-imagine-image-pro" and n_images > 1:
-            raise ValueError("The pro model supports only 1 input image.")
-        if model_id != "grok-imagine-image-pro" and n_images > 3:
-            raise ValueError("A maximum of 3 input images is supported.")
-        if aspect_ratio != "auto" and n_images == 1:
+        if n_images > max_images:
+            raise ValueError(
+                f"The {model_id} model supports at most {max_images} input "
+                f"image{'s' if max_images > 1 else ''}; {n_images} are connected."
+            )
+        if aspect_ratio != "auto" and model_id in _GROK_IMAGE_EDIT_ASPECT_RATIO_NEEDS_MULTIPLE and n_images == 1:
             raise ValueError(
                 "Custom aspect ratio is only allowed when multiple images are connected to the image input."
             )
@@ -547,7 +603,7 @@ class GrokImageEditNodeV2(IO.ComfyNode):
             cls,
             ApiEndpoint(path="/proxy/xai/v1/images/edits", method="POST"),
             data=ImageEditRequest(
-                model=model_id,
+                model=_GROK_IMAGE_MODEL_API_IDS.get(model_id, model_id),
                 images=[
                     InputUrlObject(url=f"data:image/png;base64,{tensor_to_base64_string(i)}") for i in flat_tensors
                 ],
@@ -556,6 +612,7 @@ class GrokImageEditNodeV2(IO.ComfyNode):
                 n=number_of_images,
                 seed=seed,
                 aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio,
+                quality=model.get("quality") if model_id in _GROK_IMAGE_QUALITY_MODELS else None,
             ),
             response_model=ImageGenerationResponse,
         )

comfy_api_nodes/nodes_ltxv.py

@@ -6,8 +6,12 @@ from typing_extensions import override
 from comfy_api.latest import IO, ComfyExtension, Input, InputImpl
 from comfy_api_nodes.util import (
     ApiEndpoint,
+    download_url_to_video_output,
     get_number_of_images,
+    poll_op,
+    sync_op,
     sync_op_raw,
+    upload_audio_to_comfyapi,
     upload_images_to_comfyapi,
     validate_string,
 )
@@ -17,6 +21,11 @@ MODELS_MAP = {
     "LTX-2 (Fast)": "ltx-2-fast",
 }
 
+V25_MODELS_MAP = {
+    "LTX-2.5 (Fast)": "ltx-2-5-fast",
+    "LTX-2.5 (Pro)": "ltx-2-5-pro",
+}
+
 
 class ExecuteTaskRequest(BaseModel):
     prompt: str = Field(...)
@@ -26,6 +35,48 @@ class ExecuteTaskRequest(BaseModel):
     fps: int | None = Field(25)
     generate_audio: bool | None = Field(True)
     image_uri: str | None = Field(None)
+    last_frame_uri: str | None = Field(None)
+
+
+class AudioToVideoRequest(BaseModel):
+    prompt: str = Field(...)
+    model: str = Field(...)
+    resolution: str = Field(...)
+    audio_uri: str = Field(...)
+    image_uri: str | None = Field(None)
+
+
+class Ltx25SubmitResponse(BaseModel):
+    id: str = Field(...)
+
+
+class Ltx25JobResult(BaseModel):
+    video_url: str | None = Field(None)
+
+
+class Ltx25JobStatusResponse(BaseModel):
+    id: str = Field(...)
+    status: str = Field(...)
+    result: Ltx25JobResult | None = Field(None)
+
+
+async def _v25_submit_and_poll(cls: type[IO.ComfyNode], route: str, data: BaseModel) -> IO.NodeOutput:
+    submit = await sync_op(
+        cls,
+        ApiEndpoint(f"/proxy/ltx/v2/{route}", "POST"),
+        response_model=Ltx25SubmitResponse,
+        data=data,
+        max_retries=1,
+    )
+    job = await poll_op(
+        cls,
+        ApiEndpoint(f"/proxy/ltx/v2/{route}/{submit.id}"),
+        response_model=Ltx25JobStatusResponse,
+        status_extractor=lambda r: r.status,
+    )
+    if not job.result or not job.result.video_url:
+        raise RuntimeError(f"LTX job {job.id} completed without a video URL.")
+    return IO.NodeOutput(await download_url_to_video_output(job.result.video_url, cls=cls))
 
 
 PRICE_BADGE = IO.PriceBadge(
@@ -43,6 +94,128 @@ PRICE_BADGE = IO.PriceBadge(
     """,
 )
 
+V25_PRICE_BADGE = IO.PriceBadge(
+    depends_on=IO.PriceBadgeDepends(widgets=["model", "model.duration", "model.resolution"]),
+    expr="""
+    (
+      $prices := {
+        "ltx-2.5 (fast)": {
+          "1280x720":0.1287,"720x1280":0.1287,
+          "1920x1080":0.1859,"1080x1920":0.1859,
+          "2560x1440":0.2717,"1440x2560":0.2717,
+          "3840x2160":0.429,"2160x3840":0.429
+        },
+        "ltx-2.5 (pro)": {
+          "1280x720":0.1716,"720x1280":0.1716,
+          "1920x1080":0.2431,"1080x1920":0.2431
+        }
+      };
+      $model := $lookup(widgets, "model");
+      $table := $type($model) = "string" ? $lookup($prices, $model) : undefined;
+      $res := $lookup(widgets, "model.resolution");
+      $pps := $type($table) = "object" and $type($res) = "string" ? $lookup($table, $res) : undefined;
+      $durRaw := $lookup(widgets, "model.duration");
+      $dur := $type($durRaw) in ["string", "number"] ? $number($durRaw) : undefined;
+      $type($pps) = "number" and $type($dur) = "number"
+        ? {"type":"usd","usd": $pps * $dur}
+        : undefined
+    )
+    """,
+)
+
+V25_A2V_PRICE_BADGE = IO.PriceBadge(
+    depends_on=IO.PriceBadgeDepends(widgets=["model"]),
+    expr="""
+    (
+      $rates := {"ltx-2.5 (fast)":0.1859, "ltx-2.5 (pro)":0.2431};
+      $model := $lookup(widgets, "model");
+      $rate := $type($model) = "string" ? $lookup($rates, $model) : undefined;
+      $type($rate) = "number"
+        ? {"type":"usd","usd": $rate, "format":{"suffix":"/second"}}
+        : undefined
+    )
+    """,
+)
+
+
+def _v25_generation_inputs(
+    durations: list[str], resolutions: list[str], fps_options: list[str], tooltip: str | None
+) -> list:
+    return [
+        IO.Combo.Input(
+            "duration",
+            options=durations,
+            default="8",
+            tooltip=tooltip,
+        ),
+        IO.Combo.Input(
+            "resolution",
+            options=resolutions,
+            default="1920x1080",
+        ),
+        IO.Combo.Input("fps", options=fps_options, default="25"),
+        IO.Boolean.Input(
+            "generate_audio",
+            default=True,
+            tooltip="When true, the generated video will include AI-generated audio matching the scene.",
+            advanced=True,
+        ),
+    ]
+
+
+def _v25_model_combo() -> IO.DynamicCombo.Input:
+    return IO.DynamicCombo.Input(
+        "model",
+        options=[
+            IO.DynamicCombo.Option(
+                "LTX-2.5 (Fast)",
+                _v25_generation_inputs(
+                    ["2", "3", "4", "5", "6", "8", "10", "12", "14", "16", "18", "20"],
+                    [
+                        "1280x720",
+                        "720x1280",
+                        "1920x1080",
+                        "1080x1920",
+                        "2560x1440",
+                        "1440x2560",
+                        "3840x2160",
+                        "2160x3840",
+                    ],
+                    ["24", "25", "48", "50"],
+                    "Video duration in seconds. Durations over 10s require a 720p/1080p resolution and 24/25 FPS.",
+                ),
+            ),
+            IO.DynamicCombo.Option(
+                "LTX-2.5 (Pro)",
+                _v25_generation_inputs(
+                    ["2", "3", "4", "5", "6", "8", "10"],
+                    ["1280x720", "720x1280", "1920x1080", "1080x1920"],
+                    ["24", "25", "50"],
+                    "Video duration in seconds.",
+                ),
+            ),
+        ],
+    )
+
+
+def _v25_seed_input() -> IO.Int.Input:
+    return IO.Int.Input(
+        "seed",
+        default=42,
+        min=0,
+        max=0xFFFFFFFF,
+        control_after_generate=True,
+        tooltip="Seed to determine if node should re-run; "
+        "actual results are nondeterministic regardless of seed.",
+    )
+
+
+def _v25_validate_settings(model: dict) -> None:
+    if int(model["duration"]) > 10 and (
+        int(model["fps"]) > 25 or model["resolution"] in ("2560x1440", "1440x2560", "3840x2160", "2160x3840")
+    ):
+        raise ValueError("Durations over 10s require a 720p or 1080p resolution and 24/25 FPS.")
+
 
 class TextToVideoNode(IO.ComfyNode):
     @classmethod
@@ -86,6 +259,7 @@ class TextToVideoNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=PRICE_BADGE,
         )
 
@@ -164,6 +338,7 @@ class ImageToVideoNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=PRICE_BADGE,
         )
 
@@ -203,12 +378,217 @@ class ImageToVideoNode(IO.ComfyNode):
         return IO.NodeOutput(InputImpl.VideoFromFile(BytesIO(response)))
 
 
+class Ltx25TextToVideoNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="LtxApi25TextToVideo",
+            display_name="LTX 2.5 Text To Video",
+            category="partner/video/LTXV",
+            description="Professional-quality videos with customizable duration and resolution.",
+            inputs=[
+                _v25_model_combo(),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                ),
+                _v25_seed_input(),
+            ],
+            outputs=[
+                IO.Video.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=V25_PRICE_BADGE,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        prompt: str,
+        seed: int = 42,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1, max_length=10000)
+        _v25_validate_settings(model)
+        return await _v25_submit_and_poll(
+            cls,
+            "text-to-video",
+            ExecuteTaskRequest(
+                prompt=prompt,
+                model=V25_MODELS_MAP[model["model"]],
+                duration=int(model["duration"]),
+                resolution=model["resolution"],
+                fps=int(model["fps"]),
+                generate_audio=model["generate_audio"],
+            ),
+        )
+
+
+class Ltx25ImageToVideoNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="LtxApi25ImageToVideo",
+            display_name="LTX 2.5 Image To Video",
+            category="partner/video/LTXV",
+            description="Professional-quality videos with customizable duration and resolution based on start image.",
+            inputs=[
+                IO.Image.Input("image", tooltip="First frame to be used for the video."),
+                _v25_model_combo(),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                ),
+                _v25_seed_input(),
+                IO.Image.Input(
+                    "last_frame",
+                    optional=True,
+                    tooltip="Last frame to be used for the video.",
+                ),
+            ],
+            outputs=[
+                IO.Video.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=V25_PRICE_BADGE,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        model: dict,
+        prompt: str,
+        seed: int = 42,
+        last_frame: Input.Image | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1, max_length=10000)
+        _v25_validate_settings(model)
+        if get_number_of_images(image) != 1:
+            raise ValueError("Currently only one input image is supported.")
+        last_frame_uri = None
+        if last_frame is not None:
+            if get_number_of_images(last_frame) != 1:
+                raise ValueError("Currently only one last frame image is supported.")
+            last_frame_uri = (await upload_images_to_comfyapi(cls, last_frame, max_images=1, mime_type="image/png"))[0]
+        return await _v25_submit_and_poll(
+            cls,
+            "image-to-video",
+            ExecuteTaskRequest(
+                image_uri=(await upload_images_to_comfyapi(cls, image, max_images=1, mime_type="image/png"))[0],
+                last_frame_uri=last_frame_uri,
+                prompt=prompt,
+                model=V25_MODELS_MAP[model["model"]],
+                duration=int(model["duration"]),
+                resolution=model["resolution"],
+                fps=int(model["fps"]),
+                generate_audio=model["generate_audio"],
+            ),
+        )
+
+
+class Ltx25AudioToVideoNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="LtxApi25AudioToVideo",
+            display_name="LTX 2.5 Audio To Video",
+            category="partner/video/LTXV",
+            description="Generate a video driven by an audio track, with an optional first frame image.",
+            inputs=[
+                IO.Audio.Input(
+                    "audio",
+                    tooltip="Audio track driving the video. Its length (2-20 seconds) sets the video duration.",
+                ),
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "LTX-2.5 (Fast)",
+                            [IO.Combo.Input("resolution", options=["1920x1080", "1080x1920"])],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "LTX-2.5 (Pro)",
+                            [IO.Combo.Input("resolution", options=["1920x1080", "1080x1920"])],
+                        ),
+                    ],
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                ),
+                _v25_seed_input(),
+                IO.Image.Input(
+                    "image",
+                    optional=True,
+                    tooltip="Optional first frame to be used for the video.",
+                ),
+            ],
+            outputs=[
+                IO.Video.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=V25_A2V_PRICE_BADGE,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        audio: Input.Audio,
+        model: dict,
+        prompt: str,
+        seed: int = 42,
+        image: Input.Image | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1, max_length=10000)
+        audio_duration = audio["waveform"].shape[-1] / audio["sample_rate"]
+        if not 2 <= audio_duration <= 20:
+            raise ValueError(f"Audio duration must be between 2 and 20 seconds, got {audio_duration:.1f}s.")
+        image_uri = None
+        if image is not None:
+            if get_number_of_images(image) != 1:
+                raise ValueError("Currently only one input image is supported.")
+            image_uri = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type="image/png"))[0]
+        return await _v25_submit_and_poll(
+            cls,
+            "audio-to-video",
+            AudioToVideoRequest(
+                prompt=prompt,
+                model=V25_MODELS_MAP[model["model"]],
+                resolution=model["resolution"],
+                audio_uri=await upload_audio_to_comfyapi(cls, audio),
+                image_uri=image_uri,
+            ),
+        )
+
+
 class LtxvApiExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
         return [
             TextToVideoNode,
             ImageToVideoNode,
+            Ltx25TextToVideoNode,
+            Ltx25ImageToVideoNode,
+            Ltx25AudioToVideoNode,
         ]
 
 

comfy_api_nodes/nodes_qwen.py

@@ -0,0 +1,442 @@
+import math
+import re
+
+import torch
+from typing_extensions import override
+
+from comfy_api.latest import IO, ComfyExtension
+from comfy_api_nodes.apis.qwen import (
+    QwenImageContentItem,
+    QwenImageGenerationRequest,
+    QwenImageGenerationResponse,
+    QwenImageInputField,
+    QwenImageMessage,
+    QwenImageParametersField,
+)
+from comfy_api_nodes.util import (
+    ApiEndpoint,
+    download_url_to_image_tensor,
+    sync_op,
+    tensor_to_base64_string,
+    validate_string,
+)
+
+GENERATION_PATH = "/proxy/qwen/api/v1/services/aigc/multimodal-generation/generation"
+QWEN_IMAGE_MODELS = ["qwen-image-3.0-pro", "qwen-image-3.0"]
+MIN_AREA = 262144  # 512*512
+MAX_AREA = 6553600  # 2560*2560
+MAX_ASPECT = 8  # the API allows aspect ratios from 1:8 to 8:1
+MAX_INPUT_BYTES = 10 * 1024 * 1024  # the API rejects decoded input images over 10MB
+
+_IMAGE_REF_RE = re.compile(r"@image(?P<idx>\d*)(?!\w)", re.IGNORECASE | re.ASCII)
+
+
+def _resolve_image_refs(prompt: str, total_images: int) -> str:
+    """Rewrite @Image1-style references (shared partner-node syntax, 1-based; an unnumbered
+    @image means the first image) into the plain 'Image N' wording the model resolves
+    natively. A tag counts only at a word boundary or right after a previous tag, so
+    adjacent tags like '@Image1@Image2' all resolve while addresses like user@image1.com
+    pass through untouched."""
+    parts = []
+    pos = 0
+    prev_end = -1
+    for match in _IMAGE_REF_RE.finditer(prompt):
+        start = match.start()
+        if start > 0 and start != prev_end and (prompt[start - 1].isalnum() or prompt[start - 1] == "_"):
+            continue
+        idx = int(match.group("idx") or 1)
+        if not 1 <= idx <= total_images:
+            raise ValueError(
+                f"The prompt references @Image{idx}, but only {total_images} reference images "
+                f"are connected (a batched input counts once per image)."
+            )
+        parts.append(prompt[pos:start])
+        parts.append(f"Image {idx}")
+        pos = match.end()
+        prev_end = match.end()
+    parts.append(prompt[pos:])
+    return "".join(parts)
+
+
+def _validate_size(width: int, height: int) -> None:
+    if not MIN_AREA <= width * height <= MAX_AREA:
+        raise ValueError(
+            f"Image area must be between {MIN_AREA} (512x512) and {MAX_AREA} (2560x2560) pixels; "
+            f"got {width}x{height} = {width * height}."
+        )
+    if width > MAX_ASPECT * height or height > MAX_ASPECT * width:
+        raise ValueError(f"Aspect ratio must be between 1:8 and 8:1; got {width}x{height}.")
+
+
+def _fit_to_size(width: int, height: int) -> tuple[int, int]:
+    """Scale dimensions into the supported pixel area and 1:8..8:1 aspect range, preserving
+    the aspect ratio where possible."""
+    if width > MAX_ASPECT * height:
+        height = math.ceil(width / MAX_ASPECT)
+    elif height > MAX_ASPECT * width:
+        width = math.ceil(height / MAX_ASPECT)
+    area = width * height
+    if area < MIN_AREA:
+        scale = math.sqrt(MIN_AREA / area)
+        width, height = math.ceil(width * scale), math.ceil(height * scale)
+    elif area > MAX_AREA:
+        scale = math.sqrt(MAX_AREA / area)
+        width, height = math.floor(width * scale), math.floor(height * scale)
+    # rounding can push the ratio a hair past the limit; trimming only ever shrinks the area
+    return min(width, MAX_ASPECT * height), min(height, MAX_ASPECT * width)
+
+
+def _image_data_uri(image: torch.Tensor) -> str:
+    """PNG data URI of an RGB view of the image, downscaled to <=2048x2048; falls back to
+    JPEG when the PNG exceeds the API's decoded-size cap (e.g. noisy, incompressible images)."""
+    image = image[..., :3]
+    b64 = tensor_to_base64_string(image, total_pixels=2048 * 2048)
+    if len(b64) * 3 > MAX_INPUT_BYTES * 4:
+        return "data:image/jpeg;base64," + tensor_to_base64_string(
+            image, total_pixels=2048 * 2048, mime_type="image/jpeg"
+        )
+    return "data:image/png;base64," + b64
+
+
+async def _download_result_images(response: QwenImageGenerationResponse) -> torch.Tensor:
+    if not response.output:
+        raise Exception(f"An unknown error occurred: {response.code} - {response.message}")
+    urls = [
+        item.image
+        for choice in response.output.choices
+        if choice.message
+        for item in choice.message.content
+        if item.image
+    ]
+    if not urls:
+        raise Exception(f"The response contains no images: {response.code} - {response.message}")
+    return torch.cat([await download_url_to_image_tensor(url) for url in urls])
+
+
+def _size_inputs() -> list[IO.Int.Input]:
+    return [
+        IO.Int.Input(
+            "width",
+            default=1024,
+            min=256,
+            max=2560,
+            step=16,
+            tooltip="The total pixel area must be between 512x512 and 2560x2560; "
+            "any aspect ratio within that area works.",
+        ),
+        IO.Int.Input(
+            "height",
+            default=1024,
+            min=256,
+            max=2560,
+            step=16,
+            tooltip="The total pixel area must be between 512x512 and 2560x2560; "
+            "any aspect ratio within that area works.",
+        ),
+    ]
+
+
+def _t2i_model_option(model_id: str) -> IO.DynamicCombo.Option:
+    return IO.DynamicCombo.Option(
+        model_id,
+        [
+            IO.String.Input(
+                "prompt",
+                multiline=True,
+                default="",
+                tooltip="Prompt describing the image. Supports English and Chinese.",
+            ),
+            IO.String.Input(
+                "negative_prompt",
+                multiline=True,
+                default="",
+                tooltip="Negative prompt describing what to avoid.",
+            ),
+            *_size_inputs(),
+        ],
+    )
+
+
+def _edit_model_option(model_id: str) -> IO.DynamicCombo.Option:
+    return IO.DynamicCombo.Option(
+        model_id,
+        [
+            IO.Autogrow.Input(
+                "images",
+                template=IO.Autogrow.TemplateNames(
+                    IO.Image.Input("image"),
+                    names=["image_1", "image_2", "image_3"],
+                    min=1,
+                ),
+                tooltip="1-3 reference images. Refer to them in the prompt as @Image1, @Image2, "
+                "@Image3, numbered in input order; a batched input counts once per image.",
+            ),
+            IO.String.Input(
+                "prompt",
+                multiline=True,
+                default="",
+                tooltip="Editing instructions. Supports English and Chinese, "
+                "and @Image1-style references to the input images.",
+            ),
+            IO.String.Input(
+                "negative_prompt",
+                multiline=True,
+                default="",
+                tooltip="Negative prompt describing what to avoid.",
+            ),
+        ],
+    )
+
+
+class QwenImageTextToImageApi(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="QwenImageTextToImageApi",
+            display_name="Qwen Image 3 Text to Image",
+            category="partner/image/Qwen",
+            description="Generates images from a text prompt using the Qwen-Image 3.0 models.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[_t2i_model_option(model_id) for model_id in QWEN_IMAGE_MODELS],
+                    tooltip="Model to use.",
+                ),
+                IO.Int.Input(
+                    "n",
+                    default=1,
+                    min=1,
+                    max=6,
+                    display_mode=IO.NumberDisplay.number,
+                    tooltip="Number of images to generate, returned as a batch.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation.",
+                ),
+                IO.Boolean.Input(
+                    "prompt_extend",
+                    default=True,
+                    tooltip="Whether to enhance the prompt with AI assistance.",
+                    advanced=True,
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add an AI-generated watermark to the result.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Image.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "model.width", "model.height", "n"]),
+                expr="""
+                (
+                  $isPro := widgets.model = "qwen-image-3.0-pro";
+                  $area := $lookup(widgets, "model.width") * $lookup(widgets, "model.height");
+                  $rate := $isPro ? ($area > 2250000 ? 0.10725 : 0.0572) : 0.0429;
+                  {"type":"usd","usd": $rate * widgets.n}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        n: int = 1,
+        seed: int = 42,
+        prompt_extend: bool = True,
+        watermark: bool = False,
+    ):
+        validate_string(model["prompt"], strip_whitespace=False, min_length=1)
+        width, height = model["width"], model["height"]
+        _validate_size(width, height)
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=GENERATION_PATH, method="POST"),
+            response_model=QwenImageGenerationResponse,
+            data=QwenImageGenerationRequest(
+                model=model["model"],
+                input=QwenImageInputField(
+                    messages=[QwenImageMessage(content=[QwenImageContentItem(text=model["prompt"])])],
+                ),
+                parameters=QwenImageParametersField(
+                    size=f"{width}*{height}",
+                    n=n,
+                    seed=seed,
+                    prompt_extend=prompt_extend,
+                    watermark=watermark,
+                    negative_prompt=model["negative_prompt"] or None,
+                ),
+            ),
+        )
+        return IO.NodeOutput(await _download_result_images(response))
+
+
+class QwenImageEditApi(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="QwenImageEditApi",
+            display_name="Qwen Image 3 Edit",
+            category="partner/image/Qwen",
+            description="Edits or combines up to 3 reference images guided by a text prompt "
+            "using the Qwen-Image 3.0 models.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[_edit_model_option(model_id) for model_id in QWEN_IMAGE_MODELS],
+                    tooltip="Model to use.",
+                ),
+                IO.DynamicCombo.Input(
+                    "size",
+                    options=[
+                        IO.DynamicCombo.Option("match input", []),
+                        IO.DynamicCombo.Option("auto", []),
+                        IO.DynamicCombo.Option("custom", _size_inputs()),
+                    ],
+                    tooltip="Output resolution. 'match input' reuses the first reference image's size, "
+                    "'auto' lets the model pick a size with the same aspect ratio, "
+                    "'custom' sets an explicit width and height.",
+                ),
+                IO.Int.Input(
+                    "n",
+                    default=1,
+                    min=1,
+                    max=6,
+                    display_mode=IO.NumberDisplay.number,
+                    tooltip="Number of images to generate, returned as a batch.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation.",
+                ),
+                IO.Boolean.Input(
+                    "prompt_extend",
+                    default=True,
+                    tooltip="Whether to enhance the prompt with AI assistance.",
+                    advanced=True,
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add an AI-generated watermark to the result.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Image.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                depends_on=IO.PriceBadgeDepends(
+                    widgets=["model", "size", "size.width", "size.height", "n"],
+                    input_groups=["model.images"],
+                ),
+                expr="""
+                (
+                  $isPro := widgets.model = "qwen-image-3.0-pro";
+                  $mode := widgets.size;
+                  $count := $max([$lookup(inputGroups, "model.images"), 1]);
+                  $inputCost := 0.00429 * $count;
+                  $area := $mode = "custom"
+                    ? $lookup(widgets, "size.width") * $lookup(widgets, "size.height") : 0;
+                  $customRate := $area > 2250000 ? 0.10725 : 0.0572;
+                  $isPro and $mode != "custom"
+                    ? {"type":"range_usd",
+                       "min_usd": 0.0572 * widgets.n + $inputCost,
+                       "max_usd": 0.10725 * widgets.n + $inputCost}
+                    : {"type":"usd",
+                       "usd": ($isPro ? $customRate : 0.0429) * widgets.n + $inputCost}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        size: dict,
+        n: int = 1,
+        seed: int = 42,
+        prompt_extend: bool = True,
+        watermark: bool = False,
+    ):
+        validate_string(model["prompt"], strip_whitespace=False, min_length=1)
+        reference_images = [image for key in model["images"] for image in model["images"][key]]
+        if len(reference_images) > 3:
+            raise ValueError(
+                f"A maximum of 3 reference images is supported; got {len(reference_images)} "
+                f"(a batched input counts once per image)."
+            )
+        prompt = _resolve_image_refs(model["prompt"], len(reference_images))
+        if size["size"] == "custom":
+            _validate_size(size["width"], size["height"])
+            size_str = f"{size['width']}*{size['height']}"
+        elif size["size"] == "match input":
+            height, width = reference_images[0].shape[0], reference_images[0].shape[1]
+            width, height = _fit_to_size(width, height)
+            size_str = f"{width}*{height}"
+        else:  # auto: the API picks a size preserving the input aspect ratio (1.9-4.2 MP)
+            size_str = None
+        content = [QwenImageContentItem(image=_image_data_uri(image)) for image in reference_images]
+        content.append(QwenImageContentItem(text=prompt))
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=GENERATION_PATH, method="POST"),
+            response_model=QwenImageGenerationResponse,
+            data=QwenImageGenerationRequest(
+                model=model["model"],
+                input=QwenImageInputField(messages=[QwenImageMessage(content=content)]),
+                parameters=QwenImageParametersField(
+                    size=size_str,
+                    n=n,
+                    seed=seed,
+                    prompt_extend=prompt_extend,
+                    watermark=watermark,
+                    negative_prompt=model["negative_prompt"] or None,
+                ),
+            ),
+        )
+        return IO.NodeOutput(await _download_result_images(response))
+
+
+class QwenApiExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            QwenImageTextToImageApi,
+            QwenImageEditApi,
+        ]
+
+
+async def comfy_entrypoint() -> QwenApiExtension:
+    return QwenApiExtension()

comfy_extras/nodes_custom_sampler.py

@@ -718,15 +718,7 @@ class Noise_EmptyNoise:
         self.seed = 0
 
     def generate_noise(self, input_latent):
-        latent_image = input_latent["samples"]
-        if latent_image.is_nested:
-            tensors = latent_image.unbind()
-            zeros = []
-            for t in tensors:
-                zeros.append(torch.zeros(t.shape, dtype=t.dtype, layout=t.layout, device="cpu"))
-            return comfy.nested_tensor.NestedTensor(zeros)
-        else:
-            return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+        return comfy.sample.prepare_empty_noise(input_latent["samples"])
 
 
 class Noise_RandomNoise:

comfy_extras/nodes_lt.py

@@ -2,11 +2,14 @@ import nodes
 import node_helpers
 import torch
 import torchaudio
+import comfy.ldm.lightricks.duration_head
 import comfy.model_management
 import comfy.model_sampling
 import comfy.samplers
 import comfy.utils
+import logging
 import math
+import re
 import numpy as np
 import av
 from io import BytesIO
@@ -934,6 +937,243 @@ class LTXVReferenceAudio(io.ComfyNode):
         return io.NodeOutput(m, positive, negative)
 
 
+class LTXVSpatioTemporalGuidance(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVSpatioTemporalGuidance",
+            display_name="LTXV Spatio-Temporal Guidance (STG)",
+            category="advanced/guidance",
+            description="Runs one extra pass per step with the self-attention of the selected blocks degraded to a value-passthrough, "
+                        "then guides away from it - improving spatial detail and motion coherence.",
+            inputs=[
+                io.Model.Input("model"),
+                io.Float.Input("scale", default=1.0, min=0.0, max=100.0, step=0.01, round=0.01),
+                io.String.Input("blocks", default="29", tooltip="Comma-separated transformer block indices to perturb."),
+                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001, advanced=True),
+                io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001, advanced=True),
+            ],
+            outputs=[io.Model.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, scale, blocks, start_percent, end_percent) -> io.NodeOutput:
+        block_set = frozenset(int(b) for b in re.findall(r"\d+", blocks))
+
+        m = model.clone()
+        model_sampling = m.get_model_object("model_sampling")
+        sigma_start = model_sampling.percent_to_sigma(start_percent)
+        sigma_end = model_sampling.percent_to_sigma(end_percent)
+
+        def post_cfg_function(args):
+            if scale == 0 or not block_set:
+                return args["denoised"]
+
+            sigma_ = args["sigma"][0].item()
+            if sigma_ > sigma_start or sigma_ < sigma_end:
+                return args["denoised"]
+
+            cond_pred = args["cond_denoised"]
+            cond = args["cond"]
+            cfg_result = args["denoised"]
+            x = args["input"]
+
+            model_options = args["model_options"].copy()
+            transformer_options = model_options.get("transformer_options", {}).copy()
+            transformer_options["stg_self_attn_blocks"] = block_set
+            model_options["transformer_options"] = transformer_options
+
+            (perturbed,) = comfy.samplers.calc_cond_batch(args["model"], [cond], x, args["sigma"], model_options)
+
+            return cfg_result + (cond_pred - perturbed) * scale
+
+        m.set_model_sampler_post_cfg_function(post_cfg_function)
+        return io.NodeOutput(m)
+
+
+class LTXVModalityGuidance(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVModalityGuidance",
+            display_name="LTXV Modality Guidance (A/V coupling)",
+            category="advanced/guidance",
+            description="Cross-modal (audio-video) guidance for LTXV-AV. Runs one extra forward "
+                        "pass per step with the a2v/v2a cross-attention severed, then pushes the "
+                        "result toward the coupled prediction - strengthening audio-visual sync "
+                        "(e.g. lip-sync). Reference default modality_scale is 3.0. Stacks with the "
+                        "dual-CFG guider and STG. Set to 1.0 to disable (no extra pass).",
+            inputs=[
+                io.Model.Input("model"),
+                io.Float.Input("modality_scale", default=3.0, min=1.0, max=100.0, step=0.1, round=0.01),
+                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001, advanced=True),
+                io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001, advanced=True),
+            ],
+            outputs=[io.Model.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, modality_scale, start_percent, end_percent) -> io.NodeOutput:
+        m = model.clone()
+        model_sampling = m.get_model_object("model_sampling")
+        sigma_start = model_sampling.percent_to_sigma(start_percent)
+        sigma_end = model_sampling.percent_to_sigma(end_percent)
+
+        def post_cfg_function(args):
+            if math.isclose(modality_scale, 1.0):
+                return args["denoised"]
+
+            sigma_ = args["sigma"][0].item()
+            if sigma_ > sigma_start or sigma_ < sigma_end:
+                return args["denoised"]
+
+            cond_pred = args["cond_denoised"]
+            cond = args["cond"]
+            cfg_result = args["denoised"]
+            x = args["input"]
+
+            # Extra pass with audio-video cross-attention severed (both directions)
+            model_options = args["model_options"].copy()
+            transformer_options = model_options.get("transformer_options", {}).copy()
+            transformer_options["a2v_cross_attn"] = False
+            transformer_options["v2a_cross_attn"] = False
+            model_options["transformer_options"] = transformer_options
+
+            (mod_pred,) = comfy.samplers.calc_cond_batch(
+                args["model"], [cond], x, args["sigma"], model_options
+            )
+
+            # (modality_scale - 1) * (cond - uncond_modality), per the reference guider.
+            return cfg_result + (cond_pred - mod_pred) * (modality_scale - 1.0)
+
+        m.set_model_sampler_post_cfg_function(post_cfg_function)
+        return io.NodeOutput(m)
+
+
+class Guider_LTXAVDualCFG(comfy.samplers.CFGGuider):
+    """CFG guider that applies separate guidance scales to the video and audio
+    modalities of a packed LTXV-AV latent.
+    """
+
+    def set_conds(self, positive, negative):
+        self.inner_set_conds({"positive": positive, "negative": negative})
+
+    def set_cfg(self, video_cfg, audio_cfg):
+        self.video_cfg = video_cfg
+        self.audio_cfg = audio_cfg
+        self.cfg = max(video_cfg, audio_cfg)
+
+    def sample(self, noise, latent_image, *args, **kwargs):
+        # Capture the video/audio split from the nested latent before it is packed.
+        self._v_numel = None
+        if getattr(latent_image, "is_nested", False):
+            parts = latent_image.unbind()
+            if len(parts) >= 2:
+                self._v_numel = math.prod(parts[0].shape[1:])
+        return super().sample(noise, latent_image, *args, **kwargs)
+
+    def predict_noise(self, x, timestep, model_options={}, seed=None):
+        v = getattr(self, "_v_numel", None)
+        if v is None or math.isclose(self.video_cfg, self.audio_cfg):
+            # Not an AV latent, or equal scales: fall back to standard single-CFG.
+            self.cfg = self.video_cfg
+            return super().predict_noise(x, timestep, model_options, seed)
+
+        video_cfg, audio_cfg = self.video_cfg, self.audio_cfg
+
+        def dual_cfg(args):
+            # Noise-space: cond = x - cond_pred, uncond = x - uncond_pred; the
+            # returned tensor is subtracted from x by cfg_function.
+            cond, uncond = args["cond"], args["uncond"]
+            out = uncond + (cond - uncond) * video_cfg
+            out[..., v:] = uncond[..., v:] + (cond[..., v:] - uncond[..., v:]) * audio_cfg
+            return out
+
+        # disable_cfg1_optimization so the uncond pass always runs even if one of the two scales is 1.0.
+        model_options = {**model_options, "sampler_cfg_function": dual_cfg, "disable_cfg1_optimization": True}
+        return super().predict_noise(x, timestep, model_options, seed)
+
+
+class LTXVDualCFGGuider(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVDualCFGGuider",
+            display_name="LTXV Dual CFG Guider",
+            category="model/sampling/guiders",
+            description="Separate CFG scales for the video and audio modalities of a packed LTXV-AV latent.",
+            inputs=[
+                io.Model.Input("model"),
+                io.Conditioning.Input("positive"),
+                io.Conditioning.Input("negative"),
+                io.Float.Input("video_cfg", default=3.0, min=0.0, max=100.0, step=0.1, round=0.01),
+                io.Float.Input("audio_cfg", default=7.0, min=0.0, max=100.0, step=0.1, round=0.01),
+            ],
+            outputs=[io.Guider.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, positive, negative, video_cfg, audio_cfg) -> io.NodeOutput:
+        guider = Guider_LTXAVDualCFG(model)
+        guider.set_conds(positive, negative)
+        guider.set_cfg(video_cfg, audio_cfg)
+        return io.NodeOutput(guider)
+
+
+class LTXVDurationPredictor(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVDurationPredictor",
+            display_name="LTXV Duration Predictor",
+            category="conditioning/video_models",
+            description="Predicts the natural shot duration for a prompt using the LTX 2.4 duration "
+                        "head (loaded with ModelPatchLoader), and snaps it to the VAE's 8k+1 frame grid.",
+            search_aliases=["auto duration", "duration head", "num_frames"],
+            inputs=[
+                io.Model.Input("model"),
+                io.Conditioning.Input("positive"),
+                io.Custom("MODEL_PATCH").Input("duration_head",
+                                               tooltip="LTX 2.4 duration head loaded with ModelPatchLoader."),
+                io.Float.Input("frame_rate", default=24.0, min=1.0, max=120.0, step=0.01),
+                io.Float.Input("min_seconds", default=1.0, min=0.5, max=120.0, step=0.1),
+                io.Float.Input("max_seconds", default=20.0, min=0.5, max=120.0, step=0.1),
+            ],
+            outputs=[
+                io.Int.Output(display_name="num_frames"),
+                io.Float.Output(display_name="seconds", tooltip="Raw (unclamped) predicted duration."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model, positive, duration_head, frame_rate, min_seconds, max_seconds) -> io.NodeOutput:
+        dm = model.model.diffusion_model
+        head = duration_head.model
+        if not isinstance(head, comfy.ldm.lightricks.duration_head.DurationHead):
+            raise ValueError("The connected model_patch is not an LTX duration head.")
+
+        context = positive[0][0]
+        meta = positive[0][1]
+        if context.shape[0] != 1:
+            context = context[:1]
+
+        # Run the caption connectors exactly the way sampling does.
+        comfy.model_management.load_models_gpu([model, duration_head])
+        device = model.load_device
+        head = head.to(device)
+        with torch.no_grad():
+            context = context.to(device=device, dtype=model.model.get_dtype_inference())
+            processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False))
+            video_tokens = processed[..., :dm.cross_attention_dim].float()
+            audio_tokens = processed[..., dm.cross_attention_dim:].float()
+            seconds = float(head(video_tokens, audio_tokens)[0])
+
+        num_frames = comfy.ldm.lightricks.duration_head.seconds_to_num_frames(
+            seconds, frame_rate, min_seconds, max_seconds)
+        logging.info("LTXV duration head predicted %.2fs -> %d frames @ %.2f fps", seconds, num_frames, frame_rate)
+        return io.NodeOutput(num_frames, seconds)
+
+
 class LtxvExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
@@ -951,6 +1191,10 @@ class LtxvExtension(ComfyExtension):
             LTXVConcatAVLatent,
             LTXVSeparateAVLatent,
             LTXVReferenceAudio,
+            LTXVDualCFGGuider,
+            LTXVModalityGuidance,
+            LTXVSpatioTemporalGuidance,
+            LTXVDurationPredictor,
         ]
 
 

comfy_extras/nodes_lt_audio.py

@@ -173,7 +173,7 @@ class LTXAVTextEncoderLoader(io.ComfyNode):
             node_id="LTXAVTextEncoderLoader",
             display_name="Load LTXV Audio Text Encoder",
             category="model/loaders",
-            description="Recipes:\nltxav: gemma 3 12B",
+            description="Recipes:\nltxav: gemma 3 12B or matching gemma 4 model",
             inputs=[
                 io.Combo.Input(
                     "text_encoder",

comfy_extras/nodes_model_advanced.py

@@ -1,6 +1,9 @@
+import logging
+
 import comfy.sd
 import comfy.model_sampling
 import comfy.latent_formats
+import comfy.ldm.modules.attention
 import nodes
 import torch
 import node_helpers
@@ -346,6 +349,39 @@ class ModelComputeDtype:
         return (m, )
 
 
+class ModelAttentionBackend:
+    @classmethod
+    def INPUT_TYPES(s):
+        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,),
+                             }}
+
+    @classmethod
+    def VALIDATE_INPUTS(s, attention):
+        return True
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "model/patch"
+
+    def patch(self, model, attention):
+        attention_name = {
+            "comfy kitchen attention": "comfy_kitchen_int8",
+            "pytorch attention": "pytorch",
+        }.get(attention)
+        attention_function = comfy.ldm.modules.attention.get_attention_function(attention_name, None)
+        if attention_function is None:
+            logging.warning("Attention backend '%s' is unavailable; using PyTorch attention.", attention)
+            attention_function = comfy.ldm.modules.attention.get_attention_function("pytorch")
+        m = model.clone()
+        m.set_model_optimized_attention(attention_function)
+        return (m, )
+
+
 NODE_CLASS_MAPPINGS = {
     "ModelSamplingDiscrete": ModelSamplingDiscrete,
     "ModelSamplingContinuousEDM": ModelSamplingContinuousEDM,
@@ -357,4 +393,5 @@ NODE_CLASS_MAPPINGS = {
     "ModelNoiseScale": ModelNoiseScale,
     "RescaleCFG": RescaleCFG,
     "ModelComputeDtype": ModelComputeDtype,
+    "ModelAttentionBackend": ModelAttentionBackend,
 }

comfy_extras/nodes_model_patch.py

@@ -10,6 +10,7 @@ import comfy.ldm.lumina.controlnet
 import comfy.ldm.supir.supir_modules
 import comfy.ldm.anima.lllite
 import comfy.ldm.wan.uni3c
+import comfy.ldm.lightricks.duration_head
 from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
 from comfy_api.latest import io
 from comfy.ldm.supir.supir_patch import SUPIRPatch
@@ -296,6 +297,10 @@ class ModelPatchLoader:
                     device=comfy.model_management.unet_offload_device(),
                     dtype=dtype,
                     operations=comfy.ops.manual_cast)
+        elif any(k.endswith("duration_head.attention_pooler.query_tokens") for k in sd) or "attention_pooler.query_tokens" in sd:
+            sd = comfy.ldm.lightricks.duration_head.normalize_state_dict(sd)
+            sd = {k: v.float() for k, v in sd.items()}  # tiny head, keep fp32
+            model = comfy.ldm.lightricks.duration_head.DurationHead()
         elif "audio_proj.proj1.weight" in sd:
             model = MultiTalkModelPatch(
                     audio_window=5, context_tokens=32, vae_scale=4,

comfy_extras/nodes_textgen.py

@@ -1,3 +1,4 @@
+import re
 from comfy_api.latest import ComfyExtension, io
 from typing_extensions import override
 
@@ -152,6 +153,64 @@ You are a Creative Assistant writing concise, action-focused image-to-video prom
 Style: realistic - cinematic - The woman glances at her watch and smiles warmly. She speaks in a cheerful, friendly voice, "I think we're right on time!" In the background, a café barista prepares drinks at the counter. The barista calls out in a clear, upbeat tone, "Two cappuccinos ready!" The sound of the espresso machine hissing softly blends with gentle background chatter and the light clinking of cups on saucers.
 """
 
+LTX24_T2V_SYSTEM_PROMPT = """You are given a user's short text-to-video request. Write a single, highly detailed audio-visual caption describing the video that best fulfills that request, in the EXACT style of the training captions used for this video model. The generated video is scored against the user's ORIGINAL request, so preserve every element the user stated; expand faithfully into the full caption style without contradicting or dropping anything they asked for.
+
+Match this captioning style precisely:
+
+1. Begin immediately with the action or visual detail. Do NOT use "The scene opens…", "We see…", "There is…".
+
+2. Objective, observable description only. Do not infer emotions or intentions — describe what is visible and audible (e.g. not "he looks sad" but "his eyebrows angle downward and his lips are pressed together").
+
+3. Full visual detail: environment (materials, textures, lighting, colors), character appearance (clothing, posture, facial details), and the spatial positioning of all elements. When a human appears, identify them specifically (gendered terms when clearly implied; differentiate multiple people consistently) and describe visible physical attributes — apparent gender presentation, skin tone, estimated age group, hair color/length/style, build, clothing and accessories. Do not infer ethnicity, nationality, religion, or culture.
+
+4. Precise motion and cinematic description. For every shot you MUST include, woven naturally into the prose (never as tags or labels):
+   - Shot type (exactly one: extreme wide shot / wide shot / medium shot / medium close-up / close-up / extreme close-up)
+   - Camera motion (always stated; if none, explicitly say the camera remains static). Camera movement is expected and good — match the user if they specified it, otherwise choose the treatment that best presents the requested scene.
+   - Camera viewpoint relative to subject (front-facing / back-facing / side view / over-the-shoulder / top-down / low-angle / high-angle).
+   Express these as flowing prose: "a medium shot frames…, captured from a front-facing angle as the camera slowly pans…". Never as "medium shot, static camera —".
+
+5. Complete soundscape, integrated naturally: any dialogue (quote it exactly, in the original language), tone of voice, background music (type, mood, volume changes), and environmental sounds (footsteps, wind, traffic, animals). If the request implies sound, describe it plausibly.
+
+6. Strict chronological, real-time flow using transitions like "Initially…", "A moment later…", "Simultaneously…". Keep every stated action in motion.
+
+7. One single continuous paragraph. No bullet points, no section headers, no labels like "Audio:" or "Visual:". Exhaustive and lossless — include background elements, subtle movements, lighting, secondary sounds — detailed enough to reconstruct the scene. Aim for a rich, complete paragraph (roughly 150–220 words).
+
+If the user wrote in another language, produce the English caption of the same content. Output ONLY the caption text — no JSON, no preamble.
+
+AESTHETIC QUALITY (in addition to the above, without breaking the objective caption style): render the described scene with strong visual production value — cinematic, film-grade color and contrast, beautiful natural lighting, crisp fine detail and texture, pleasing composition and depth. Weave these quality descriptors naturally into the same observable prose (e.g. "warm cinematic lighting", "richly saturated film-grade color", "crisp high-resolution detail") — describe how the exact requested scene LOOKS at its most visually striking, never adding new objects or actions. Keep everything else (framing triple, soundscape, chronological single paragraph, faithfulness) exactly as specified.
+"""
+
+
+LTX24_I2V_SYSTEM_PROMPT = """You are given a REFERENCE IMAGE (the exact first frame of the video) and a user's short image-to-video request. Write a single, highly detailed audio-visual caption describing the video that BEGINS from this exact reference image and best fulfills that request, in the EXACT style of the training captions used for this video model. The generated video is scored against the user's ORIGINAL request, so preserve every element the user stated; expand faithfully into the full caption style without contradicting or dropping anything they asked for.
+
+FIRST-FRAME / IMAGE GROUNDING (do this first): the opening of your caption must match the reference image exactly — same subject(s), identity, appearance, clothing, setting, lighting, and composition as shown. The video starts on this frame; describe it faithfully, then narrate chronologically as the user's requested action unfolds from it. Never contradict, replace, or invent things not consistent with the image. Single continuous take — no hard cuts.
+
+Match this captioning style precisely:
+
+1. Begin immediately with the action or visual detail. Do NOT use "The scene opens…", "We see…", "There is…".
+
+2. Objective, observable description only. Do not infer emotions or intentions — describe what is visible and audible (e.g. not "he looks sad" but "his eyebrows angle downward and his lips are pressed together").
+
+3. Full visual detail: environment (materials, textures, lighting, colors), character appearance (clothing, posture, facial details), and the spatial positioning of all elements — grounded in and consistent with the reference image. When a human appears, identify them specifically (gendered terms when clearly implied; differentiate multiple people consistently) and describe visible physical attributes — apparent gender presentation, skin tone, estimated age group, hair color/length/style, build, clothing and accessories. Do not infer ethnicity, nationality, religion, or culture.
+
+4. Precise motion and cinematic description. For every shot you MUST include, woven naturally into the prose (never as tags or labels):
+   - Shot type (exactly one: extreme wide shot / wide shot / medium shot / medium close-up / close-up / extreme close-up) — consistent with how the reference image is framed at the start.
+   - Camera motion (always stated; if none, explicitly say the camera remains static). Camera movement is expected and good — match the user if they specified it, otherwise choose the treatment that best presents the requested scene starting from this frame.
+   - Camera viewpoint relative to subject (front-facing / back-facing / side view / over-the-shoulder / top-down / low-angle / high-angle) — matching the reference image's viewpoint at the opening.
+   Express these as flowing prose: "a medium shot frames…, captured from a front-facing angle as the camera slowly pans…". Never as "medium shot, static camera —".
+
+5. Complete soundscape, integrated naturally: any dialogue (quote it exactly, in the original language), tone of voice, background music (type, mood, volume changes), and environmental sounds (footsteps, wind, traffic, animals). If the request implies sound, describe it plausibly.
+
+6. Strict chronological, real-time flow using transitions like "Initially…", "A moment later…", "Simultaneously…". Keep the user's requested motion/action central and in motion throughout.
+
+7. One single continuous paragraph. No bullet points, no section headers, no labels like "Audio:" or "Visual:". Exhaustive and lossless — include background elements, subtle movements, lighting, secondary sounds — detailed enough to reconstruct the scene. Aim for a rich, complete paragraph (roughly 150–220 words).
+
+If the user wrote in another language, produce the English caption of the same content. Output ONLY the caption text — no JSON, no preamble.
+
+AESTHETIC QUALITY (in addition to the above, without breaking the objective caption style or contradicting the reference image): render the described scene with strong visual production value — cinematic, film-grade color and contrast, beautiful natural lighting, crisp fine detail and texture, pleasing composition and depth. Weave these quality descriptors naturally into the same observable prose (e.g. "warm cinematic lighting", "richly saturated film-grade color", "crisp high-resolution detail") — describe how the exact requested scene, starting from this frame, LOOKS at its most visually striking, never adding new objects or actions and never contradicting the first frame. Keep everything else (first-frame grounding, framing triple, soundscape, chronological single paragraph, faithfulness) exactly as specified.
+"""
+
+
 class TextGenerateLTX2Prompt(TextGenerate):
     @classmethod
     def define_schema(cls):
@@ -167,11 +226,42 @@ class TextGenerateLTX2Prompt(TextGenerate):
 
     @classmethod
     def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None) -> io.NodeOutput:
-        if image is None:
-            formatted_prompt = f"<start_of_turn>system\n{LTX2_T2V_SYSTEM_PROMPT.strip()}<end_of_turn>\n<start_of_turn>user\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n<start_of_turn>model\n"
+        # Gemma 3 and Gemma 4 use different chat-turn markers and image tokens.
+        # The Gemma 4 text encoder is the LTX 2.4 path; Gemma 3 is LTX 2.0.
+        is_gemma4 = "gemma4" in getattr(clip.tokenizer, "clip_name", "")
+
+        if is_gemma4:
+            if image is not None:
+                system = LTX24_I2V_SYSTEM_PROMPT.strip()
+                user_text = f"User Raw Input Prompt: {prompt}."
+            else:
+                system = LTX24_T2V_SYSTEM_PROMPT.strip()
+                user_text = f"user prompt: {prompt}"
+            think_prefix = "<|think|>\n" if thinking else ""
+            model_open = "" if thinking else "<|channel>final\n"
+            media = "<|image><|image|><image|>\n\n" if image is not None else ""
+            formatted_prompt = (
+                f"<|turn>system\n{think_prefix}{system}<turn|>\n"
+                f"<|turn>user\n{media}{user_text}<turn|>\n"
+                f"<|turn>model\n{model_open}"
+            )
         else:
-            formatted_prompt = f"<start_of_turn>system\n{LTX2_I2V_SYSTEM_PROMPT.strip()}<end_of_turn>\n<start_of_turn>user\n\n<image_soft_token>\n\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n<start_of_turn>model\n"
-        return super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
+            system = (LTX2_I2V_SYSTEM_PROMPT if image is not None else LTX2_T2V_SYSTEM_PROMPT).strip()
+            media = "\n<image_soft_token>\n" if image is not None else ""
+            formatted_prompt = (
+                f"<start_of_turn>system\n{system}<end_of_turn>\n"
+                f"<start_of_turn>user\n{media}\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n"
+                f"<start_of_turn>model\n"
+            )
+
+        out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
+
+        text = out.args[0]
+        text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
+        if "</think>" in text:  # unclosed/truncated reasoning: keep what follows the last close
+            text = text.rsplit("</think>", 1)[-1]
+        text = re.sub(r"</?think>|<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?", "", text).strip()
+        return io.NodeOutput(text)
 
 
 class TextgenExtension(ComfyExtension):

comfyui_version.py

@@ -1,3 +1,3 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
-__version__ = "0.31.0"
+__version__ = "0.32.0"

nodes.py

@@ -364,8 +364,12 @@ class VAEDecodeTiled:
             temporal_size = None
             temporal_overlap = None
 
+        latent = samples["samples"]
+        if latent.is_nested:
+            latent = latent.unbind()[0]
+
         compression = vae.spacial_compression_decode()
-        images = vae.decode_tiled(samples["samples"], tile_x=tile_size // compression, tile_y=tile_size // compression, overlap=overlap // compression, tile_t=temporal_size, overlap_t=temporal_overlap)
+        images = vae.decode_tiled(latent, tile_x=tile_size // compression, tile_y=tile_size // compression, overlap=overlap // compression, tile_t=temporal_size, overlap_t=temporal_overlap)
         if len(images.shape) == 5: #Combine batches
             images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
         return (images, )
@@ -1566,7 +1570,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None))
 
     if disable_noise:
-        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+        noise = comfy.sample.prepare_empty_noise(latent_image)
     else:
         batch_inds = latent["batch_index"] if "batch_index" in latent else None
         noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)

openapi.yaml

@@ -1521,7 +1521,8 @@ paths:
                 Supports filtering by tags, name, metadata, and sorting options.
             operationId: listAssets
             parameters:
-                - description: Filter assets that have ALL of these tags
+                - deprecated: true
+                  description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
                   explode: false
                   in: query
                   name: include_tags
@@ -1530,7 +1531,8 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - deprecated: true
+                  description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
                   explode: false
                   in: query
                   name: exclude_tags
@@ -1539,6 +1541,33 @@ paths:
                         type: string
                     type: array
                   style: form
+                - description: Filter assets that have ALL of these tags
+                  explode: false
+                  in: query
+                  name: tags_all
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: Filter assets that have AT LEAST ONE of these tags
+                  explode: false
+                  in: query
+                  name: tags_any
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: Exclude assets that have ANY of these tags
+                  explode: false
+                  in: query
+                  name: tags_none
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
                 - description: Filter assets where name contains this substring (case-insensitive)
                   in: query
                   name: name_contains
@@ -2312,7 +2341,8 @@ paths:
                 Only returns tags with non-zero counts (tags that exist on matching assets).
             operationId: getAssetTagHistogram
             parameters:
-                - description: Filter assets that have ALL of these tags
+                - deprecated: true
+                  description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
                   explode: false
                   in: query
                   name: include_tags
@@ -2321,7 +2351,8 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - deprecated: true
+                  description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
                   explode: false
                   in: query
                   name: exclude_tags
@@ -2330,6 +2361,33 @@ paths:
                         type: string
                     type: array
                   style: form
+                - description: Filter assets that have ALL of these tags
+                  explode: false
+                  in: query
+                  name: tags_all
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: Filter assets that have AT LEAST ONE of these tags
+                  explode: false
+                  in: query
+                  name: tags_any
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: Exclude assets that have ANY of these tags
+                  explode: false
+                  in: query
+                  name: tags_none
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
                 - description: Filter assets where name contains this substring (case-insensitive)
                   in: query
                   name: name_contains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.31.0"
+version = "0.32.0"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.48.7
-comfyui-workflow-templates==0.11.37
+comfyui-workflow-templates==0.11.39
 comfyui-embedded-docs==0.5.9
 torch
 torchsde
@@ -22,7 +22,7 @@ alembic
 SQLAlchemy>=2.0.0
 filelock
 av>=16.0.0
-comfy-kitchen==0.2.28
+comfy-kitchen==0.2.30
 comfy-aimdo==0.4.13
 requests
 simpleeval>=1.0.0

tests-unit/assets_test/test_list_filter.py

@@ -1,10 +1,14 @@
 import time
 import uuid
+import warnings
 
 import pytest
 import requests
 from helpers import assert_hash_fields_consistent
 
+from app.assets.api import routes as assets_routes
+from app.assets.api import schemas_in
+
 
 def test_list_assets_paging_and_sort(http: requests.Session, api_base: str, asset_factory, make_asset_bytes):
     names = ["a1_u.safetensors", "a2_u.safetensors", "a3_u.safetensors"]
@@ -337,3 +341,418 @@ def test_list_assets_name_contains_literal_underscore(
     assert b["name"] not in names, "Underscore must be escaped — should not match 'fooxbar'"
     assert c["name"] not in names, "Underscore must be escaped — should not match 'foobar'"
     assert body["total"] == 1
+
+
+def test_list_assets_tags_any_alone(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-any-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    a = asset_factory("any_a.safetensors", [*t, f"{scope}-alpha"], {}, make_asset_bytes("any_a"))
+    b = asset_factory("any_b.safetensors", [*t, f"{scope}-beta"], {}, make_asset_bytes("any_b"))
+    c = asset_factory("any_c.safetensors", [*t, f"{scope}-gamma"], {}, make_asset_bytes("any_c"))
+
+    r = http.get(
+        api_base + "/api/assets",
+        params={"tags_any": f"{scope}-alpha,{scope}-beta", "limit": "50"},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    names = [x["name"] for x in body["assets"]]
+    assert a["name"] in names
+    assert b["name"] in names
+    assert c["name"] not in names
+
+
+def test_list_assets_tags_any_with_tags_all(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-anyall-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    x = asset_factory("aa_x.safetensors", [*t, alpha], {}, make_asset_bytes("aa_x"))
+    y = asset_factory("aa_y.safetensors", [*t, beta], {}, make_asset_bytes("aa_y"))
+    w = asset_factory("aa_w.safetensors", t, {}, make_asset_bytes("aa_w"))
+    d = asset_factory(
+        "aa_d.safetensors",
+        ["models", "model_type:checkpoints", "unit-tests", f"{scope}-other", alpha],
+        {},
+        make_asset_bytes("aa_d"),
+    )
+
+    r = http.get(
+        api_base + "/api/assets",
+        params={"tags_all": f"unit-tests,{scope}", "tags_any": f"{alpha},{beta}", "limit": "50"},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    names = [a["name"] for a in body["assets"]]
+    assert x["name"] in names
+    assert y["name"] in names
+    assert w["name"] not in names, "asset matching tags_all but not tags_any must be excluded"
+    assert d["name"] not in names, "asset matching tags_any but not tags_all must be excluded"
+
+
+def test_list_assets_tags_none_wins_over_tags_any(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-nonewins-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    x = asset_factory("nw_x.safetensors", [*t, alpha], {}, make_asset_bytes("nw_x"))
+    y = asset_factory("nw_y.safetensors", [*t, alpha, beta], {}, make_asset_bytes("nw_y"))
+
+    r = http.get(
+        api_base + "/api/assets",
+        params={"tags_any": alpha, "tags_none": beta, "limit": "50"},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    names = [a["name"] for a in body["assets"]]
+    assert x["name"] in names
+    assert y["name"] not in names, "tags_none must exclude an asset even when it matches tags_any"
+
+
+def test_list_assets_empty_tag_filter_lists_behave_as_absent(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-empty-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    a = asset_factory("em_a.safetensors", t, {}, make_asset_bytes("em_a"))
+    b = asset_factory("em_b.safetensors", t, {}, make_asset_bytes("em_b"))
+    expected = {a["name"], b["name"]}
+
+    # Empty new-name lists impose no constraint.
+    r1 = http.get(
+        api_base + "/api/assets",
+        params={"tags_all": f"unit-tests,{scope}", "tags_any": "", "tags_none": ""},
+        timeout=120,
+    )
+    b1 = r1.json()
+    assert r1.status_code == 200, b1
+    assert {x["name"] for x in b1["assets"]} == expected
+
+    # An empty new-name param alongside old names must not trigger validation.
+    r2 = http.get(
+        api_base + "/api/assets",
+        params={"include_tags": f"unit-tests,{scope}", "tags_any": ""},
+        timeout=120,
+    )
+    b2 = r2.json()
+    assert r2.status_code == 200, b2
+    assert {x["name"] for x in b2["assets"]} == expected
+
+    # An empty tags_all next to include_tags is not a mixed-spelling conflict.
+    r3 = http.get(
+        api_base + "/api/assets",
+        params={"include_tags": f"unit-tests,{scope}", "tags_all": ""},
+        timeout=120,
+    )
+    b3 = r3.json()
+    assert r3.status_code == 200, b3
+    assert {x["name"] for x in b3["assets"]} == expected
+
+
+def test_list_assets_old_names_match_new_names(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-alias-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    asset_factory("al_a.safetensors", [*t, alpha], {}, make_asset_bytes("al_a"))
+    asset_factory("al_b.safetensors", [*t, beta], {}, make_asset_bytes("al_b"))
+
+    def names_for(params: dict) -> tuple[list, int]:
+        r = http.get(api_base + "/api/assets", params={**params, "sort": "name", "order": "asc"}, timeout=120)
+        body = r.json()
+        assert r.status_code == 200, body
+        return [x["name"] for x in body["assets"]], body["total"]
+
+    # include_tags ≡ tags_all
+    old_names, old_total = names_for({"include_tags": f"unit-tests,{scope}"})
+    new_names, new_total = names_for({"tags_all": f"unit-tests,{scope}"})
+    assert old_names == new_names
+    assert old_total == new_total
+
+    # exclude_tags ≡ tags_none (and old/new spellings mix across slots)
+    old_names, old_total = names_for({"include_tags": f"unit-tests,{scope}", "exclude_tags": alpha})
+    new_names, new_total = names_for({"tags_all": f"unit-tests,{scope}", "tags_none": alpha})
+    mixed_names, mixed_total = names_for({"include_tags": f"unit-tests,{scope}", "tags_none": alpha})
+    assert old_names == new_names == mixed_names == ["al_b.safetensors"]
+    assert old_total == new_total == mixed_total == 1
+
+
+@pytest.mark.parametrize(
+    "params,expected_parameters",
+    [
+        ({"include_tags": "mx-x", "tags_all": "mx-y"}, ["include_tags", "tags_all"]),
+        ({"exclude_tags": "mx-x", "tags_none": "mx-y"}, ["exclude_tags", "tags_none"]),
+    ],
+    ids=["include_tags_with_tags_all", "exclude_tags_with_tags_none"],
+)
+def test_list_assets_mixed_tag_spellings_rejected(http, api_base, params, expected_parameters):
+    r = http.get(api_base + "/api/assets", params=params, timeout=120)
+    body = r.json()
+    assert r.status_code == 400, body
+    assert body["error"]["code"] == "INVALID_TAG_FILTER"
+    assert body["error"]["details"]["parameters"] == expected_parameters
+
+
+@pytest.mark.parametrize(
+    "params,conflicting,parameters",
+    [
+        (
+            {"tags_all": "cf-x", "tags_none": "cf-x"},
+            ["cf-x"],
+            ["tags_all", "tags_none"],
+        ),
+        (
+            {"include_tags": "cf-x", "tags_none": "cf-x"},
+            ["cf-x"],
+            ["include_tags", "tags_none"],
+        ),
+        (
+            {"tags_all": "cf-a,cf-b", "tags_none": "cf-b,cf-c"},
+            ["cf-b"],
+            ["tags_all", "tags_none"],
+        ),
+    ],
+    ids=["new_names", "include_tags_remapped", "partial_overlap"],
+)
+def test_list_assets_all_none_conflict_rejected(http, api_base, params, conflicting, parameters):
+    r = http.get(api_base + "/api/assets", params=params, timeout=120)
+    body = r.json()
+    assert r.status_code == 400, body
+    assert body["error"]["code"] == "INVALID_TAG_FILTER"
+    assert body["error"]["details"]["conflicting_tags"] == conflicting
+    assert body["error"]["details"]["parameters"] == parameters
+
+
+def test_list_assets_any_none_overlap_accepted(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-deadterm-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    x = asset_factory("dt_x.safetensors", [*t, alpha], {}, make_asset_bytes("dt_x"))
+    y = asset_factory("dt_y.safetensors", [*t, beta], {}, make_asset_bytes("dt_y"))
+
+    # alpha is a dead term (in both tags_any and tags_none) but the query is valid.
+    r = http.get(
+        api_base + "/api/assets",
+        params={"tags_any": f"{alpha},{beta}", "tags_none": alpha, "limit": "50"},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    names = [a["name"] for a in body["assets"]]
+    assert y["name"] in names
+    assert x["name"] not in names
+
+
+def test_list_assets_legacy_include_exclude_conflict_still_200(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-legacy-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    asset_factory("lg_a.safetensors", t, {}, make_asset_bytes("lg_a"))
+
+    # Old names only: the self-contradictory query stays an empty 200, never a 400.
+    r = http.get(
+        api_base + "/api/assets",
+        params={"include_tags": scope, "exclude_tags": scope},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    assert body["assets"] == []
+
+
+def test_tags_refine_new_tag_filters(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"rf-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    asset_factory("rf_a.safetensors", [*t, alpha], {}, make_asset_bytes("rf_a"))
+    asset_factory("rf_b.safetensors", [*t, beta], {}, make_asset_bytes("rf_b"))
+
+    r = http.get(
+        api_base + "/api/assets/tags/refine",
+        params={"tags_any": f"{alpha},{beta}", "tags_none": alpha},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    counts = body["tag_counts"]
+    assert counts.get(beta) == 1
+    assert alpha not in counts
+
+    r2 = http.get(
+        api_base + "/api/assets/tags/refine",
+        params={"tags_all": "rf-x", "tags_none": "rf-x"},
+        timeout=120,
+    )
+    body2 = r2.json()
+    assert r2.status_code == 400, body2
+    assert body2["error"]["code"] == "INVALID_TAG_FILTER"
+    assert body2["error"]["details"]["conflicting_tags"] == ["rf-x"]
+
+
+def test_list_assets_cross_slot_old_new_combinations(http, api_base, asset_factory, make_asset_bytes):
+    """Old and new spellings of *different* slots combine freely; only
+    same-slot mixing is rejected."""
+    scope = f"lf-cross-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    a = asset_factory("cs_a.safetensors", [*t, alpha], {}, make_asset_bytes("cs_a"))
+    b = asset_factory("cs_b.safetensors", [*t, beta], {}, make_asset_bytes("cs_b"))
+
+    def names_for(params: dict) -> set:
+        r = http.get(api_base + "/api/assets", params=params, timeout=120)
+        body = r.json()
+        assert r.status_code == 200, body
+        return {x["name"] for x in body["assets"]}
+
+    assert names_for(
+        {"include_tags": f"unit-tests,{scope}", "tags_any": alpha}
+    ) == {a["name"]}
+    assert names_for(
+        {"tags_all": f"unit-tests,{scope}", "exclude_tags": alpha}
+    ) == {b["name"]}
+    assert names_for(
+        {"tags_any": f"{alpha},{beta}", "exclude_tags": alpha}
+    ) == {b["name"]}
+
+
+def test_list_assets_repeated_query_keys_concatenate(http, api_base, asset_factory, make_asset_bytes):
+    """Repeated occurrences of a tag param concatenate before the CSV split
+    (Core-local behavior, not a cross-platform guarantee)."""
+    scope = f"lf-repeat-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha, beta = f"{scope}-alpha", f"{scope}-beta"
+    a = asset_factory("rp_a.safetensors", [*t, alpha], {}, make_asset_bytes("rp_a"))
+    b = asset_factory("rp_b.safetensors", [*t, beta], {}, make_asset_bytes("rp_b"))
+
+    # requests encodes a list value as repeated keys: tags_any=<alpha>&tags_any=<beta>
+    r = http.get(
+        api_base + "/api/assets",
+        params={"tags_any": [alpha, beta], "limit": "50"},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 200, body
+    names = {x["name"] for x in body["assets"]}
+    assert {a["name"], b["name"]} <= names
+
+
+def test_list_assets_tags_any_cursor_pagination_consistent(http, api_base, asset_factory, make_asset_bytes):
+    scope = f"lf-anypage-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    alpha = f"{scope}-alpha"
+    expected = set()
+    for i in range(3):
+        made = asset_factory(f"pg_{i}.safetensors", [*t, alpha], {}, make_asset_bytes(f"pg_{i}"))
+        expected.add(made["name"])
+
+    r1 = http.get(
+        api_base + "/api/assets",
+        params={"tags_any": alpha, "limit": "2", "sort": "name", "order": "asc"},
+        timeout=120,
+    )
+    b1 = r1.json()
+    assert r1.status_code == 200, b1
+    assert b1["total"] == 3
+    assert b1["has_more"] is True
+    assert b1.get("next_cursor"), "expected a keyset cursor on the first page"
+
+    r2 = http.get(
+        api_base + "/api/assets",
+        params={
+            "tags_any": alpha,
+            "limit": "2",
+            "sort": "name",
+            "order": "asc",
+            "after": b1["next_cursor"],
+        },
+        timeout=120,
+    )
+    b2 = r2.json()
+    assert r2.status_code == 200, b2
+    assert b2["has_more"] is False
+
+    page1 = {x["name"] for x in b1["assets"]}
+    page2 = {x["name"] for x in b2["assets"]}
+    assert not page1 & page2, "cursor pages must not overlap"
+    assert page1 | page2 == expected
+
+
+def test_tags_refine_mixed_spellings_rejected_and_legacy_conflict_kept(http, api_base):
+    r = http.get(
+        api_base + "/api/assets/tags/refine",
+        params={"include_tags": "rfmx-x", "tags_all": "rfmx-y"},
+        timeout=120,
+    )
+    body = r.json()
+    assert r.status_code == 400, body
+    assert body["error"]["code"] == "INVALID_TAG_FILTER"
+    assert body["error"]["details"]["parameters"] == ["include_tags", "tags_all"]
+
+    # Old names only: the refine route keeps legacy behaviour too.
+    r2 = http.get(
+        api_base + "/api/assets/tags/refine",
+        params={"include_tags": "rfmx-z", "exclude_tags": "rfmx-z"},
+        timeout=120,
+    )
+    body2 = r2.json()
+    assert r2.status_code == 200, body2
+    assert body2["tag_counts"] == {}
+
+
+def test_list_assets_tag_values_case_sensitive(http, api_base, asset_factory, make_asset_bytes):
+    """Case-distinct tags are distinct; the all/none conflict check is byte-exact."""
+    scope = f"lf-case-{uuid.uuid4().hex[:6]}"
+    t = ["models", "model_type:checkpoints", "unit-tests", scope]
+    upper, lower = f"{scope}-ALPHA", f"{scope}-alpha"
+    a = asset_factory("cx_a.safetensors", [*t, upper], {}, make_asset_bytes("cx_a"))
+    b = asset_factory("cx_b.safetensors", [*t, lower], {}, make_asset_bytes("cx_b"))
+
+    def names_for(params: dict) -> set:
+        r = http.get(api_base + "/api/assets", params=params, timeout=120)
+        body = r.json()
+        assert r.status_code == 200, body
+        return {x["name"] for x in body["assets"]}
+
+    assert names_for({"tags_all": f"unit-tests,{scope},{upper}"}) == {a["name"]}
+    assert names_for({"tags_any": lower, "limit": "50"}) == {b["name"]}
+    # Case-distinct all/none pair is NOT a conflict — byte-exact comparison.
+    assert names_for({"tags_all": f"unit-tests,{scope},{upper}", "tags_none": lower}) == {a["name"]}
+
+
+def test_tag_list_cap_applies_to_all_spellings(http, api_base):
+    """The cap covers the legacy spellings too."""
+    big = ",".join(f"cap-{i}" for i in range(101))
+    for param in ("tags_any", "include_tags"):
+        r = http.get(api_base + "/api/assets", params={param: big}, timeout=120)
+        body = r.json()
+        assert r.status_code == 400, body
+        assert body["error"]["code"] == "INVALID_TAG_FILTER"
+        assert body["error"]["details"]["parameter"] == param
+        assert body["error"]["details"]["max"] == 100
+
+    exact = ",".join(f"cap-{i}" for i in range(100))
+    r = http.get(api_base + "/api/assets", params={"tags_any": exact}, timeout=120)
+    assert r.status_code == 200, r.json()
+
+    # The cap counts normalized (deduped) tags, not raw CSV items.
+    dups = ",".join("cap-dup" for _ in range(150))
+    r = http.get(api_base + "/api/assets", params={"tags_any": dups}, timeout=120)
+    assert r.status_code == 200, r.json()
+
+
+def test_resolve_tag_filters_no_deprecation_warning():
+    """The deprecated-field warning is for API clients; the server's own remap
+    shim must not fire it on every request."""
+    for q in (
+        schemas_in.ListAssetsQuery(tags_all="a", tags_none="b"),
+        schemas_in.TagsRefineQuery(tags_any="c"),
+    ):
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", DeprecationWarning)
+            assets_routes._resolve_tag_filters(q)
+
+
+def test_tag_filter_alias_fields_marked_deprecated():
+    for model in (schemas_in.ListAssetsQuery, schemas_in.TagsRefineQuery):
+        props = model.model_json_schema()["properties"]
+        for field in ("include_tags", "exclude_tags"):
+            assert props[field].get("deprecated") is True, (model.__name__, field)
+        for field in ("tags_all", "tags_any", "tags_none"):
+            assert "deprecated" not in props[field], (model.__name__, field)

tests-unit/comfy_test/test_vae_decode_tiled_nested.py

@@ -0,0 +1,27 @@
+from unittest.mock import MagicMock
+
+import torch
+
+from comfy.cli_args import args as cli_args
+
+if not torch.cuda.is_available():
+    cli_args.cpu = True
+
+import comfy.nested_tensor  # noqa: E402
+import nodes  # noqa: E402
+
+
+def test_vae_decode_tiled_unwraps_nested_tensor():
+    video = torch.zeros(1, 4, 2, 8, 8)
+    audio = torch.zeros(1, 2, 2, 40)
+    samples = {"samples": comfy.nested_tensor.NestedTensor((video, audio))}
+
+    vae = MagicMock()
+    vae.temporal_compression_decode.return_value = None
+    vae.spacial_compression_decode.return_value = 8
+    vae.decode_tiled.return_value = torch.zeros(1, 3, 2, 8, 8)
+
+    nodes.VAEDecodeTiled().decode(vae, samples, tile_size=512)
+
+    decoded_arg = vae.decode_tiled.call_args[0][0]
+    assert decoded_arg is video