Compare commits

..
6 Commits
Author SHA1 Message Date
Yousef R. GamaleldinandGitHub a3ca7ef350 Merge eda4ba84f2 into 51bf508a0b 2026-07-08 01:24:18 +08:00
Terry Jia eda4ba84f2 fix(jobs): treat text file outputs as previewable and stop counting text previews 2026-07-05 22:54:19 -04:00
Yousef RafatandTerry Jia 949c4b8d07 savedResult 2026-07-05 22:54:19 -04:00
cd53aff8d8 Update comfy_extras/nodes_text.py
Co-authored-by: Alexis Rolland <alexis@comfy.org>
2026-07-05 22:54:19 -04:00
Yousef RafatandTerry Jia 84f96a56ca . 2026-07-05 22:54:19 -04:00
Yousef RafatandTerry Jia 978895e1a9 Add Save Text Node 2026-07-05 22:54:19 -04:00
30 changed files with 239 additions and 1230 deletions
-2
View File
@@ -127,8 +127,6 @@
- Do not add unnecessary `try`/`except` blocks. Use them for optional dependency,
platform, or backend capability detection only when the program has a useful
fallback. Prefer specific exception types when changing new code.
- If a library version is pinned in `requirements.txt`, do not add code to
ComfyUI to handle older versions of that library.
- Remove any workarounds for PyTorch versions that ComfyUI no longer officially
supports. Deprecated workarounds include catching an exception and rerunning
the same op with the input cast to float. If a workaround does not have a
+1 -1
View File
@@ -229,7 +229,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia
Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
torch 2.4 and above is supported but some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old.
### Instructions:
-1
View File
@@ -225,7 +225,6 @@ parser.add_argument(
)
parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path. Overrides --base-directory.")
parser.add_argument("--models-directory", type=is_valid_directory, default=None, help="Set the ComfyUI models directory. Overrides the models folder in --base-directory.")
parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.")
-3
View File
@@ -540,9 +540,6 @@ class Wan21(LatentFormat):
latents_std = self.latents_std.to(latent.device, latent.dtype)
return latent * latents_std / self.scale_factor + latents_mean
class LingBotVideo(Wan21):
pass
class Wan22(Wan21):
latent_channels = 48
latent_dimensions = 3
+5 -2
View File
@@ -217,7 +217,10 @@ class AceStepAttention(nn.Module):
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
n_rep = self.num_heads // self.num_kv_heads
if n_rep > 1:
key_states = key_states.repeat_interleave(n_rep, dim=1)
value_states = value_states.repeat_interleave(n_rep, dim=1)
attn_bias = None
if self.sliding_window is not None and not self.is_cross_attention:
@@ -241,7 +244,7 @@ class AceStepAttention(nn.Module):
else:
attn_bias = window_bias
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False, **gqa_kwargs)
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False)
attn_output = self.o_proj(attn_output)
return attn_output
+7 -4
View File
@@ -425,16 +425,19 @@ class Attention(nn.Module):
if n == 1 and causal:
causal = False
gqa_kwargs = {"enable_gqa": True} if h != kv_h else {}
if h != kv_h:
# Repeat interleave kv_heads to match q_heads
heads_per_kv_head = h // kv_h
k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
if self.differential:
q, q_diff = q.unbind(dim=1)
k, k_diff = k.unbind(dim=1)
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
out = out - out_diff
else:
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
out = self.to_out(out)
+5 -2
View File
@@ -74,8 +74,11 @@ class BooguDoubleStreamProcessor(nn.Module):
key = key.transpose(1, 2)
value = value.transpose(1, 2)
gqa_kwargs = {"enable_gqa": True} if attn.kv_heads < attn.heads else {}
hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs)
if attn.kv_heads < attn.heads:
key = key.repeat_interleave(attn.heads // attn.kv_heads, dim=1)
value = value.repeat_interleave(attn.heads // attn.kv_heads, dim=1)
hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
# Split back to instruction/image, apply per-stream output projections, recombine.
instruct_hidden_states = self.instruct_out(hidden_states[:, :L_instruct])
-518
View File
@@ -1,518 +0,0 @@
import math
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from comfy.ldm.flux.math import apply_rope1, rope
from comfy.ldm.flux.layers import timestep_embedding
from comfy.ldm.modules.attention import optimized_attention
class LingBotVideoRotaryEmbedding(nn.Module):
def __init__(self, axes_dims: Tuple[int, ...], axes_lens: Tuple[int, ...], theta: float):
super().__init__()
self.axes_dims = tuple(axes_dims)
self.theta = theta
def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
return torch.cat(
[rope(position_ids[None, :, i], self.axes_dims[i], self.theta) for i in range(len(self.axes_dims))],
dim=-3,
).squeeze(0)
def make_joint_position_ids(
text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device, padded_text_len: Optional[int] = None
) -> torch.Tensor:
"""3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1.
Matches patchify_and_embed: cap start (1,0,0); vision start (cap_len+1,0,0);
freqs ordered with x first and cap second (same order as cat_interleave).
"""
tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
hh = torch.arange(grid_h, device=device, dtype=torch.int32)
ww = torch.arange(grid_w, device=device, dtype=torch.int32)
grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
if padded_text_len is None:
padded_text_len = text_len
text_t = torch.arange(padded_text_len, device=device, dtype=torch.int32) + 1
text_pos = torch.stack(
[text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1
)
return torch.cat([grid, text_pos], dim=0) # (Nx + L, 3)
class LingBotVideoTimestepEmbedding(nn.Module):
def __init__(self, in_channels, time_embed_dim, bias=True, device=None, dtype=None, operations=None):
super().__init__()
self.linear_1 = operations.Linear(in_channels, time_embed_dim, bias=bias, device=device, dtype=dtype)
self.act = nn.SiLU()
self.linear_2 = operations.Linear(time_embed_dim, time_embed_dim, bias=bias, device=device, dtype=dtype)
def forward(self, sample):
return self.linear_2(self.act(self.linear_1(sample)))
class LingBotVideoTextEmbedder(nn.Module):
"""Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
def __init__(self, text_dim: int, hidden_size: int, device=None, dtype=None, operations=None):
super().__init__()
self.norm = operations.RMSNorm(text_dim, eps=1e-6, elementwise_affine=True, device=device, dtype=dtype)
self.linear_1 = operations.Linear(text_dim, hidden_size, bias=True, device=device, dtype=dtype)
self.linear_2 = operations.Linear(hidden_size, hidden_size, bias=True, device=device, dtype=dtype)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.norm(x)
return self.linear_2(F.silu(self.linear_1(x)))
class LingBotVideoAttention(nn.Module):
def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias, device=None, dtype=None, operations=None):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.to_q = operations.Linear(hidden_size, hidden_size, bias=qkv_bias, device=device, dtype=dtype)
self.to_k = operations.Linear(hidden_size, hidden_size, bias=qkv_bias, device=device, dtype=dtype)
self.to_v = operations.Linear(hidden_size, hidden_size, bias=qkv_bias, device=device, dtype=dtype)
self.norm_q = operations.RMSNorm(self.head_dim, eps=norm_eps, elementwise_affine=True, device=device, dtype=dtype)
self.norm_k = operations.RMSNorm(self.head_dim, eps=norm_eps, elementwise_affine=True, device=device, dtype=dtype)
self.to_out = operations.Linear(hidden_size, hidden_size, bias=out_bias, device=device, dtype=dtype)
def forward(
self,
x,
rotary_emb,
attention_mask=None,
transformer_options={},
):
q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
q = apply_rope1(self.norm_q(q), rotary_emb)
k = apply_rope1(self.norm_k(k), rotary_emb)
out = optimized_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
heads=self.num_heads,
mask=attention_mask,
skip_reshape=True,
transformer_options=transformer_options,
)
return self.to_out(out)
class LingBotVideoMLP(nn.Module):
def __init__(self, hidden_size, intermediate_size, device=None, dtype=None, operations=None):
super().__init__()
self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, device=device, dtype=dtype)
def forward(self, x):
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class LingBotVideoRouter(nn.Module):
"""Matches the TokenChoiceTopKRouter inference path (no capacity/jitter/load stats).
The asymmetry must be preserved: selection uses the bias-added score, while gating
weights gather the bias-free score.
"""
def __init__(self, hidden_size, num_experts, top_k, score_func, norm_topk_prob,
n_group, topk_group, route_scale, device=None, dtype=None, operations=None):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.score_func = score_func
self.norm_topk_prob = norm_topk_prob
self.n_group = n_group
self.topk_group = topk_group
self.route_scale = route_scale
self.weight = nn.Parameter(torch.empty(num_experts, hidden_size, device=device, dtype=dtype))
self.register_buffer("e_score_correction_bias", torch.zeros(num_experts, device=device, dtype=dtype), persistent=True)
def _group_limited_topk(self, scores_for_choice):
seq_len = scores_for_choice.shape[0]
experts_per_group = self.num_experts // self.n_group
grouped = scores_for_choice.view(seq_len, self.n_group, experts_per_group)
group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1)
group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
group_mask = torch.zeros_like(group_scores)
group_mask.scatter_(1, group_idx, 1)
score_mask = (
group_mask.unsqueeze(-1)
.expand(seq_len, self.n_group, experts_per_group)
.reshape(seq_len, -1)
)
masked = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
return torch.topk(masked, k=self.top_k, dim=-1, sorted=False)[1]
def forward(self, tokens: torch.Tensor):
logits = F.linear(tokens, self.weight)
if self.score_func == "softmax":
scores = F.softmax(logits, dim=-1)
else:
scores = logits.sigmoid()
scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
if self.n_group is not None and self.n_group > 1:
top_indices = self._group_limited_topk(scores_for_choice)
else:
top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
top_scores = scores.gather(1, top_indices)
if self.top_k > 1 and self.norm_topk_prob:
top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
top_scores = top_scores * self.route_scale
return top_indices, top_scores.to(tokens.dtype), logits, scores, scores_for_choice
class LingBotVideoGroupedExperts(nn.Module):
"""Weight layout matches GroupedExperts: w1 [E,I,H], w2 [E,H,I], w3 [E,I,H]. Eager per-expert compute."""
def __init__(self, num_experts, hidden_size, intermediate_size, device=None, dtype=None):
super().__init__()
self.num_experts = num_experts
self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype))
self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size, device=device, dtype=dtype))
self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype))
class LingBotVideoSparseMoeBlock(nn.Module):
def __init__(self, hidden_size, intermediate_size, num_experts, top_k,
moe_intermediate_size, score_func, norm_topk_prob, n_group, topk_group,
routed_scaling_factor, n_shared_experts, device=None, dtype=None, operations=None):
super().__init__()
self.hidden_size = hidden_size
self.num_experts = num_experts
self.router = LingBotVideoRouter(
hidden_size, num_experts, top_k, score_func, norm_topk_prob,
n_group, topk_group, routed_scaling_factor, device=device, dtype=dtype, operations=operations,
)
self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size, device=device, dtype=dtype)
self.shared_experts = None
if n_shared_experts is not None and n_shared_experts > 0:
self.shared_experts = LingBotVideoMLP(
hidden_size, moe_intermediate_size * n_shared_experts, device=device, dtype=dtype, operations=operations
)
def _run_expert(self, expert_idx: int, tokens: torch.Tensor) -> torch.Tensor:
h = F.silu(tokens @ self.experts.w1[expert_idx].transpose(-2, -1))
h = h * (tokens @ self.experts.w3[expert_idx].transpose(-2, -1))
return h @ self.experts.w2[expert_idx].transpose(-2, -1)
def _run_selected_experts(
self,
tokens: torch.Tensor,
top_scores: torch.Tensor,
top_indices: torch.Tensor,
) -> torch.Tensor:
out = tokens.new_zeros(tokens.shape)
for expert_idx in range(self.num_experts):
selected = top_indices == expert_idx
if not bool(selected.any()):
continue
token_indices, choice_indices = torch.where(selected)
expert_tokens = tokens[token_indices]
expert_output = self._run_expert(expert_idx, expert_tokens)
expert_output = expert_output * top_scores[token_indices, choice_indices].unsqueeze(-1)
out.index_add_(0, token_indices, expert_output)
return out
def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
# hidden_states: (B, S, H); padding_mask: (B*S,) with 1=valid (only needed when B>1)
B = hidden_states.shape[0]
tokens = hidden_states.view(-1, self.hidden_size)
top_indices, top_scores, logits, scores, scores_for_choice = self.router(tokens)
del logits, scores, scores_for_choice
if padding_mask is not None:
pm = padding_mask.unsqueeze(-1).to(top_scores.dtype)
top_scores = top_scores * pm
top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9)
top_scores = top_scores * self.router.route_scale
out = self._run_selected_experts(tokens, top_scores, top_indices)
out = out.view(B, -1, self.hidden_size)
if self.shared_experts is not None:
shared_output = self.shared_experts(hidden_states)
out = out + shared_output
return out
class LingBotVideoBlock(nn.Module):
def __init__(
self,
hidden_size,
num_attention_heads,
intermediate_size,
norm_eps,
qkv_bias,
out_bias,
num_experts,
num_experts_per_tok,
moe_intermediate_size,
decoder_sparse_step,
mlp_only_layers,
n_shared_experts,
score_func,
norm_topk_prob,
n_group,
topk_group,
routed_scaling_factor,
layer_idx: int,
device=None,
dtype=None,
operations=None,
):
super().__init__()
self.layer_idx = layer_idx
h = hidden_size
self.scale_shift_table = nn.Parameter(torch.empty(1, 6 * h, device=device, dtype=dtype))
self.norm1 = operations.RMSNorm(h, eps=norm_eps, elementwise_affine=True, device=device, dtype=dtype)
self.attn = LingBotVideoAttention(
h, num_attention_heads, norm_eps, qkv_bias, out_bias, device=device, dtype=dtype, operations=operations
)
self.norm_post_attn = operations.RMSNorm(h, eps=norm_eps, elementwise_affine=True, device=device, dtype=dtype)
self.norm2 = operations.RMSNorm(h, eps=norm_eps, elementwise_affine=True, device=device, dtype=dtype)
# Sparsity decision matches MoEBlock: mlp_only_layers + decoder_sparse_step + num_experts
if layer_idx not in mlp_only_layers and (
num_experts > 0 and (layer_idx + 1) % decoder_sparse_step == 0
):
self.ffn = LingBotVideoSparseMoeBlock(
h, intermediate_size, num_experts, num_experts_per_tok,
moe_intermediate_size, score_func, norm_topk_prob,
n_group, topk_group, routed_scaling_factor,
n_shared_experts, device=device, dtype=dtype, operations=operations,
)
else:
self.ffn = LingBotVideoMLP(h, intermediate_size, device=device, dtype=dtype, operations=operations)
self.norm_post_ffn = operations.RMSNorm(h, eps=norm_eps, elementwise_affine=True, device=device, dtype=dtype)
def forward(
self,
x,
temb6,
rotary_emb,
attention_mask=None,
moe_padding_mask=None,
transformer_options={},
):
expected_tokens = x.shape[0] * x.shape[1]
if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
raise ValueError(
"LingBotVideoBlock expects token-level temb6 with shape "
f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
)
mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
attn_in = self.norm1(x) * scale_msa + shift_msa
attn_out = self.attn(
attn_in,
rotary_emb,
attention_mask,
transformer_options=transformer_options,
)
x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
ffn_in = self.norm2(x) * scale_mlp + shift_mlp
if isinstance(self.ffn, LingBotVideoSparseMoeBlock):
ffn_out = self.ffn(ffn_in, padding_mask=moe_padding_mask)
else:
ffn_out = self.ffn(ffn_in)
ffn_normed = self.norm_post_ffn(ffn_out)
x = x + (gate_mlp * ffn_normed).to(x.dtype)
return x
class LingBotVideo(nn.Module):
_no_split_modules = ["LingBotVideoBlock"]
def __init__(
self,
image_model=None,
patch_size: Tuple[int, int, int] = (1, 2, 2),
in_channels: int = 16,
out_channels: int = 16,
hidden_size: int = 2048,
num_attention_heads: int = 16,
depth: int = 24,
intermediate_size: int = 6144,
text_dim: int = 2560,
freq_dim: int = 256,
norm_eps: float = 1e-6,
rope_theta: float = 256.0,
axes_dims: Tuple[int, int, int] = (32, 48, 48),
axes_lens: Tuple[int, int, int] = (8192, 1024, 1024),
qkv_bias: bool = False,
out_bias: bool = True,
patch_embed_bias: bool = True,
timestep_mlp_bias: bool = True,
num_experts: int = 0,
num_experts_per_tok: int = 8,
moe_intermediate_size: int = 512,
decoder_sparse_step: int = 1,
mlp_only_layers: Tuple[int, ...] = (),
n_shared_experts: Optional[int] = None,
score_func: str = "sigmoid",
norm_topk_prob: bool = True,
n_group: Optional[int] = None,
topk_group: Optional[int] = None,
routed_scaling_factor: float = 1.0,
dtype=None,
device=None,
operations=None,
):
super().__init__()
self.dtype = dtype
self.patch_size = tuple(patch_size)
self.out_channels = out_channels
head_dim = hidden_size // num_attention_heads
assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
mlp_only_layers = tuple(mlp_only_layers)
self.patch_embedder = operations.Linear(
in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias, device=device, dtype=dtype
)
self.freq_dim = freq_dim
self.time_embedder = LingBotVideoTimestepEmbedding(
freq_dim, hidden_size, bias=timestep_mlp_bias, device=device, dtype=dtype, operations=operations
)
self.time_modulation = nn.Sequential(
nn.SiLU(),
operations.Linear(hidden_size, 6 * hidden_size, device=device, dtype=dtype),
)
self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size, device=device, dtype=dtype, operations=operations)
self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
self.blocks = nn.ModuleList(
[
LingBotVideoBlock(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
norm_eps=norm_eps,
qkv_bias=qkv_bias,
out_bias=out_bias,
num_experts=num_experts,
num_experts_per_tok=num_experts_per_tok,
moe_intermediate_size=moe_intermediate_size,
decoder_sparse_step=decoder_sparse_step,
mlp_only_layers=mlp_only_layers,
n_shared_experts=n_shared_experts,
score_func=score_func,
norm_topk_prob=norm_topk_prob,
n_group=n_group,
topk_group=topk_group,
routed_scaling_factor=routed_scaling_factor,
layer_idx=i,
device=device,
dtype=dtype,
operations=operations,
)
for i in range(depth)
]
)
self.norm_out = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps, device=device, dtype=dtype)
self.norm_out_modulation = nn.Sequential(
nn.SiLU(),
operations.Linear(hidden_size, 2 * hidden_size, device=device, dtype=dtype),
)
self.proj_out = operations.Linear(hidden_size, math.prod(patch_size) * out_channels, device=device, dtype=dtype)
def forward(
self,
hidden_states: torch.Tensor, # (B, C, T, H, W)
timestep: torch.Tensor, # (B,) ∈ [0, 1000](= sigma*1000)
context: torch.Tensor = None, # (B, L, text_dim)
encoder_attention_mask: Optional[torch.Tensor] = None, # (B, L) 1=valid
attention_mask: Optional[torch.Tensor] = None,
transformer_options={},
**kwargs,
):
encoder_hidden_states = context
if encoder_hidden_states is None:
raise ValueError("LingBotVideo requires text conditioning.")
if encoder_attention_mask is None:
encoder_attention_mask = attention_mask
B, C, T, H, W = hidden_states.shape
pF, pH, pW = self.patch_size
gt, gh, gw = T // pF, H // pH, W // pW
n_video = gt * gh * gw
L = encoder_hidden_states.shape[1]
device = hidden_states.device
if encoder_attention_mask is not None:
text_lens = encoder_attention_mask.sum(dim=-1).long()
else:
text_lens = torch.full((B,), L, dtype=torch.long, device=device)
text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
# patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed
patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(
B,
n_video,
pF * pH * pW * C,
)
x = self.patch_embedder(patch_tokens)
text = self.text_embedder(encoder_hidden_states)
joint = torch.cat([x, text], dim=1) # [video; text]
joint_seq_len = joint.shape[1]
# Per-sample RoPE: video t-axis start = real text length of this sample + 1
rotary_parts = [
self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device, L))
for i in range(B)
]
rotary = torch.stack(rotary_parts, dim=0).unsqueeze(2) # (B, S, 1, head_dim/2, 2, 2)
attention_mask = None
moe_padding_mask = None
has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
if has_padding:
key_mask = torch.cat(
[torch.ones(B, n_video, dtype=torch.bool, device=device),
encoder_attention_mask.bool()],
dim=1,
)
attention_mask = key_mask[:, None, None, :] # (B,1,1,S) → SDPA broadcast
moe_padding_mask = key_mask.reshape(-1) # (B*S,)
timestep_proj = timestep_embedding(timestep.to(hidden_states.dtype), self.freq_dim, time_factor=1.0)
t_emb = self.time_embedder(timestep_proj) # (B, D)
temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1) # (B, S, D)
temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1))
temb6 = temb6.reshape(B, joint_seq_len, -1) # (B, S, 6D)
temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1)
for block in self.blocks:
joint = block(
joint,
temb6,
rotary,
attention_mask,
moe_padding_mask,
transformer_options=transformer_options,
)
final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
projected = self.proj_out(final_hidden)
x = projected[:, :n_video]
# unpatchify (matches the rearrange in postprocess)
Cout = self.out_channels
x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
return x
LingBotVideoTransformer3DModel = LingBotVideo
+65 -98
View File
@@ -1,6 +1,5 @@
import math
import sys
import inspect
import torch
import torch.nn.functional as F
@@ -15,16 +14,16 @@ from .sub_quadratic_attention import efficient_dot_product_attention
from comfy import model_management
TORCH_HAS_GQA = model_management.torch_version_numeric >= (2, 5)
if model_management.xformers_enabled():
import xformers
import xformers.ops
SAGE_ATTENTION_IS_AVAILABLE = False
SAGE_ATTENTION_SUPPORTS_MASK = False
try:
from sageattention import sageattn
SAGE_ATTENTION_IS_AVAILABLE = True
SAGE_ATTENTION_SUPPORTS_MASK = "attn_mask" in inspect.signature(sageattn).parameters
except ImportError as e:
if model_management.sage_attention_enabled():
if e.name == "sageattention":
@@ -90,44 +89,6 @@ def default(val, d):
return val
return d
def _gqa_repeat_factor(query_heads, key_heads, value_heads):
if key_heads != value_heads:
raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
if query_heads == key_heads:
return 1
if query_heads % key_heads != 0:
raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
return query_heads // key_heads
def _repeat_kv_for_gqa(k, v, query_heads, head_dim):
n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
if n_rep > 1:
k = k.repeat_interleave(n_rep, dim=head_dim)
v = v.repeat_interleave(n_rep, dim=head_dim)
return k, v
def _heads_from_dim(tensor, dim_head, name):
inner_dim = tensor.shape[-1]
if inner_dim % dim_head != 0:
raise ValueError(f"{name} inner dimension {inner_dim} is not divisible by head dimension {dim_head}")
return inner_dim // dim_head
def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_kv=True):
q = q.unsqueeze(3).reshape(b, -1, heads, dim_head)
if enable_gqa:
key_heads = _heads_from_dim(k, dim_head, "Key")
value_heads = _heads_from_dim(v, dim_head, "Value")
else:
key_heads = heads
value_heads = heads
k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head)
v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head)
if enable_gqa:
_gqa_repeat_factor(heads, key_heads, value_heads)
if expand_kv:
k, v = _repeat_kv_for_gqa(k, v, heads, -2)
return q, k, v
# feedforward
class GEGLU(nn.Module):
@@ -191,19 +152,28 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
b, _, dim_head = q.shape
dim_head //= heads
if kwargs.get("enable_gqa", False) and q.shape[-3] != k.shape[-3]:
n_rep = q.shape[-3] // k.shape[-3]
k = k.repeat_interleave(n_rep, dim=-3)
v = v.repeat_interleave(n_rep, dim=-3)
scale = kwargs.get("scale", dim_head ** -0.5)
h = heads
if skip_reshape:
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
q, k, v = map(
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head),
(q, k, v),
)
else:
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v))
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, -1, heads, dim_head)
.permute(0, 2, 1, 3)
.reshape(b * heads, -1, dim_head)
.contiguous(),
(q, k, v),
)
# force cast to fp32 to avoid overflowing
if attn_precision == torch.float32:
@@ -261,16 +231,13 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
query = query * (kwargs["scale"] * dim_head ** 0.5)
if skip_reshape:
if kwargs.get("enable_gqa", False):
key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3)
query = query.reshape(b * heads, -1, dim_head)
value = value.reshape(b * heads, -1, dim_head)
key = key.reshape(b * heads, -1, dim_head).movedim(1, 2)
else:
query, key, value = _reshape_qkv_to_heads(query, key, value, b, heads, dim_head, kwargs.get("enable_gqa", False))
query = query.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
value = value.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
key = key.permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
query = query.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
value = value.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
dtype = query.dtype
@@ -337,15 +304,19 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
scale = kwargs.get("scale", dim_head ** -0.5)
if skip_reshape:
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
q, k, v = map(
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head),
(q, k, v),
)
else:
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v))
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, -1, heads, dim_head)
.permute(0, 2, 1, 3)
.reshape(b * heads, -1, dim_head)
.contiguous(),
(q, k, v),
)
r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
@@ -467,7 +438,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
disabled_xformers = True
if disabled_xformers:
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs)
if skip_reshape:
# b h k d -> b k h d
@@ -475,12 +446,13 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
lambda t: t.permute(0, 2, 1, 3),
(q, k, v),
)
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2)
# actually do the reshaping
else:
dim_head //= heads
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
q, k, v = map(
lambda t: t.reshape(b, -1, heads, dim_head),
(q, k, v),
)
if mask is not None:
# add a singleton batch dimension
@@ -502,7 +474,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
mask = mask_out[..., :mask.shape[-1]]
mask = mask.expand(b, heads, -1, -1)
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, scale=kwargs.get("scale", None))
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
if skip_output_reshape:
out = out.permute(0, 2, 1, 3)
@@ -526,8 +498,10 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
else:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False)
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
q, k, v = map(
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
(q, k, v),
)
if mask is not None:
# add a batch dimension if there isn't already one
@@ -537,7 +511,9 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
if mask.ndim == 3:
mask = mask.unsqueeze(1)
sdpa_keys = ("scale", "enable_gqa")
# Pass through extra SDPA kwargs (scale, enable_gqa) if provided
# enable_gqa requires PyTorch 2.5+; older versions use manual KV expansion above
sdpa_keys = ("scale", "enable_gqa") if TORCH_HAS_GQA else ("scale",)
sdpa_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys}
if SDP_BATCH_LIMIT >= b:
@@ -565,19 +541,20 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
@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):
if kwargs.get("low_precision_attention", True) is False:
return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
exception_fallback = False
if skip_reshape:
b, _, _, dim_head = q.shape
tensor_layout = "HND"
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
else:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
q, k, v = map(
lambda t: t.view(b, -1, heads, dim_head),
(q, k, v),
)
tensor_layout = "NHD"
if mask is not None:
@@ -588,12 +565,8 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
if mask.ndim == 3:
mask = mask.unsqueeze(1)
sage_kwargs = {"is_causal": False, "tensor_layout": tensor_layout, "sm_scale": kwargs.get("scale", None), "smooth_k": False}
if mask is not None:
sage_kwargs["attn_mask"] = mask
try:
out = sageattn(q, k, v, **sage_kwargs)
out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout)
except Exception as e:
logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e))
exception_fallback = True
@@ -643,6 +616,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
skip_output_reshape=skip_output_reshape,
**kwargs
)
q_s, k_s, v_s = q, k, v
N = q.shape[2]
dim_head = D
else:
@@ -668,15 +642,11 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
**kwargs
)
if skip_reshape:
q_s = q
if kwargs.get("enable_gqa", False):
k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3)
else:
k_s, v_s = k, v
else:
q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False))
q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s))
if not skip_reshape:
q_s, k_s, v_s = map(
lambda t: t.view(B, -1, heads, dim_head).permute(0, 2, 1, 3).contiguous(),
(q, k, v),
)
B, H, L, D = q_s.shape
try:
@@ -692,7 +662,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
q, k, v, heads,
mask=mask,
attn_precision=attn_precision,
skip_reshape=skip_reshape,
skip_reshape=False,
skip_output_reshape=skip_output_reshape,
**kwargs
)
@@ -711,20 +681,19 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
try:
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal, softmax_scale=softmax_scale_arg)
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
@flash_attn_wrapper.register_fake
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False, softmax_scale=-1.0):
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False):
# Output shape is the same as q
return q.new_empty(q.shape)
except AttributeError as error:
FLASH_ATTN_ERROR = error
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}"
@wrap_attn
@@ -734,8 +703,10 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
else:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False)
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
q, k, v = map(
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
(q, k, v),
)
if mask is not None:
# add a batch dimension if there isn't already one
@@ -754,16 +725,10 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
v.transpose(1, 2),
dropout_p=0.0,
causal=False,
softmax_scale=kwargs.get("scale", -1.0),
).transpose(1, 2)
except Exception as e:
logging.warning(f"Flash Attention failed, using default SDPA: {e}")
sdpa_extra = {}
if kwargs.get("enable_gqa", False):
sdpa_extra["enable_gqa"] = True
if "scale" in kwargs:
sdpa_extra["scale"] = kwargs["scale"]
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra)
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
if not skip_output_reshape:
out = (
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
@@ -1244,3 +1209,5 @@ class SpatialVideoTransformer(SpatialTransformer):
x = self.proj_out(x)
out = x + x_in
return out
+5 -2
View File
@@ -141,8 +141,11 @@ class Attention(nn.Module):
key = key.transpose(1, 2)
value = value.transpose(1, 2)
gqa_kwargs = {"enable_gqa": True} if self.kv_heads < self.heads else {}
hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs)
if self.kv_heads < self.heads:
key = key.repeat_interleave(self.heads // self.kv_heads, dim=1)
value = value.repeat_interleave(self.heads // self.kv_heads, dim=1)
hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
hidden_states = self.to_out[0](hidden_states)
return hidden_states
-15
View File
@@ -63,7 +63,6 @@ import comfy.ldm.kandinsky5.model
import comfy.ldm.anima.model
import comfy.ldm.ace.ace_step15
import comfy.ldm.cogvideo.model
import comfy.ldm.lingbot_video.model
import comfy.ldm.rt_detr.rtdetr_v4
import comfy.ldm.ernie.model
import comfy.ldm.sam3.detector
@@ -1377,20 +1376,6 @@ class HunyuanVideoSkyreelsI2V(HunyuanVideo):
def scale_latent_inpaint(self, latent_image, **kwargs):
return super().scale_latent_inpaint(latent_image=latent_image, **kwargs)
class LingBotVideo(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.lingbot_video.model.LingBotVideo)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
cross_attn = kwargs.get("cross_attn", None)
if cross_attn is not None:
out["c_crossattn"] = comfy.conds.CONDRegular(cross_attn)
return out
def scale_latent_inpaint(self, latent_image, **kwargs):
return latent_image
class CosmosVideo(BaseModel):
def __init__(self, model_config, model_type=ModelType.EDM, image_to_video=False, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.cosmos.model.GeneralDIT)
-48
View File
@@ -232,54 +232,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["meanflow_sum"] = False
return dit_config
if '{}patch_embedder.weight'.format(key_prefix) in state_dict_keys and '{}text_embedder.norm.weight'.format(key_prefix) in state_dict_keys and '{}blocks.0.attn.to_q.weight'.format(key_prefix) in state_dict_keys: # LingBot Video
dit_config = {}
patch_size = (1, 2, 2)
patch_prod = math.prod(patch_size)
patch_embed = state_dict['{}patch_embedder.weight'.format(key_prefix)]
proj_out = state_dict['{}proj_out.weight'.format(key_prefix)]
hidden_size = patch_embed.shape[0]
dit_config["image_model"] = "lingbot_video"
dit_config["patch_size"] = patch_size
dit_config["in_channels"] = 16
dit_config["out_channels"] = proj_out.shape[0] // patch_prod
dit_config["hidden_size"] = hidden_size
dit_config["num_attention_heads"] = hidden_size // 128
dit_config["depth"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
dit_config["text_dim"] = state_dict['{}text_embedder.norm.weight'.format(key_prefix)].shape[0]
dit_config["freq_dim"] = 256
dit_config["qkv_bias"] = '{}blocks.0.attn.to_q.bias'.format(key_prefix) in state_dict_keys
dit_config["out_bias"] = '{}blocks.0.attn.to_out.bias'.format(key_prefix) in state_dict_keys
dit_config["patch_embed_bias"] = '{}patch_embedder.bias'.format(key_prefix) in state_dict_keys
dit_config["timestep_mlp_bias"] = '{}time_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys
moe_key = '{}blocks.0.ffn.experts.w1'.format(key_prefix)
if moe_key in state_dict_keys:
experts = state_dict[moe_key]
dit_config["num_experts"] = experts.shape[0]
dit_config["moe_intermediate_size"] = experts.shape[1]
if experts.shape[0] == 128:
dit_config["n_group"] = 4
dit_config["topk_group"] = 2
dit_config["routed_scaling_factor"] = 2.5
mlp_only_layers = []
for i in range(dit_config["depth"]):
if '{}blocks.{}.ffn.gate_proj.weight'.format(key_prefix, i) in state_dict_keys:
mlp_only_layers.append(i)
dit_config["mlp_only_layers"] = tuple(mlp_only_layers)
if len(mlp_only_layers) > 0:
dit_config["intermediate_size"] = state_dict['{}blocks.{}.ffn.gate_proj.weight'.format(key_prefix, mlp_only_layers[0])].shape[0]
if '{}blocks.0.ffn.shared_experts.gate_proj.weight'.format(key_prefix) in state_dict_keys:
shared = state_dict['{}blocks.0.ffn.shared_experts.gate_proj.weight'.format(key_prefix)]
dit_config["n_shared_experts"] = shared.shape[0] // experts.shape[1]
else:
dit_config["num_experts"] = 0
dit_config["intermediate_size"] = state_dict['{}blocks.0.ffn.gate_proj.weight'.format(key_prefix)].shape[0]
if metadata is not None and "config" in metadata:
dit_config.update(json.loads(metadata["config"]).get("transformer", {}))
return dit_config
if any_suffix_in(state_dict_keys, key_prefix, 'double_blocks.0.img_attn.norm.key_norm.', ["weight", "scale"]) and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.norms.0.', ["weight", "scale"])): #Flux, Chroma or Chroma Radiance (has no img_in.weight)
dit_config = {}
if '{}double_stream_modulation_img.lin.weight'.format(key_prefix) in state_dict_keys:
-2
View File
@@ -174,8 +174,6 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
elif xfer_dest2 is not None:
xfer_source.prepare(xfer_dest2, stream, copy=True, commit=False)
return
else:
return
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream, r2=xfer_dest2)
def handle_pin(m, pin, source, dest, subset="weights", size=None):
+1 -10
View File
@@ -69,7 +69,6 @@ import comfy.text_encoders.ace15
import comfy.text_encoders.longcat_image
import comfy.text_encoders.qwen35
import comfy.text_encoders.qwen3vl
import comfy.text_encoders.lingbot_video
import comfy.text_encoders.boogu
import comfy.text_encoders.ernie
import comfy.text_encoders.gemma4
@@ -1256,10 +1255,7 @@ class VAE:
return None
def is_dynamic(self):
# A VAE built from a state dict with no detectable VAE weights returns early
# from __init__ ("No VAE weights detected") before self.patcher is assigned.
patcher = getattr(self, "patcher", None)
return patcher is not None and patcher.is_dynamic()
return self.patcher.is_dynamic()
class StyleModel:
def __init__(self, model, device="cpu"):
@@ -1314,7 +1310,6 @@ class CLIPType(Enum):
IDEOGRAM4 = 30
BOOGU = 31
KREA2 = 32
LINGBOT_VIDEO = 33
@@ -1648,10 +1643,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
clip_target.tokenizer = comfy.text_encoders.flux.KleinTokenizer8B if te_model == TEModel.QWEN3VL_8B else comfy.text_encoders.flux.KleinTokenizer
elif clip_type == CLIPType.LINGBOT_VIDEO and te_model == TEModel.QWEN3VL_4B:
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
clip_target.clip = comfy.text_encoders.lingbot_video.te(**llama_detect(clip_data), model_type="qwen3vl_4b")
clip_target.tokenizer = comfy.text_encoders.lingbot_video.tokenizer(model_type="qwen3vl_4b")
else:
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
qwen3vl_type = {TEModel.QWEN3VL_4B: "qwen3vl_4b", TEModel.QWEN3VL_8B: "qwen3vl_8b"}[te_model]
-37
View File
@@ -27,8 +27,6 @@ import comfy.text_encoders.z_image
import comfy.text_encoders.ideogram4
import comfy.text_encoders.boogu
import comfy.text_encoders.krea2
import comfy.text_encoders.qwen3vl
import comfy.text_encoders.lingbot_video
import comfy.text_encoders.anima
import comfy.text_encoders.ace15
import comfy.text_encoders.longcat_image
@@ -1025,40 +1023,6 @@ class HunyuanVideoSkyreelsI2V(HunyuanVideo):
out = model_base.HunyuanVideoSkyreelsI2V(self, device=device)
return out
class LingBotVideo(supported_models_base.BASE):
unet_config = {
"image_model": "lingbot_video",
}
sampling_settings = {
"shift": 3.0,
}
unet_extra_config = {}
latent_format = latent_formats.LingBotVideo
memory_usage_factor = 1.8
supported_inference_dtypes = [torch.bfloat16, torch.float32]
vae_key_prefix = ["vae."]
text_encoder_key_prefix = ["text_encoders."]
def __init__(self, unet_config):
super().__init__(unet_config)
self.memory_usage_factor = self.memory_usage_factor * (unet_config.get("hidden_size", 2048) / 2048)
def get_model(self, state_dict, prefix="", device=None):
return model_base.LingBotVideo(self, device=device)
def clip_target(self, state_dict={}):
pref = self.text_encoder_key_prefix[0]
qwen3vl_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref))
if len(qwen3vl_detect) > 0:
qwen3vl_detect["model_type"] = "qwen3vl_4b"
return supported_models_base.ClipTarget(comfy.text_encoders.lingbot_video.tokenizer(model_type="qwen3vl_4b"), comfy.text_encoders.lingbot_video.te(**qwen3vl_detect))
return None
class CosmosT2V(supported_models_base.BASE):
unet_config = {
"image_model": "cosmos",
@@ -2353,7 +2317,6 @@ models = [
HunyuanVideoSkyreelsI2V,
HunyuanVideoI2V,
HunyuanVideo,
LingBotVideo,
CosmosT2V,
CosmosI2V,
CosmosT2IPredict2,
+5 -1
View File
@@ -12,7 +12,7 @@ import torch.nn.functional as F
import comfy.ops
from comfy import sd1_clip
from comfy.ldm.modules.attention import optimized_attention_for_device
from comfy.ldm.modules.attention import TORCH_HAS_GQA, optimized_attention_for_device
from comfy.text_encoders.llama import RMSNorm, apply_rope
@@ -110,6 +110,10 @@ def _attention_with_sinks(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sin
putting the sink logit in the mask at that column.
"""
if num_kv_groups > 1 and not TORCH_HAS_GQA:
k = k.repeat_interleave(num_kv_groups, dim=1)
v = v.repeat_interleave(num_kv_groups, dim=1)
B, _, S_q, D = q.shape
H_kv = k.shape[1]
S_kv = k.shape[-2]
-149
View File
@@ -1,149 +0,0 @@
import comfy.text_encoders.qwen3vl
from comfy import sd1_clip
PAD_TOKEN = 151643
CROP_MARKER = {"type": "lingbot_video_crop_start"}
PROMPT_TEMPLATE = (
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
"a text prompt with an image reference, or a text prompt with a video reference "
"or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
"visual descriptions suitable for video generation. Evaluate the level of detail "
"in the user's input: if it is simple, enrich it by adding specifics about colors, "
"shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
"progression, and spatial relationships to create vivid, concrete, and temporally "
"coherent scenes to create vivid and concrete scenes. Please generate only the "
"enhanced description for the prompt below and avoid including any additional "
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
"<|im_start|>assistant\n"
)
IMAGE_PROMPT_TEMPLATE = "<|vision_start|><|image_pad|><|vision_end|>"
def _marker_tuple(example=None):
if example is not None and len(example) > 2:
return (CROP_MARKER, 1.0, None)
return (CROP_MARKER, 1.0)
class LingBotVideoTokenizer(comfy.text_encoders.qwen3vl.Qwen3VLTokenizer):
def __init__(self, embedding_directory=None, tokenizer_data={}, model_type="qwen3vl_4b"):
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type=model_type)
self._lingbot_crop_start = None
def lingbot_crop_start(self):
if self._lingbot_crop_start is not None:
return self._lingbot_crop_start
prefix = PROMPT_TEMPLATE.split("{}")[0]
tokens = super().tokenize_with_weights(prefix, thinking=True)
key = next(iter(tokens))
count = 0
for token in tokens[key][0]:
token_id = token[0]
if isinstance(token_id, int) and token_id == PAD_TOKEN:
continue
count += 1
self._lingbot_crop_start = count
return count
def tokenize_with_weights(self, text, return_word_ids=False, images=[], prevent_empty_text=False, thinking=True, **kwargs):
image = kwargs.get("image", None)
if image is not None and len(images) == 0:
images = [image[i:i + 1] for i in range(image.shape[0])]
prompt_text = text
if len(images) > 0 and not prompt_text.startswith("<|vision_start|>"):
prompt_text = IMAGE_PROMPT_TEMPLATE + prompt_text
tokens = super().tokenize_with_weights(
prompt_text,
return_word_ids=return_word_ids,
llama_template=PROMPT_TEMPLATE,
images=images,
prevent_empty_text=prevent_empty_text,
thinking=thinking,
**kwargs,
)
crop_start = self.lingbot_crop_start()
for key in tokens:
for row in tokens[key]:
example = row[0] if len(row) > 0 else None
row.insert(crop_start, _marker_tuple(example))
return tokens
class LingBotVideoClipModel(comfy.text_encoders.qwen3vl.Qwen3VLClipModel):
def __init__(self, device="cpu", dtype=None, attention_mask=True, model_options={}, model_type="qwen3vl_4b"):
super().__init__(
device=device,
layer="last",
layer_idx=None,
dtype=dtype,
attention_mask=attention_mask,
model_options=model_options,
model_type=model_type,
)
self.return_attention_masks = False
@staticmethod
def _strip_crop_markers(tokens):
clean_tokens = []
crop_starts = []
for row in tokens:
clean_row = []
crop_start = None
for token in row:
token_value = token[0] if isinstance(token, tuple) else token
if isinstance(token_value, dict) and token_value.get("type") == CROP_MARKER["type"]:
crop_start = len(clean_row)
continue
clean_row.append(token)
clean_tokens.append(clean_row)
crop_starts.append(crop_start)
return clean_tokens, crop_starts
def forward(self, tokens):
clean_tokens, crop_starts = self._strip_crop_markers(tokens)
out = super().forward(clean_tokens)
crop_starts = [c for c in crop_starts if c is not None]
if len(crop_starts) == 0:
return out
crop_start = min(crop_starts)
z, pooled_output = out[:2]
z = z[:, crop_start:]
return z, pooled_output
class LingBotVideoTEModel(comfy.text_encoders.qwen3vl.Qwen3VLTEModel):
def __init__(self, device="cpu", dtype=None, model_options={}, model_type="qwen3vl_4b"):
clip_model = lambda **kw: LingBotVideoClipModel(**kw, model_type=model_type)
sd1_clip.SD1ClipModel.__init__(
self,
device=device,
dtype=dtype,
name=model_type,
clip_model=clip_model,
model_options=model_options,
)
def tokenizer(model_type="qwen3vl_4b"):
class LingBotVideoTokenizer_(LingBotVideoTokenizer):
def __init__(self, embedding_directory=None, tokenizer_data={}):
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type=model_type)
return LingBotVideoTokenizer_
def te(dtype_llama=None, llama_quantization_metadata=None, model_type="qwen3vl_4b"):
class LingBotVideoTEModel_(LingBotVideoTEModel):
def __init__(self, device="cpu", dtype=None, model_options={}):
if dtype_llama is not None:
dtype = dtype_llama
if llama_quantization_metadata is not None:
model_options = model_options.copy()
model_options["quantization_metadata"] = llama_quantization_metadata
super().__init__(device=device, dtype=dtype, model_options=model_options, model_type=model_type)
return LingBotVideoTEModel_
+4 -2
View File
@@ -550,8 +550,10 @@ class Attention(nn.Module):
xv = xv[:, :, -sliding_window:]
attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
return self.o_proj(output), present_key_value
class MLP(nn.Module):
+6 -2
View File
@@ -366,8 +366,12 @@ class GatedAttention(nn.Module):
xv = torch.cat((past_value[:, :, :index], xv), dim=2)
present_key_value = (xk, xv, index + num_tokens)
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
# Expand KV heads for GQA
if self.num_heads != self.num_kv_heads:
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
output = output * gate.sigmoid()
return self.o_proj(output), present_key_value
-21
View File
@@ -90,27 +90,6 @@ class Qwen3VL(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
deepstack = [torch.cat([deepstack[i], ds[i]], dim=0) for i in range(len(ds))]
return position_ids, visual_pos_masks, deepstack
def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[], **kwargs):
position_ids = kwargs.pop("position_ids", None)
visual_pos_masks = kwargs.pop("visual_pos_masks", None)
deepstack_embeds = kwargs.pop("deepstack_embeds", None)
if embeds is not None and position_ids is None:
position_ids, visual_pos_masks, deepstack_embeds = self.build_image_inputs(embeds, embeds_info)
return self.model(
input_ids,
attention_mask=attention_mask,
embeds=embeds,
num_tokens=num_tokens,
intermediate_output=intermediate_output,
final_layer_norm_intermediate=final_layer_norm_intermediate,
dtype=dtype,
position_ids=position_ids,
embeds_info=embeds_info,
visual_pos_masks=visual_pos_masks,
deepstack_embeds=deepstack_embeds,
**kwargs,
)
def _make_qwen3vl_model(model_type):
class Qwen3VL_(Qwen3VL):
+2 -15
View File
@@ -24,8 +24,8 @@ class Seedream4TaskCreationRequest(BaseModel):
image: list[str] | None = Field(None, description="Image URLs")
size: str = Field(...)
seed: int = Field(..., ge=0, le=2147483647)
sequential_image_generation: str | None = Field("disabled")
sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15))
sequential_image_generation: str = Field("disabled")
sequential_image_generation_options: Seedream4Options = Field(Seedream4Options(max_images=15))
watermark: bool = Field(False)
output_format: str | None = None
@@ -261,19 +261,6 @@ _PRESETS_SEEDREAM_4K = [
_CUSTOM_PRESET = [("Custom", None, None)]
_PRESETS_SEEDREAM_2K_PRO = [
("(2K) 2048x2048 (1:1)", 2048, 2048),
("(2K) 1728x2304 (3:4)", 1728, 2304),
("(2K) 2304x1728 (4:3)", 2304, 1728),
# ("(2K) 2848x1600 (16:9)", 2848, 1600), # 4,556,800 px - temporarily unavailable
# ("(2K) 1600x2848 (9:16)", 1600, 2848), # 4,556,800 px - temporarily unavailable
("(2K) 1664x2496 (2:3)", 1664, 2496),
("(2K) 2496x1664 (3:2)", 2496, 1664),
# ("(2K) 3136x1344 (21:9)", 3136, 1344), # 4,214,784 px - temporarily unavailable
]
RECOMMENDED_PRESETS_SEEDREAM_5_PRO = (
_PRESETS_SEEDREAM_1K + _PRESETS_SEEDREAM_2K_PRO + _CUSTOM_PRESET
)
RECOMMENDED_PRESETS_SEEDREAM_5_LITE = (
_PRESETS_SEEDREAM_2K + _PRESETS_SEEDREAM_3K + _PRESETS_SEEDREAM_4K + _CUSTOM_PRESET
)
+44 -96
View File
@@ -16,7 +16,6 @@ from comfy_api_nodes.apis.bytedance import (
RECOMMENDED_PRESETS_SEEDREAM_4_0,
RECOMMENDED_PRESETS_SEEDREAM_4_5,
RECOMMENDED_PRESETS_SEEDREAM_5_LITE,
RECOMMENDED_PRESETS_SEEDREAM_5_PRO,
SEEDANCE2_REF_VIDEO_PIXEL_LIMITS,
VIDEO_TASKS_EXECUTION_TIME,
GetAssetResponse,
@@ -81,14 +80,12 @@ _VERIFICATION_POLL_TIMEOUT_SEC = 120
_VERIFICATION_POLL_INTERVAL_SEC = 3
SEEDREAM_MODELS = {
"seedream 5.0 pro": "seedream-5-0-pro-260628",
"seedream 5.0 lite": "seedream-5-0-260128",
"seedream-4-5-251128": "seedream-4-5-251128",
"seedream-4-0-250828": "seedream-4-0-250828",
}
SEEDREAM_PRESETS = {
"seedream-5-0-pro-260628": RECOMMENDED_PRESETS_SEEDREAM_5_PRO,
"seedream-5-0-260128": RECOMMENDED_PRESETS_SEEDREAM_5_LITE,
"seedream-4-5-251128": RECOMMENDED_PRESETS_SEEDREAM_4_5,
"seedream-4-0-250828": RECOMMENDED_PRESETS_SEEDREAM_4_0,
@@ -746,15 +743,8 @@ class ByteDanceSeedreamNode(IO.ComfyNode):
return IO.NodeOutput(torch.cat([await download_url_to_image_tensor(i) for i in urls]))
def _seedream_model_inputs(
*,
max_ref_images: int,
presets: list,
max_width: int = 6240,
max_height: int = 4992,
supports_batch: bool = True,
):
inputs = [
def _seedream_model_inputs(*, max_ref_images: int, presets: list):
return [
IO.Combo.Input(
"size_preset",
options=[label for label, _, _ in presets],
@@ -764,7 +754,7 @@ def _seedream_model_inputs(
"width",
default=2048,
min=1024,
max=max_width,
max=6240,
step=2,
tooltip="Custom width for image. Value is working only if `size_preset` is set to `Custom`",
),
@@ -772,27 +762,22 @@ def _seedream_model_inputs(
"height",
default=2048,
min=1024,
max=max_height,
max=4992,
step=2,
tooltip="Custom height for image. Value is working only if `size_preset` is set to `Custom`",
),
]
if supports_batch:
inputs.append(
IO.Int.Input(
"max_images",
default=1,
min=1,
max=max_ref_images,
step=1,
display_mode=IO.NumberDisplay.number,
tooltip="Maximum number of images to generate. With 1, exactly one image is produced. "
"With >1, the model generates between 1 and max_images related images "
"(e.g., story scenes, character variations). "
"Total images (input + generated) cannot exceed 15.",
)
)
inputs.append(
IO.Int.Input(
"max_images",
default=1,
min=1,
max=max_ref_images,
step=1,
display_mode=IO.NumberDisplay.number,
tooltip="Maximum number of images to generate. With 1, exactly one image is produced. "
"With >1, the model generates between 1 and max_images related images "
"(e.g., story scenes, character variations). "
"Total images (input + generated) cannot exceed 15.",
),
IO.Autogrow.Input(
"images",
template=IO.Autogrow.TemplateNames(
@@ -802,18 +787,14 @@ def _seedream_model_inputs(
),
tooltip=f"Optional reference image(s) for image-to-image or multi-reference generation. "
f"Up to {max_ref_images} images.",
)
)
if supports_batch:
inputs.append(
IO.Boolean.Input(
"fail_on_partial",
default=False,
tooltip="If enabled, abort execution if any requested images are missing or return an error.",
advanced=True,
)
)
return inputs
),
IO.Boolean.Input(
"fail_on_partial",
default=False,
tooltip="If enabled, abort execution if any requested images are missing or return an error.",
advanced=True,
),
]
class ByteDanceSeedreamNodeV2(IO.ComfyNode):
@@ -835,16 +816,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option(
"seedream 5.0 pro",
_seedream_model_inputs(
max_ref_images=10,
presets=RECOMMENDED_PRESETS_SEEDREAM_5_PRO,
max_width=3136,
max_height=2496,
supports_batch=False,
),
),
IO.DynamicCombo.Option(
"seedream 5.0 lite",
_seedream_model_inputs(max_ref_images=14, presets=RECOMMENDED_PRESETS_SEEDREAM_5_LITE),
@@ -886,27 +857,15 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
],
is_api_node=True,
price_badge=IO.PriceBadge(
depends_on=IO.PriceBadgeDepends(
widgets=["model", "model.size_preset", "model.width", "model.height"]
),
depends_on=IO.PriceBadgeDepends(widgets=["model"]),
expr="""
(
$sp := $lookup(widgets, "model.size_preset");
$px := $lookup(widgets, "model.width") * $lookup(widgets, "model.height");
$isPro := $contains(widgets.model, "5.0 pro");
$price := $isPro
? (
$contains($sp, "custom")
? ($px <= 2360000 ? 0.045 : 0.09)
: ($contains($sp, "1k") ? 0.045 : 0.09)
)
: $contains(widgets.model, "5.0 lite") ? 0.035
: $contains(widgets.model, "4-5") ? 0.04
: 0.03;
$price := $contains(widgets.model, "5.0 lite") ? 0.035 :
$contains(widgets.model, "4-5") ? 0.04 : 0.03;
{
"type": "usd",
"type":"usd",
"usd": $price,
"format": { "suffix": $isPro ? "/Image" : " x images/Run", "approximate": true }
"format": { "suffix":" x images/Run", "approximate": true }
}
)
""",
@@ -924,7 +883,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
validate_string(prompt, strip_whitespace=True, min_length=1)
model_id = SEEDREAM_MODELS[model["model"]]
presets = SEEDREAM_PRESETS[model_id]
is_pro = "seedream-5-0-pro" in model_id
size_preset = model.get("size_preset", presets[0][0])
width = model.get("width", 2048)
@@ -944,29 +902,19 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
out_num_pixels = w * h
mp_provided = out_num_pixels / 1_000_000.0
if is_pro:
if out_num_pixels < 921_600:
raise ValueError(
f"Minimum image resolution for the selected model is 0.92MP, but {mp_provided:.2f}MP provided."
)
if out_num_pixels > 4_194_304:
raise ValueError(
f"Maximum image resolution for the selected model is 4.19MP, but {mp_provided:.2f}MP provided."
)
else:
if ("seedream-4-5" in model_id or "seedream-5-0" in model_id) and out_num_pixels < 3_686_400:
raise ValueError(
f"Minimum image resolution for the selected model is 3.68MP, but {mp_provided:.2f}MP provided."
)
if "seedream-4-0" in model_id and out_num_pixels < 921_600:
raise ValueError(
f"Minimum image resolution that the selected model can generate is 0.92MP, "
f"but {mp_provided:.2f}MP provided."
)
if out_num_pixels > 16_777_216:
raise ValueError(
f"Maximum image resolution for the selected model is 16.78MP, but {mp_provided:.2f}MP provided."
)
if ("seedream-4-5" in model_id or "seedream-5-0" in model_id) and out_num_pixels < 3686400:
raise ValueError(
f"Minimum image resolution for the selected model is 3.68MP, but {mp_provided:.2f}MP provided."
)
if "seedream-4-0" in model_id and out_num_pixels < 921600:
raise ValueError(
f"Minimum image resolution that the selected model can generate is 0.92MP, "
f"but {mp_provided:.2f}MP provided."
)
if out_num_pixels > 16_777_216:
raise ValueError(
f"Maximum image resolution for the selected model is 16.78MP, but {mp_provided:.2f}MP provided."
)
image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None]
n_input_images = sum(get_number_of_images(t) for t in image_tensors)
@@ -1002,8 +950,8 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
image=reference_images_urls,
size=f"{w}x{h}",
seed=seed,
sequential_image_generation=None if is_pro else sequential_image_generation,
sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images),
sequential_image_generation=sequential_image_generation,
sequential_image_generation_options=Seedream4Options(max_images=max_images),
watermark=watermark,
),
)
+12 -3
View File
@@ -56,6 +56,9 @@ PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'})
# 3D file extensions for preview fallback (no dedicated media_type exists)
THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb', '.usdz'})
# Text file extensions for preview fallback (the formats SaveText can produce)
TEXT_EXTENSIONS = frozenset({'.txt', '.md', '.json'})
def has_3d_extension(filename: str) -> bool:
lower = filename.lower()
@@ -143,9 +146,10 @@ def is_previewable(media_type: str, item: dict) -> bool:
Maintains backwards compatibility with existing logic.
Priority:
1. media_type is 'images', 'video', 'audio', or '3d'
1. media_type is 'images', 'video', 'audio', '3d', or 'text'
2. format field starts with 'video/' or 'audio/'
3. filename has a 3D extension (.obj, .fbx, .gltf, .glb, .usdz)
4. filename has a text extension (.txt, .md, .json, ...)
"""
if media_type in PREVIEWABLE_MEDIA_TYPES:
return True
@@ -156,10 +160,12 @@ def is_previewable(media_type: str, item: dict) -> bool:
if fmt and (fmt.startswith('video/') or fmt.startswith('audio/')):
return True
# Check for 3D files by extension
# Check for 3D and text files by extension
filename = item.get('filename', '').lower()
if any(filename.endswith(ext) for ext in THREE_D_EXTENSIONS):
return True
if any(filename.endswith(ext) for ext in TEXT_EXTENSIONS):
return True
return False
@@ -255,6 +261,10 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
Preview priority (matching frontend):
1. type="output" with previewable media
2. Any previewable media
Text content entries (strings under 'text') are preview-only metadata,
matching the frontend's METADATA_KEYS: they can serve as the fallback
preview but are not counted as outputs.
"""
count = 0
preview_output = None
@@ -275,7 +285,6 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
if normalized is None:
# Not a 3D file string — check for text preview
if media_type == 'text':
count += 1
if preview_output is None:
if isinstance(item, tuple):
text_value = item[0] if item else ''
-167
View File
@@ -1,167 +0,0 @@
import math
import nodes
import numpy as np
import torch
from PIL import Image
import comfy.latent_formats
import comfy.model_management
import comfy.utils
from comfy_api.latest import ComfyExtension, io
from typing_extensions import override
IMAGE_MIN_TOKEN_NUM = 4
IMAGE_MAX_TOKEN_NUM = 16384
MAX_RATIO = 200
SPATIAL_MERGE_SIZE = 2
VISION_PATCH_SIZE = 16
def _crop_image(image, width, height):
image = image[:1].movedim(-1, 1)
image = comfy.utils.common_upscale(image, width, height, "bilinear", "center")
return image.movedim(1, -1)[:, :, :, :3]
def _round_by_factor(number, factor):
return round(number / factor) * factor
def _ceil_by_factor(number, factor):
return math.ceil(number / factor) * factor
def _floor_by_factor(number, factor):
return math.floor(number / factor) * factor
def _smart_resize(height, width, factor, min_pixels=None, max_pixels=None):
max_pixels = max_pixels if max_pixels is not None else IMAGE_MAX_TOKEN_NUM * factor ** 2
min_pixels = min_pixels if min_pixels is not None else IMAGE_MIN_TOKEN_NUM * factor ** 2
if max_pixels < min_pixels:
raise ValueError("max_pixels must be greater than or equal to min_pixels.")
if max(height, width) / min(height, width) > MAX_RATIO:
raise ValueError(f"LingBotVideo image aspect ratio must be smaller than {MAX_RATIO}.")
resized_height = max(factor, _round_by_factor(height, factor))
resized_width = max(factor, _round_by_factor(width, factor))
if resized_height * resized_width > max_pixels:
beta = math.sqrt((height * width) / max_pixels)
resized_height = _floor_by_factor(height / beta, factor)
resized_width = _floor_by_factor(width / beta, factor)
elif resized_height * resized_width < min_pixels:
beta = math.sqrt(min_pixels / (height * width))
resized_height = _ceil_by_factor(height * beta, factor)
resized_width = _ceil_by_factor(width * beta, factor)
return resized_height, resized_width
def _vlm_image(image):
factor = VISION_PATCH_SIZE * SPATIAL_MERGE_SIZE
height, width = image.shape[1:3]
resized_height, resized_width = _smart_resize(height, width, factor)
array = image[0].detach().cpu().clamp(0, 1).mul(255).byte().numpy()
pil_image = Image.fromarray(array, mode="RGB").resize((resized_width, resized_height))
array = np.asarray(pil_image).astype(np.float32) / 255.0
return torch.from_numpy(array).unsqueeze(0)
class TextEncodeLingBotVideoI2V(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncodeLingBotVideoI2V",
category="model/conditioning/lingbot_video",
inputs=[
io.Clip.Input("clip"),
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
io.String.Input("negative_prompt", multiline=True, dynamic_prompts=True, default=""),
io.Int.Input("width", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Image.Input("image", optional=True),
],
outputs=[
io.Conditioning.Output(display_name="positive"),
io.Conditioning.Output(display_name="negative"),
],
)
@classmethod
def execute(cls, clip, prompt, negative_prompt, width, height, image=None) -> io.NodeOutput:
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"LingBotVideo width and height must be multiples of 16, got {width}x{height}.")
if image is None:
images = []
else:
image = _crop_image(image, width, height)
images = [_vlm_image(image)]
positive = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=images))
negative = clip.encode_from_tokens_scheduled(clip.tokenize(negative_prompt, images=images))
return io.NodeOutput(positive, negative)
class LingBotImageToVideo(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LingBotImageToVideo",
category="model/conditioning/lingbot_video",
inputs=[
io.Conditioning.Input("positive"),
io.Conditioning.Input("negative"),
io.Vae.Input("vae"),
io.Int.Input("width", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Int.Input("length", default=81, min=1, max=nodes.MAX_RESOLUTION, step=4),
io.Int.Input("batch_size", default=1, min=1, max=4096),
io.Image.Input("start_image", optional=True),
],
outputs=[
io.Conditioning.Output(display_name="positive"),
io.Conditioning.Output(display_name="negative"),
io.Latent.Output(display_name="latent"),
],
)
@classmethod
def execute(cls, positive, negative, vae, width, height, length, batch_size, start_image=None) -> io.NodeOutput:
if length != 1 and (length - 1) % 4 != 0:
raise ValueError(f"LingBotVideo length must be 1 or 4n+1, got {length}.")
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"LingBotVideo width and height must be multiples of 16, got {width}x{height}.")
latent_frames = ((length - 1) // 4) + 1
latent = torch.zeros(
[batch_size, 16, latent_frames, height // 8, width // 8],
device=comfy.model_management.intermediate_device(),
)
out_latent = {"samples": latent}
if start_image is not None:
start_image = _crop_image(start_image, width, height)
cond_latent = comfy.latent_formats.LingBotVideo().process_in(vae.encode(start_image))
cond_latent = comfy.utils.resize_to_batch_size(cond_latent, batch_size)
cond_t = min(cond_latent.shape[2], latent.shape[2])
latent[:, :, :cond_t] = cond_latent[:, :, :cond_t].to(device=latent.device, dtype=latent.dtype)
noise_mask = torch.ones(
(batch_size, 1, latent.shape[2], latent.shape[3], latent.shape[4]),
device=latent.device,
dtype=latent.dtype,
)
noise_mask[:, :, :cond_t] = 0.0
out_latent["noise_mask"] = noise_mask
return io.NodeOutput(positive, negative, out_latent)
class LingBotVideoExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
TextEncodeLingBotVideoI2V,
LingBotImageToVideo,
]
async def comfy_entrypoint() -> LingBotVideoExtension:
return LingBotVideoExtension()
+73
View File
@@ -0,0 +1,73 @@
import os
import json
from typing_extensions import override
from comfy_api.latest import io, ComfyExtension, ui
import folder_paths
import logging
class SaveTextNode(io.ComfyNode):
"""Save text content to .txt, .md, or .json."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SaveText",
search_aliases=["save text", "write text", "export text"],
display_name="Save Text",
category="text",
description="Save text content to a file in the output directory.",
inputs=[
io.String.Input("text", force_input=True),
io.String.Input("filename_prefix", default="ComfyUI"),
io.Combo.Input("format", options=["txt", "md", "json"], default="txt"),
],
outputs=[io.String.Output(display_name="text")],
is_output_node=True,
)
@classmethod
def execute(cls, text, filename_prefix, format):
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
filename_prefix,
folder_paths.get_output_directory(),
1,
1,
)
file = f"{filename}_{counter:05}.{format}"
filepath = os.path.join(full_output_folder, file)
if format == "json":
# tries to pretty print otherwise saves normally
try:
data = json.loads(text)
with open(filepath, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2, ensure_ascii=False)
except json.JSONDecodeError:
logging.warning("Saved JSON as a raw text")
with open(filepath, "w", encoding="utf-8") as f:
f.write(text)
else:
with open(filepath, "w", encoding="utf-8") as f:
f.write(text)
return io.NodeOutput(
text,
ui={
"text": (text,),
"files": [
ui.SavedResult(file, subfolder, io.FolderType.output)
]
}
)
class TextExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
SaveTextNode
]
async def comfy_entrypoint() -> TextExtension:
return TextExtension()
+1 -5
View File
@@ -17,11 +17,7 @@ if args.base_directory:
else:
base_path = os.path.dirname(os.path.realpath(__file__))
if args.models_directory:
models_dir = os.path.abspath(args.models_directory)
else:
models_dir = os.path.join(base_path, "models")
models_dir = os.path.join(base_path, "models")
folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])
-4
View File
@@ -131,10 +131,6 @@ def apply_custom_paths():
if args.base_directory:
logging.info(f"Setting base directory to: {folder_paths.base_path}")
# --models-directory
if args.models_directory:
logging.info(f"Setting models directory to: {folder_paths.models_dir}")
# --output-directory, --input-directory, --user-directory
if args.output_directory:
output_dir = os.path.abspath(args.output_directory)
+2 -2
View File
@@ -992,7 +992,7 @@ class CLIPLoader:
@classmethod
def INPUT_TYPES(s):
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "lingbot_video"], ),
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2"], ),
},
"optional": {
"device": (["default", "cpu"], {"advanced": True}),
@@ -2460,7 +2460,6 @@ async def init_builtin_extra_nodes():
"nodes_tcfg.py",
"nodes_context_windows.py",
"nodes_qwen.py",
"nodes_lingbot_video.py",
"nodes_boogu.py",
"nodes_chroma_radiance.py",
"nodes_pid.py",
@@ -2504,6 +2503,7 @@ async def init_builtin_extra_nodes():
"nodes_triposplat.py",
"nodes_depth_anything_3.py",
"nodes_seed.py",
"nodes_text.py",
]
import_failed = []
+1 -1
View File
@@ -1,5 +1,5 @@
comfyui-frontend-package==1.45.20
comfyui-workflow-templates==0.11.6
comfyui-workflow-templates==0.11.2
comfyui-embedded-docs==0.5.7
torch
torchsde
-17
View File
@@ -163,20 +163,3 @@ def test_base_path_change_clears_old(set_base_dir):
for name in ["controlnet", "diffusion_models", "text_encoders"]:
assert len(folder_paths.get_folder_paths(name)) == 2
def test_models_directory_cli_and_getters(temp_dir):
try:
with patch.object(sys, 'argv', ["main.py", "--models-directory", temp_dir]):
reload(comfy.cli_args)
reload(folder_paths)
assert folder_paths.models_dir == os.path.abspath(temp_dir)
with pytest.raises(Exception):
comfy.cli_args.is_valid_directory(os.path.join(temp_dir, "non_existent_folder_path"))
finally:
with patch.object(sys, 'argv', ["main.py"]):
reload(comfy.cli_args)
reload(folder_paths)