Merge remote-tracking branch 'upstream/master' into fix-minimax-h3-decode-oom-dynamic-vram

# Conflicts:
#	comfy/sd.py
This commit is contained in:
chelsealong 2026-08-13 13:12:46 +00:00
commit 6008a10429
53 changed files with 3512 additions and 220 deletions

View File

@ -37,7 +37,7 @@
ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
- ComfyUI natively supports the latest open-source state of the art models.
- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
- [Partner nodes](https://docs.comfy.org/tutorials/partner-nodes/overview#partner-nodes) provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
- It is available on Windows, Linux, and macOS, locally with our [desktop application](https://www.comfy.org/download), our [portable install](#installing) or on our [cloud](https://www.comfy.org/cloud).
- The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
- It integrates seamlessly into production pipelines with our API endpoints.

View File

@ -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,

View File

@ -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:

View File

@ -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)

View File

@ -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(

View File

@ -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()

View File

@ -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,

View File

@ -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,

View File

@ -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.")

View File

@ -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])

View File

@ -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,

View File

@ -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

View File

@ -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,

View File

@ -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.

View File

@ -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

View File

@ -0,0 +1,520 @@
"""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
import comfy.model_management
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):
out_device = device
if not comfy.model_management.supports_fp64(device):
device = torch.device("cpu")
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(dtype=torch.float32, device=out_device)
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)

View File

@ -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
@ -165,9 +165,10 @@ class Attention(nn.Module):
else:
q = self.q_norm(q.view(s, self.heads, self.head_dim))
k = self.k_norm(k.view(s, self.heads, self.head_dim))
q = q.transpose(0, 1).unsqueeze(0)
k = k.transpose(0, 1).unsqueeze(0)
v = v.transpose(0, 1).unsqueeze(0)
v = v.clone()
q = AttentionTensorContainer(q.transpose(0, 1).unsqueeze(0))
k = AttentionTensorContainer(k.transpose(0, 1).unsqueeze(0))
v = AttentionTensorContainer(v.transpose(0, 1).unsqueeze(0))
out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)
return self.out_proj(out.squeeze(0))

View File

@ -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:

View File

@ -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):

View File

@ -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
@ -829,11 +830,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["use_adaln_lora"] = True
dit_config["adaln_lora_dim"] = 256
dit_config["num_blocks"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
if dit_config["model_channels"] == 2048:
dit_config["num_blocks"] = 28
dit_config["num_heads"] = 16
elif dit_config["model_channels"] == 5120:
dit_config["num_blocks"] = 36
dit_config["num_heads"] = 40
if dit_config["in_channels"] == 16:

View File

@ -490,28 +490,36 @@ try:
except:
rocm_version = (6, -1)
def aotriton_supported(gpu_arch):
path = torch.__path__[0]
path = os.path.join(os.path.join(path, "lib"), "aotriton.images")
gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path))))
if gpu_arch in gfx:
return True
if "{}x".format(gpu_arch[:-1]) in gfx:
return True
if "{}xx".format(gpu_arch[:-2]) in gfx:
return True
return False
def aotriton_supported():
"""Whether pytorch reports flash attention as usable on this gpu.
can_use_flash_attention() evaluates runtime eligibility for the given
parameters; on a ROCm build that includes checking the gpu arch against the
kernel images AOTriton was compiled for. Querying it avoids assuming where
those images live inside the torch install. The probe tensor is shaped and
typed to pass the unrelated SDPA checks, so False means no hardware support
rather than a rejected shape.
"""
try:
if not torch.backends.cuda.is_flash_attention_available(): # not built with flash attention
return False
q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device())
params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False)
return torch.backends.cuda.can_use_flash_attention(params, False)
except (AttributeError, RuntimeError, TypeError) as e:
logging.warning("Could not query aotriton support: {}".format(e))
return False
logging.info("AMD arch: {}".format(arch))
logging.info("ROCm version: {}".format(rocm_version))
if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
if aotriton_supported(arch): # AMD efficient attention implementation depends on aotriton.
if aotriton_supported(): # AMD efficient attention implementation depends on aotriton.
if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much
if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]): # TODO: more arches, TODO: gfx950
ENABLE_PYTORCH_ATTENTION = True
if rocm_version >= (7, 0):
if any((a in arch) for a in ["gfx1200", "gfx1201"]):
ENABLE_PYTORCH_ATTENTION = True
if any((a in arch) for a in ["gfx1200", "gfx1201"]):
ENABLE_PYTORCH_ATTENTION = True
if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4):
if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
SUPPORT_FP8_OPS = True
@ -1658,6 +1666,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

View File

@ -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")

View File

@ -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

View File

@ -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

View File

@ -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}
@ -1228,16 +1245,48 @@ class VAE:
# decode_tiled is identical to decode for these VAEs, so freeing the
# memory other models hold onto is the only thing that can make the retry succeed.
model_management.free_memory(1e30, self.device, keep_loaded=[model_management.LoadedModel(self.patcher)])
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 = {}
@ -1711,12 +1760,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:
@ -1884,9 +1942,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

View File

@ -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)

View File

@ -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={}):

View File

@ -1183,6 +1183,7 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po
class Gemma4_Tokenizer():
tokenizer_json_data = None
prime_empty_thought = False
def state_dict(self):
if self.tokenizer_json_data is not None:
@ -1333,8 +1334,8 @@ class Gemma4_Tokenizer():
num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
n_audio_tokens = self._audio_token_count(num_samples)
media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
# Non-thinking mode primes an empty thought channel so the model answers directly.
model_open = "" if thinking else "<|channel>thought\n<channel|>"
# 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline.
model_open = "<|channel>thought\n<channel|>" if self.prime_empty_thought and not thinking else ""
llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\n<|turn>model\n{model_open}"
text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
@ -1418,6 +1419,7 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
"""Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
embedding_size = 3840
prime_empty_thought = True
def _extract_audio_features(self, waveform, sample_rate):
audio = self._resample_16k(waveform, sample_rate)
@ -1443,14 +1445,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 +1476,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:
@ -1489,7 +1502,7 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No
# Variants
def _make_variant(config_cls):
def _make_variant(config_cls, prime_empty_thought=False):
audio = config_cls.audio_config is not None
bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,)
class Variant(*bases):
@ -1499,8 +1512,8 @@ def _make_variant(config_cls):
if audio:
self._init_audio(self.model.config, dtype, device, operations)
embedding_size = config_cls.hidden_size
if embedding_size != Gemma4SDTokenizer.embedding_size:
tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size})
if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought:
tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought})
class Tokenizer(Gemma4Tokenizer):
tokenizer_class = tok_cls
Variant.tokenizer = Tokenizer
@ -1510,7 +1523,7 @@ def _make_variant(config_cls):
Gemma4_E4B = _make_variant(Gemma4Config)
Gemma4_E2B = _make_variant(Gemma4_E2B_Config)
Gemma4_31B = _make_variant(Gemma4_31B_Config)
Gemma4_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True)
# Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).

View File

@ -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

View File

@ -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 {}

View File

@ -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)

View File

@ -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):

View File

@ -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.")

View File

@ -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,
)

View File

@ -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,
]

View File

@ -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()

View File

@ -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:

View File

@ -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,
]

View File

@ -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",

View File

@ -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,
}

View File

@ -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,

View File

@ -29,7 +29,7 @@ class PreviewAny():
value = str(source)
elif source is not None:
try:
value = json.dumps(source, indent=4)
value = json.dumps(source, indent=4, ensure_ascii=False)
except Exception:
try:
value = str(source)

View File

@ -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 150220 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 150220 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):

View File

@ -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"

View File

@ -248,7 +248,7 @@ import hook_breaker_ac10a0
import comfy.memory_management
import comfy.model_patcher
if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()):
if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia()):
if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
else:

View File

@ -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)

View File

@ -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

View File

@ -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"

View File

@ -1,5 +1,5 @@
comfyui-frontend-package==1.48.7
comfyui-workflow-templates==0.11.37
comfyui-workflow-templates==0.11.40
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.31
comfy-aimdo==0.4.13
requests
simpleeval>=1.0.0

View File

@ -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)

View File

@ -0,0 +1,30 @@
from unittest.mock import patch, MagicMock
mock_nodes = MagicMock()
mock_nodes.MAX_RESOLUTION = 16384
mock_server = MagicMock()
with patch.dict("sys.modules", {"nodes": mock_nodes, "server": mock_server}):
from comfy_extras.nodes_preview_any import PreviewAny
class TestPreviewAnyMain:
@staticmethod
def _exec(source) -> dict:
return PreviewAny().main(source)
def test_dict_keeps_non_ascii(self):
result = self._exec({"greeting": "你好"})
assert "你好" in result["ui"]["text"][0]
assert "\\u" not in result["ui"]["text"][0]
assert result["result"][0] == result["ui"]["text"][0]
def test_list_keeps_non_ascii(self):
result = self._exec(["你好", "こんにちは"])
assert "こんにちは" in result["result"][0]
assert "\\u" not in result["result"][0]
def test_string_passthrough(self):
result = self._exec("你好")
assert result["ui"]["text"][0] == "你好"
assert result["result"][0] == "你好"

View File

@ -0,0 +1,61 @@
"""Gemma4 chat template regression tests."""
import pytest
import torch
from comfy.cli_args import args
if not torch.cuda.is_available():
args.cpu = True
import comfy.text_encoders.gemma4 as gemma4 # noqa: E402
PROMPT = "describe a cute anime girl with fennec ears"
THOUGHT_BLOCK = "<|channel>thought\n<channel|>"
# E2B/E4B and 12B/31B ship different canonical chat templates: only the latter prime a
# closed thought block when thinking is off.
NO_PRIMING = [gemma4.Gemma4_E2B, gemma4.Gemma4_E4B]
PRIMING = [gemma4.Gemma4_31B, gemma4.Gemma4_12B]
class _CaptureTemplate:
"""Stands in for SDTokenizer.tokenize_with_weights so the built template is checked without model files."""
llama_text = ""
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
self.llama_text = text
return {}
def build_template(variant, **kwargs):
prime = variant.tokenizer.tokenizer_class.prime_empty_thought
probe = type("Probe", (gemma4.Gemma4_Tokenizer, _CaptureTemplate), {"prime_empty_thought": prime})()
probe.tokenize_with_weights(PROMPT, **kwargs)
return probe.llama_text
@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
def test_thinking_enabled_only_asks_via_the_system_turn(variant):
template = build_template(variant, skip_template=False, thinking=True)
assert template == f"<|turn>system\n<|think|>\n<turn|>\n<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
@pytest.mark.parametrize("variant", NO_PRIMING)
def test_thinking_disabled_does_not_prime_a_thought_channel(variant):
template = build_template(variant, skip_template=False, thinking=False)
assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
assert "channel" not in template
assert "<|think|>" not in template
@pytest.mark.parametrize("variant", PRIMING)
def test_thinking_disabled_primes_a_thought_channel(variant):
template = build_template(variant, skip_template=False, thinking=False)
assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n{THOUGHT_BLOCK}"
@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
@pytest.mark.parametrize("thinking", [False, True])
def test_skip_template_passes_text_through_unchanged(variant, thinking):
assert build_template(variant, skip_template=True, thinking=thinking) == PROMPT

View File

@ -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