Merge branch 'Comfy-Org:master' into logc4-pr
This commit is contained in:
commit
3332e0949e
41
AGENTS.md
41
AGENTS.md
|
|
@ -19,6 +19,9 @@
|
||||||
better to remove a broken feature path than keep a complicated partial fix.
|
better to remove a broken feature path than keep a complicated partial fix.
|
||||||
- Preserve existing APIs, node names, model-loading behavior, file layout, and
|
- Preserve existing APIs, node names, model-loading behavior, file layout, and
|
||||||
workflow compatibility unless the change is explicitly about replacing them.
|
workflow compatibility unless the change is explicitly about replacing them.
|
||||||
|
- When compatibility is explicitly out of scope, remove compatibility-only
|
||||||
|
aliases, duplicate nodes, legacy entry points, and preset wrappers instead of
|
||||||
|
retaining parallel ways to perform the same operation.
|
||||||
- Code must look hand-written for this repository. Changes that read like
|
- Code must look hand-written for this repository. Changes that read like
|
||||||
generic AI-generated code will be rejected automatically: unnecessary helper
|
generic AI-generated code will be rejected automatically: unnecessary helper
|
||||||
layers, vague names, boilerplate comments, defensive branches without a real
|
layers, vague names, boilerplate comments, defensive branches without a real
|
||||||
|
|
@ -96,6 +99,13 @@
|
||||||
unless they are read by current code and change current behavior. Remove
|
unless they are read by current code and change current behavior. Remove
|
||||||
pass-through or stored-but-unused values instead of preserving upstream or
|
pass-through or stored-but-unused values instead of preserving upstream or
|
||||||
deprecated API baggage.
|
deprecated API baggage.
|
||||||
|
- Do not add a model-specific option to a shared helper when only one caller
|
||||||
|
needs it. Keep one-off behavior at the model integration boundary, or extend
|
||||||
|
the shared helper only when the option is a coherent reusable capability.
|
||||||
|
- Implementations of shared model interfaces should accept the standard caller
|
||||||
|
contract without model-specific rejection branches for optional capabilities
|
||||||
|
they do not consume. Let supported behavior be determined by implementation
|
||||||
|
paths that actually use those inputs.
|
||||||
- If an implementation needs auxiliary values for its own workflow, expose them
|
- If an implementation needs auxiliary values for its own workflow, expose them
|
||||||
through a private helper or a clearly named implementation-specific method
|
through a private helper or a clearly named implementation-specific method
|
||||||
instead of overloading the public method's return contract.
|
instead of overloading the public method's return contract.
|
||||||
|
|
@ -154,6 +164,10 @@
|
||||||
`comfy-kitchen` helpers where they already solve the problem.
|
`comfy-kitchen` helpers where they already solve the problem.
|
||||||
- Use optimized comfy-kitchen ops in places where they improve performance
|
- Use optimized comfy-kitchen ops in places where they improve performance
|
||||||
without changing the expected dtype, device, memory, or interface behavior.
|
without changing the expected dtype, device, memory, or interface behavior.
|
||||||
|
- Prefer ComfyUI's shared optimized kernels and backend dispatchers over
|
||||||
|
handwritten implementations of the same operation. Remove duplicate local
|
||||||
|
kernels and adapt inputs to the shared operation's documented layout while
|
||||||
|
preserving the model's original math and output contract.
|
||||||
- All models should use the optimized attention function selected by ComfyUI.
|
- All models should use the optimized attention function selected by ComfyUI.
|
||||||
Treat optimized backend functions, dispatch helpers, and capability-selected
|
Treat optimized backend functions, dispatch helpers, and capability-selected
|
||||||
callables as opaque. Higher-level code must not inspect function identity,
|
callables as opaque. Higher-level code must not inspect function identity,
|
||||||
|
|
@ -176,6 +190,12 @@
|
||||||
- Model detection code that inspects linear weight shapes should only use the
|
- Model detection code that inspects linear weight shapes should only use the
|
||||||
first dimension. The second dimension may be half the original size for
|
first dimension. The second dimension may be half the original size for
|
||||||
NVFP4 or other 4-bit quantized models.
|
NVFP4 or other 4-bit quantized models.
|
||||||
|
- A model-detection signature must guard every state-dict key it dereferences.
|
||||||
|
Do not partially match a format and then raise an incidental `KeyError` while
|
||||||
|
extracting its configuration.
|
||||||
|
- Order model-detection checks from established or more-specific signatures to
|
||||||
|
newer or broader signatures. Put a broad new detector near the generic
|
||||||
|
fallback when giving it higher precedence could steal another model family.
|
||||||
- Avoid adding `einops` usage in core inference code. Use native torch tensor
|
- Avoid adding `einops` usage in core inference code. Use native torch tensor
|
||||||
ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
|
ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
|
||||||
`unsqueeze`, and `squeeze` instead.
|
`unsqueeze`, and `squeeze` instead.
|
||||||
|
|
@ -192,11 +212,23 @@
|
||||||
methods for scalar or structural calculations.
|
methods for scalar or structural calculations.
|
||||||
- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
|
- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
|
||||||
storage dtype, bias dtype, and original tensor shape metadata.
|
storage dtype, bias dtype, and original tensor shape metadata.
|
||||||
|
- Do not cast the result of an optimized backend operation back to its input
|
||||||
|
dtype unless that backend's documented result contract requires normalization.
|
||||||
|
In particular, trust the selected optimized-attention implementation to honor
|
||||||
|
its dtype contract.
|
||||||
- Keep model-native latent layout handling inside the model or latent-format
|
- Keep model-native latent layout handling inside the model or latent-format
|
||||||
owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
|
owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
|
||||||
dimensions in nodes or other caller-side adapters just to satisfy a model
|
dimensions in nodes or other caller-side adapters just to satisfy a model
|
||||||
forward; the model path should consume and return the native latent shape for
|
forward; the model path should consume and return the native latent shape for
|
||||||
that model family.
|
that model family.
|
||||||
|
- DiT models should accept latent dimensions that are not exact patch-size
|
||||||
|
multiples. Use `comfy.ldm.common_dit.pad_to_patch_size` on every patchified
|
||||||
|
target or reference input, then crop only the target output back to its
|
||||||
|
original dimensions.
|
||||||
|
- Avoid defensive shape and configuration checks that merely replace the clear
|
||||||
|
failure from the tensor operation immediately below them. Add explicit
|
||||||
|
validation only when it provides materially better context at a real boundary
|
||||||
|
or prevents silent incorrect output.
|
||||||
- Assume inputs to the main model forward are already in the compute dtype by
|
- Assume inputs to the main model forward are already in the compute dtype by
|
||||||
default, except integer inputs such as some model timestep tensors. Do not add
|
default, except integer inputs such as some model timestep tensors. Do not add
|
||||||
defensive or convenience casts in model code; it is better for invalid dtype
|
defensive or convenience casts in model code; it is better for invalid dtype
|
||||||
|
|
@ -260,6 +292,15 @@
|
||||||
- Model implementations should add the minimal number of ComfyUI nodes required
|
- Model implementations should add the minimal number of ComfyUI nodes required
|
||||||
to run the model. Reuse existing nodes as much as possible; adapting the model
|
to run the model. Reuse existing nodes as much as possible; adapting the model
|
||||||
to work with existing nodes is strongly preferred over creating new nodes.
|
to work with existing nodes is strongly preferred over creating new nodes.
|
||||||
|
- Use `io.Autogrow` for a variable number of repeated inputs instead of a fixed
|
||||||
|
series of numbered optional sockets. Set its minimum to zero when the model
|
||||||
|
has a valid no-item path, and cap it only when the model has a real limit.
|
||||||
|
- Mark inputs optional when execution has a valid path that does not read them.
|
||||||
|
If one optional input is needed only to process another optional input, do not
|
||||||
|
force users on the path that supplies neither to connect it.
|
||||||
|
- Conditioning nodes should normally output conditioning only. Do not expose
|
||||||
|
input or intermediate images as convenience outputs for downstream sizing or
|
||||||
|
routing; use the existing image path or a dedicated image operation instead.
|
||||||
- Nodes should output only values they own. Do not add pass-through outputs for
|
- Nodes should output only values they own. Do not add pass-through outputs for
|
||||||
workflow convenience unless the node is explicitly an output node. Existing
|
workflow convenience unless the node is explicitly an output node. Existing
|
||||||
models, latents, conditioning, or other inputs should flow directly to the
|
models, latents, conditioning, or other inputs should flow directly to the
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,7 @@ from torchvision import transforms
|
||||||
import comfy.patcher_extension
|
import comfy.patcher_extension
|
||||||
from comfy.ldm.modules.attention import optimized_attention
|
from comfy.ldm.modules.attention import optimized_attention
|
||||||
import comfy.ldm.common_dit
|
import comfy.ldm.common_dit
|
||||||
|
import comfy.ops
|
||||||
import comfy.quant_ops
|
import comfy.quant_ops
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -161,11 +162,16 @@ class Attention(nn.Module):
|
||||||
def apply_norm_and_rotary_pos_emb(
|
def apply_norm_and_rotary_pos_emb(
|
||||||
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
|
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||||
q = self.q_norm(q)
|
|
||||||
k = self.k_norm(k)
|
|
||||||
v = self.v_norm(v)
|
v = self.v_norm(v)
|
||||||
if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
|
if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
|
||||||
q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rope_emb)
|
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, q, offloadable=True)
|
||||||
|
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, k, offloadable=True)
|
||||||
|
q, k = comfy.quant_ops.ck.rms_rope_split_half(q, k, rope_emb, q_scale, k_scale, self.q_norm.eps)
|
||||||
|
comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
|
||||||
|
comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
|
||||||
|
else:
|
||||||
|
q = self.q_norm(q)
|
||||||
|
k = self.k_norm(k)
|
||||||
return q, k, v
|
return q, k, v
|
||||||
|
|
||||||
q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
|
q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
|
||||||
|
|
|
||||||
|
|
@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None):
|
||||||
The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
|
The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def two_pass_attention(q, k, v, heads, **kwargs):
|
def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs):
|
||||||
B, H, T, D = q.shape
|
B, H, T, D = q.shape
|
||||||
|
|
||||||
if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
|
if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
|
||||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||||
elif ar_len >= T:
|
elif ar_len >= T:
|
||||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
|
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
|
||||||
elif ar_len <= 0:
|
elif ar_len <= 0:
|
||||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||||
else:
|
else:
|
||||||
out_ar = comfy.ops.scaled_dot_product_attention(
|
out_ar = comfy.ops.scaled_dot_product_attention(
|
||||||
q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
|
q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
|
||||||
attn_mask=None, dropout_p=0.0, is_causal=True,
|
attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
|
||||||
)
|
)
|
||||||
out_gen = optimized_attention(
|
out_gen = optimized_attention(
|
||||||
q[:, :, ar_len:], k, v, heads,
|
q[:, :, ar_len:], k, v, heads,
|
||||||
mask=None, skip_reshape=True, skip_output_reshape=True,
|
mask=None, skip_reshape=True, skip_output_reshape=True,
|
||||||
transformer_options=transformer_options,
|
transformer_options=transformer_options, enable_gqa=enable_gqa,
|
||||||
)
|
)
|
||||||
out = torch.cat([out_ar, out_gen], dim=2)
|
out = torch.cat([out_ar, out_gen], dim=2)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,445 @@
|
||||||
|
# https://github.com/jdopensource/JoyAI-Image-Edit (Apache 2.0)
|
||||||
|
import math
|
||||||
|
from typing import Optional, Tuple
|
||||||
|
|
||||||
|
import comfy_kitchen
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
import comfy.ldm.common_dit
|
||||||
|
import comfy.ops
|
||||||
|
import comfy.patcher_extension
|
||||||
|
from comfy.ldm.lightricks.model import GELU_approx, PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
|
||||||
|
from comfy.ldm.modules.attention import optimized_attention
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageModulate(nn.Module):
|
||||||
|
def __init__(self, hidden_size: int, factor: int, dtype=None, device=None):
|
||||||
|
super().__init__()
|
||||||
|
self.factor = factor
|
||||||
|
self.modulate_table = nn.Parameter(
|
||||||
|
torch.empty(1, factor, hidden_size, dtype=dtype, device=device)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> list:
|
||||||
|
if x.ndim != 3:
|
||||||
|
x = x.unsqueeze(1)
|
||||||
|
table = comfy.ops.cast_to_input(self.modulate_table, x)
|
||||||
|
return [o.squeeze(1) for o in (table + x).chunk(self.factor, dim=1)]
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageFeedForward(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
inner_dim: int,
|
||||||
|
dtype=None,
|
||||||
|
device=None,
|
||||||
|
operations=None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.net = nn.ModuleList([
|
||||||
|
GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations),
|
||||||
|
nn.Identity(),
|
||||||
|
operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device),
|
||||||
|
])
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
for module in self.net:
|
||||||
|
x = module(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageAttention(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
num_attention_heads: int,
|
||||||
|
attention_head_dim: int,
|
||||||
|
eps: float = 1e-6,
|
||||||
|
dtype=None,
|
||||||
|
device=None,
|
||||||
|
operations=None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.num_attention_heads = num_attention_heads
|
||||||
|
inner_dim = num_attention_heads * attention_head_dim
|
||||||
|
|
||||||
|
self.img_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device)
|
||||||
|
self.img_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.img_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.img_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
self.txt_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device)
|
||||||
|
self.txt_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.txt_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.txt_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
img: torch.Tensor,
|
||||||
|
txt: torch.Tensor,
|
||||||
|
image_rotary_emb: torch.Tensor,
|
||||||
|
transformer_options=None,
|
||||||
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
heads = self.num_attention_heads
|
||||||
|
|
||||||
|
img_q, img_k, img_v = self.img_attn_qkv(img).chunk(3, dim=-1)
|
||||||
|
txt_q, txt_k, txt_v = self.txt_attn_qkv(txt).chunk(3, dim=-1)
|
||||||
|
|
||||||
|
img_q = img_q.unflatten(-1, (heads, -1))
|
||||||
|
img_k = img_k.unflatten(-1, (heads, -1))
|
||||||
|
img_v = img_v.unflatten(-1, (heads, -1))
|
||||||
|
txt_q = txt_q.unflatten(-1, (heads, -1))
|
||||||
|
txt_k = txt_k.unflatten(-1, (heads, -1))
|
||||||
|
txt_v = txt_v.unflatten(-1, (heads, -1))
|
||||||
|
|
||||||
|
img_q = self.img_attn_q_norm(img_q)
|
||||||
|
img_k = self.img_attn_k_norm(img_k)
|
||||||
|
txt_q = self.txt_attn_q_norm(txt_q)
|
||||||
|
txt_k = self.txt_attn_k_norm(txt_k)
|
||||||
|
|
||||||
|
img_q, img_k = comfy_kitchen.apply_rope(img_q, img_k, image_rotary_emb)
|
||||||
|
|
||||||
|
joint_q = torch.cat([img_q, txt_q], dim=1)
|
||||||
|
joint_k = torch.cat([img_k, txt_k], dim=1)
|
||||||
|
joint_v = torch.cat([img_v, txt_v], dim=1)
|
||||||
|
|
||||||
|
joint_q = joint_q.flatten(2, 3)
|
||||||
|
joint_k = joint_k.flatten(2, 3)
|
||||||
|
joint_v = joint_v.flatten(2, 3)
|
||||||
|
|
||||||
|
joint_out = optimized_attention(joint_q, joint_k, joint_v, heads=heads, transformer_options=transformer_options)
|
||||||
|
|
||||||
|
seq_img = img.shape[1]
|
||||||
|
img_out = joint_out[:, :seq_img, :]
|
||||||
|
txt_out = joint_out[:, seq_img:, :]
|
||||||
|
|
||||||
|
img_out = self.img_attn_proj(img_out)
|
||||||
|
txt_out = self.txt_attn_proj(txt_out)
|
||||||
|
return img_out, txt_out
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageTransformerBlock(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
num_attention_heads: int,
|
||||||
|
attention_head_dim: int,
|
||||||
|
mlp_width_ratio: float = 4.0,
|
||||||
|
eps: float = 1e-6,
|
||||||
|
dtype=None,
|
||||||
|
device=None,
|
||||||
|
operations=None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
mlp_hidden_dim = int(dim * mlp_width_ratio)
|
||||||
|
|
||||||
|
self.img_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device)
|
||||||
|
self.img_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.img_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.img_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
|
||||||
|
|
||||||
|
self.txt_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device)
|
||||||
|
self.txt_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.txt_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||||
|
self.txt_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
|
||||||
|
|
||||||
|
self.attn = JoyImageAttention(
|
||||||
|
dim=dim,
|
||||||
|
num_attention_heads=num_attention_heads,
|
||||||
|
attention_head_dim=attention_head_dim,
|
||||||
|
eps=eps,
|
||||||
|
dtype=dtype,
|
||||||
|
device=device,
|
||||||
|
operations=operations,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
encoder_hidden_states: torch.Tensor,
|
||||||
|
temb: torch.Tensor,
|
||||||
|
image_rotary_emb: torch.Tensor,
|
||||||
|
transformer_options=None,
|
||||||
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
(
|
||||||
|
img_mod1_shift,
|
||||||
|
img_mod1_scale,
|
||||||
|
img_mod1_gate,
|
||||||
|
img_mod2_shift,
|
||||||
|
img_mod2_scale,
|
||||||
|
img_mod2_gate,
|
||||||
|
) = self.img_mod(temb)
|
||||||
|
(
|
||||||
|
txt_mod1_shift,
|
||||||
|
txt_mod1_scale,
|
||||||
|
txt_mod1_gate,
|
||||||
|
txt_mod2_shift,
|
||||||
|
txt_mod2_scale,
|
||||||
|
txt_mod2_gate,
|
||||||
|
) = self.txt_mod(temb)
|
||||||
|
|
||||||
|
img_normed = self.img_norm1(hidden_states)
|
||||||
|
txt_normed = self.txt_norm1(encoder_hidden_states)
|
||||||
|
img_modulated = img_normed * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1)
|
||||||
|
txt_modulated = txt_normed * (1 + txt_mod1_scale.unsqueeze(1)) + txt_mod1_shift.unsqueeze(1)
|
||||||
|
|
||||||
|
img_attn, txt_attn = self.attn(img_modulated, txt_modulated, image_rotary_emb, transformer_options=transformer_options)
|
||||||
|
|
||||||
|
hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1)
|
||||||
|
encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1)
|
||||||
|
|
||||||
|
img_ffn_normed = self.img_norm2(hidden_states)
|
||||||
|
txt_ffn_normed = self.txt_norm2(encoder_hidden_states)
|
||||||
|
img_ffn_input = img_ffn_normed * (1 + img_mod2_scale.unsqueeze(1)) + img_mod2_shift.unsqueeze(1)
|
||||||
|
txt_ffn_input = txt_ffn_normed * (1 + txt_mod2_scale.unsqueeze(1)) + txt_mod2_shift.unsqueeze(1)
|
||||||
|
hidden_states = hidden_states + self.img_mlp(img_ffn_input) * img_mod2_gate.unsqueeze(1)
|
||||||
|
encoder_hidden_states = encoder_hidden_states + self.txt_mlp(txt_ffn_input) * txt_mod2_gate.unsqueeze(1)
|
||||||
|
|
||||||
|
return hidden_states, encoder_hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageTimeTextImageEmbedding(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
time_freq_dim: int,
|
||||||
|
time_proj_dim: int,
|
||||||
|
text_embed_dim: int,
|
||||||
|
dtype=None,
|
||||||
|
device=None,
|
||||||
|
operations=None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||||
|
self.time_embedder = TimestepEmbedding(
|
||||||
|
in_channels=time_freq_dim,
|
||||||
|
time_embed_dim=dim,
|
||||||
|
dtype=dtype,
|
||||||
|
device=device,
|
||||||
|
operations=operations,
|
||||||
|
)
|
||||||
|
self.act_fn = nn.SiLU()
|
||||||
|
self.time_proj = operations.Linear(dim, time_proj_dim, bias=True, dtype=dtype, device=device)
|
||||||
|
self.text_embedder = PixArtAlphaTextProjection(
|
||||||
|
text_embed_dim, dim, act_fn="gelu_tanh", dtype=dtype, device=device, operations=operations,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor):
|
||||||
|
timestep = self.timesteps_proj(timestep)
|
||||||
|
temb = self.time_embedder(timestep.to(dtype=encoder_hidden_states.dtype)).type_as(encoder_hidden_states)
|
||||||
|
timestep_proj = self.time_proj(self.act_fn(temb))
|
||||||
|
encoder_hidden_states = self.text_embedder(encoder_hidden_states)
|
||||||
|
return temb, timestep_proj, encoder_hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageTransformer3DModel(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
patch_size: list = [1, 2, 2],
|
||||||
|
in_channels: int = 16,
|
||||||
|
out_channels: Optional[int] = None,
|
||||||
|
hidden_size: int = 3072,
|
||||||
|
num_attention_heads: int = 24,
|
||||||
|
text_dim: int = 4096,
|
||||||
|
mlp_width_ratio: float = 4.0,
|
||||||
|
num_layers: int = 20,
|
||||||
|
rope_dim_list: list = [16, 56, 56],
|
||||||
|
theta: int = 256,
|
||||||
|
image_model=None,
|
||||||
|
dtype=None,
|
||||||
|
device=None,
|
||||||
|
operations=None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.dtype = dtype
|
||||||
|
self.out_channels = out_channels or in_channels
|
||||||
|
self.patch_size = list(patch_size)
|
||||||
|
self.rope_dim_list = list(rope_dim_list)
|
||||||
|
self.theta = theta
|
||||||
|
|
||||||
|
attention_head_dim = hidden_size // num_attention_heads
|
||||||
|
|
||||||
|
self.img_in = operations.Conv3d(
|
||||||
|
in_channels,
|
||||||
|
hidden_size,
|
||||||
|
kernel_size=tuple(self.patch_size),
|
||||||
|
stride=tuple(self.patch_size),
|
||||||
|
dtype=dtype,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.condition_embedder = JoyImageTimeTextImageEmbedding(
|
||||||
|
dim=hidden_size,
|
||||||
|
time_freq_dim=256,
|
||||||
|
time_proj_dim=hidden_size * 6,
|
||||||
|
text_embed_dim=text_dim,
|
||||||
|
dtype=dtype,
|
||||||
|
device=device,
|
||||||
|
operations=operations,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.double_blocks = nn.ModuleList([
|
||||||
|
JoyImageTransformerBlock(
|
||||||
|
dim=hidden_size,
|
||||||
|
num_attention_heads=num_attention_heads,
|
||||||
|
attention_head_dim=attention_head_dim,
|
||||||
|
mlp_width_ratio=mlp_width_ratio,
|
||||||
|
dtype=dtype,
|
||||||
|
device=device,
|
||||||
|
operations=operations,
|
||||||
|
)
|
||||||
|
for _ in range(num_layers)
|
||||||
|
])
|
||||||
|
|
||||||
|
self.norm_out = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
|
||||||
|
self.proj_out = operations.Linear(
|
||||||
|
hidden_size,
|
||||||
|
self.out_channels * math.prod(self.patch_size),
|
||||||
|
bias=True,
|
||||||
|
dtype=dtype,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _get_rotary_pos_embed_for_range(
|
||||||
|
self,
|
||||||
|
start: Tuple[int, int, int],
|
||||||
|
stop: Tuple[int, int, int],
|
||||||
|
device=None,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
# 3D RoPE for the patch grid range [start, stop) over (t, h, w). Token order after
|
||||||
|
# reshape(-1) is (t, h, w), matching the img_in Conv3d flatten.
|
||||||
|
rope_dim_list = self.rope_dim_list
|
||||||
|
|
||||||
|
grids = [torch.arange(start[i], stop[i], dtype=torch.float32, device=device) for i in range(3)]
|
||||||
|
mesh = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=0)
|
||||||
|
|
||||||
|
angles_parts = []
|
||||||
|
for i, dim in enumerate(rope_dim_list):
|
||||||
|
pos = mesh[i].reshape(-1)
|
||||||
|
freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device)[: (dim // 2)] / dim))
|
||||||
|
angles_parts.append(torch.outer(pos, freqs))
|
||||||
|
|
||||||
|
angles = torch.cat(angles_parts, dim=1)
|
||||||
|
cos = angles.cos()
|
||||||
|
sin = angles.sin()
|
||||||
|
return torch.stack((cos, -sin, sin, cos), dim=-1).unflatten(-1, (2, 2))
|
||||||
|
|
||||||
|
def get_rotary_pos_embed_for_components(
|
||||||
|
self,
|
||||||
|
component_sizes,
|
||||||
|
device=None,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
# Per-component 3D RoPE. component_sizes is a list of (t, h, w) patch grid sizes in
|
||||||
|
# sequence order [target, ref0, ref1, ...]; h/w restart at 0 for each component while t
|
||||||
|
# continues from the running offset, giving every image its own temporal position band.
|
||||||
|
freqs_parts = []
|
||||||
|
t_offset = 0
|
||||||
|
for (t, h, w) in component_sizes:
|
||||||
|
freqs = self._get_rotary_pos_embed_for_range(
|
||||||
|
start=(t_offset, 0, 0),
|
||||||
|
stop=(t_offset + t, h, w),
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
freqs_parts.append(freqs)
|
||||||
|
t_offset += t
|
||||||
|
return torch.cat(freqs_parts, dim=0).unsqueeze(0).unsqueeze(2)
|
||||||
|
|
||||||
|
def unpatchify(self, x: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
|
||||||
|
c = self.out_channels
|
||||||
|
pt, ph, pw = self.patch_size
|
||||||
|
x = x.reshape(x.shape[0], t, h, w, pt, ph, pw, c)
|
||||||
|
x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)
|
||||||
|
return x.reshape(x.shape[0], c, t * pt, h * ph, w * pw)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
timestep: torch.Tensor,
|
||||||
|
context: torch.Tensor = None,
|
||||||
|
ref_latents=None,
|
||||||
|
control=None,
|
||||||
|
transformer_options=None,
|
||||||
|
**kwargs,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
transformer_options = {} if transformer_options is None else transformer_options.copy()
|
||||||
|
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||||
|
self._forward,
|
||||||
|
self,
|
||||||
|
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
|
||||||
|
).execute(hidden_states, timestep, context, ref_latents, transformer_options, **kwargs)
|
||||||
|
|
||||||
|
def _forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
timestep: torch.Tensor,
|
||||||
|
context: torch.Tensor,
|
||||||
|
ref_latents=None,
|
||||||
|
transformer_options=None,
|
||||||
|
**kwargs,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
pt, ph, pw = self.patch_size
|
||||||
|
_, _, ot, oh, ow = hidden_states.shape
|
||||||
|
|
||||||
|
components = [hidden_states, *(ref_latents or [])]
|
||||||
|
component_sizes = []
|
||||||
|
img_tokens = []
|
||||||
|
for comp in components:
|
||||||
|
comp = comfy.ldm.common_dit.pad_to_patch_size(comp, self.patch_size)
|
||||||
|
_, _, ct, ch, cw = comp.shape
|
||||||
|
component_sizes.append((ct // pt, ch // ph, cw // pw))
|
||||||
|
tokens = self.img_in(comp).flatten(2).transpose(1, 2) # (B, n_i, D)
|
||||||
|
img_tokens.append(tokens)
|
||||||
|
|
||||||
|
img = torch.cat(img_tokens, dim=1)
|
||||||
|
|
||||||
|
_, vec, txt = self.condition_embedder(timestep, context)
|
||||||
|
vec = vec.unflatten(1, (6, -1))
|
||||||
|
|
||||||
|
image_rotary_emb = self.get_rotary_pos_embed_for_components(
|
||||||
|
component_sizes,
|
||||||
|
device=hidden_states.device,
|
||||||
|
)
|
||||||
|
|
||||||
|
patches_replace = transformer_options.get("patches_replace", {})
|
||||||
|
blocks_replace = patches_replace.get("dit", {})
|
||||||
|
transformer_options["total_blocks"] = len(self.double_blocks)
|
||||||
|
transformer_options["block_type"] = "double"
|
||||||
|
for i, block in enumerate(self.double_blocks):
|
||||||
|
transformer_options["block_index"] = i
|
||||||
|
if ("double_block", i) in blocks_replace:
|
||||||
|
def block_wrap(args):
|
||||||
|
out = {}
|
||||||
|
out["img"], out["txt"] = block(
|
||||||
|
hidden_states=args["img"],
|
||||||
|
encoder_hidden_states=args["txt"],
|
||||||
|
temb=args["vec"],
|
||||||
|
image_rotary_emb=args["pe"],
|
||||||
|
transformer_options=args.get("transformer_options"),
|
||||||
|
)
|
||||||
|
return out
|
||||||
|
|
||||||
|
out = blocks_replace[("double_block", i)]({"img": img,
|
||||||
|
"txt": txt,
|
||||||
|
"vec": vec,
|
||||||
|
"pe": image_rotary_emb,
|
||||||
|
"transformer_options": transformer_options},
|
||||||
|
{"original_block": block_wrap})
|
||||||
|
txt = out["txt"]
|
||||||
|
img = out["img"]
|
||||||
|
else:
|
||||||
|
img, txt = block(
|
||||||
|
hidden_states=img,
|
||||||
|
encoder_hidden_states=txt,
|
||||||
|
temb=vec,
|
||||||
|
image_rotary_emb=image_rotary_emb,
|
||||||
|
transformer_options=transformer_options,
|
||||||
|
)
|
||||||
|
|
||||||
|
tt, th, tw = component_sizes[0]
|
||||||
|
target_tokens = tt * th * tw
|
||||||
|
img = img[:, :target_tokens, :]
|
||||||
|
img = self.proj_out(self.norm_out(img))
|
||||||
|
img = self.unpatchify(img, tt, th, tw)
|
||||||
|
return img[:, :, :ot, :oh, :ow]
|
||||||
|
|
@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||||
return out
|
return out
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
|
@torch.library.custom_op("comfy::flash_attn", mutates_args=())
|
||||||
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
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, softmax_scale: float = -1.0) -> torch.Tensor:
|
||||||
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale
|
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale
|
||||||
|
|
|
||||||
|
|
@ -197,6 +197,9 @@ class PixDiT_T2I(nn.Module):
|
||||||
"""Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
|
"""Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
|
||||||
return s
|
return s
|
||||||
|
|
||||||
|
def _pre_pixel_blocks(self, s, **kwargs):
|
||||||
|
return s
|
||||||
|
|
||||||
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
|
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
|
||||||
H_orig, W_orig = x.shape[2], x.shape[3]
|
H_orig, W_orig = x.shape[2], x.shape[3]
|
||||||
x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
|
x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
|
||||||
|
|
@ -226,6 +229,7 @@ class PixDiT_T2I(nn.Module):
|
||||||
s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
|
s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
|
||||||
s = F.silu(t_emb + s)
|
s = F.silu(t_emb + s)
|
||||||
|
|
||||||
|
s = self._pre_pixel_blocks(s, **kwargs)
|
||||||
s_cond = s.view(B * L, self.hidden_size)
|
s_cond = s.view(B * L, self.hidden_size)
|
||||||
x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
|
x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
|
||||||
for blk in self.pixel_blocks:
|
for blk in self.pixel_blocks:
|
||||||
|
|
|
||||||
|
|
@ -13,15 +13,15 @@ from .model import PixDiT_T2I
|
||||||
from .modules import precompute_freqs_cis_2d
|
from .modules import precompute_freqs_cis_2d
|
||||||
|
|
||||||
|
|
||||||
class SigmaAwareGatePerTokenPerDim(nn.Module):
|
class SigmaAwareGate(nn.Module):
|
||||||
"""gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq.
|
"""gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq.
|
||||||
|
|
||||||
Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1.
|
Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, dim: int, dtype=None, device=None, operations=None):
|
def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device)
|
self.content_proj = operations.Linear(dim * 2, 1 if per_token else dim, dtype=dtype, device=device)
|
||||||
self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device))
|
self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device))
|
||||||
|
|
||||||
def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
|
def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
|
||||||
|
|
@ -36,15 +36,15 @@ class SigmaAwareGatePerTokenPerDim(nn.Module):
|
||||||
class ResBlock(nn.Module):
|
class ResBlock(nn.Module):
|
||||||
"""Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip."""
|
"""Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip."""
|
||||||
|
|
||||||
def __init__(self, channels: int, num_groups: int = 4, dtype=None, device=None, operations=None):
|
def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", dtype=None, device=None, operations=None):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.block = nn.Sequential(
|
self.block = nn.Sequential(
|
||||||
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
||||||
nn.SiLU(),
|
nn.SiLU(),
|
||||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
|
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||||
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
||||||
nn.SiLU(),
|
nn.SiLU(),
|
||||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
|
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||||
)
|
)
|
||||||
|
|
||||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
|
@ -62,9 +62,13 @@ class LQProjection2D(nn.Module):
|
||||||
patch_size: int = 16,
|
patch_size: int = 16,
|
||||||
sr_scale: int = 4,
|
sr_scale: int = 4,
|
||||||
latent_spatial_down_factor: int = 8,
|
latent_spatial_down_factor: int = 8,
|
||||||
|
latent_unpatchify_factor: int = 1,
|
||||||
num_res_blocks: int = 4,
|
num_res_blocks: int = 4,
|
||||||
num_outputs: int = 7,
|
num_outputs: int = 7,
|
||||||
interval: int = 2,
|
interval: int = 2,
|
||||||
|
conv_padding_mode: str = "zeros",
|
||||||
|
gate_per_token: bool = False,
|
||||||
|
pit_output: bool = False,
|
||||||
dtype=None, device=None, operations=None,
|
dtype=None, device=None, operations=None,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
@ -74,34 +78,38 @@ class LQProjection2D(nn.Module):
|
||||||
self.patch_size = patch_size
|
self.patch_size = patch_size
|
||||||
self.sr_scale = sr_scale
|
self.sr_scale = sr_scale
|
||||||
self.latent_spatial_down_factor = latent_spatial_down_factor
|
self.latent_spatial_down_factor = latent_spatial_down_factor
|
||||||
|
self.latent_unpatchify_factor = latent_unpatchify_factor
|
||||||
self.num_outputs = num_outputs
|
self.num_outputs = num_outputs
|
||||||
self.interval = interval
|
self.interval = interval
|
||||||
|
|
||||||
z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size
|
effective_latent_channels = latent_channels // (latent_unpatchify_factor * latent_unpatchify_factor)
|
||||||
|
effective_spatial_down_factor = latent_spatial_down_factor // latent_unpatchify_factor
|
||||||
|
z_to_patch_ratio = (sr_scale * effective_spatial_down_factor) / patch_size
|
||||||
self.z_to_patch_ratio = z_to_patch_ratio
|
self.z_to_patch_ratio = z_to_patch_ratio
|
||||||
if z_to_patch_ratio >= 1:
|
if z_to_patch_ratio >= 1:
|
||||||
self.latent_fold_factor = 0
|
self.latent_fold_factor = 0
|
||||||
latent_proj_in_ch = latent_channels
|
latent_proj_in_ch = effective_latent_channels
|
||||||
else:
|
else:
|
||||||
fold_factor = int(1 / z_to_patch_ratio)
|
fold_factor = int(1 / z_to_patch_ratio)
|
||||||
assert fold_factor * z_to_patch_ratio == 1.0
|
assert fold_factor * z_to_patch_ratio == 1.0
|
||||||
self.latent_fold_factor = fold_factor
|
self.latent_fold_factor = fold_factor
|
||||||
latent_proj_in_ch = latent_channels * fold_factor * fold_factor
|
latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor
|
||||||
|
|
||||||
layers = [
|
layers = [
|
||||||
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
|
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||||
nn.SiLU(),
|
nn.SiLU(),
|
||||||
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
|
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||||
]
|
]
|
||||||
for _ in range(num_res_blocks):
|
for _ in range(num_res_blocks):
|
||||||
layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations))
|
layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, dtype=dtype, device=device, operations=operations))
|
||||||
self.latent_proj = nn.Sequential(*layers)
|
self.latent_proj = nn.Sequential(*layers)
|
||||||
|
|
||||||
self.output_heads = nn.ModuleList(
|
self.output_heads = nn.ModuleList(
|
||||||
[operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)]
|
[operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)]
|
||||||
)
|
)
|
||||||
|
self.pit_head = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) if pit_output else None
|
||||||
self.gate_modules = nn.ModuleList(
|
self.gate_modules = nn.ModuleList(
|
||||||
[SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations)
|
[SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations)
|
||||||
for _ in range(num_outputs)]
|
for _ in range(num_outputs)]
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -115,6 +123,11 @@ class LQProjection2D(nn.Module):
|
||||||
return self.gate_modules[out_idx](x, lq_feature, sigma)
|
return self.gate_modules[out_idx](x, lq_feature, sigma)
|
||||||
|
|
||||||
def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor:
|
def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor:
|
||||||
|
f = self.latent_unpatchify_factor
|
||||||
|
if f > 1:
|
||||||
|
B, C, H, W = lq_latent.shape
|
||||||
|
lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W)
|
||||||
|
lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f)
|
||||||
B, z_dim = lq_latent.shape[:2]
|
B, z_dim = lq_latent.shape[:2]
|
||||||
if self.z_to_patch_ratio >= 1:
|
if self.z_to_patch_ratio >= 1:
|
||||||
if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW:
|
if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW:
|
||||||
|
|
@ -134,7 +147,10 @@ class LQProjection2D(nn.Module):
|
||||||
feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
|
feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
|
||||||
B, C, H, W = feat.shape
|
B, C, H, W = feat.shape
|
||||||
tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
|
tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
|
||||||
return [head(tokens) for head in self.output_heads]
|
outputs = [head(tokens) for head in self.output_heads]
|
||||||
|
if self.pit_head is not None:
|
||||||
|
outputs.append(self.pit_head(tokens))
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
|
||||||
class PidNet(PixDiT_T2I):
|
class PidNet(PixDiT_T2I):
|
||||||
|
|
@ -148,6 +164,10 @@ class PidNet(PixDiT_T2I):
|
||||||
lq_interval: int = 2,
|
lq_interval: int = 2,
|
||||||
sr_scale: int = 4,
|
sr_scale: int = 4,
|
||||||
latent_spatial_down_factor: int = 8,
|
latent_spatial_down_factor: int = 8,
|
||||||
|
lq_latent_unpatchify_factor: int = 1,
|
||||||
|
lq_conv_padding_mode: str = "zeros",
|
||||||
|
lq_gate_per_token: bool = False,
|
||||||
|
pit_lq_inject: bool = False,
|
||||||
rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
|
rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
|
||||||
rope_ref_w: int = 1024,
|
rope_ref_w: int = 1024,
|
||||||
image_model=None,
|
image_model=None,
|
||||||
|
|
@ -165,6 +185,8 @@ class PidNet(PixDiT_T2I):
|
||||||
for blk in self.pixel_blocks:
|
for blk in self.pixel_blocks:
|
||||||
blk._rope_fn = _pit_rope_fn
|
blk._rope_fn = _pit_rope_fn
|
||||||
|
|
||||||
|
self.pit_lq_inject = pit_lq_inject
|
||||||
|
|
||||||
num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
|
num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
|
||||||
self.lq_proj = LQProjection2D(
|
self.lq_proj = LQProjection2D(
|
||||||
latent_channels=lq_latent_channels,
|
latent_channels=lq_latent_channels,
|
||||||
|
|
@ -173,13 +195,20 @@ class PidNet(PixDiT_T2I):
|
||||||
patch_size=self.patch_size,
|
patch_size=self.patch_size,
|
||||||
sr_scale=sr_scale,
|
sr_scale=sr_scale,
|
||||||
latent_spatial_down_factor=latent_spatial_down_factor,
|
latent_spatial_down_factor=latent_spatial_down_factor,
|
||||||
|
latent_unpatchify_factor=lq_latent_unpatchify_factor,
|
||||||
num_res_blocks=lq_num_res_blocks,
|
num_res_blocks=lq_num_res_blocks,
|
||||||
num_outputs=num_lq_outputs,
|
num_outputs=num_lq_outputs,
|
||||||
interval=lq_interval,
|
interval=lq_interval,
|
||||||
|
conv_padding_mode=lq_conv_padding_mode,
|
||||||
|
gate_per_token=lq_gate_per_token,
|
||||||
|
pit_output=pit_lq_inject,
|
||||||
dtype=dtype,
|
dtype=dtype,
|
||||||
device=device,
|
device=device,
|
||||||
operations=operations,
|
operations=operations,
|
||||||
)
|
)
|
||||||
|
self.pit_lq_gate = SigmaAwareGate(
|
||||||
|
self.hidden_size, per_token=lq_gate_per_token, dtype=dtype, device=device, operations=operations
|
||||||
|
) if pit_lq_inject else None
|
||||||
|
|
||||||
def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
|
def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
|
||||||
return precompute_freqs_cis_2d(
|
return precompute_freqs_cis_2d(
|
||||||
|
|
@ -197,6 +226,11 @@ class PidNet(PixDiT_T2I):
|
||||||
return s
|
return s
|
||||||
return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx)
|
return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx)
|
||||||
|
|
||||||
|
def _pre_pixel_blocks(self, s, pid_pit_lq_feature=None, pid_degrade_sigma=None, **kwargs):
|
||||||
|
if pid_pit_lq_feature is None:
|
||||||
|
return s
|
||||||
|
return self.pit_lq_gate(s, pid_pit_lq_feature, pid_degrade_sigma)
|
||||||
|
|
||||||
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
|
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
|
||||||
if lq_latent is None:
|
if lq_latent is None:
|
||||||
raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
|
raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
|
||||||
|
|
@ -216,12 +250,14 @@ class PidNet(PixDiT_T2I):
|
||||||
degrade_sigma = degrade_sigma.expand(B).contiguous()
|
degrade_sigma = degrade_sigma.expand(B).contiguous()
|
||||||
|
|
||||||
lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws)
|
lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws)
|
||||||
|
pit_lq_feature = lq_features.pop() if self.pit_lq_inject else None
|
||||||
|
|
||||||
return super()._forward(
|
return super()._forward(
|
||||||
x, timesteps,
|
x, timesteps,
|
||||||
context=context, attention_mask=attention_mask,
|
context=context, attention_mask=attention_mask,
|
||||||
transformer_options=transformer_options,
|
transformer_options=transformer_options,
|
||||||
pid_lq_features=lq_features,
|
pid_lq_features=lq_features,
|
||||||
|
pid_pit_lq_feature=pit_lq_feature,
|
||||||
pid_degrade_sigma=degrade_sigma,
|
pid_degrade_sigma=degrade_sigma,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -58,6 +58,7 @@ import comfy.ldm.omnigen.omnigen2
|
||||||
import comfy.ldm.seedvr.model
|
import comfy.ldm.seedvr.model
|
||||||
import comfy.ldm.boogu.model
|
import comfy.ldm.boogu.model
|
||||||
import comfy.ldm.qwen_image.model
|
import comfy.ldm.qwen_image.model
|
||||||
|
import comfy.ldm.joyimage.model
|
||||||
import comfy.ldm.ideogram4.model
|
import comfy.ldm.ideogram4.model
|
||||||
import comfy.ldm.krea2.model
|
import comfy.ldm.krea2.model
|
||||||
import comfy.ldm.kandinsky5.model
|
import comfy.ldm.kandinsky5.model
|
||||||
|
|
@ -2276,6 +2277,28 @@ class QwenImage(BaseModel):
|
||||||
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
class JoyImage(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel)
|
||||||
|
self.memory_usage_factor_conds = ("ref_latents",)
|
||||||
|
|
||||||
|
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)
|
||||||
|
ref_latents = kwargs.get("reference_latents", None)
|
||||||
|
if ref_latents is not None:
|
||||||
|
out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
|
||||||
|
return out
|
||||||
|
|
||||||
|
def extra_conds_shapes(self, **kwargs):
|
||||||
|
out = {}
|
||||||
|
ref_latents = kwargs.get("reference_latents", None)
|
||||||
|
if ref_latents is not None:
|
||||||
|
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||||
|
return out
|
||||||
|
|
||||||
class Ideogram4(BaseModel):
|
class Ideogram4(BaseModel):
|
||||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
|
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
|
||||||
|
|
|
||||||
|
|
@ -470,15 +470,46 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||||
# PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
|
# PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
|
||||||
_lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
|
_lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
|
||||||
if _lq_w_key in state_dict_keys:
|
if _lq_w_key in state_dict_keys:
|
||||||
in_ch = int(state_dict[_lq_w_key].shape[1])
|
latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])
|
||||||
|
hidden_dim = int(state_dict[_lq_w_key].shape[0])
|
||||||
_gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
|
_gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
|
||||||
num_gates = len({k[len(_gate_prefix):].split('.')[0]
|
num_gates = len({k[len(_gate_prefix):].split('.')[0]
|
||||||
for k in state_dict_keys if k.startswith(_gate_prefix)})
|
for k in state_dict_keys if k.startswith(_gate_prefix)})
|
||||||
|
pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys
|
||||||
dit_config = {"image_model": "pid",
|
dit_config = {"image_model": "pid",
|
||||||
"lq_latent_channels": in_ch,
|
"lq_hidden_dim": hidden_dim}
|
||||||
"latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
|
|
||||||
if num_gates > 0:
|
if num_gates > 0:
|
||||||
dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
|
dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
|
||||||
|
if pid_v1_5:
|
||||||
|
pid_v1_5_variants = {
|
||||||
|
16: { # Flux and QwenImage
|
||||||
|
"lq_latent_channels": 16,
|
||||||
|
"latent_spatial_down_factor": 8,
|
||||||
|
"lq_latent_unpatchify_factor": 1,
|
||||||
|
},
|
||||||
|
32: { # Flux2 after 2x latent unpatchify
|
||||||
|
"lq_latent_channels": 128,
|
||||||
|
"latent_spatial_down_factor": 16,
|
||||||
|
"lq_latent_unpatchify_factor": 2,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
variant = pid_v1_5_variants.get(latent_proj_in_channels)
|
||||||
|
if variant is None:
|
||||||
|
raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels")
|
||||||
|
gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)]
|
||||||
|
dit_config.update(variant)
|
||||||
|
dit_config.update({
|
||||||
|
"lq_conv_padding_mode": "replicate",
|
||||||
|
"lq_gate_per_token": gate_weight.shape[0] == 1,
|
||||||
|
"pit_lq_inject": True,
|
||||||
|
"rope_ref_h": 2048,
|
||||||
|
"rope_ref_w": 2048,
|
||||||
|
})
|
||||||
|
else:
|
||||||
|
dit_config.update({
|
||||||
|
"lq_latent_channels": latent_proj_in_channels,
|
||||||
|
"latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8,
|
||||||
|
})
|
||||||
return dit_config
|
return dit_config
|
||||||
|
|
||||||
if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I
|
if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I
|
||||||
|
|
@ -1027,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||||
dit_config["image_model"] = "SAM31"
|
dit_config["image_model"] = "SAM31"
|
||||||
return dit_config
|
return dit_config
|
||||||
|
|
||||||
|
if (
|
||||||
|
'{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and '{}img_in.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5
|
||||||
|
):
|
||||||
|
img_in = state_dict['{}img_in.weight'.format(key_prefix)]
|
||||||
|
head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
|
||||||
|
return {
|
||||||
|
"image_model": "joyimage",
|
||||||
|
"in_channels": img_in.shape[1],
|
||||||
|
"hidden_size": img_in.shape[0],
|
||||||
|
"patch_size": list(img_in.shape[2:]),
|
||||||
|
"num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'),
|
||||||
|
"num_attention_heads": img_in.shape[0] // head_dim,
|
||||||
|
"text_dim": 4096,
|
||||||
|
}
|
||||||
|
|
||||||
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -76,6 +76,7 @@ import comfy.text_encoders.gemma4
|
||||||
import comfy.text_encoders.cogvideo
|
import comfy.text_encoders.cogvideo
|
||||||
import comfy.text_encoders.sa3
|
import comfy.text_encoders.sa3
|
||||||
import comfy.text_encoders.gpt_oss
|
import comfy.text_encoders.gpt_oss
|
||||||
|
import comfy.text_encoders.joyimage
|
||||||
|
|
||||||
import comfy.model_patcher
|
import comfy.model_patcher
|
||||||
import comfy.lora
|
import comfy.lora
|
||||||
|
|
@ -1377,6 +1378,7 @@ class CLIPType(Enum):
|
||||||
IDEOGRAM4 = 30
|
IDEOGRAM4 = 30
|
||||||
BOOGU = 31
|
BOOGU = 31
|
||||||
KREA2 = 32
|
KREA2 = 32
|
||||||
|
JOYIMAGE = 33
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -1706,6 +1708,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
|
||||||
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_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.krea2.te(**llama_detect(clip_data))
|
clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
|
||||||
clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
|
clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
|
||||||
|
elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B: # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx.
|
||||||
|
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.joyimage.te(**llama_detect(clip_data))
|
||||||
|
clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer
|
||||||
elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
|
elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
|
||||||
klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
|
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.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
|
||||||
|
|
|
||||||
|
|
@ -27,6 +27,7 @@ import comfy.text_encoders.z_image
|
||||||
import comfy.text_encoders.ideogram4
|
import comfy.text_encoders.ideogram4
|
||||||
import comfy.text_encoders.boogu
|
import comfy.text_encoders.boogu
|
||||||
import comfy.text_encoders.krea2
|
import comfy.text_encoders.krea2
|
||||||
|
import comfy.text_encoders.joyimage
|
||||||
import comfy.text_encoders.anima
|
import comfy.text_encoders.anima
|
||||||
import comfy.text_encoders.ace15
|
import comfy.text_encoders.ace15
|
||||||
import comfy.text_encoders.longcat_image
|
import comfy.text_encoders.longcat_image
|
||||||
|
|
@ -1911,6 +1912,38 @@ class QwenImage(supported_models_base.BASE):
|
||||||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
||||||
return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
|
return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
|
||||||
|
|
||||||
|
class JoyImage(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"image_model": "joyimage",
|
||||||
|
}
|
||||||
|
|
||||||
|
sampling_settings = {
|
||||||
|
"multiplier": 1000,
|
||||||
|
"shift": 1.5,
|
||||||
|
}
|
||||||
|
|
||||||
|
memory_usage_factor = 1.8
|
||||||
|
|
||||||
|
unet_extra_config = {
|
||||||
|
"theta": 10000,
|
||||||
|
"rope_dim_list": [16, 56, 56],
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.Wan21
|
||||||
|
|
||||||
|
supported_inference_dtypes = [torch.bfloat16, torch.float32]
|
||||||
|
|
||||||
|
vae_key_prefix = ["vae."]
|
||||||
|
text_encoder_key_prefix = ["text_encoders."]
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
return model_base.JoyImage(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.transformer.".format(pref))
|
||||||
|
return supported_models_base.ClipTarget(comfy.text_encoders.joyimage.JoyImageTokenizer, comfy.text_encoders.joyimage.te(**qwen3vl_detect))
|
||||||
|
|
||||||
class HunyuanImage21(HunyuanVideo):
|
class HunyuanImage21(HunyuanVideo):
|
||||||
unet_config = {
|
unet_config = {
|
||||||
"image_model": "hunyuan_video",
|
"image_model": "hunyuan_video",
|
||||||
|
|
@ -2389,6 +2422,7 @@ models = [
|
||||||
Omnigen2,
|
Omnigen2,
|
||||||
Boogu,
|
Boogu,
|
||||||
QwenImage,
|
QwenImage,
|
||||||
|
JoyImage,
|
||||||
Ideogram4,
|
Ideogram4,
|
||||||
Krea2,
|
Krea2,
|
||||||
Flux2,
|
Flux2,
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,97 @@
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from comfy import sd1_clip
|
||||||
|
import comfy.text_encoders.qwen_vl
|
||||||
|
from comfy.text_encoders.qwen3vl import Qwen3VL, Qwen3VLTokenizer
|
||||||
|
|
||||||
|
JOYIMAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||||
|
JOYIMAGE_TEMPLATE_TEXT = (
|
||||||
|
"<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
|
||||||
|
"quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
|
||||||
|
"<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
|
||||||
|
)
|
||||||
|
JOYIMAGE_TEMPLATE_IMAGE = (
|
||||||
|
"<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
|
||||||
|
"quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
|
||||||
|
f"<|im_start|>user\n{JOYIMAGE_VISION_BLOCK}{{}}<|im_end|>\n<|im_start|>assistant\n"
|
||||||
|
)
|
||||||
|
# The DiT was trained without the leading system-prompt tokens.
|
||||||
|
JOYIMAGE_DROP_IDX = 34
|
||||||
|
PAD_TOKEN = 151643
|
||||||
|
|
||||||
|
|
||||||
|
class Qwen3VL8B_JoyImage(Qwen3VL):
|
||||||
|
model_type = "qwen3vl_8b"
|
||||||
|
|
||||||
|
def preprocess_embed(self, embed, device):
|
||||||
|
if embed["type"] == "image":
|
||||||
|
image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(
|
||||||
|
embed["data"], min_pixels=65536, max_pixels=16777216, patch_size=16,
|
||||||
|
image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5],
|
||||||
|
interpolation="bicubic",
|
||||||
|
)
|
||||||
|
merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)
|
||||||
|
return merged, {"grid": grid, "deepstack": deepstack}
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageTokenizer(Qwen3VLTokenizer):
|
||||||
|
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||||
|
super().__init__(
|
||||||
|
embedding_directory=embedding_directory, tokenizer_data=tokenizer_data,
|
||||||
|
model_type="qwen3vl_8b",
|
||||||
|
)
|
||||||
|
self.llama_template = JOYIMAGE_TEMPLATE_TEXT
|
||||||
|
self.llama_template_images = JOYIMAGE_TEMPLATE_IMAGE
|
||||||
|
|
||||||
|
def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=None, **kwargs):
|
||||||
|
kwargs.pop("thinking", None)
|
||||||
|
return super().tokenize_with_weights(
|
||||||
|
text, return_word_ids=return_word_ids, llama_template=llama_template,
|
||||||
|
images=images or [], thinking=True, **kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _JoyImageClipModel(sd1_clip.SDClipModel):
|
||||||
|
def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None,
|
||||||
|
attention_mask=True, model_options={}):
|
||||||
|
super().__init__(
|
||||||
|
device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={},
|
||||||
|
# JoyImage conditions on the pre-final-norm output of the last decoder layer.
|
||||||
|
dtype=dtype, special_tokens={"pad": PAD_TOKEN}, layer_norm_hidden_state=False,
|
||||||
|
model_class=Qwen3VL8B_JoyImage, enable_attention_masks=attention_mask,
|
||||||
|
return_attention_masks=attention_mask, model_options=model_options,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageTEModel(sd1_clip.SD1ClipModel):
|
||||||
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||||
|
super().__init__(
|
||||||
|
device=device, dtype=dtype, name="qwen3vl_8b",
|
||||||
|
clip_model=_JoyImageClipModel, model_options=model_options,
|
||||||
|
)
|
||||||
|
|
||||||
|
def encode_token_weights(self, token_weight_pairs):
|
||||||
|
out, pooled, extra = super().encode_token_weights(token_weight_pairs)
|
||||||
|
if out.shape[1] <= JOYIMAGE_DROP_IDX:
|
||||||
|
raise ValueError(
|
||||||
|
f"JoyImageTEModel: encoded sequence length {out.shape[1]} is shorter "
|
||||||
|
f"than drop_idx={JOYIMAGE_DROP_IDX}; the prompt did not include the "
|
||||||
|
f"template prefix."
|
||||||
|
)
|
||||||
|
out = out[:, JOYIMAGE_DROP_IDX:]
|
||||||
|
if "attention_mask" in extra:
|
||||||
|
extra["attention_mask"] = extra["attention_mask"][:, JOYIMAGE_DROP_IDX:]
|
||||||
|
return out, pooled, extra
|
||||||
|
|
||||||
|
|
||||||
|
def te(dtype_llama=None, llama_quantization_metadata=None):
|
||||||
|
class JoyImageTEModel_(JoyImageTEModel):
|
||||||
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||||
|
if llama_quantization_metadata is not None:
|
||||||
|
model_options = model_options.copy()
|
||||||
|
model_options["quantization_metadata"] = llama_quantization_metadata
|
||||||
|
if dtype_llama is not None:
|
||||||
|
dtype = dtype_llama
|
||||||
|
super().__init__(device=device, dtype=dtype, model_options=model_options)
|
||||||
|
return JoyImageTEModel_
|
||||||
|
|
@ -15,6 +15,7 @@ def process_qwen2vl_images(
|
||||||
merge_size: int = 2,
|
merge_size: int = 2,
|
||||||
image_mean: list = None,
|
image_mean: list = None,
|
||||||
image_std: list = None,
|
image_std: list = None,
|
||||||
|
interpolation: str = "bilinear",
|
||||||
):
|
):
|
||||||
if image_mean is None:
|
if image_mean is None:
|
||||||
image_mean = [0.48145466, 0.4578275, 0.40821073]
|
image_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||||
|
|
@ -47,10 +48,9 @@ def process_qwen2vl_images(
|
||||||
img_resized = F.interpolate(
|
img_resized = F.interpolate(
|
||||||
img.unsqueeze(0),
|
img.unsqueeze(0),
|
||||||
size=(h_bar, w_bar),
|
size=(h_bar, w_bar),
|
||||||
mode='bilinear',
|
mode=interpolation,
|
||||||
align_corners=False
|
align_corners=False
|
||||||
).squeeze(0)
|
).squeeze(0)
|
||||||
|
|
||||||
normalized = img_resized.clone()
|
normalized = img_resized.clone()
|
||||||
for c in range(3):
|
for c in range(3):
|
||||||
normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]
|
normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,6 @@
|
||||||
from av.container import InputContainer
|
from av.container import InputContainer
|
||||||
from av.subtitles.stream import SubtitleStream
|
from av.subtitles.stream import SubtitleStream
|
||||||
|
from av.video.reformatter import ColorRange
|
||||||
from fractions import Fraction
|
from fractions import Fraction
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from .._input import AudioInput, VideoInput
|
from .._input import AudioInput, VideoInput
|
||||||
|
|
@ -9,6 +10,7 @@ import itertools
|
||||||
import json
|
import json
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import math
|
import math
|
||||||
|
import os
|
||||||
import torch
|
import torch
|
||||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||||
import logging
|
import logging
|
||||||
|
|
@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
|
||||||
return max(component.bits for component in stream.format.components)
|
return max(component.bits for component in stream.format.components)
|
||||||
|
|
||||||
|
|
||||||
|
def last_decodable_audio_stream(container: InputContainer):
|
||||||
|
"""Streams FFmpeg has no decoder for have no codec context, and decoding their
|
||||||
|
packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
|
||||||
|
stream = next(
|
||||||
|
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
if stream is None and len(container.streams.audio):
|
||||||
|
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||||
|
return stream
|
||||||
|
|
||||||
|
|
||||||
|
def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
|
||||||
|
"""Containers probed only up to a window (mpegts) leave audio codec parameters unset when
|
||||||
|
audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
|
||||||
|
Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
|
||||||
|
for i, packet in enumerate(container.demux(audio_stream)):
|
||||||
|
try:
|
||||||
|
frames = packet.decode()
|
||||||
|
except av.error.FFmpegError:
|
||||||
|
frames = ()
|
||||||
|
if frames:
|
||||||
|
return frames[0].sample_rate, frames[0].layout.nb_channels
|
||||||
|
if i >= max_packets:
|
||||||
|
break
|
||||||
|
return 0, 0
|
||||||
|
|
||||||
|
|
||||||
|
def write_output_metadata(container: InputContainer, output, metadata: dict | None):
|
||||||
|
"""Copy the source container's metadata, then overlay the caller's tags."""
|
||||||
|
for key, value in container.metadata.items():
|
||||||
|
if metadata is None or key not in metadata:
|
||||||
|
output.metadata[key] = value
|
||||||
|
if metadata is not None:
|
||||||
|
for key, value in metadata.items():
|
||||||
|
output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
|
||||||
|
|
||||||
|
|
||||||
|
def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
|
||||||
|
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
||||||
|
raise ValueError("Only MP4 format is supported for now")
|
||||||
|
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
||||||
|
raise ValueError("Only H264 codec is supported for now")
|
||||||
|
open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
|
||||||
|
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
||||||
|
open_kwargs["format"] = format.value
|
||||||
|
elif isinstance(path, io.BytesIO):
|
||||||
|
open_kwargs["format"] = "mp4" # no file extension to infer the format from
|
||||||
|
return open_kwargs
|
||||||
|
|
||||||
|
|
||||||
class VideoFromFile(VideoInput):
|
class VideoFromFile(VideoInput):
|
||||||
"""
|
"""
|
||||||
Class representing video input from a file.
|
Class representing video input from a file.
|
||||||
|
|
@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
|
||||||
return estimated_frames
|
return estimated_frames
|
||||||
|
|
||||||
# 3. Last resort: decode frames and count them (streaming)
|
# 3. Last resort: decode frames and count them (streaming)
|
||||||
if self.__start_time < 0:
|
start_time, duration = self.get_active_trim_window()
|
||||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
|
||||||
else:
|
|
||||||
start_time = self.__start_time
|
|
||||||
frame_count = 1
|
frame_count = 1
|
||||||
start_pts = int(start_time / video_stream.time_base)
|
start_pts = int(start_time / video_stream.time_base)
|
||||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||||
container.seek(start_pts, stream=video_stream)
|
container.seek(start_pts, stream=video_stream)
|
||||||
frame_iterator = (
|
frame_iterator = (
|
||||||
container.decode(video_stream)
|
container.decode(video_stream)
|
||||||
|
|
@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
|
||||||
|
|
||||||
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
||||||
video_stream = self._get_first_video_stream(container)
|
video_stream = self._get_first_video_stream(container)
|
||||||
if self.__start_time < 0:
|
start_time, duration = self.get_active_trim_window()
|
||||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
|
||||||
else:
|
|
||||||
start_time = self.__start_time
|
|
||||||
|
|
||||||
# Get video frames
|
# Get video frames
|
||||||
frames = []
|
frames = []
|
||||||
audio_frames = []
|
audio_frames = []
|
||||||
alphas = None
|
alphas = None
|
||||||
start_pts = int(start_time / video_stream.time_base)
|
start_pts = int(start_time / video_stream.time_base)
|
||||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||||
|
|
||||||
if start_pts != 0:
|
if start_pts != 0:
|
||||||
container.seek(start_pts, stream=video_stream)
|
container.seek(start_pts, stream=video_stream)
|
||||||
|
|
@ -281,18 +328,11 @@ class VideoFromFile(VideoInput):
|
||||||
video_done = False
|
video_done = False
|
||||||
audio_done = True
|
audio_done = True
|
||||||
|
|
||||||
# Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
|
audio_stream = last_decodable_audio_stream(container)
|
||||||
# and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
|
|
||||||
audio_stream = next(
|
|
||||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
if audio_stream is not None:
|
if audio_stream is not None:
|
||||||
streams += [audio_stream]
|
streams += [audio_stream]
|
||||||
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
||||||
audio_done = False
|
audio_done = False
|
||||||
elif len(container.streams.audio):
|
|
||||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
|
||||||
|
|
||||||
for packet in container.demux(*streams):
|
for packet in container.demux(*streams):
|
||||||
if video_done and audio_done:
|
if video_done and audio_done:
|
||||||
|
|
@ -305,7 +345,7 @@ class VideoFromFile(VideoInput):
|
||||||
for frame in packet.decode():
|
for frame in packet.decode():
|
||||||
if frame.pts < start_pts:
|
if frame.pts < start_pts:
|
||||||
continue
|
continue
|
||||||
if self.__duration and frame.pts >= end_pts:
|
if duration and frame.pts >= end_pts:
|
||||||
video_done = True
|
video_done = True
|
||||||
break
|
break
|
||||||
|
|
||||||
|
|
@ -372,7 +412,7 @@ class VideoFromFile(VideoInput):
|
||||||
map(resampler.resample, packet.decode())
|
map(resampler.resample, packet.decode())
|
||||||
)
|
)
|
||||||
for frame in aframes:
|
for frame in aframes:
|
||||||
if self.__duration and frame.time > start_time + self.__duration:
|
if duration and frame.time > start_time + duration:
|
||||||
audio_done = True
|
audio_done = True
|
||||||
break
|
break
|
||||||
|
|
||||||
|
|
@ -394,8 +434,8 @@ class VideoFromFile(VideoInput):
|
||||||
|
|
||||||
if len(audio_frames) > 0:
|
if len(audio_frames) > 0:
|
||||||
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
||||||
if self.__duration:
|
if duration:
|
||||||
audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
|
audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
|
||||||
|
|
||||||
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
|
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
|
||||||
audio = AudioInput({
|
audio = AudioInput({
|
||||||
|
|
@ -441,28 +481,14 @@ class VideoFromFile(VideoInput):
|
||||||
if not reuse_streams:
|
if not reuse_streams:
|
||||||
if bit_depth is None:
|
if bit_depth is None:
|
||||||
bit_depth = source_bit_depth
|
bit_depth = source_bit_depth
|
||||||
components = self.get_components_internal(container)
|
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
|
||||||
video = VideoFromComponents(components)
|
|
||||||
return video.save_to(
|
|
||||||
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
|
|
||||||
)
|
|
||||||
|
|
||||||
streams = container.streams
|
streams = container.streams
|
||||||
|
|
||||||
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
||||||
with av.open(path, **open_kwargs) as output_container:
|
with av.open(path, **open_kwargs) as output_container:
|
||||||
# Copy over the original metadata
|
# Add metadata before writing any streams
|
||||||
for key, value in container.metadata.items():
|
write_output_metadata(container, output_container, metadata)
|
||||||
if metadata is None or key not in metadata:
|
|
||||||
output_container.metadata[key] = value
|
|
||||||
|
|
||||||
# Add our new metadata
|
|
||||||
if metadata is not None:
|
|
||||||
for key, value in metadata.items():
|
|
||||||
if isinstance(value, str):
|
|
||||||
output_container.metadata[key] = value
|
|
||||||
else:
|
|
||||||
output_container.metadata[key] = json.dumps(value)
|
|
||||||
|
|
||||||
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
||||||
stream_map = {}
|
stream_map = {}
|
||||||
|
|
@ -480,6 +506,282 @@ class VideoFromFile(VideoInput):
|
||||||
packet.stream = stream_map[packet.stream]
|
packet.stream = stream_map[packet.stream]
|
||||||
output_container.mux(packet)
|
output_container.mux(packet)
|
||||||
|
|
||||||
|
def _save_transcoded(
|
||||||
|
self,
|
||||||
|
container: InputContainer,
|
||||||
|
path: str | io.BytesIO,
|
||||||
|
format: VideoContainer,
|
||||||
|
codec: VideoCodec,
|
||||||
|
metadata: dict | None,
|
||||||
|
bit_depth: int,
|
||||||
|
):
|
||||||
|
"""Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
|
||||||
|
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||||
|
video_stream = self._get_first_video_stream(container)
|
||||||
|
start_time, duration = self.get_active_trim_window()
|
||||||
|
start_pts = int(start_time / video_stream.time_base)
|
||||||
|
end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
|
||||||
|
stream_end_pts = None
|
||||||
|
if video_stream.duration is not None:
|
||||||
|
stream_end_pts = (video_stream.start_time or 0) + video_stream.duration
|
||||||
|
output_end_pts = end_pts
|
||||||
|
if stream_end_pts is not None and (output_end_pts is None or stream_end_pts < output_end_pts):
|
||||||
|
output_end_pts = stream_end_pts
|
||||||
|
if start_pts != 0:
|
||||||
|
container.seek(start_pts, stream=video_stream)
|
||||||
|
|
||||||
|
audio_stream = last_decodable_audio_stream(container)
|
||||||
|
pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p"
|
||||||
|
rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||||
|
|
||||||
|
resampler = None
|
||||||
|
sample_rate = 0
|
||||||
|
audio_time_base = None
|
||||||
|
duration_cap = None
|
||||||
|
if audio_stream is not None:
|
||||||
|
sample_rate = audio_stream.codec_context.sample_rate
|
||||||
|
channels = audio_stream.codec_context.channels
|
||||||
|
if not sample_rate:
|
||||||
|
sample_rate, channels = probe_audio_params(container, audio_stream)
|
||||||
|
container.seek(start_pts, stream=video_stream)
|
||||||
|
if sample_rate:
|
||||||
|
audio_stream.codec_context.flush_buffers()
|
||||||
|
else:
|
||||||
|
logging.warning("Audio stream parameters could not be determined; ignoring audio.")
|
||||||
|
audio_stream = None
|
||||||
|
if audio_stream is not None:
|
||||||
|
audio_time_base = Fraction(1, sample_rate)
|
||||||
|
layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
|
||||||
|
resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
|
||||||
|
if duration:
|
||||||
|
duration_cap = math.ceil(duration * sample_rate)
|
||||||
|
|
||||||
|
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
|
||||||
|
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
|
||||||
|
video_done = False
|
||||||
|
audio_done = audio_stream is None
|
||||||
|
video_pts_offset = None
|
||||||
|
last_video_pts = None
|
||||||
|
last_video_end = None
|
||||||
|
# rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise
|
||||||
|
video_frame_durations = {}
|
||||||
|
source_size = None
|
||||||
|
rotation_k = 0
|
||||||
|
rotation_filter = None
|
||||||
|
audio_started = False
|
||||||
|
samples_written = 0
|
||||||
|
pending_audio = []
|
||||||
|
# The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims),
|
||||||
|
# and never seeking back keeps webm/mkv leading audio intact.
|
||||||
|
output = None
|
||||||
|
out_video = None
|
||||||
|
out_audio = None
|
||||||
|
|
||||||
|
def audio_frame_from_ndarray(nd_planar):
|
||||||
|
frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout)
|
||||||
|
frame.sample_rate = sample_rate
|
||||||
|
return frame
|
||||||
|
|
||||||
|
def drain_audio(final=False):
|
||||||
|
# Audio may cover the pts span of the video written so far, capped by the requested duration
|
||||||
|
nonlocal samples_written, audio_done
|
||||||
|
if last_video_end is None:
|
||||||
|
cap = 0
|
||||||
|
else:
|
||||||
|
cap = math.ceil(last_video_end * video_stream.time_base * sample_rate)
|
||||||
|
if duration_cap is not None:
|
||||||
|
cap = min(cap, duration_cap)
|
||||||
|
while pending_audio and not audio_done:
|
||||||
|
frame = pending_audio[0]
|
||||||
|
if samples_written + frame.samples <= cap:
|
||||||
|
frame.pts = samples_written
|
||||||
|
frame.time_base = audio_time_base
|
||||||
|
output.mux(out_audio.encode(frame))
|
||||||
|
samples_written += frame.samples
|
||||||
|
pending_audio.pop(0)
|
||||||
|
continue
|
||||||
|
if final:
|
||||||
|
keep = frame.to_ndarray()[..., :cap - samples_written]
|
||||||
|
if keep.shape[-1] > 0:
|
||||||
|
tail = audio_frame_from_ndarray(keep)
|
||||||
|
tail.pts = samples_written
|
||||||
|
tail.time_base = audio_time_base
|
||||||
|
output.mux(out_audio.encode(tail))
|
||||||
|
samples_written += keep.shape[-1]
|
||||||
|
pending_audio.clear()
|
||||||
|
break
|
||||||
|
if duration_cap is not None and samples_written >= duration_cap:
|
||||||
|
audio_done = True
|
||||||
|
return cap
|
||||||
|
|
||||||
|
try:
|
||||||
|
for packet in container.demux(*streams):
|
||||||
|
if video_done and audio_done:
|
||||||
|
break
|
||||||
|
|
||||||
|
if packet.stream == video_stream and not video_done:
|
||||||
|
try:
|
||||||
|
frames = packet.decode()
|
||||||
|
except av.error.InvalidDataError:
|
||||||
|
logging.info("pyav decode error")
|
||||||
|
continue
|
||||||
|
for frame in frames:
|
||||||
|
if frame.pts is not None and frame.pts < start_pts:
|
||||||
|
continue
|
||||||
|
if end_pts is not None and frame.pts is not None and frame.pts >= end_pts:
|
||||||
|
video_done = True
|
||||||
|
if last_video_pts is not None:
|
||||||
|
# the source continues past the window: hold the last kept frame to the window end
|
||||||
|
end_offset = video_pts_offset if video_pts_offset is not None else start_pts
|
||||||
|
last_video_end = max(last_video_end, end_pts - end_offset)
|
||||||
|
break
|
||||||
|
# the source's true display duration of this frame; average_rate is not a
|
||||||
|
# frame duration (sparse/VFR sources), so it is only the fallback
|
||||||
|
frame_duration = frame.duration if frame.duration else pts_step
|
||||||
|
if end_pts is not None and frame.pts is not None:
|
||||||
|
frame_duration = min(frame_duration, end_pts - frame.pts)
|
||||||
|
if output is None:
|
||||||
|
rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
|
||||||
|
if rotation_k % 2:
|
||||||
|
out_width, out_height = frame.height, frame.width
|
||||||
|
else:
|
||||||
|
out_width, out_height = frame.width, frame.height
|
||||||
|
if out_width % 2 or out_height % 2:
|
||||||
|
raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
|
||||||
|
source_size = (frame.width, frame.height)
|
||||||
|
output = av.open(path, **open_kwargs)
|
||||||
|
# Add metadata before writing any streams
|
||||||
|
write_output_metadata(container, output, metadata)
|
||||||
|
out_video = output.add_stream("h264", rate=rate)
|
||||||
|
# no B-frames: reordering makes mp4 sample durations follow decode order,
|
||||||
|
# so irregular-VFR spans and trim windows land wrong
|
||||||
|
out_video.codec_context.max_b_frames = 0
|
||||||
|
out_video.width = out_width
|
||||||
|
out_video.height = out_height
|
||||||
|
out_video.pix_fmt = pix_fmt
|
||||||
|
# source pts pass through (rebased to 0), so variable frame rate survives
|
||||||
|
out_video.codec_context.time_base = video_stream.time_base
|
||||||
|
if audio_stream is not None:
|
||||||
|
out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
|
||||||
|
if (frame.width, frame.height) != source_size:
|
||||||
|
# encoding would silently rescale the new geometry into the old one
|
||||||
|
raise ValueError(
|
||||||
|
f"Video resolution changes mid-stream "
|
||||||
|
f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode"
|
||||||
|
)
|
||||||
|
if rotation_k:
|
||||||
|
if rotation_filter is None:
|
||||||
|
g = av.filter.Graph()
|
||||||
|
g_src = g.add_buffer(width=frame.width, height=frame.height,
|
||||||
|
format=frame.format.name, time_base=video_stream.time_base)
|
||||||
|
tail = g_src
|
||||||
|
for filter_name, filter_args in {1: [("transpose", "cclock")],
|
||||||
|
2: [("hflip", None), ("vflip", None)],
|
||||||
|
3: [("transpose", "clock")]}[rotation_k]:
|
||||||
|
step = g.add(filter_name, filter_args)
|
||||||
|
tail.link_to(step)
|
||||||
|
tail = step
|
||||||
|
g_sink = g.add("buffersink")
|
||||||
|
tail.link_to(g_sink)
|
||||||
|
g.configure()
|
||||||
|
rotation_filter = (g_src, g_sink)
|
||||||
|
rotation_filter[0].push(frame)
|
||||||
|
frame = rotation_filter[1].pull()
|
||||||
|
if frame.color_range == ColorRange.JPEG:
|
||||||
|
# compress full-range sources (yuvj/MJPEG) to limited range
|
||||||
|
frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
|
||||||
|
else:
|
||||||
|
frame = frame.reformat(format=pix_fmt)
|
||||||
|
frame_output_end = None
|
||||||
|
if frame.pts is not None:
|
||||||
|
if video_pts_offset is None:
|
||||||
|
video_pts_offset = frame.pts
|
||||||
|
frame.pts -= video_pts_offset
|
||||||
|
if output_end_pts is not None:
|
||||||
|
frame_output_end = output_end_pts - video_pts_offset
|
||||||
|
if frame.pts + frame_duration > frame_output_end:
|
||||||
|
clamped_pts = frame_output_end - frame_duration
|
||||||
|
if clamped_pts >= 0 and (last_video_pts is None or clamped_pts > last_video_pts):
|
||||||
|
frame.pts = min(frame.pts, clamped_pts)
|
||||||
|
elif frame.pts < frame_output_end:
|
||||||
|
frame_duration = frame_output_end - frame.pts
|
||||||
|
else:
|
||||||
|
continue
|
||||||
|
if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts):
|
||||||
|
# broken sources emit missing/backward timestamps mid-stream, which the
|
||||||
|
# muxer rejects; nudge them forward by one nominal frame interval
|
||||||
|
frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step
|
||||||
|
if frame_output_end is not None and frame.pts + frame_duration > frame_output_end:
|
||||||
|
if frame.pts >= frame_output_end:
|
||||||
|
continue
|
||||||
|
frame_duration = frame_output_end - frame.pts
|
||||||
|
last_video_pts = frame.pts
|
||||||
|
last_video_end = frame.pts + frame_duration
|
||||||
|
video_frame_durations[frame.pts] = frame_duration
|
||||||
|
# the decoded pict_type would force x264's frame types (intra-only
|
||||||
|
# sources like MJPEG/ProRes would come out all-keyframe)
|
||||||
|
frame.pict_type = 0
|
||||||
|
for out_packet in out_video.encode(frame):
|
||||||
|
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||||
|
output.mux(out_packet)
|
||||||
|
drain_audio()
|
||||||
|
|
||||||
|
elif packet.stream == audio_stream and not audio_done:
|
||||||
|
for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
|
||||||
|
frame_start = None
|
||||||
|
if resampled.pts is not None:
|
||||||
|
# passthrough frames keep the source stream's time base
|
||||||
|
tb = resampled.time_base if resampled.time_base else audio_time_base
|
||||||
|
frame_start = float(resampled.pts * tb)
|
||||||
|
if duration and not audio_started and frame_start >= start_time + duration:
|
||||||
|
audio_done = True
|
||||||
|
break
|
||||||
|
if not audio_started:
|
||||||
|
if frame_start is None:
|
||||||
|
frame_start = 0.0
|
||||||
|
to_skip = max(0, int((start_time - frame_start) * sample_rate))
|
||||||
|
if to_skip >= resampled.samples:
|
||||||
|
continue
|
||||||
|
audio_started = True
|
||||||
|
if duration and frame_start > start_time:
|
||||||
|
duration_cap = min(duration_cap, math.ceil((start_time + duration - frame_start) * sample_rate))
|
||||||
|
if to_skip:
|
||||||
|
pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
|
||||||
|
continue
|
||||||
|
pending_audio.append(resampled)
|
||||||
|
if video_done:
|
||||||
|
# the video window is complete so the cap is final, but containers
|
||||||
|
# that interleave audio behind video (fragmented mp4) still owe most
|
||||||
|
# of it: stop only once the demuxed audio covers the cap
|
||||||
|
cap = drain_audio()
|
||||||
|
if pending_audio or samples_written >= cap:
|
||||||
|
drain_audio(final=True)
|
||||||
|
audio_done = True
|
||||||
|
break
|
||||||
|
|
||||||
|
if output is None:
|
||||||
|
raise ValueError(f"No decodable video frames found in file '{self.__file}'")
|
||||||
|
if out_audio is not None and not audio_done:
|
||||||
|
drain_audio(final=True)
|
||||||
|
window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
|
||||||
|
for out_packet in out_video.encode(None):
|
||||||
|
duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||||
|
if out_packet.pts == last_video_pts:
|
||||||
|
duration = max(duration, window_fill)
|
||||||
|
out_packet.duration = duration
|
||||||
|
output.mux(out_packet)
|
||||||
|
if out_audio is not None:
|
||||||
|
output.mux(out_audio.encode(None))
|
||||||
|
except BaseException:
|
||||||
|
if output is not None:
|
||||||
|
output.close()
|
||||||
|
if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
|
||||||
|
os.remove(path)
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
if output is not None:
|
||||||
|
output.close()
|
||||||
|
|
||||||
def _get_first_video_stream(self, container: InputContainer):
|
def _get_first_video_stream(self, container: InputContainer):
|
||||||
if len(container.streams.video):
|
if len(container.streams.video):
|
||||||
return container.streams.video[0]
|
return container.streams.video[0]
|
||||||
|
|
@ -527,22 +829,12 @@ class VideoFromComponents(VideoInput):
|
||||||
bit_depth: int | None = None,
|
bit_depth: int | None = None,
|
||||||
):
|
):
|
||||||
"""Save the video to a file path or BytesIO buffer."""
|
"""Save the video to a file path or BytesIO buffer."""
|
||||||
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||||
raise ValueError("Only MP4 format is supported for now")
|
|
||||||
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
|
||||||
raise ValueError("Only H264 codec is supported for now")
|
|
||||||
# None means "use the depth this video was created with" (CreateVideo's choice).
|
# None means "use the depth this video was created with" (CreateVideo's choice).
|
||||||
if bit_depth is None:
|
if bit_depth is None:
|
||||||
bit_depth = self.__bit_depth
|
bit_depth = self.__bit_depth
|
||||||
is_10bit = bit_depth >= 10
|
is_10bit = bit_depth >= 10
|
||||||
extra_kwargs = {}
|
with av.open(path, **open_kwargs) as output:
|
||||||
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
|
||||||
extra_kwargs["format"] = format.value
|
|
||||||
elif isinstance(path, io.BytesIO):
|
|
||||||
# BytesIO has no file extension, so av.open can't infer the format.
|
|
||||||
# Default to mp4 since that's the only supported format anyway.
|
|
||||||
extra_kwargs["format"] = "mp4"
|
|
||||||
with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
|
|
||||||
# Add metadata before writing any streams
|
# Add metadata before writing any streams
|
||||||
if metadata is not None:
|
if metadata is not None:
|
||||||
for key, value in metadata.items():
|
for key, value in metadata.items():
|
||||||
|
|
|
||||||
|
|
@ -128,7 +128,7 @@ class OpenAIResponse(ModelResponseProperties, ResponseProperties):
|
||||||
parallel_tool_calls: bool | None = Field(True)
|
parallel_tool_calls: bool | None = Field(True)
|
||||||
status: str | None = Field(
|
status: str | None = Field(
|
||||||
None,
|
None,
|
||||||
description="One of `completed`, `failed`, `in_progress`, or `incomplete`.",
|
description="One of `completed`, `failed`, `in_progress`, `incomplete`, `queued`, or `cancelled`.",
|
||||||
)
|
)
|
||||||
usage: ResponseUsage | None = Field(None)
|
usage: ResponseUsage | None = Field(None)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,49 @@
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
|
||||||
|
class SyncInputItem(BaseModel):
|
||||||
|
type: str = Field(..., description="Input kind: 'video', 'image' or 'audio'.")
|
||||||
|
url: str = Field(...)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncActiveSpeakerDetection(BaseModel):
|
||||||
|
auto_detect: bool | None = Field(
|
||||||
|
None, description="Detect the active speaker automatically. Video input only; rejected for images."
|
||||||
|
)
|
||||||
|
frame_number: int | None = Field(
|
||||||
|
None, description="Frame used for manual speaker selection. Must be 0 for image inputs."
|
||||||
|
)
|
||||||
|
coordinates: list[int] | None = Field(
|
||||||
|
None, description="Pixel [x, y] of the speaker's face in the frame selected by frame_number."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncGenerationOptions(BaseModel):
|
||||||
|
sync_mode: str | None = Field(
|
||||||
|
None,
|
||||||
|
description="How to resolve an audio/video duration mismatch: "
|
||||||
|
"cut_off, bounce, loop, silence or remap. Ignored for image inputs.",
|
||||||
|
)
|
||||||
|
i2v_prompt: str | None = Field(
|
||||||
|
None, description="Motion prompt for image-to-video generation. Image input only."
|
||||||
|
)
|
||||||
|
active_speaker_detection: SyncActiveSpeakerDetection | None = Field(None)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncGenerationRequest(BaseModel):
|
||||||
|
model: str = Field(..., description="Generation model, e.g. 'sync-3'.")
|
||||||
|
input: list[SyncInputItem] = Field(
|
||||||
|
..., description="Exactly one visual input (video or image) plus one audio input."
|
||||||
|
)
|
||||||
|
options: SyncGenerationOptions | None = Field(None)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncGeneration(BaseModel):
|
||||||
|
"""Subset of the Generation object returned by POST /v2/generate and GET /v2/generate/{id}."""
|
||||||
|
|
||||||
|
id: str = Field(...)
|
||||||
|
status: str = Field(..., description="PENDING | PROCESSING | COMPLETED | FAILED | REJECTED")
|
||||||
|
outputUrl: str | None = Field(None)
|
||||||
|
outputDuration: float | None = Field(None)
|
||||||
|
error: str | None = Field(None, description="Human-readable failure message.")
|
||||||
|
errorCode: str | None = Field(None, description="Stable machine-readable code from the GET /v2/errors catalog.")
|
||||||
|
|
@ -1133,7 +1133,9 @@ class GeminiImage2(IO.ComfyNode):
|
||||||
) -> IO.NodeOutput:
|
) -> IO.NodeOutput:
|
||||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||||
if model == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
if model == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||||
model = "gemini-3.1-flash-image-preview"
|
model = "gemini-3.1-flash-image"
|
||||||
|
elif model == "gemini-3-pro-image-preview":
|
||||||
|
model = "gemini-3-pro-image"
|
||||||
|
|
||||||
parts: list[GeminiPart] = [GeminiPart(text=prompt)]
|
parts: list[GeminiPart] = [GeminiPart(text=prompt)]
|
||||||
if images is not None:
|
if images is not None:
|
||||||
|
|
@ -1507,7 +1509,7 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
||||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||||
model_choice = model["model"]
|
model_choice = model["model"]
|
||||||
if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||||
model_id = "gemini-3.1-flash-image-preview"
|
model_id = "gemini-3.1-flash-image"
|
||||||
elif model_choice == "Nano Banana 2 Lite":
|
elif model_choice == "Nano Banana 2 Lite":
|
||||||
model_id = "gemini-3.1-flash-lite-image"
|
model_id = "gemini-3.1-flash-lite-image"
|
||||||
else:
|
else:
|
||||||
|
|
|
||||||
|
|
@ -41,6 +41,9 @@ STARTING_POINT_ID_PATTERN = r"<starting_point_id:(.*)>"
|
||||||
|
|
||||||
|
|
||||||
class SupportedOpenAIModel(str, Enum):
|
class SupportedOpenAIModel(str, Enum):
|
||||||
|
gpt_5_6_sol = "gpt-5.6-sol"
|
||||||
|
gpt_5_6_terra = "gpt-5.6-terra"
|
||||||
|
gpt_5_6_luna = "gpt-5.6-luna"
|
||||||
gpt_5_5_pro = "gpt-5.5-pro"
|
gpt_5_5_pro = "gpt-5.5-pro"
|
||||||
gpt_5_5 = "gpt-5.5"
|
gpt_5_5 = "gpt-5.5"
|
||||||
gpt_5 = "gpt-5"
|
gpt_5 = "gpt-5"
|
||||||
|
|
@ -1063,6 +1066,21 @@ class OpenAIChatNode(IO.ComfyNode):
|
||||||
"usd": [0.002, 0.008],
|
"usd": [0.002, 0.008],
|
||||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
}
|
}
|
||||||
|
: $contains($m, "gpt-5.6-terra") ? {
|
||||||
|
"type": "list_usd",
|
||||||
|
"usd": [0.0025, 0.015],
|
||||||
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
|
}
|
||||||
|
: $contains($m, "gpt-5.6-luna") ? {
|
||||||
|
"type": "list_usd",
|
||||||
|
"usd": [0.001, 0.006],
|
||||||
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
|
}
|
||||||
|
: $contains($m, "gpt-5.6") ? {
|
||||||
|
"type": "list_usd",
|
||||||
|
"usd": [0.005, 0.03],
|
||||||
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
|
}
|
||||||
: $contains($m, "gpt-5.5-pro") ? {
|
: $contains($m, "gpt-5.5-pro") ? {
|
||||||
"type": "list_usd",
|
"type": "list_usd",
|
||||||
"usd": [0.03, 0.18],
|
"usd": [0.03, 0.18],
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,391 @@
|
||||||
|
from typing_extensions import override
|
||||||
|
|
||||||
|
from comfy_api.latest import IO, ComfyExtension, Input
|
||||||
|
from comfy_api_nodes.apis.sync_so import (
|
||||||
|
SyncActiveSpeakerDetection,
|
||||||
|
SyncGeneration,
|
||||||
|
SyncGenerationOptions,
|
||||||
|
SyncGenerationRequest,
|
||||||
|
SyncInputItem,
|
||||||
|
)
|
||||||
|
from comfy_api_nodes.util import (
|
||||||
|
ApiEndpoint,
|
||||||
|
download_url_to_video_output,
|
||||||
|
downscale_image_tensor,
|
||||||
|
downscale_image_tensor_by_max_side,
|
||||||
|
get_image_dimensions,
|
||||||
|
get_number_of_images,
|
||||||
|
poll_op,
|
||||||
|
sync_op,
|
||||||
|
upload_audio_to_comfyapi,
|
||||||
|
upload_image_to_comfyapi,
|
||||||
|
upload_video_to_comfyapi,
|
||||||
|
validate_audio_duration,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncLipSyncNode(IO.ComfyNode):
|
||||||
|
@classmethod
|
||||||
|
def define_schema(cls) -> IO.Schema:
|
||||||
|
return IO.Schema(
|
||||||
|
node_id="SyncLipSyncNode",
|
||||||
|
display_name="sync.so Lip Sync",
|
||||||
|
category="partner/video/sync.so",
|
||||||
|
description=(
|
||||||
|
"Re-sync mouth movement in a video to new speech audio using sync.so. "
|
||||||
|
"Handles close-ups, profiles and obstructions automatically while preserving "
|
||||||
|
"the speaker's expression. Cost scales with output duration."
|
||||||
|
),
|
||||||
|
inputs=[
|
||||||
|
IO.Video.Input(
|
||||||
|
"video",
|
||||||
|
tooltip="Footage of the speaker to re-sync. Up to 4K (4096x2160); "
|
||||||
|
"a constant frame rate of 24/25/30 fps works best.",
|
||||||
|
),
|
||||||
|
IO.Audio.Input(
|
||||||
|
"audio",
|
||||||
|
tooltip="Speech audio to sync the mouth to.",
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"seed",
|
||||||
|
default=42,
|
||||||
|
min=0,
|
||||||
|
max=2147483647,
|
||||||
|
control_after_generate=True,
|
||||||
|
tooltip="Seed controls whether the node should re-run; "
|
||||||
|
"results are non-deterministic regardless of seed.",
|
||||||
|
),
|
||||||
|
IO.DynamicCombo.Input(
|
||||||
|
"model",
|
||||||
|
options=[
|
||||||
|
IO.DynamicCombo.Option(
|
||||||
|
"sync-3",
|
||||||
|
[
|
||||||
|
IO.Combo.Input(
|
||||||
|
"sync_mode",
|
||||||
|
options=["bounce", "cut_off", "loop", "silence", "remap"],
|
||||||
|
default="bounce",
|
||||||
|
tooltip=(
|
||||||
|
"How to handle a duration mismatch between video and audio; "
|
||||||
|
"this also sets the output length. "
|
||||||
|
"bounce: video plays forward then backward until the audio ends "
|
||||||
|
"(output = audio length). "
|
||||||
|
"loop: video restarts until the audio ends (output = audio length). "
|
||||||
|
"remap: video is time-stretched to match the audio (output = audio length). "
|
||||||
|
"cut_off: the longer track is trimmed (output = shorter length). "
|
||||||
|
"silence: nothing is trimmed; the shorter track is padded "
|
||||||
|
"(output = longer length)."
|
||||||
|
),
|
||||||
|
),
|
||||||
|
IO.Combo.Input(
|
||||||
|
"speaker_selection",
|
||||||
|
options=["default", "auto-detect", "coordinates"],
|
||||||
|
default="default",
|
||||||
|
tooltip=(
|
||||||
|
"Which face to lipsync when several people are visible. "
|
||||||
|
"default: let the model decide. "
|
||||||
|
"auto-detect: detect and follow the active speaker. "
|
||||||
|
"coordinates: target the face at pixel (speaker_x, speaker_y) "
|
||||||
|
"in the frame chosen by speaker_frame."
|
||||||
|
),
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"speaker_frame",
|
||||||
|
default=0,
|
||||||
|
min=0,
|
||||||
|
max=1_000_000,
|
||||||
|
advanced=True,
|
||||||
|
tooltip="Video frame used to locate the speaker. "
|
||||||
|
"Only used when speaker_selection is 'coordinates'.",
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"speaker_x",
|
||||||
|
default=0,
|
||||||
|
min=0,
|
||||||
|
max=4096,
|
||||||
|
advanced=True,
|
||||||
|
tooltip="X pixel coordinate of the speaker's face. "
|
||||||
|
"Only used when speaker_selection is 'coordinates'.",
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"speaker_y",
|
||||||
|
default=0,
|
||||||
|
min=0,
|
||||||
|
max=4096,
|
||||||
|
advanced=True,
|
||||||
|
tooltip="Y pixel coordinate of the speaker's face. "
|
||||||
|
"Only used when speaker_selection is 'coordinates'.",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
],
|
||||||
|
tooltip="sync.so generation model.",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
outputs=[IO.Video.Output()],
|
||||||
|
hidden=[
|
||||||
|
IO.Hidden.auth_token_comfy_org,
|
||||||
|
IO.Hidden.api_key_comfy_org,
|
||||||
|
IO.Hidden.unique_id,
|
||||||
|
],
|
||||||
|
is_api_node=True,
|
||||||
|
price_badge=IO.PriceBadge(
|
||||||
|
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def execute(
|
||||||
|
cls,
|
||||||
|
video: Input.Video,
|
||||||
|
audio: Input.Audio,
|
||||||
|
seed: int,
|
||||||
|
model: dict,
|
||||||
|
) -> IO.NodeOutput:
|
||||||
|
try:
|
||||||
|
width, height = video.get_dimensions()
|
||||||
|
except Exception:
|
||||||
|
width = height = None
|
||||||
|
if width and height and (max(width, height) > 4096 or width * height > 4096 * 2160):
|
||||||
|
raise ValueError(
|
||||||
|
f"sync.so rejects videos above 4K (4096x2160); got {width}x{height}. Downscale the video first."
|
||||||
|
)
|
||||||
|
validate_audio_duration(audio, max_duration=600)
|
||||||
|
|
||||||
|
if model["speaker_selection"] == "auto-detect":
|
||||||
|
speaker_detection = SyncActiveSpeakerDetection(auto_detect=True)
|
||||||
|
elif model["speaker_selection"] == "coordinates":
|
||||||
|
speaker_detection = SyncActiveSpeakerDetection(
|
||||||
|
frame_number=model["speaker_frame"],
|
||||||
|
coordinates=[model["speaker_x"], model["speaker_y"]],
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
speaker_detection = None
|
||||||
|
|
||||||
|
video_url = await upload_video_to_comfyapi(cls, video, max_duration=600)
|
||||||
|
audio_url = await upload_audio_to_comfyapi(cls, audio)
|
||||||
|
|
||||||
|
generation = await sync_op(
|
||||||
|
cls,
|
||||||
|
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
|
||||||
|
response_model=SyncGeneration,
|
||||||
|
data=SyncGenerationRequest(
|
||||||
|
model=model["model"],
|
||||||
|
input=[
|
||||||
|
SyncInputItem(type="video", url=video_url),
|
||||||
|
SyncInputItem(type="audio", url=audio_url),
|
||||||
|
],
|
||||||
|
options=SyncGenerationOptions(
|
||||||
|
sync_mode=model["sync_mode"],
|
||||||
|
active_speaker_detection=speaker_detection,
|
||||||
|
),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
generation = await poll_op(
|
||||||
|
cls,
|
||||||
|
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
|
||||||
|
response_model=SyncGeneration,
|
||||||
|
status_extractor=lambda g: g.status,
|
||||||
|
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
|
||||||
|
failed_statuses=[],
|
||||||
|
queued_statuses=["PENDING"],
|
||||||
|
poll_interval=10.0,
|
||||||
|
)
|
||||||
|
if generation.status != "COMPLETED":
|
||||||
|
code = f" [{generation.errorCode}]" if generation.errorCode else ""
|
||||||
|
raise ValueError(
|
||||||
|
f"sync.so generation {generation.status.lower()}{code}: "
|
||||||
|
f"{generation.error or 'no error details provided'}"
|
||||||
|
)
|
||||||
|
if not generation.outputUrl:
|
||||||
|
raise ValueError("sync.so generation completed but no output URL was returned.")
|
||||||
|
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
|
||||||
|
|
||||||
|
|
||||||
|
class SyncTalkingImageNode(IO.ComfyNode):
|
||||||
|
@classmethod
|
||||||
|
def define_schema(cls) -> IO.Schema:
|
||||||
|
return IO.Schema(
|
||||||
|
node_id="SyncTalkingImageNode",
|
||||||
|
display_name="sync.so Talking Image",
|
||||||
|
category="partner/video/sync.so",
|
||||||
|
description=(
|
||||||
|
"Animate a still portrait into a talking video driven by speech audio, "
|
||||||
|
"using sync.so's sync-3 model. The output duration matches the audio. "
|
||||||
|
"Cost scales with output duration."
|
||||||
|
),
|
||||||
|
inputs=[
|
||||||
|
IO.Image.Input(
|
||||||
|
"image",
|
||||||
|
tooltip="A single image with a clearly visible face, up to 4K (4096x2160).",
|
||||||
|
),
|
||||||
|
IO.Audio.Input(
|
||||||
|
"audio",
|
||||||
|
tooltip="Speech audio driving the talking video; the output duration matches it. "
|
||||||
|
"Chain any TTS node here to drive the animation from text.",
|
||||||
|
),
|
||||||
|
IO.String.Input(
|
||||||
|
"prompt",
|
||||||
|
multiline=True,
|
||||||
|
default="",
|
||||||
|
tooltip="Optional guidance for how the portrait comes to life, e.g. "
|
||||||
|
"'make the subject smile and look at the camera'. "
|
||||||
|
"Leave empty for natural talking motion.",
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"seed",
|
||||||
|
default=0,
|
||||||
|
min=0,
|
||||||
|
max=2147483647,
|
||||||
|
control_after_generate=True,
|
||||||
|
tooltip="Seed controls whether the node should re-run; "
|
||||||
|
"results are non-deterministic regardless of seed.",
|
||||||
|
),
|
||||||
|
IO.DynamicCombo.Input(
|
||||||
|
"model",
|
||||||
|
options=[
|
||||||
|
IO.DynamicCombo.Option(
|
||||||
|
"sync-3",
|
||||||
|
[
|
||||||
|
IO.Combo.Input(
|
||||||
|
"speaker_selection",
|
||||||
|
options=["default", "coordinates"],
|
||||||
|
default="default",
|
||||||
|
tooltip=(
|
||||||
|
"Which face to animate when several people are visible. "
|
||||||
|
"default: let the model decide. "
|
||||||
|
"coordinates: target the face at pixel (speaker_x, speaker_y) "
|
||||||
|
"in the image. Auto-detection is not supported for images."
|
||||||
|
),
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"speaker_x",
|
||||||
|
default=0,
|
||||||
|
min=0,
|
||||||
|
max=4096,
|
||||||
|
advanced=True,
|
||||||
|
tooltip="X pixel coordinate of the speaker's face. "
|
||||||
|
"Only used when speaker_selection is 'coordinates'.",
|
||||||
|
),
|
||||||
|
IO.Int.Input(
|
||||||
|
"speaker_y",
|
||||||
|
default=0,
|
||||||
|
min=0,
|
||||||
|
max=4096,
|
||||||
|
advanced=True,
|
||||||
|
tooltip="Y pixel coordinate of the speaker's face. "
|
||||||
|
"Only used when speaker_selection is 'coordinates'.",
|
||||||
|
),
|
||||||
|
IO.Boolean.Input(
|
||||||
|
"auto_downscale",
|
||||||
|
default=True,
|
||||||
|
advanced=True,
|
||||||
|
tooltip="Automatically downscale the image if it exceeds the 4K "
|
||||||
|
"(4096x2160) input limit; speaker coordinates are scaled to match. "
|
||||||
|
"When disabled, an oversized image raises an error instead.",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
],
|
||||||
|
tooltip="sync.so generation model. Image input is exclusive to sync-3.",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
outputs=[IO.Video.Output()],
|
||||||
|
hidden=[
|
||||||
|
IO.Hidden.auth_token_comfy_org,
|
||||||
|
IO.Hidden.api_key_comfy_org,
|
||||||
|
IO.Hidden.unique_id,
|
||||||
|
],
|
||||||
|
is_api_node=True,
|
||||||
|
price_badge=IO.PriceBadge(
|
||||||
|
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
async def execute(
|
||||||
|
cls,
|
||||||
|
image: Input.Image,
|
||||||
|
audio: Input.Audio,
|
||||||
|
prompt: str,
|
||||||
|
seed: int,
|
||||||
|
model: dict,
|
||||||
|
) -> IO.NodeOutput:
|
||||||
|
if get_number_of_images(image) != 1:
|
||||||
|
raise ValueError("Exactly one image is required; got a batch. Pick one frame first.")
|
||||||
|
validate_audio_duration(audio, max_duration=600)
|
||||||
|
|
||||||
|
height, width = get_image_dimensions(image)
|
||||||
|
speaker_x, speaker_y = model["speaker_x"], model["speaker_y"]
|
||||||
|
if max(width, height) > 4096 or width * height > 4096 * 2160:
|
||||||
|
if not model["auto_downscale"]:
|
||||||
|
raise ValueError(
|
||||||
|
f"sync.so rejects images above 4K (4096x2160); got {width}x{height}. "
|
||||||
|
"Downscale the image first or enable auto_downscale."
|
||||||
|
)
|
||||||
|
image = downscale_image_tensor(image, total_pixels=4096 * 2160)
|
||||||
|
image = downscale_image_tensor_by_max_side(image, max_side=4096)
|
||||||
|
new_height, new_width = get_image_dimensions(image)
|
||||||
|
# speaker coordinates are given in the original image's pixel space
|
||||||
|
speaker_x = min(new_width - 1, round(speaker_x * new_width / width))
|
||||||
|
speaker_y = min(new_height - 1, round(speaker_y * new_height / height))
|
||||||
|
|
||||||
|
if model["speaker_selection"] == "coordinates":
|
||||||
|
speaker_detection = SyncActiveSpeakerDetection(
|
||||||
|
frame_number=0, # images have a single frame; auto_detect is rejected by the API
|
||||||
|
coordinates=[speaker_x, speaker_y],
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
speaker_detection = None
|
||||||
|
|
||||||
|
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||||
|
audio_url = await upload_audio_to_comfyapi(cls, audio)
|
||||||
|
|
||||||
|
generation = await sync_op(
|
||||||
|
cls,
|
||||||
|
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
|
||||||
|
response_model=SyncGeneration,
|
||||||
|
data=SyncGenerationRequest(
|
||||||
|
model=model["model"],
|
||||||
|
input=[
|
||||||
|
SyncInputItem(type="image", url=image_url),
|
||||||
|
SyncInputItem(type="audio", url=audio_url),
|
||||||
|
],
|
||||||
|
options=SyncGenerationOptions(
|
||||||
|
i2v_prompt=prompt.strip() or None,
|
||||||
|
active_speaker_detection=speaker_detection,
|
||||||
|
),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
generation = await poll_op(
|
||||||
|
cls,
|
||||||
|
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
|
||||||
|
response_model=SyncGeneration,
|
||||||
|
status_extractor=lambda g: g.status,
|
||||||
|
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
|
||||||
|
failed_statuses=[],
|
||||||
|
queued_statuses=["PENDING"],
|
||||||
|
poll_interval=10.0,
|
||||||
|
)
|
||||||
|
if generation.status != "COMPLETED":
|
||||||
|
code = f" [{generation.errorCode}]" if generation.errorCode else ""
|
||||||
|
raise ValueError(
|
||||||
|
f"sync.so generation {generation.status.lower()}{code}: "
|
||||||
|
f"{generation.error or 'no error details provided'}"
|
||||||
|
)
|
||||||
|
if not generation.outputUrl:
|
||||||
|
raise ValueError("sync.so generation completed but no output URL was returned.")
|
||||||
|
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
|
||||||
|
|
||||||
|
|
||||||
|
class SyncExtension(ComfyExtension):
|
||||||
|
@override
|
||||||
|
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||||
|
return [
|
||||||
|
SyncLipSyncNode,
|
||||||
|
SyncTalkingImageNode,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
async def comfy_entrypoint() -> SyncExtension:
|
||||||
|
return SyncExtension()
|
||||||
|
|
@ -15,6 +15,8 @@ from comfy.comfy_api_env import normalize_comfy_api_base
|
||||||
from comfy.deploy_environment import get_deploy_environment
|
from comfy.deploy_environment import get_deploy_environment
|
||||||
from comfy.model_management import processing_interrupted
|
from comfy.model_management import processing_interrupted
|
||||||
from comfy_api.latest import IO
|
from comfy_api.latest import IO
|
||||||
|
from comfy_execution.utils import get_executing_context
|
||||||
|
from comfyui_version import __version__ as comfyui_version
|
||||||
|
|
||||||
from .common_exceptions import ProcessingInterrupted
|
from .common_exceptions import ProcessingInterrupted
|
||||||
|
|
||||||
|
|
@ -56,11 +58,16 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
|
||||||
relative/cloud URLs resolved against ``default_base_url()``; because the result
|
relative/cloud URLs resolved against ``default_base_url()``; because the result
|
||||||
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
||||||
"""
|
"""
|
||||||
return {
|
headers = {
|
||||||
**get_auth_header(node_cls),
|
**get_auth_header(node_cls),
|
||||||
"Comfy-Env": get_deploy_environment(),
|
"Comfy-Env": get_deploy_environment(),
|
||||||
"Comfy-Usage-Source": get_usage_source(node_cls),
|
"Comfy-Usage-Source": get_usage_source(node_cls),
|
||||||
|
"Comfy-Core-Version": comfyui_version,
|
||||||
}
|
}
|
||||||
|
ctx = get_executing_context()
|
||||||
|
if ctx is not None:
|
||||||
|
headers["Comfy-Job-Id"] = ctx.prompt_id
|
||||||
|
return headers
|
||||||
|
|
||||||
|
|
||||||
def default_base_url() -> str:
|
def default_base_url() -> str:
|
||||||
|
|
|
||||||
|
|
@ -844,15 +844,18 @@ class ImageMergeTileList(IO.ComfyNode):
|
||||||
# Format specifications
|
# Format specifications
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format,
|
# Maps (file_format, bit_depth, num_channels) -> (quantization scale, numpy dtype,
|
||||||
# stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
# av frame pix_fmt, stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
||||||
_FORMAT_SPECS = {
|
_FORMAT_SPECS = {
|
||||||
("png", "8-bit", False): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
|
("png", "8-bit", 1): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "gray", "stream_fmt": "gray"},
|
||||||
("png", "8-bit", True): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
|
("png", "8-bit", 3): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
|
||||||
("png", "16-bit", False): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
|
("png", "8-bit", 4): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
|
||||||
("png", "16-bit", True): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"},
|
("png", "16-bit", 1): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "gray16le", "stream_fmt": "gray16be"},
|
||||||
("exr", "32-bit float", False): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"},
|
("png", "16-bit", 3): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
|
||||||
("exr", "32-bit float", True): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
("png", "16-bit", 4): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"},
|
||||||
|
("exr", "32-bit float", 1): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "grayf32le", "stream_fmt": "grayf32le"},
|
||||||
|
("exr", "32-bit float", 3): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"},
|
||||||
|
("exr", "32-bit float", 4): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -891,10 +894,11 @@ def hlg_to_linear(t: torch.Tensor) -> torch.Tensor:
|
||||||
return torch.cat([hlg_to_linear(rgb), alpha], dim=-1)
|
return torch.cat([hlg_to_linear(rgb), alpha], dim=-1)
|
||||||
|
|
||||||
# Piecewise: sqrt branch below 0.5, log branch above.
|
# Piecewise: sqrt branch below 0.5, log branch above.
|
||||||
# Clamp inside the log branch so negative / out-of-range values don't blow up;
|
# Clamp the log branch at the 0.5 branch point (not above it) so the
|
||||||
|
# unselected lane stays finite in exp() without altering selected values;
|
||||||
# values above 1.0 are allowed and extrapolate naturally.
|
# values above 1.0 are allowed and extrapolate naturally.
|
||||||
low = (t ** 2) / 3.0
|
low = (t ** 2) / 3.0
|
||||||
high = (torch.exp((t.clamp(min=_HLG_C) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
high = (torch.exp((t.clamp(min=0.5) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||||
return torch.where(t <= 0.5, low, high)
|
return torch.where(t <= 0.5, low, high)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -1111,7 +1115,8 @@ def _encode_image(
|
||||||
bit_depth: str,
|
bit_depth: str,
|
||||||
colorspace: str,
|
colorspace: str,
|
||||||
) -> bytes:
|
) -> bytes:
|
||||||
"""Encode a single HxWxC tensor to PNG or EXR bytes in memory.
|
"""Encode a single HxWxC (or channel-less HxW grayscale) tensor to PNG or
|
||||||
|
EXR bytes in memory. Grayscale is written as single-channel PNG / Y-only EXR.
|
||||||
|
|
||||||
For EXR the input is interpreted according to `colorspace` and converted
|
For EXR the input is interpreted according to `colorspace` and converted
|
||||||
to scene-linear (EXR's convention) before writing:
|
to scene-linear (EXR's convention) before writing:
|
||||||
|
|
@ -1128,10 +1133,16 @@ def _encode_image(
|
||||||
For PNG, colorspace selection does not modify pixels — PNG is delivered
|
For PNG, colorspace selection does not modify pixels — PNG is delivered
|
||||||
sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
|
sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
|
||||||
"""
|
"""
|
||||||
|
if img_tensor.ndim == 2:
|
||||||
|
img_tensor = img_tensor.unsqueeze(-1) # Some nodes emit grayscale as (H, W) with no channel dim, mask-style.
|
||||||
height, width, num_channels = img_tensor.shape
|
height, width, num_channels = img_tensor.shape
|
||||||
has_alpha = num_channels == 4
|
|
||||||
|
|
||||||
spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)]
|
spec = _FORMAT_SPECS.get((file_format, bit_depth, num_channels))
|
||||||
|
if spec is None:
|
||||||
|
raise ValueError(
|
||||||
|
f"No {file_format}/{bit_depth} encoder for {num_channels}-channel images: "
|
||||||
|
"supported channel counts are 1 (grayscale), 3 (RGB) and 4 (RGBA)."
|
||||||
|
)
|
||||||
|
|
||||||
if spec["dtype"] == np.float32:
|
if spec["dtype"] == np.float32:
|
||||||
# EXR path: preserve full range, no clamp.
|
# EXR path: preserve full range, no clamp.
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,102 @@
|
||||||
|
from typing_extensions import override
|
||||||
|
|
||||||
|
import comfy.utils
|
||||||
|
import node_helpers
|
||||||
|
from comfy_api.latest import ComfyExtension, io
|
||||||
|
|
||||||
|
|
||||||
|
# fmt: off
|
||||||
|
BUCKETS_1024 = [
|
||||||
|
(512, 1792), (512, 1856), (512, 1920), (512, 1984), (512, 2048),
|
||||||
|
(576, 1600), (576, 1664), (576, 1728), (576, 1792),
|
||||||
|
(640, 1472), (640, 1536), (640, 1600),
|
||||||
|
(704, 1344), (704, 1408), (704, 1472),
|
||||||
|
(768, 1216), (768, 1280), (768, 1344),
|
||||||
|
(832, 1152), (832, 1216),
|
||||||
|
(896, 1088), (896, 1152),
|
||||||
|
(960, 1024), (960, 1088),
|
||||||
|
(1024, 960), (1024, 1024),
|
||||||
|
(1088, 896), (1088, 960),
|
||||||
|
(1152, 832), (1152, 896),
|
||||||
|
(1216, 768), (1216, 832),
|
||||||
|
(1280, 768),
|
||||||
|
(1344, 704), (1344, 768),
|
||||||
|
(1408, 704),
|
||||||
|
(1472, 640), (1472, 704),
|
||||||
|
(1536, 640),
|
||||||
|
(1600, 576), (1600, 640),
|
||||||
|
(1664, 576),
|
||||||
|
(1728, 576),
|
||||||
|
(1792, 512), (1792, 576),
|
||||||
|
(1856, 512),
|
||||||
|
(1920, 512),
|
||||||
|
(1984, 512),
|
||||||
|
(2048, 512),
|
||||||
|
]
|
||||||
|
# fmt: on
|
||||||
|
|
||||||
|
|
||||||
|
def _find_best_bucket(height: int, width: int) -> tuple[int, int]:
|
||||||
|
target_ratio = height / width
|
||||||
|
return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio))
|
||||||
|
|
||||||
|
|
||||||
|
def _resize_reference(image):
|
||||||
|
if image.shape[0] != 1:
|
||||||
|
raise ValueError("JoyImage reference inputs must contain one image each")
|
||||||
|
samples = image.movedim(-1, 1)
|
||||||
|
bucket_h, bucket_w = _find_best_bucket(samples.shape[2], samples.shape[3])
|
||||||
|
resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center")
|
||||||
|
return resized.movedim(1, -1)[:, :, :, :3]
|
||||||
|
|
||||||
|
|
||||||
|
def _encode(clip, prompt, vae, images):
|
||||||
|
resized_images = [_resize_reference(image) for image in images]
|
||||||
|
conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=resized_images))
|
||||||
|
if vae is not None and resized_images:
|
||||||
|
ref_latents = [vae.encode(image) for image in resized_images]
|
||||||
|
conditioning = node_helpers.conditioning_set_values(
|
||||||
|
conditioning, {"reference_latents": ref_latents}, append=True,
|
||||||
|
)
|
||||||
|
return conditioning
|
||||||
|
|
||||||
|
|
||||||
|
class TextEncodeJoyImageEdit(io.ComfyNode):
|
||||||
|
@classmethod
|
||||||
|
def define_schema(cls):
|
||||||
|
image_template = io.Autogrow.TemplatePrefix(
|
||||||
|
io.Image.Input("image"),
|
||||||
|
prefix="image",
|
||||||
|
min=0,
|
||||||
|
max=6,
|
||||||
|
)
|
||||||
|
return io.Schema(
|
||||||
|
node_id="TextEncodeJoyImageEdit",
|
||||||
|
category="model/conditioning/joyimage",
|
||||||
|
inputs=[
|
||||||
|
io.Clip.Input("clip"),
|
||||||
|
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
||||||
|
io.Vae.Input("vae", optional=True),
|
||||||
|
io.Autogrow.Input("images", template=image_template, optional=True),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
io.Conditioning.Output(),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput:
|
||||||
|
images = images or {}
|
||||||
|
return io.NodeOutput(_encode(clip, prompt, vae, list(images.values())))
|
||||||
|
|
||||||
|
|
||||||
|
class JoyImageExtension(ComfyExtension):
|
||||||
|
@override
|
||||||
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||||
|
return [
|
||||||
|
TextEncodeJoyImageEdit,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
async def comfy_entrypoint() -> JoyImageExtension:
|
||||||
|
return JoyImageExtension()
|
||||||
|
|
@ -174,8 +174,9 @@ class Preview3DAdvanced(IO.ComfyNode):
|
||||||
filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
|
filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
|
||||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||||
|
|
||||||
|
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||||
camera_info_input = kwargs.get("camera_info", None)
|
camera_info_input = kwargs.get("camera_info", None)
|
||||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||||
return IO.NodeOutput(
|
return IO.NodeOutput(
|
||||||
|
|
@ -243,8 +244,9 @@ class PreviewGaussianSplat(IO.ComfyNode):
|
||||||
filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
|
filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
|
||||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||||
|
|
||||||
|
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||||
camera_info_input = kwargs.get("camera_info", None)
|
camera_info_input = kwargs.get("camera_info", None)
|
||||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||||
return IO.NodeOutput(
|
return IO.NodeOutput(
|
||||||
|
|
@ -303,8 +305,9 @@ class PreviewPointCloud(IO.ComfyNode):
|
||||||
filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
|
filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
|
||||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||||
|
|
||||||
|
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||||
camera_info_input = kwargs.get("camera_info", None)
|
camera_info_input = kwargs.get("camera_info", None)
|
||||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||||
return IO.NodeOutput(
|
return IO.NodeOutput(
|
||||||
|
|
@ -375,8 +378,9 @@ class Load3DAdvanced(IO.ComfyNode):
|
||||||
file_3d = None
|
file_3d = None
|
||||||
if model_file and model_file != "none":
|
if model_file and model_file != "none":
|
||||||
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
|
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
|
||||||
|
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||||
model_3d_info = viewport_state.get('model_3d_info', [])
|
model_3d_info = viewport_state.get('model_3d_info', [])
|
||||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height)
|
return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height)
|
||||||
|
|
||||||
|
|
||||||
class Load3DExtension(ComfyExtension):
|
class Load3DExtension(ComfyExtension):
|
||||||
|
|
|
||||||
|
|
@ -418,8 +418,9 @@ def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str:
|
||||||
|
|
||||||
def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
|
def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
|
||||||
model_file = _save_file3d_to_output(model_3d, filename_prefix)
|
model_file = _save_file3d_to_output(model_3d, filename_prefix)
|
||||||
|
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||||
camera_info_input = kwargs.get("camera_info", None)
|
camera_info_input = kwargs.get("camera_info", None)
|
||||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||||
return IO.NodeOutput(
|
return IO.NodeOutput(
|
||||||
|
|
|
||||||
|
|
@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
|
||||||
display_name="Save Video",
|
display_name="Save Video",
|
||||||
category="video",
|
category="video",
|
||||||
essentials_category="Basics",
|
essentials_category="Basics",
|
||||||
description="Saves the input images to your ComfyUI output directory.",
|
description="Saves the input videos to your ComfyUI output directory.",
|
||||||
inputs=[
|
inputs=[
|
||||||
io.Video.Input("video", tooltip="The video to save."),
|
io.Video.Input("video", tooltip="The video to save."),
|
||||||
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),
|
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,3 @@
|
||||||
# This file is automatically generated by the build process when version is
|
# This file is automatically generated by the build process when version is
|
||||||
# updated in pyproject.toml.
|
# updated in pyproject.toml.
|
||||||
__version__ = "0.27.0"
|
__version__ = "0.28.0"
|
||||||
|
|
|
||||||
|
|
@ -426,12 +426,12 @@ def _is_intermediate_output(dynprompt, node_id):
|
||||||
|
|
||||||
|
|
||||||
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
|
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
|
||||||
|
if cached.ui is not None:
|
||||||
|
ui_outputs[node_id] = cached.ui
|
||||||
if server.client_id is None:
|
if server.client_id is None:
|
||||||
return
|
return
|
||||||
cached_ui = cached.ui or {}
|
cached_ui = cached.ui or {}
|
||||||
server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
|
server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
|
||||||
if cached.ui is not None:
|
|
||||||
ui_outputs[node_id] = cached.ui
|
|
||||||
|
|
||||||
async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
|
async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
|
||||||
unique_id = current_item
|
unique_id = current_item
|
||||||
|
|
|
||||||
5
nodes.py
5
nodes.py
|
|
@ -992,7 +992,7 @@ class CLIPLoader:
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def INPUT_TYPES(s):
|
||||||
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
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"], ),
|
"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", "joyimage"], ),
|
||||||
},
|
},
|
||||||
"optional": {
|
"optional": {
|
||||||
"device": (["default", "cpu"], {"advanced": True}),
|
"device": (["default", "cpu"], {"advanced": True}),
|
||||||
|
|
@ -1002,7 +1002,7 @@ class CLIPLoader:
|
||||||
|
|
||||||
CATEGORY = "model/loaders"
|
CATEGORY = "model/loaders"
|
||||||
|
|
||||||
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
||||||
|
|
||||||
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
|
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
|
||||||
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||||
|
|
@ -2462,6 +2462,7 @@ async def init_builtin_extra_nodes():
|
||||||
"nodes_seedvr.py",
|
"nodes_seedvr.py",
|
||||||
"nodes_context_windows.py",
|
"nodes_context_windows.py",
|
||||||
"nodes_qwen.py",
|
"nodes_qwen.py",
|
||||||
|
"nodes_joyimage.py",
|
||||||
"nodes_boogu.py",
|
"nodes_boogu.py",
|
||||||
"nodes_chroma_radiance.py",
|
"nodes_chroma_radiance.py",
|
||||||
"nodes_pid.py",
|
"nodes_pid.py",
|
||||||
|
|
|
||||||
45
openapi.yaml
45
openapi.yaml
|
|
@ -7,18 +7,18 @@ components:
|
||||||
description: Timestamp when the asset was created
|
description: Timestamp when the asset was created
|
||||||
format: date-time
|
format: date-time
|
||||||
type: string
|
type: string
|
||||||
|
display_name:
|
||||||
|
description: Display name of the asset. Mirrors name for backwards compatibility.
|
||||||
|
nullable: true
|
||||||
|
type: string
|
||||||
|
file_path:
|
||||||
|
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
|
||||||
|
nullable: true
|
||||||
|
type: string
|
||||||
hash:
|
hash:
|
||||||
description: Blake3 hash of the asset content.
|
description: Blake3 hash of the asset content.
|
||||||
pattern: ^blake3:[a-f0-9]{64}$
|
pattern: ^blake3:[a-f0-9]{64}$
|
||||||
type: string
|
type: string
|
||||||
loader_path:
|
|
||||||
description: The value a loader consumes to load this asset. Null when no loader can resolve the file.
|
|
||||||
nullable: true
|
|
||||||
type: string
|
|
||||||
display_name:
|
|
||||||
description: Human-facing label for the asset. Not unique.
|
|
||||||
nullable: true
|
|
||||||
type: string
|
|
||||||
id:
|
id:
|
||||||
description: Unique identifier for the asset
|
description: Unique identifier for the asset
|
||||||
format: uuid
|
format: uuid
|
||||||
|
|
@ -144,6 +144,14 @@ components:
|
||||||
AssetUpdated:
|
AssetUpdated:
|
||||||
description: Response returned when an existing asset is successfully updated.
|
description: Response returned when an existing asset is successfully updated.
|
||||||
properties:
|
properties:
|
||||||
|
display_name:
|
||||||
|
description: Display name of the asset. Mirrors name for backwards compatibility.
|
||||||
|
nullable: true
|
||||||
|
type: string
|
||||||
|
file_path:
|
||||||
|
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
|
||||||
|
nullable: true
|
||||||
|
type: string
|
||||||
hash:
|
hash:
|
||||||
description: Blake3 hash of the asset content.
|
description: Blake3 hash of the asset content.
|
||||||
pattern: ^blake3:[a-f0-9]{64}$
|
pattern: ^blake3:[a-f0-9]{64}$
|
||||||
|
|
@ -775,14 +783,6 @@ components:
|
||||||
ModelFolder:
|
ModelFolder:
|
||||||
description: Represents a folder containing models
|
description: Represents a folder containing models
|
||||||
properties:
|
properties:
|
||||||
extensions:
|
|
||||||
description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all).
|
|
||||||
example:
|
|
||||||
- .ckpt
|
|
||||||
- .safetensors
|
|
||||||
items:
|
|
||||||
type: string
|
|
||||||
type: array
|
|
||||||
folders:
|
folders:
|
||||||
description: List of paths where models of this type are stored
|
description: List of paths where models of this type are stored
|
||||||
example:
|
example:
|
||||||
|
|
@ -1644,7 +1644,7 @@ paths:
|
||||||
format: uuid
|
format: uuid
|
||||||
type: string
|
type: string
|
||||||
tags:
|
tags:
|
||||||
description: JSON-encoded array of tag strings. For new byte uploads, include exactly one destination role (`input`, `output`, or `models`); `models` uploads also require exactly one `model_type:<folder_name>` tag. Extra tags are stored as labels and do not create path components.
|
description: JSON-encoded array of freeform tag strings, e.g. '["models","checkpoint"]'. Common types include "models", "input", "output", and "temp", but any tag can be used in any order.
|
||||||
type: string
|
type: string
|
||||||
user_metadata:
|
user_metadata:
|
||||||
description: Custom JSON metadata as a string
|
description: Custom JSON metadata as a string
|
||||||
|
|
@ -1829,7 +1829,7 @@ paths:
|
||||||
content:
|
content:
|
||||||
application/json:
|
application/json:
|
||||||
schema:
|
schema:
|
||||||
$ref: '#/components/schemas/Asset'
|
$ref: '#/components/schemas/AssetUpdated'
|
||||||
description: Asset updated successfully
|
description: Asset updated successfully
|
||||||
"400":
|
"400":
|
||||||
content:
|
content:
|
||||||
|
|
@ -2470,9 +2470,6 @@ paths:
|
||||||
supports_preview_metadata:
|
supports_preview_metadata:
|
||||||
description: Whether the server supports preview metadata
|
description: Whether the server supports preview metadata
|
||||||
type: boolean
|
type: boolean
|
||||||
supports_model_type_tags:
|
|
||||||
description: Whether the server supports namespaced model type asset tags
|
|
||||||
type: boolean
|
|
||||||
type: object
|
type: object
|
||||||
description: Success
|
description: Success
|
||||||
headers:
|
headers:
|
||||||
|
|
@ -3300,6 +3297,12 @@ paths:
|
||||||
schema:
|
schema:
|
||||||
$ref: '#/components/schemas/ErrorResponse'
|
$ref: '#/components/schemas/ErrorResponse'
|
||||||
description: Invalid request parameters
|
description: Invalid request parameters
|
||||||
|
"401":
|
||||||
|
content:
|
||||||
|
application/json:
|
||||||
|
schema:
|
||||||
|
$ref: '#/components/schemas/ErrorResponse'
|
||||||
|
description: Unauthorized - Authentication required
|
||||||
"500":
|
"500":
|
||||||
content:
|
content:
|
||||||
application/json:
|
application/json:
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
[project]
|
[project]
|
||||||
name = "ComfyUI"
|
name = "ComfyUI"
|
||||||
version = "0.27.0"
|
version = "0.28.0"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
license = { file = "LICENSE" }
|
license = { file = "LICENSE" }
|
||||||
requires-python = ">=3.10"
|
requires-python = ">=3.10"
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
comfyui-frontend-package==1.45.20
|
comfyui-frontend-package==1.45.21
|
||||||
comfyui-workflow-templates==0.11.6
|
comfyui-workflow-templates==0.11.9
|
||||||
comfyui-embedded-docs==0.5.7
|
comfyui-embedded-docs==0.5.8
|
||||||
torch
|
torch
|
||||||
torchsde
|
torchsde
|
||||||
torchvision
|
torchvision
|
||||||
|
|
@ -22,7 +22,7 @@ alembic
|
||||||
SQLAlchemy>=2.0.0
|
SQLAlchemy>=2.0.0
|
||||||
filelock
|
filelock
|
||||||
av>=16.0.0
|
av>=16.0.0
|
||||||
comfy-kitchen==0.2.18
|
comfy-kitchen==0.2.21
|
||||||
comfy-aimdo==0.4.10
|
comfy-aimdo==0.4.10
|
||||||
requests
|
requests
|
||||||
simpleeval>=1.0.0
|
simpleeval>=1.0.0
|
||||||
|
|
|
||||||
|
|
@ -2,11 +2,12 @@ import pytest
|
||||||
import torch
|
import torch
|
||||||
import tempfile
|
import tempfile
|
||||||
import os
|
import os
|
||||||
|
import sys
|
||||||
import av
|
import av
|
||||||
import io
|
import io
|
||||||
from fractions import Fraction
|
from fractions import Fraction
|
||||||
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
|
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
|
||||||
from comfy_api.util.video_types import VideoComponents
|
from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
|
||||||
from comfy_api.input.basic_types import AudioInput
|
from comfy_api.input.basic_types import AudioInput
|
||||||
from av.error import InvalidDataError
|
from av.error import InvalidDataError
|
||||||
|
|
||||||
|
|
@ -237,3 +238,526 @@ def test_duration_consistency(video_components):
|
||||||
manual_duration = float(components.images.shape[0] / components.frame_rate)
|
manual_duration = float(components.images.shape[0] / components.frame_rate)
|
||||||
|
|
||||||
assert duration == pytest.approx(manual_duration)
|
assert duration == pytest.approx(manual_duration)
|
||||||
|
|
||||||
|
|
||||||
|
def create_transcode_source(
|
||||||
|
width=64, height=64, frames=30, fps=30, audio_streams=1, undecodable_audio=0, rotation=False,
|
||||||
|
container_format="mov", audio_codec="pcm_s16le",
|
||||||
|
):
|
||||||
|
"""Create a temp video that save_to must transcode (mpeg4 video, so codec != h264).
|
||||||
|
|
||||||
|
``undecodable_audio`` trailing PCM streams get their fourcc corrupted so no decoder exists
|
||||||
|
(``codec_context is None``), like the APAC track in iPhone spatial-audio recordings.
|
||||||
|
``rotation`` patches a 90-degree display matrix into the video track header.
|
||||||
|
"""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format=container_format) as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||||
|
video_stream.width = width
|
||||||
|
video_stream.height = height
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio = []
|
||||||
|
for _ in range(audio_streams + undecodable_audio):
|
||||||
|
stream = container.add_stream(audio_codec, rate=44100)
|
||||||
|
stream.sample_rate = 44100
|
||||||
|
audio.append(stream)
|
||||||
|
|
||||||
|
for i in range(frames):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((height, width, 3), (i * 7) % 256, dtype=torch.uint8).numpy(),
|
||||||
|
format="rgb24",
|
||||||
|
)
|
||||||
|
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||||
|
# write audio in 1024-sample frames, like real decoders produce, so the
|
||||||
|
# per-frame skip/cap logic in the transcode path actually runs
|
||||||
|
for stream in audio:
|
||||||
|
for offset in range(0, 44100 * frames // fps, 1024):
|
||||||
|
n = min(1024, 44100 * frames // fps - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, n, dtype=torch.int16).numpy(), format="s16", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = 44100
|
||||||
|
audio_frame.pts = offset
|
||||||
|
container.mux(stream.encode(audio_frame))
|
||||||
|
for stream in [video_stream, *audio]:
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
data = bytearray(buffer.getvalue())
|
||||||
|
end = len(data)
|
||||||
|
for _ in range(undecodable_audio):
|
||||||
|
end = data.rindex(b"sowt", 0, end)
|
||||||
|
data[end:end + 4] = b"Xpac"
|
||||||
|
if rotation:
|
||||||
|
# the 3x3 display matrix sits 40 bytes into the version-0 tkhd payload; first tkhd
|
||||||
|
# inside moov = video track (search from moov so mdat bytes can't false-match)
|
||||||
|
matrix_offset = data.index(b"tkhd", data.rindex(b"moov")) + 4 + 40
|
||||||
|
values = [0, 1 << 16, 0, -(1 << 16), 0, 0, 0, 0, 1 << 30]
|
||||||
|
data[matrix_offset:matrix_offset + 36] = b"".join(v.to_bytes(4, "big", signed=True) for v in values)
|
||||||
|
|
||||||
|
tmp = tempfile.NamedTemporaryFile(suffix=f".{container_format}", delete=False)
|
||||||
|
tmp.write(bytes(data))
|
||||||
|
tmp.close()
|
||||||
|
return tmp.name
|
||||||
|
|
||||||
|
|
||||||
|
def transcode_and_probe(video):
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
video.save_to(buffer, format=VideoContainer.MP4, codec=VideoCodec.H264)
|
||||||
|
buffer.seek(0)
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
video_stream = container.streams.video[0]
|
||||||
|
audio_stream = container.streams.audio[0] if container.streams.audio else None
|
||||||
|
frames = 0
|
||||||
|
first_pts = None
|
||||||
|
for packet in container.demux(video_stream):
|
||||||
|
for frame in packet.decode():
|
||||||
|
if first_pts is None:
|
||||||
|
first_pts = frame.pts
|
||||||
|
frames += 1
|
||||||
|
return {
|
||||||
|
"codec": video_stream.codec_context.name,
|
||||||
|
"width": video_stream.codec_context.width,
|
||||||
|
"height": video_stream.codec_context.height,
|
||||||
|
"frames": frames,
|
||||||
|
"first_pts": first_pts,
|
||||||
|
"video_seconds": float(video_stream.duration * video_stream.time_base) if video_stream.duration else None,
|
||||||
|
"audio_seconds": float(audio_stream.duration * audio_stream.time_base)
|
||||||
|
if audio_stream and audio_stream.duration else None,
|
||||||
|
"audio_codecs": [s.codec_context.name for s in container.streams.audio],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_streams_without_buffering_frames():
|
||||||
|
"""Transcoding must not decode the whole video into memory first (~2 GiB for this source)"""
|
||||||
|
resource = pytest.importorskip("resource") # no getrusage on Windows
|
||||||
|
rss_scale = 1 if sys.platform == "darwin" else 1024 # ru_maxrss: bytes on macOS, KiB elsewhere
|
||||||
|
# ru_maxrss is a lifetime peak: a heavier test running earlier would shrink the measured
|
||||||
|
# delta and quietly defang this canary, so keep this source the biggest thing in the suite
|
||||||
|
file_path = create_transcode_source(width=640, height=480, frames=300)
|
||||||
|
try:
|
||||||
|
rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path))
|
||||||
|
rss_delta = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale - rss_before
|
||||||
|
|
||||||
|
assert result["codec"] == "h264"
|
||||||
|
assert result["frames"] == 300
|
||||||
|
assert rss_delta < 500 * 2**20, f"transcode buffered frames in RAM (peak grew {rss_delta / 2**20:.0f} MiB)"
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_honors_trim_window():
|
||||||
|
"""start_time/duration trim applies to both video and audio on the streaming path"""
|
||||||
|
file_path = create_transcode_source(frames=90) # 3s @ 30fps
|
||||||
|
try:
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
|
||||||
|
assert result["frames"] == pytest.approx(30, abs=2)
|
||||||
|
assert result["first_pts"] == 0 # trimmed output is rebased to start at zero
|
||||||
|
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_keeps_audio_of_sparse_video():
|
||||||
|
"""Audio that runs ahead of a sparse video track (slideshows, timelapses) must be
|
||||||
|
kept in full — it is only clamped to the video's end, never to the video cursor."""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
|
||||||
|
for t in (0, 30, 60): # 3 frames spread over 60 seconds
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), t * 4, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = t * 15360
|
||||||
|
frame.time_base = Fraction(1, 15360)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
for offset in range(0, 48000 * 60, 1024):
|
||||||
|
n = min(1024, 48000 * 60 - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = 48000
|
||||||
|
audio_frame.pts = offset
|
||||||
|
audio_frame.time_base = Fraction(1, 48000)
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
container.mux(audio_stream.encode(None))
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["audio_seconds"] == pytest.approx(60.0, abs=1.0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_vfr_audio_covers_video_span():
|
||||||
|
"""A trim window in the sparse region of a VFR file keeps audio for the true pts span
|
||||||
|
of the kept frames. Deriving the span as frames/average_rate undercuts it badly: the
|
||||||
|
average is dominated by the dense region (and can be plain wrong on MediaRecorder files)."""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
|
||||||
|
# 10 frames inside the first second, then one every 1.25 s
|
||||||
|
for i, t in enumerate([x / 10 for x in range(10)] + [1.0, 2.25, 3.5, 4.75]):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 16) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = int(t * 15360)
|
||||||
|
frame.time_base = Fraction(1, 15360)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
for offset in range(0, 48000 * 6, 1024):
|
||||||
|
n = min(1024, 48000 * 6 - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = 48000
|
||||||
|
audio_frame.pts = offset
|
||||||
|
audio_frame.time_base = Fraction(1, 48000)
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
container.mux(audio_stream.encode(None))
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer, start_time=1, duration=5))
|
||||||
|
# kept frames: 1.0/2.25/3.5/4.75 s -> rebased span 3.75 s + one nominal interval
|
||||||
|
assert result["frames"] == 4
|
||||||
|
assert result["audio_seconds"] == pytest.approx(4.0, abs=0.45)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_trims_audio_in_stream_time_base_units():
|
||||||
|
"""Matroska audio timestamps tick in 1/1000, not 1/sample_rate; trim and audio timing
|
||||||
|
must convert through the frame's time base instead of assuming sample units. AAC audio,
|
||||||
|
because it decodes straight to the encoder's format and hits the resampler passthrough
|
||||||
|
that keeps the source time base on the frames."""
|
||||||
|
file_path = create_transcode_source(frames=90, container_format="matroska", audio_codec="aac")
|
||||||
|
try:
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
|
||||||
|
assert result["audio_codecs"] == ["aac"]
|
||||||
|
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_learns_unprobed_audio_params():
|
||||||
|
"""mpegts is only probed a few seconds deep at open, so an audio stream whose first
|
||||||
|
packet comes later (live captures where audio kicks in late) still has sample_rate 0
|
||||||
|
when the transcode starts; the parameters must be learned from the stream itself."""
|
||||||
|
sample_rate, fps, video_seconds, audio_start = 48000, 30, 13, 12
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mpegts") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||||
|
for i in range(video_seconds * fps):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
)
|
||||||
|
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||||
|
for offset in range(0, (video_seconds - audio_start) * sample_rate, 1024):
|
||||||
|
n = min(1024, (video_seconds - audio_start) * sample_rate - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = sample_rate
|
||||||
|
audio_frame.pts = audio_start * sample_rate + offset
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
for stream in (video_stream, audio_stream):
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
# the scenario requires unprobed parameters; if a future FFmpeg probes deeper,
|
||||||
|
# push audio_start/video_seconds further out to restore it
|
||||||
|
assert container.streams.audio[0].codec_context.sample_rate == 0
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["frames"] == video_seconds * fps
|
||||||
|
assert result["audio_codecs"] == ["aac"]
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
trimmed_before_audio = transcode_and_probe(VideoFromFile(buffer, duration=1))
|
||||||
|
assert trimmed_before_audio["frames"] == fps
|
||||||
|
assert trimmed_before_audio["audio_codecs"] == []
|
||||||
|
assert trimmed_before_audio["audio_seconds"] is None
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
trimmed_crossing_audio = transcode_and_probe(VideoFromFile(buffer, start_time=11.5, duration=1))
|
||||||
|
assert trimmed_crossing_audio["frames"] == fps
|
||||||
|
assert trimmed_crossing_audio["audio_codecs"] == ["aac"]
|
||||||
|
assert trimmed_crossing_audio["video_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||||
|
assert trimmed_crossing_audio["audio_seconds"] == pytest.approx(0.5, abs=0.1)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_trimmed_fragmented_mp4_keeps_audio():
|
||||||
|
"""Fragmented mp4 (MediaRecorder, DASH/HLS-derived files) delivers audio well behind
|
||||||
|
video, so when the trim window's last video frame arrives the audio demuxed so far
|
||||||
|
does not cover the window yet; the transcode must keep demuxing audio until it does
|
||||||
|
instead of finalizing on the first audio frame it sees afterwards."""
|
||||||
|
sample_rate, fps, seconds = 48000, 30, 6
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4", options={"movflags": "frag_keyframe+empty_moov"}) as container:
|
||||||
|
video_stream = container.add_stream("h264", rate=fps)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||||
|
next_audio_pts = 0
|
||||||
|
for i in range(seconds * fps):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
)
|
||||||
|
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||||
|
while next_audio_pts / sample_rate <= i / fps: # feed audio alongside, like a live pipeline
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, 1024, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = sample_rate
|
||||||
|
audio_frame.pts = next_audio_pts
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
next_audio_pts += 1024
|
||||||
|
for stream in (video_stream, audio_stream):
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer, start_time=0.5, duration=1.0))
|
||||||
|
assert result["video_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_sparse_video_keeps_true_duration():
|
||||||
|
"""average_rate is not a frame duration: a 3-frame video spanning 60 s averages
|
||||||
|
0.05 fps, and padding the last frame with 1/average_rate used to extend the
|
||||||
|
output — and the audio kept with it — about 20 s past the source span."""
|
||||||
|
sample_rate = 48000
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||||
|
for i, second in enumerate((0, 30, 60)):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), i * 80, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = second * 30
|
||||||
|
frame.time_base = Fraction(1, 30)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
for offset in range(0, 90 * sample_rate, 1024):
|
||||||
|
n = min(1024, 90 * sample_rate - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = sample_rate
|
||||||
|
audio_frame.pts = offset
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
for stream in (video_stream, audio_stream):
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["frames"] == 3
|
||||||
|
# the last frame keeps its true stts duration (1/30 s), not 1/average_rate (~20 s)
|
||||||
|
assert result["video_seconds"] == pytest.approx(60.03, abs=0.05)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(60.03, abs=0.1)
|
||||||
|
|
||||||
|
trimmed = transcode_and_probe(VideoFromFile(buffer, duration=45))
|
||||||
|
assert trimmed["frames"] == 2
|
||||||
|
# a kept frame whose source duration crosses the window end is clamped to it
|
||||||
|
assert trimmed["video_seconds"] == pytest.approx(45.0, abs=0.05)
|
||||||
|
assert trimmed["audio_seconds"] == pytest.approx(45.0, abs=0.1)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_clamps_final_pts_to_declared_stream_duration():
|
||||||
|
"""Some iPhone MOVs report a video stream duration that ends before the final
|
||||||
|
decoded frame's nominal duration. A transcode must not turn that trailing
|
||||||
|
timestamp quirk into an extra frame interval compared to the source/remux path."""
|
||||||
|
fps = 30
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
for i, pts in enumerate([*range(31), 32]):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = pts
|
||||||
|
frame.time_base = Fraction(1, fps)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
|
||||||
|
class _StreamProxy:
|
||||||
|
def __init__(self, stream, duration):
|
||||||
|
self._stream = stream
|
||||||
|
self.duration = duration
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._stream, name)
|
||||||
|
|
||||||
|
class _StreamsProxy:
|
||||||
|
def __init__(self, video_stream):
|
||||||
|
self.video = [video_stream]
|
||||||
|
self.audio = []
|
||||||
|
|
||||||
|
class _PacketProxy:
|
||||||
|
def __init__(self, packet, stream):
|
||||||
|
self._packet = packet
|
||||||
|
self.stream = stream
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._packet, name)
|
||||||
|
|
||||||
|
class _ContainerProxy:
|
||||||
|
def __init__(self, container, stream):
|
||||||
|
self._container = container
|
||||||
|
self._stream = stream
|
||||||
|
self.streams = _StreamsProxy(stream)
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._container, name)
|
||||||
|
|
||||||
|
def demux(self, *streams):
|
||||||
|
for packet in self._container.demux(self._stream._stream):
|
||||||
|
yield _PacketProxy(packet, self._stream)
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
output = io.BytesIO()
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
real_stream = container.streams.video[0]
|
||||||
|
declared_duration = 32 * int(round((1 / fps) / real_stream.time_base))
|
||||||
|
stream = _StreamProxy(real_stream, declared_duration)
|
||||||
|
VideoFromFile(buffer)._save_transcoded(
|
||||||
|
_ContainerProxy(container, stream), output, VideoContainer.MP4, VideoCodec.H264, None, 8
|
||||||
|
)
|
||||||
|
|
||||||
|
output.seek(0)
|
||||||
|
with av.open(output) as container:
|
||||||
|
video_stream = container.streams.video[0]
|
||||||
|
frames = [f for p in container.demux(video_stream) for f in p.decode()]
|
||||||
|
assert len(frames) == 32
|
||||||
|
assert float(video_stream.duration * video_stream.time_base) == pytest.approx(32 / fps, abs=0.01)
|
||||||
|
assert float(frames[-1].pts * frames[-1].time_base) == pytest.approx(31 / fps, abs=0.01)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_irregular_vfr_keeps_span():
|
||||||
|
"""B-frames reorder packets, and mp4 sample durations follow decode order: the dts
|
||||||
|
timeline ends before the pts timeline, so an irregular-VFR source's tail holds fell
|
||||||
|
out of the container (this 20.23 s span used to come out as 15.27 s, and the 10 s
|
||||||
|
trim as 6.03 s). The transcode encodes without B-frames so every sample keeps its
|
||||||
|
true display duration."""
|
||||||
|
durations = [1, 1, 60, 1, 1, 120, 1, 180, 1, 1, 150, 90] # 1/30 s ticks, span 20.2333 s
|
||||||
|
generator = torch.Generator().manual_seed(7)
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
pts = 0
|
||||||
|
for duration in durations:
|
||||||
|
# textured frames, so an encoder with default settings has B-frames to gain from
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.randint(0, 255, (64, 64, 3), generator=generator, dtype=torch.uint8).numpy(),
|
||||||
|
format="rgb24",
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = pts
|
||||||
|
frame.time_base = Fraction(1, 30)
|
||||||
|
pts += duration
|
||||||
|
for packet in video_stream.encode(frame):
|
||||||
|
packet.duration = duration # exact stts in the source
|
||||||
|
container.mux(packet)
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["frames"] == len(durations)
|
||||||
|
assert result["video_seconds"] == pytest.approx(sum(durations) / 30, abs=0.05)
|
||||||
|
|
||||||
|
trimmed = transcode_and_probe(VideoFromFile(buffer, duration=10))
|
||||||
|
assert trimmed["frames"] == 8 # frames at 12.167 s+ fall outside the window
|
||||||
|
assert trimmed["video_seconds"] == pytest.approx(10.0, abs=0.05)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_trim_survives_missing_leading_pts():
|
||||||
|
"""A trim should survive pts-less kept frames followed by a real-pts frame past the window."""
|
||||||
|
nulled_frames = 0
|
||||||
|
|
||||||
|
class _PacketProxy:
|
||||||
|
def __init__(self, packet):
|
||||||
|
self._packet = packet
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._packet, name)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def stream(self):
|
||||||
|
return self._packet.stream
|
||||||
|
|
||||||
|
def decode(self):
|
||||||
|
nonlocal nulled_frames
|
||||||
|
frames = self._packet.decode()
|
||||||
|
for frame in frames:
|
||||||
|
if nulled_frames < 2:
|
||||||
|
frame.pts = None
|
||||||
|
nulled_frames += 1
|
||||||
|
return frames
|
||||||
|
|
||||||
|
class _ContainerProxy:
|
||||||
|
def __init__(self, real):
|
||||||
|
self._real = real
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._real, name)
|
||||||
|
|
||||||
|
def demux(self, *streams):
|
||||||
|
for packet in self._real.demux(*streams):
|
||||||
|
yield _PacketProxy(packet)
|
||||||
|
|
||||||
|
file_path = create_transcode_source(frames=10, audio_streams=0)
|
||||||
|
try:
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(file_path) as container:
|
||||||
|
# 0.05 s window: both pts-less frames are kept (synthesized pts 0 and 512),
|
||||||
|
# and the first real-pts frame (1024 ticks) already lies past end_pts (768)
|
||||||
|
VideoFromFile(file_path, duration=0.05)._save_transcoded(
|
||||||
|
_ContainerProxy(container), buffer, VideoContainer.MP4, VideoCodec.H264, None, 8
|
||||||
|
)
|
||||||
|
assert nulled_frames == 2
|
||||||
|
buffer.seek(0)
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
video_stream = container.streams.video[0]
|
||||||
|
frames = [f for p in container.demux(video_stream) for f in p.decode()]
|
||||||
|
assert len(frames) == 2
|
||||||
|
assert float(video_stream.duration * video_stream.time_base) == pytest.approx(2 / 30, abs=0.01)
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_bakes_rotation():
|
||||||
|
"""A 90-degree display-matrix rotation swaps the output dimensions (portrait video)"""
|
||||||
|
file_path = create_transcode_source(width=64, height=32, rotation=True)
|
||||||
|
try:
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path))
|
||||||
|
assert (result["width"], result["height"]) == (32, 64)
|
||||||
|
assert result["frames"] == 30
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_skips_undecodable_audio():
|
||||||
|
"""Streaming transcode keeps the decodable audio track and drops undecodable ones;
|
||||||
|
with no decodable audio at all the output is video-only instead of crashing."""
|
||||||
|
mixed = all_bad = None
|
||||||
|
try:
|
||||||
|
mixed = create_transcode_source(audio_streams=1, undecodable_audio=1)
|
||||||
|
all_bad = create_transcode_source(audio_streams=0, undecodable_audio=2)
|
||||||
|
result = transcode_and_probe(VideoFromFile(mixed))
|
||||||
|
assert result["audio_codecs"] == ["aac"]
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
assert transcode_and_probe(VideoFromFile(all_bad))["audio_codecs"] == []
|
||||||
|
finally:
|
||||||
|
for path in (mixed, all_bad):
|
||||||
|
if path:
|
||||||
|
os.unlink(path)
|
||||||
|
|
|
||||||
|
|
@ -97,6 +97,32 @@ def _make_seedvr2_3b_shared_mm_sd():
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _make_pid_v1_5_sd(latent_proj_channels=16):
|
||||||
|
sd = {
|
||||||
|
"pixel_embedder.proj.weight": torch.empty(16, 3, device="meta"),
|
||||||
|
"lq_proj.latent_proj.0.weight": torch.empty(1024, latent_proj_channels, 3, 3, device="meta"),
|
||||||
|
"lq_proj.pit_head.weight": torch.empty(1536, 1024, device="meta"),
|
||||||
|
"lq_proj.gate_modules.0.content_proj.weight": torch.empty(1, 3072, device="meta"),
|
||||||
|
"pixel_blocks.0.attn.q_norm.weight": torch.empty(72, device="meta"),
|
||||||
|
"pixel_blocks.0.adaLN_modulation.0.weight": torch.empty(24576, 1536, device="meta"),
|
||||||
|
"pixel_blocks.0.adaLN_modulation.0.bias": torch.empty(24576, device="meta"),
|
||||||
|
}
|
||||||
|
for i in range(7):
|
||||||
|
sd[f"lq_proj.gate_modules.{i}.log_alpha"] = torch.empty((), device="meta")
|
||||||
|
return sd
|
||||||
|
|
||||||
|
|
||||||
|
def _make_joyimage_edit_plus_sd():
|
||||||
|
sd = {
|
||||||
|
"img_in.weight": torch.empty(4096, 16, 1, 2, 2, device="meta"),
|
||||||
|
"condition_embedder.time_embedder.linear_1.weight": torch.empty(1, device="meta"),
|
||||||
|
"double_blocks.0.attn.img_attn_q_norm.weight": torch.empty(128, device="meta"),
|
||||||
|
}
|
||||||
|
for i in range(40):
|
||||||
|
sd[f"double_blocks.{i}.attn.img_attn_qkv.weight"] = torch.empty(1, device="meta")
|
||||||
|
return sd
|
||||||
|
|
||||||
|
|
||||||
def _add_model_diffusion_prefix(sd):
|
def _add_model_diffusion_prefix(sd):
|
||||||
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
|
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
|
||||||
|
|
||||||
|
|
@ -206,6 +232,63 @@ class TestModelDetection:
|
||||||
|
|
||||||
assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2"
|
assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2"
|
||||||
|
|
||||||
|
def test_pid_v1_5_detection(self):
|
||||||
|
sd = _make_pid_v1_5_sd()
|
||||||
|
unet_config = detect_unet_config(sd, "")
|
||||||
|
|
||||||
|
assert unet_config == {
|
||||||
|
"image_model": "pid",
|
||||||
|
"lq_latent_channels": 16,
|
||||||
|
"lq_hidden_dim": 1024,
|
||||||
|
"latent_spatial_down_factor": 8,
|
||||||
|
"lq_interval": 2,
|
||||||
|
"lq_latent_unpatchify_factor": 1,
|
||||||
|
"lq_conv_padding_mode": "replicate",
|
||||||
|
"lq_gate_per_token": True,
|
||||||
|
"pit_lq_inject": True,
|
||||||
|
"rope_ref_h": 2048,
|
||||||
|
"rope_ref_w": 2048,
|
||||||
|
}
|
||||||
|
assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "PiD"
|
||||||
|
|
||||||
|
def test_pid_v1_5_flux2_detection(self):
|
||||||
|
unet_config = detect_unet_config(_make_pid_v1_5_sd(latent_proj_channels=32), "")
|
||||||
|
|
||||||
|
assert unet_config["lq_latent_channels"] == 128
|
||||||
|
assert unet_config["latent_spatial_down_factor"] == 16
|
||||||
|
assert unet_config["lq_latent_unpatchify_factor"] == 2
|
||||||
|
|
||||||
|
def test_pid_v1_5_pixel_adaln_conversion(self):
|
||||||
|
sd = _make_pid_v1_5_sd()
|
||||||
|
model_config = model_config_from_unet_config(detect_unet_config(sd, ""), sd)
|
||||||
|
processed = model_config.process_unet_state_dict(sd)
|
||||||
|
|
||||||
|
assert processed["pixel_blocks.0.attn.q_norm.weight"].shape == (72,)
|
||||||
|
assert processed["pixel_blocks.0.adaLN_modulation_msa.weight"].shape == (12288, 1536)
|
||||||
|
assert processed["pixel_blocks.0.adaLN_modulation_mlp.weight"].shape == (12288, 1536)
|
||||||
|
assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,)
|
||||||
|
assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,)
|
||||||
|
|
||||||
|
def test_joyimage_edit_plus_detection(self):
|
||||||
|
sd = _make_joyimage_edit_plus_sd()
|
||||||
|
unet_config = detect_unet_config(sd, "")
|
||||||
|
|
||||||
|
assert unet_config == {
|
||||||
|
"image_model": "joyimage",
|
||||||
|
"in_channels": 16,
|
||||||
|
"hidden_size": 4096,
|
||||||
|
"patch_size": [1, 2, 2],
|
||||||
|
"num_layers": 40,
|
||||||
|
"num_attention_heads": 32,
|
||||||
|
"text_dim": 4096,
|
||||||
|
}
|
||||||
|
assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "JoyImage"
|
||||||
|
|
||||||
|
def test_incomplete_joyimage_signature_is_not_detected(self):
|
||||||
|
sd = _make_joyimage_edit_plus_sd()
|
||||||
|
del sd["double_blocks.0.attn.img_attn_q_norm.weight"]
|
||||||
|
assert detect_unet_config(sd, "") is None
|
||||||
|
|
||||||
def test_unet_config_and_required_keys_combination_is_unique(self):
|
def test_unet_config_and_required_keys_combination_is_unique(self):
|
||||||
"""Each model in the registry must have a unique combination of
|
"""Each model in the registry must have a unique combination of
|
||||||
``unet_config`` and ``required_keys``. If two models share the same
|
``unet_config`` and ``required_keys``. If two models share the same
|
||||||
|
|
|
||||||
|
|
@ -818,6 +818,30 @@ class TestExecution:
|
||||||
except urllib.error.HTTPError:
|
except urllib.error.HTTPError:
|
||||||
pass # Expected behavior
|
pass # Expected behavior
|
||||||
|
|
||||||
|
def test_cached_outputs_in_job_without_client_id(self, client: ComfyClient, builder: GraphBuilder):
|
||||||
|
g = builder
|
||||||
|
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||||
|
output = g.node("SaveImage", images=image.out(0))
|
||||||
|
|
||||||
|
# Prime the cache with a normal run.
|
||||||
|
client.run(g)
|
||||||
|
|
||||||
|
# Resubmit anonymously (no client_id) so output nodes are cache hits with no websocket client.
|
||||||
|
data = json.dumps({"prompt": g.finalize()}).encode('utf-8')
|
||||||
|
req = urllib.request.Request(f"http://{client.server_address}/prompt", data=data)
|
||||||
|
prompt_id = json.loads(urllib.request.urlopen(req).read())['prompt_id']
|
||||||
|
|
||||||
|
for _ in range(100):
|
||||||
|
job = client.get_job(prompt_id)
|
||||||
|
if job is not None and job['status'] not in ('pending', 'in_progress'):
|
||||||
|
break
|
||||||
|
time.sleep(0.1)
|
||||||
|
else:
|
||||||
|
raise AssertionError("Prompt did not complete in time")
|
||||||
|
|
||||||
|
assert job['status'] == 'completed'
|
||||||
|
assert output.id in job['outputs'], "Cached outputs must appear in job outputs without a client_id"
|
||||||
|
|
||||||
def _create_history_item(self, client, builder):
|
def _create_history_item(self, client, builder):
|
||||||
g = GraphBuilder(prefix="offset_test")
|
g = GraphBuilder(prefix="offset_test")
|
||||||
input_node = g.node(
|
input_node = g.node(
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue