mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-21 23:41:28 +08:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a3ca7ef350 | ||
|
|
eda4ba84f2 | ||
|
|
949c4b8d07 | ||
|
|
cd53aff8d8 | ||
|
|
84f96a56ca | ||
|
|
978895e1a9 |
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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.")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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_
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
@@ -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 ''
|
||||
|
||||
@@ -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()
|
||||
@@ -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
@@ -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"])
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user