mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
Compare commits
13 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4800e78518 | ||
|
|
b08e6cf35f | ||
|
|
1d1099bea0 | ||
|
|
c67f95607f | ||
|
|
8edea4a65d | ||
|
|
0f42ba5146 | ||
|
|
71b73e3b2b | ||
|
|
6a8ff7a929 | ||
|
|
285a98944c | ||
|
|
03978e1e81 | ||
|
|
678d42c90e | ||
|
|
87d23b8176 | ||
|
|
cc6b352511 |
41
AGENTS.md
41
AGENTS.md
@ -19,6 +19,9 @@
|
||||
better to remove a broken feature path than keep a complicated partial fix.
|
||||
- Preserve existing APIs, node names, model-loading behavior, file layout, and
|
||||
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
|
||||
generic AI-generated code will be rejected automatically: unnecessary helper
|
||||
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
|
||||
pass-through or stored-but-unused values instead of preserving upstream or
|
||||
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
|
||||
through a private helper or a clearly named implementation-specific method
|
||||
instead of overloading the public method's return contract.
|
||||
@ -154,6 +164,10 @@
|
||||
`comfy-kitchen` helpers where they already solve the problem.
|
||||
- Use optimized comfy-kitchen ops in places where they improve performance
|
||||
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.
|
||||
Treat optimized backend functions, dispatch helpers, and capability-selected
|
||||
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
|
||||
first dimension. The second dimension may be half the original size for
|
||||
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
|
||||
ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
|
||||
`unsqueeze`, and `squeeze` instead.
|
||||
@ -192,11 +212,23 @@
|
||||
methods for scalar or structural calculations.
|
||||
- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
|
||||
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
|
||||
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
|
||||
forward; the model path should consume and return the native latent shape for
|
||||
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
|
||||
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
|
||||
@ -260,6 +292,15 @@
|
||||
- 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 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
|
||||
workflow convenience unless the node is explicitly an output node. Existing
|
||||
models, latents, conditioning, or other inputs should flow directly to the
|
||||
|
||||
278
comfy/ldm/anima/lllite.py
Normal file
278
comfy/ldm/anima/lllite.py
Normal file
@ -0,0 +1,278 @@
|
||||
import re
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
|
||||
|
||||
MODULE_PATTERN = re.compile(r"lllite_dit_blocks_(\d+)_(self_attn_[qkv]_proj|cross_attn_q_proj|mlp_layer1)$")
|
||||
|
||||
|
||||
def _group_norm(channels, device=None, dtype=None, operations=None):
|
||||
groups = 8
|
||||
while groups > 1 and channels % groups != 0:
|
||||
groups //= 2
|
||||
return operations.GroupNorm(groups, channels, device=device, dtype=dtype)
|
||||
|
||||
|
||||
class AnimaLLLiteResBlock(nn.Module):
|
||||
def __init__(self, channels, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.norm1 = _group_norm(channels, device=device, dtype=dtype, operations=operations)
|
||||
self.conv1 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype)
|
||||
self.norm2 = _group_norm(channels, device=device, dtype=dtype, operations=operations)
|
||||
self.conv2 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.conv1(F.silu(self.norm1(x)))
|
||||
h = self.conv2(F.silu(self.norm2(h)))
|
||||
return x + h
|
||||
|
||||
|
||||
class AnimaLLLiteASPP(nn.Module):
|
||||
def __init__(self, channels, dilations, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
branches = []
|
||||
for dilation in dilations:
|
||||
if dilation == 1:
|
||||
conv = operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype)
|
||||
else:
|
||||
conv = operations.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation, device=device, dtype=dtype)
|
||||
branches.append(nn.Sequential(conv, _group_norm(channels, device=device, dtype=dtype, operations=operations), nn.SiLU()))
|
||||
self.branches = nn.ModuleList(branches)
|
||||
self.global_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.global_conv = nn.Sequential(
|
||||
operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype),
|
||||
_group_norm(channels, device=device, dtype=dtype, operations=operations),
|
||||
nn.SiLU(),
|
||||
)
|
||||
self.proj = nn.Sequential(
|
||||
operations.Conv2d(channels * (len(dilations) + 1), channels, kernel_size=1, device=device, dtype=dtype),
|
||||
_group_norm(channels, device=device, dtype=dtype, operations=operations),
|
||||
nn.SiLU(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
height, width = x.shape[-2:]
|
||||
outputs = [branch(x) for branch in self.branches]
|
||||
pooled = self.global_conv(self.global_pool(x))
|
||||
outputs.append(F.interpolate(pooled, size=(height, width), mode="bilinear", align_corners=False))
|
||||
return self.proj(torch.cat(outputs, dim=1))
|
||||
|
||||
|
||||
class AnimaLLLiteConditioning(nn.Module):
|
||||
def __init__(self, cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
half_dim = cond_dim // 2
|
||||
self.conv1 = operations.Conv2d(cond_in_channels, half_dim, kernel_size=4, stride=4, device=device, dtype=dtype)
|
||||
self.norm1 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations)
|
||||
self.conv2 = operations.Conv2d(half_dim, half_dim, kernel_size=3, padding=1, device=device, dtype=dtype)
|
||||
self.norm2 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations)
|
||||
self.conv3 = operations.Conv2d(half_dim, cond_dim, kernel_size=4, stride=4, device=device, dtype=dtype)
|
||||
self.norm3 = _group_norm(cond_dim, device=device, dtype=dtype, operations=operations)
|
||||
self.resblocks = nn.ModuleList([
|
||||
AnimaLLLiteResBlock(cond_dim, device=device, dtype=dtype, operations=operations)
|
||||
for _ in range(cond_resblocks)
|
||||
])
|
||||
self.aspp = AnimaLLLiteASPP(cond_dim, aspp_dilations, device=device, dtype=dtype, operations=operations) if aspp_dilations else None
|
||||
self.proj = operations.Conv2d(cond_dim, cond_emb_dim, kernel_size=1, device=device, dtype=dtype)
|
||||
self.out_norm = operations.LayerNorm(cond_emb_dim, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.silu(self.norm1(self.conv1(x)))
|
||||
x = F.silu(self.norm2(self.conv2(x)))
|
||||
x = F.silu(self.norm3(self.conv3(x)))
|
||||
for block in self.resblocks:
|
||||
x = block(x)
|
||||
if self.aspp is not None:
|
||||
x = self.aspp(x)
|
||||
x = self.proj(x).flatten(2).transpose(1, 2).contiguous()
|
||||
return self.out_norm(x)
|
||||
|
||||
|
||||
class AnimaLLLiteModule(nn.Module):
|
||||
def __init__(self, in_dim, cond_emb_dim, mlp_dim, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.down = operations.Linear(in_dim, mlp_dim, device=device, dtype=dtype)
|
||||
self.mid = operations.Linear(mlp_dim + cond_emb_dim, mlp_dim, device=device, dtype=dtype)
|
||||
self.cond_to_film = operations.Linear(cond_emb_dim, 2 * mlp_dim, device=device, dtype=dtype)
|
||||
self.up = operations.Linear(mlp_dim, in_dim, device=device, dtype=dtype)
|
||||
self.depth_embed = nn.Parameter(torch.empty(cond_emb_dim, device=device, dtype=dtype), requires_grad=False)
|
||||
|
||||
def forward(self, x, cond_emb, strength):
|
||||
original_shape = x.shape
|
||||
if x.ndim == 5:
|
||||
x = x.flatten(1, 3)
|
||||
|
||||
if x.shape[0] != cond_emb.shape[0]:
|
||||
if x.shape[0] % cond_emb.shape[0] != 0:
|
||||
raise ValueError(f"Anima LLLite batch mismatch: model input batch {x.shape[0]}, control batch {cond_emb.shape[0]}")
|
||||
cond_emb = cond_emb.repeat(x.shape[0] // cond_emb.shape[0], 1, 1)
|
||||
if x.shape[1] != cond_emb.shape[1]:
|
||||
raise ValueError(f"Anima LLLite sequence mismatch: model input has {x.shape[1]} tokens, control has {cond_emb.shape[1]}")
|
||||
|
||||
cond_local = cond_emb + comfy.ops.cast_to_input(self.depth_embed, cond_emb)
|
||||
hidden = F.silu(self.down(x))
|
||||
gamma, beta = self.cond_to_film(cond_local).chunk(2, dim=-1)
|
||||
hidden = self.mid(torch.cat((cond_local, hidden), dim=-1))
|
||||
hidden = F.silu(hidden * (1 + gamma) + beta)
|
||||
x = x + self.up(hidden) * strength
|
||||
|
||||
if len(original_shape) == 5:
|
||||
x = x.reshape(original_shape)
|
||||
return x
|
||||
|
||||
|
||||
class AnimaLLLite(nn.Module):
|
||||
def __init__(self, state_dict, metadata, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
metadata = metadata or {}
|
||||
version = metadata.get("lllite.version", "2")
|
||||
if version != "2":
|
||||
raise ValueError(f"Unsupported Anima LLLite version {version!r}; only named-key v2 checkpoints are supported")
|
||||
|
||||
module_names = sorted({key.split(".", 1)[0] for key in state_dict if key.startswith("lllite_dit_blocks_")})
|
||||
if not module_names:
|
||||
raise ValueError("Anima LLLite checkpoint has no lllite_dit_blocks_* modules")
|
||||
|
||||
cond_in_channels = state_dict["lllite_conditioning1.conv1.weight"].shape[1]
|
||||
cond_dim = state_dict["lllite_conditioning1.conv3.weight"].shape[0]
|
||||
cond_emb_dim = state_dict["lllite_conditioning1.proj.weight"].shape[0]
|
||||
resblock_ids = {int(key.split(".")[2]) for key in state_dict if key.startswith("lllite_conditioning1.resblocks.")}
|
||||
cond_resblocks = max(resblock_ids) + 1 if resblock_ids else 0
|
||||
use_aspp = any(key.startswith("lllite_conditioning1.aspp.") for key in state_dict)
|
||||
dilation_string = metadata.get("lllite.aspp_dilations", "1,2,4,8")
|
||||
aspp_dilations = tuple(int(value) for value in dilation_string.split(",") if value.strip()) if use_aspp else ()
|
||||
|
||||
self.cond_in_channels = cond_in_channels
|
||||
self.inpaint_masked_input = metadata.get("lllite.inpaint_masked_input", "false").lower() == "true"
|
||||
self.lllite_conditioning1 = AnimaLLLiteConditioning(
|
||||
cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations,
|
||||
device=device, dtype=dtype, operations=operations,
|
||||
)
|
||||
|
||||
self.module_names = set()
|
||||
self.block_count = 0
|
||||
self.model_dim = None
|
||||
for name in module_names:
|
||||
match = MODULE_PATTERN.fullmatch(name)
|
||||
if match is None:
|
||||
raise ValueError(f"Unsupported Anima LLLite module name: {name}")
|
||||
down_shape = state_dict[f"{name}.down.weight"].shape
|
||||
mlp_dim, in_dim = down_shape
|
||||
module_cond_dim = state_dict[f"{name}.cond_to_film.weight"].shape[1]
|
||||
if module_cond_dim != cond_emb_dim:
|
||||
raise ValueError(f"Anima LLLite conditioning dimension mismatch in {name}: {module_cond_dim} != {cond_emb_dim}")
|
||||
if self.model_dim is None:
|
||||
self.model_dim = in_dim
|
||||
elif self.model_dim != in_dim:
|
||||
raise ValueError(f"Anima LLLite model dimension mismatch in {name}: {in_dim} != {self.model_dim}")
|
||||
self.add_module(name, AnimaLLLiteModule(in_dim, cond_emb_dim, mlp_dim, device=device, dtype=dtype, operations=operations))
|
||||
self.module_names.add(name)
|
||||
self.block_count = max(self.block_count, int(match.group(1)) + 1)
|
||||
|
||||
def encode_conditioning(self, image):
|
||||
return self.lllite_conditioning1(image)
|
||||
|
||||
def apply(self, x, cond_emb, block_index, target, strength):
|
||||
name = f"lllite_dit_blocks_{block_index}_{target}"
|
||||
if name not in self.module_names:
|
||||
return x
|
||||
return self.get_submodule(name)(x, cond_emb, strength)
|
||||
|
||||
|
||||
class AnimaLLLitePatch:
|
||||
def __init__(self, model_patch, image, mask, strength, sigma_start, sigma_end):
|
||||
self.model_patch = model_patch
|
||||
self.image = image
|
||||
self.mask = mask
|
||||
self.strength = strength
|
||||
self.sigma_start = sigma_start
|
||||
self.sigma_end = sigma_end
|
||||
|
||||
def __call__(self, args):
|
||||
x = args["x"]
|
||||
transformer_options = args["transformer_options"]
|
||||
if self.strength == 0.0:
|
||||
return args
|
||||
sigmas = transformer_options.get("sigmas")
|
||||
if sigmas is not None:
|
||||
sigma = float(sigmas.max().item())
|
||||
if not self.sigma_end <= sigma <= self.sigma_start:
|
||||
return args
|
||||
if x.shape[2] != 1:
|
||||
raise ValueError(f"Anima LLLite only supports T=1, got T={x.shape[2]}")
|
||||
|
||||
target_height = x.shape[-2] * 8
|
||||
target_width = x.shape[-1] * 8
|
||||
image = comfy.utils.common_upscale(
|
||||
self.image.movedim(-1, 1), target_width, target_height, "bicubic", crop="center"
|
||||
).clamp(0.0, 1.0)
|
||||
image = image.to(device=x.device, dtype=x.dtype) * 2.0 - 1.0
|
||||
|
||||
if self.model_patch.model.cond_in_channels == 4:
|
||||
mask = self.mask
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
if mask.ndim != 4 or mask.shape[1] != 1:
|
||||
raise ValueError(f"Anima LLLite mask must have one channel, got shape {tuple(mask.shape)}")
|
||||
mask = comfy.utils.common_upscale(
|
||||
mask.float(), target_width, target_height, "nearest-exact", crop="center"
|
||||
)
|
||||
if mask.shape[0] != image.shape[0]:
|
||||
if image.shape[0] % mask.shape[0] != 0:
|
||||
raise ValueError(
|
||||
f"Anima LLLite mask batch {mask.shape[0]} cannot be broadcast to image batch {image.shape[0]}"
|
||||
)
|
||||
mask = mask.repeat(image.shape[0] // mask.shape[0], 1, 1, 1)
|
||||
mask = (mask >= 0.5).to(device=x.device, dtype=x.dtype)
|
||||
if self.model_patch.model.inpaint_masked_input:
|
||||
image = image * (mask < 0.5).to(image.dtype)
|
||||
image = torch.cat((image, mask * 2.0 - 1.0), dim=1)
|
||||
|
||||
cond_emb = self.model_patch.model.encode_conditioning(image)
|
||||
transformer_options["model_patch_data"][self] = cond_emb
|
||||
return args
|
||||
|
||||
def to(self, device_or_dtype):
|
||||
return self
|
||||
|
||||
def models(self):
|
||||
return [self.model_patch]
|
||||
|
||||
|
||||
class AnimaLLLiteAttentionPatch:
|
||||
def __init__(self, patch, targets):
|
||||
self.patch = patch
|
||||
self.targets = targets
|
||||
|
||||
def __call__(self, q, k, v, pe=None, attn_mask=None, extra_options=None):
|
||||
cond_emb = extra_options["model_patch_data"].get(self.patch)
|
||||
if cond_emb is None:
|
||||
return {"q": q, "k": k, "v": v, "pe": pe, "attn_mask": attn_mask}
|
||||
|
||||
block_index = extra_options["block_index"]
|
||||
values = {"q": q, "k": k, "v": v}
|
||||
for value_name, target in self.targets.items():
|
||||
values[value_name] = self.patch.model_patch.model.apply(
|
||||
values[value_name], cond_emb, block_index, target, self.patch.strength
|
||||
)
|
||||
|
||||
return {"q": values["q"], "k": values["k"], "v": values["v"], "pe": pe, "attn_mask": attn_mask}
|
||||
|
||||
|
||||
class AnimaLLLiteMLPPatch:
|
||||
def __init__(self, patch):
|
||||
self.patch = patch
|
||||
|
||||
def __call__(self, args):
|
||||
cond_emb = args["transformer_options"]["model_patch_data"].get(self.patch)
|
||||
if cond_emb is None:
|
||||
return args
|
||||
args["x"] = self.patch.model_patch.model.apply(
|
||||
args["x"], cond_emb, args["transformer_options"]["block_index"], "mlp_layer1", self.patch.strength
|
||||
)
|
||||
return args
|
||||
@ -14,6 +14,7 @@ from torchvision import transforms
|
||||
import comfy.patcher_extension
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.ops
|
||||
import comfy.quant_ops
|
||||
|
||||
|
||||
@ -148,11 +149,29 @@ class Attention(nn.Module):
|
||||
x: torch.Tensor,
|
||||
context: Optional[torch.Tensor] = None,
|
||||
rope_emb: Optional[torch.Tensor] = None,
|
||||
transformer_options: Optional[dict] = {},
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
q = self.q_proj(x)
|
||||
context = x if context is None else context
|
||||
k = self.k_proj(context)
|
||||
v = self.v_proj(context)
|
||||
q_input = x
|
||||
k_input = context
|
||||
v_input = context
|
||||
|
||||
transformer_patches = transformer_options.get("patches", {})
|
||||
patch_name = "attn1_patch" if self.is_selfattn else "attn2_patch"
|
||||
if patch_name in transformer_patches:
|
||||
extra_options = transformer_options.copy()
|
||||
extra_options["n_heads"] = self.n_heads
|
||||
extra_options["dim_head"] = self.head_dim
|
||||
for patch in transformer_patches[patch_name]:
|
||||
out = patch(q_input, k_input, v_input, pe=rope_emb, attn_mask=None, extra_options=extra_options)
|
||||
q_input = out.get("q", q_input)
|
||||
k_input = out.get("k", k_input)
|
||||
v_input = out.get("v", v_input)
|
||||
rope_emb = out.get("pe", rope_emb)
|
||||
|
||||
q = self.q_proj(q_input)
|
||||
k = self.k_proj(k_input)
|
||||
v = self.v_proj(v_input)
|
||||
q, k, v = map(
|
||||
lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
|
||||
(q, k, v),
|
||||
@ -161,11 +180,16 @@ class Attention(nn.Module):
|
||||
def apply_norm_and_rotary_pos_emb(
|
||||
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
v = self.v_norm(v)
|
||||
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
|
||||
|
||||
q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
|
||||
@ -188,7 +212,7 @@ class Attention(nn.Module):
|
||||
x (Tensor): The query tensor of shape [B, Mq, K]
|
||||
context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
|
||||
"""
|
||||
q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb)
|
||||
q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb, transformer_options=transformer_options)
|
||||
return self.compute_attention(q, k, v, transformer_options=transformer_options)
|
||||
|
||||
|
||||
@ -555,8 +579,14 @@ class Block(nn.Module):
|
||||
self.layer_norm_mlp,
|
||||
scale_mlp_B_T_1_1_D,
|
||||
shift_mlp_B_T_1_1_D,
|
||||
)
|
||||
result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.to(compute_dtype))
|
||||
).to(compute_dtype)
|
||||
patches = transformer_options.get("patches", {})
|
||||
if "mlp_patch" in patches:
|
||||
args = {"x": normalized_x_B_T_H_W_D, "transformer_options": transformer_options}
|
||||
for patch in patches["mlp_patch"]:
|
||||
args = patch(args)
|
||||
normalized_x_B_T_H_W_D = args["x"]
|
||||
result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
|
||||
x_B_T_H_W_D = torch.addcmul(x_B_T_H_W_D, gate_mlp_B_T_1_1_D.to(residual_dtype), result_B_T_H_W_D.to(residual_dtype))
|
||||
return x_B_T_H_W_D
|
||||
|
||||
@ -863,11 +893,22 @@ class MiniTrainDIT(nn.Module):
|
||||
x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape
|
||||
), f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
|
||||
|
||||
transformer_options = kwargs.get("transformer_options", {})
|
||||
patches = transformer_options.get("patches", {})
|
||||
if "post_input" in patches:
|
||||
transformer_options = transformer_options.copy()
|
||||
transformer_options["model_patch_data"] = {}
|
||||
|
||||
if "post_input" in patches:
|
||||
for patch in patches["post_input"]:
|
||||
out = patch({"img": x_B_T_H_W_D, "x": x_B_C_T_H_W, "transformer_options": transformer_options})
|
||||
x_B_T_H_W_D = out["img"]
|
||||
|
||||
block_kwargs = {
|
||||
"rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0),
|
||||
"adaln_lora_B_T_3D": adaln_lora_B_T_3D,
|
||||
"extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
|
||||
"transformer_options": kwargs.get("transformer_options", {}),
|
||||
"transformer_options": transformer_options,
|
||||
}
|
||||
|
||||
# The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream
|
||||
@ -877,7 +918,8 @@ class MiniTrainDIT(nn.Module):
|
||||
if x_B_T_H_W_D.dtype == torch.float16:
|
||||
x_B_T_H_W_D = x_B_T_H_W_D.float()
|
||||
|
||||
for block in self.blocks:
|
||||
for block_index, block in enumerate(self.blocks):
|
||||
transformer_options["block_index"] = block_index
|
||||
x_B_T_H_W_D = block(
|
||||
x_B_T_H_W_D,
|
||||
t_embedding_B_T_D,
|
||||
|
||||
445
comfy/ldm/joyimage/model.py
Normal file
445
comfy/ldm/joyimage/model.py
Normal file
@ -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]
|
||||
@ -58,6 +58,7 @@ import comfy.ldm.omnigen.omnigen2
|
||||
import comfy.ldm.seedvr.model
|
||||
import comfy.ldm.boogu.model
|
||||
import comfy.ldm.qwen_image.model
|
||||
import comfy.ldm.joyimage.model
|
||||
import comfy.ldm.ideogram4.model
|
||||
import comfy.ldm.krea2.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])
|
||||
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):
|
||||
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)
|
||||
|
||||
@ -1058,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
dit_config["image_model"] = "SAM31"
|
||||
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:
|
||||
return None
|
||||
|
||||
|
||||
@ -76,6 +76,7 @@ import comfy.text_encoders.gemma4
|
||||
import comfy.text_encoders.cogvideo
|
||||
import comfy.text_encoders.sa3
|
||||
import comfy.text_encoders.gpt_oss
|
||||
import comfy.text_encoders.joyimage
|
||||
|
||||
import comfy.model_patcher
|
||||
import comfy.lora
|
||||
@ -1377,6 +1378,7 @@ class CLIPType(Enum):
|
||||
IDEOGRAM4 = 30
|
||||
BOOGU = 31
|
||||
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_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
|
||||
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.
|
||||
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)
|
||||
|
||||
@ -27,6 +27,7 @@ import comfy.text_encoders.z_image
|
||||
import comfy.text_encoders.ideogram4
|
||||
import comfy.text_encoders.boogu
|
||||
import comfy.text_encoders.krea2
|
||||
import comfy.text_encoders.joyimage
|
||||
import comfy.text_encoders.anima
|
||||
import comfy.text_encoders.ace15
|
||||
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))
|
||||
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):
|
||||
unet_config = {
|
||||
"image_model": "hunyuan_video",
|
||||
@ -2389,6 +2422,7 @@ models = [
|
||||
Omnigen2,
|
||||
Boogu,
|
||||
QwenImage,
|
||||
JoyImage,
|
||||
Ideogram4,
|
||||
Krea2,
|
||||
Flux2,
|
||||
|
||||
97
comfy/text_encoders/joyimage.py
Normal file
97
comfy/text_encoders/joyimage.py
Normal file
@ -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,
|
||||
image_mean: list = None,
|
||||
image_std: list = None,
|
||||
interpolation: str = "bilinear",
|
||||
):
|
||||
if image_mean is None:
|
||||
image_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||
@ -47,10 +48,9 @@ def process_qwen2vl_images(
|
||||
img_resized = F.interpolate(
|
||||
img.unsqueeze(0),
|
||||
size=(h_bar, w_bar),
|
||||
mode='bilinear',
|
||||
mode=interpolation,
|
||||
align_corners=False
|
||||
).squeeze(0)
|
||||
|
||||
normalized = img_resized.clone()
|
||||
for c in range(3):
|
||||
normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from av.container import InputContainer
|
||||
from av.subtitles.stream import SubtitleStream
|
||||
from av.video.reformatter import ColorRange
|
||||
from fractions import Fraction
|
||||
from typing import Optional
|
||||
from .._input import AudioInput, VideoInput
|
||||
@ -9,6 +10,7 @@ import itertools
|
||||
import json
|
||||
import numpy as np
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||
import logging
|
||||
@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
|
||||
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 representing video input from a file.
|
||||
@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
|
||||
return estimated_frames
|
||||
|
||||
# 3. Last resort: decode frames and count them (streaming)
|
||||
if self.__start_time < 0:
|
||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
||||
else:
|
||||
start_time = self.__start_time
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
frame_count = 1
|
||||
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)
|
||||
frame_iterator = (
|
||||
container.decode(video_stream)
|
||||
@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
|
||||
|
||||
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
if self.__start_time < 0:
|
||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
||||
else:
|
||||
start_time = self.__start_time
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
|
||||
# Get video frames
|
||||
frames = []
|
||||
audio_frames = []
|
||||
alphas = None
|
||||
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:
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
@ -281,18 +328,11 @@ class VideoFromFile(VideoInput):
|
||||
video_done = False
|
||||
audio_done = True
|
||||
|
||||
# Use the last decodable audio stream. 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)
|
||||
audio_stream = next(
|
||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||
None,
|
||||
)
|
||||
audio_stream = last_decodable_audio_stream(container)
|
||||
if audio_stream is not None:
|
||||
streams += [audio_stream]
|
||||
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
||||
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):
|
||||
if video_done and audio_done:
|
||||
@ -305,7 +345,7 @@ class VideoFromFile(VideoInput):
|
||||
for frame in packet.decode():
|
||||
if frame.pts < start_pts:
|
||||
continue
|
||||
if self.__duration and frame.pts >= end_pts:
|
||||
if duration and frame.pts >= end_pts:
|
||||
video_done = True
|
||||
break
|
||||
|
||||
@ -372,7 +412,7 @@ class VideoFromFile(VideoInput):
|
||||
map(resampler.resample, packet.decode())
|
||||
)
|
||||
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
|
||||
break
|
||||
|
||||
@ -394,8 +434,8 @@ class VideoFromFile(VideoInput):
|
||||
|
||||
if len(audio_frames) > 0:
|
||||
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
||||
if self.__duration:
|
||||
audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
|
||||
if duration:
|
||||
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 = AudioInput({
|
||||
@ -441,28 +481,14 @@ class VideoFromFile(VideoInput):
|
||||
if not reuse_streams:
|
||||
if bit_depth is None:
|
||||
bit_depth = source_bit_depth
|
||||
components = self.get_components_internal(container)
|
||||
video = VideoFromComponents(components)
|
||||
return video.save_to(
|
||||
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
|
||||
)
|
||||
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
|
||||
|
||||
streams = container.streams
|
||||
|
||||
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
||||
with av.open(path, **open_kwargs) as output_container:
|
||||
# Copy over the original metadata
|
||||
for key, value in container.metadata.items():
|
||||
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 metadata before writing any streams
|
||||
write_output_metadata(container, output_container, metadata)
|
||||
|
||||
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
||||
stream_map = {}
|
||||
@ -480,6 +506,282 @@ class VideoFromFile(VideoInput):
|
||||
packet.stream = stream_map[packet.stream]
|
||||
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):
|
||||
if len(container.streams.video):
|
||||
return container.streams.video[0]
|
||||
@ -527,22 +829,12 @@ class VideoFromComponents(VideoInput):
|
||||
bit_depth: int | None = None,
|
||||
):
|
||||
"""Save the video to a file path or BytesIO buffer."""
|
||||
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 = mp4_output_open_kwargs(path, format, codec)
|
||||
# None means "use the depth this video was created with" (CreateVideo's choice).
|
||||
if bit_depth is None:
|
||||
bit_depth = self.__bit_depth
|
||||
is_10bit = bit_depth >= 10
|
||||
extra_kwargs = {}
|
||||
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:
|
||||
with av.open(path, **open_kwargs) as output:
|
||||
# Add metadata before writing any streams
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
|
||||
452
comfy_api_nodes/apis/heygen.py
Normal file
452
comfy_api_nodes/apis/heygen.py
Normal file
@ -0,0 +1,452 @@
|
||||
# (label, avatar_id, avatar_type, supported engines)
|
||||
HEYGEN_AVATAR_LOOKS: list[tuple[str, str, str, tuple[str, ...]]] = [
|
||||
(
|
||||
"Annie Lounge Standing Side",
|
||||
"Annie_Lounge_Standing_Side_public",
|
||||
"studio_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Yara Modern Lecture Hall",
|
||||
"fd6814ecc5e143cd899e615a80eaa2dc",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Brandon Business Sitting Front",
|
||||
"Brandon_Business_Sitting_Front_public",
|
||||
"studio_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Caroline Business Sitting Side",
|
||||
"Caroline_Business_Sitting_Side_public",
|
||||
"studio_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Ursula Lawyer Angle 4",
|
||||
"f7173d2bb8584c00bfec6905c5e9a492",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Sofia Corporate Presenter 01 Angle 3",
|
||||
"fe563971fd2d438e957372dac9e2be8c",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Seoyeon Health Nutrition Coach Angle 3",
|
||||
"fe3c5d5028d941398d064b8fc64a2dea",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Sanne Fitness Coach Angle 4",
|
||||
"d967f935a8bf4a0c8f0bccfd66c501d2",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
("Sander", "f5cd7b94056f495ca0610602d64a9aa3", "photo_avatar", ("avatar_v", "avatar_iv", "avatar_iii")),
|
||||
(
|
||||
"Rupert Personal Development Coach Angle 4",
|
||||
"f57b3e626adb4bc997b38f64884adce4",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Olivier Professor Angle 2",
|
||||
"f6659bbb094b459c87c967edbb9ee481",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Obi Health Nutrition Coach Angle 5",
|
||||
"f3dc2c38201d414382f506d2d8e8d029",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Matilda Modern Office Setting",
|
||||
"fda889ac354a440da8dbecc410981273",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Mateo Traditional Law Office",
|
||||
"ff172d6c499c4e47ba6fcc5de631e9fc",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Marlon Inviting Armchair Setting",
|
||||
"f5a57db099ab462daa3e7c604a05dacc",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Margaret Professor Angle 1",
|
||||
"fb472bc29ab04bcca576e3703978fecb",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Marek Therapy Coach Angle 3",
|
||||
"e197768703f1463a93dc25ada1f421fb",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Maeve Warm, Professional Setting",
|
||||
"faf66681d8cc48dc82c4283200b3e782",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Lorenzo Professor Angle 5",
|
||||
"fc268dc244bb40d7a554663ce723dcf0",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
("Luca", "Luca_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Bruce", "Bruce_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Nico", "Nico_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Lisa", "Lisa_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Sophie", "Sophie_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Aiko", "Aiko_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Rebecca (portrait)", "Rebecca_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Daphne in Grey blazer (portrait)", "Daphne_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Bryce in Black t-shirt", "Bryce_public_5", "studio_avatar", ("avatar_iii",)),
|
||||
("Diora in White shirt", "Diora_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Freja in White blazer", "Freja_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Albert in Blue blazer", "Albert_public_2", "studio_avatar", ("avatar_iii",)),
|
||||
("Emery in Red blazer", "Emery_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Minho in Blue shirt", "Minho_public_6", "studio_avatar", ("avatar_iii",)),
|
||||
("Aditya in Brown blazer", "Aditya_public_4", "studio_avatar", ("avatar_iii",)),
|
||||
("Nadim in Blue blazer", "Nadim_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Iker in Black blazer", "Iker_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Nour in Black blazer", "Nour_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Saskia in Blue blazer", "Saskia_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Lucien in Blue blazer", "Lucien_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Esmond in Blue suit", "Esmond_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Jinwoo in Blue suit", "Jinwoo_public_5", "studio_avatar", ("avatar_iii",)),
|
||||
("Annelore in Red sweater (portrait)", "Annelore_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Bastien in Blue shirt", "Bastien_public_4", "studio_avatar", ("avatar_iii",)),
|
||||
("Zosia in Khaki blazer", "Zosia_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Tahlia in Dark blue suit", "Tahlia_public_4", "studio_avatar", ("avatar_iii",)),
|
||||
]
|
||||
HEYGEN_AVATAR_OPTIONS = [x[0] for x in HEYGEN_AVATAR_LOOKS]
|
||||
HEYGEN_AVATAR_MAP = {x[0]: (x[1], x[2], x[3]) for x in HEYGEN_AVATAR_LOOKS}
|
||||
|
||||
# (label, voice_id) — Starfish-compatible voices for the TTS endpoint
|
||||
HEYGEN_VOICE_TTS: list[tuple[str, str]] = [
|
||||
("Chill Brian (English, male)", "d2f4f24783d04e22ab49ee8fdc3715e0"),
|
||||
("Zain (English, female)", "0047732240584155b1588455313e78ec"),
|
||||
("Narrator Mateo - Excited 🤩 (Spanish, male)", "0077225a877e457db4572ccaf245910b"),
|
||||
("Aria (English, female)", "007e1378fc454a9f976db570ba6164a7"),
|
||||
("Caryns (English, female)", "0082e70326864107823605db0d77f5e0"),
|
||||
("Klara (English, female)", "01209fdcd1c24a109c86dc24ee0f71c0"),
|
||||
("Bold Kasia - Excited 🤩 (Polish, female)", "015482a78b9a46ebae74bd0beb17765b"),
|
||||
("Shaun (English, male)", "01c42cddcfdc4665a57b8d89cba8ffc1"),
|
||||
("Senthil (English, male)", "01d674cfd32b4728a3fddd21b7e7d543"),
|
||||
("Cody (English, male)", "01f98ed43e6140349f47dbd37a416827"),
|
||||
("Saffron (English, female)", "0258bbc2cd8648cfa357adfb833f6d7b"),
|
||||
("Blanka - Lifelike (English, female)", "02880d1c6fd94b7799d91135581ed810"),
|
||||
("Rami (English, male)", "02d5366a90af4c7a87157808ff352e33"),
|
||||
("Rhodes (English, female)", "02dce0a169b3460084b6c914d18fb2c8"),
|
||||
("Michelle - Voice 1 (English, female)", "02X8sHnuxFpsq1caYWN0"),
|
||||
("Autumn - UGC 3 (English, female)", "03dca9ebfca441dba55fb14afa0791b7"),
|
||||
("Reassuring Rupert (English, male)", "03fcf8ecb0a94b6b94e9007edb7c35f8"),
|
||||
("Rose - UGC -2 (English, female)", "0495e14c2bd74eb3aeeef03583e0bce5"),
|
||||
("Derya - Lifelike - Broadcaster 🎙️ (English, female)", "04d0ae1d0af2489ca7d3bb402a39a890"),
|
||||
("Dynamic Derek (English, male)", "0516c2d857eb425c94e90b068241914e"),
|
||||
("Lotte (English, female)", "052fcfb83d1a4c2f8d0368c226fea4b9"),
|
||||
("Thanos - Broadcaster 🎙️ (English, male)", "054af44a167344d0af2722fdfef08d17"),
|
||||
("Marcia (English, female)", "05f19352e8f74b0392a8f411eba40de1"),
|
||||
("Camden (English, male)", "06468055edd4458aa131a1dfd813c1e9"),
|
||||
("Rumi (English, female)", "06672207805f41a9ad0af6797f8aa14b"),
|
||||
("Pippa (English, female)", "06b68c4dbb544935b9af984e80efa4fb"),
|
||||
("William Prescott - Broadcaster 🎙️ (English, male)", "06c816b952f14fa9b3a6c42aa151f731"),
|
||||
("Sammy (English, female)", "06e6facd99654b9dbb9308f67bf3a31c"),
|
||||
("Breezy Bagus (Indonesian, male)", "06e81a5d7c8b41818d3f0b38f7cf15a1"),
|
||||
("Ben (English, male)", "07ca39b243184dbcb82e7e0f0e524b21"),
|
||||
("Smooth Dev (English, male)", "07d2ba65847541feb97abc9b60181555"),
|
||||
("Daran inside booth (English, male)", "080d9383c0314056aef392892e009806"),
|
||||
("Peppy Stella (English, female)", "084760b4922a44599575c770070ec2d7"),
|
||||
("Silas (English, male)", "08f561403ec846dbbd8c691cc448f45a"),
|
||||
("Aditya (English, male)", "09c3d65e44e247dd8b78a97a903feb58"),
|
||||
("Christy (English, female)", "09d88c036bf449fa905900c08b235a37"),
|
||||
("Elio (English, male)", "0a0b38624ac64ec6afcd5842a977ca10"),
|
||||
("Luminous Laksh (Hindi, male)", "0adc547b76a5401c856274c379904eb7"),
|
||||
("Jeff (English, male)", "0add542e349f4ccaba6ecb3b7ced6034"),
|
||||
("Tahlia Brooks - Excited 🤩 (English, female)", "0b440d1ac2454d69a73302fc806522b1"),
|
||||
("Riya Mehta (Hindi, female)", "0b464b2f4e2249a4b5a05e60eaf41e7e"),
|
||||
("Ben Hart (English, male)", "0b47b5a637e944f9bfd49913999b344b"),
|
||||
("Skylar (English, female)", "0bbfbda5aa924a68a9d1da7b8496052a"),
|
||||
("Relaxed Reece (English, male)", "0c2151d538844c70a8b096de533f2828"),
|
||||
("Daniel (English, male)", "0c23804af39a4946ac6fda42bfff2738"),
|
||||
("Melani (English, female)", "0c54c6399ad64551a304e1a346677723"),
|
||||
("Clover (English, female)", "0ccb0bea067d4449ad367baeed7ea2e9"),
|
||||
("Pedro Lima - Serious 😐 (Portuguese, male)", "0d0e23e8170446e38b18a7380b2d30a8"),
|
||||
("Ana Carvalho (Portuguese, female)", "0d23c5b2f6004e909802a2e8bfcd52c2"),
|
||||
("Confident Connor - Excited 🤩 (English, male)", "0dd34c3eb79247238219eea35aeb58cd"),
|
||||
("Vibrant Victor (Spanish, male)", "1062976ea8bf42f4adc27c7e868b8fde"),
|
||||
("Young Olivier (French, male)", "1c5dc9a8f8cf4de0932f91d75f43a15d"),
|
||||
("Émile Noir (French, male)", "25a6a67280574d3da78e97b1935ebfc7"),
|
||||
("Steadfast Stefan (German, male)", "0eb85e6e8710473b82f7e88609ba3053"),
|
||||
("Deep Dieter (German, male)", "118949676b0a46629d1ad52981c3ef84"),
|
||||
("Serene Marco (Italian, male)", "72e922488a614041b5ab5f6ee07e3deb"),
|
||||
("Murmuring Matteo (Italian, male)", "755902b751654f30a6ef49e8bbcacfec"),
|
||||
("Gail in car (Multilingual, female)", "0214ac51f93e420f8711d568dcfbc50e"),
|
||||
("Daran outside walking (Multilingual, male)", "0ac81e725f4948dfa9638ceca216bcfa"),
|
||||
("BOB - Voice 1 (Chinese, unknown)", "dMkR1XwIkarpNqWUJLnX"),
|
||||
("Hakeem Hassan (Arabic, male)", "61a4359785664d01a59664ceb87ce6d4"),
|
||||
("Rami Idris (Arabic, male)", "a0bd2e5d41a74643be47ac75ca9171a2"),
|
||||
("Bold Kasia - Friendly 😊 (Polish, female)", "331624aec8b24a6c9287b8e16bdf54e8"),
|
||||
("Tranquil Tulin (Turkish, female)", "61646c861eb64e2d9036d8db51385356"),
|
||||
("Dynamic Derya (Turkish, female)", "664b73058b784aa89ddb2924c141d441"),
|
||||
("Quiet Dewa (Indonesian, male)", "1fa1193cf1d74f27ba58531c07ef9862"),
|
||||
("Cuong (Vietnamese, male)", "8af68d7ea38f4e7ca05cf46c3f7a590b"),
|
||||
]
|
||||
HEYGEN_VOICE_TTS_OPTIONS = [x[0] for x in HEYGEN_VOICE_TTS]
|
||||
HEYGEN_VOICE_TTS_MAP = dict(HEYGEN_VOICE_TTS)
|
||||
|
||||
# (label, voice_id) — top-ranked voices for video narration (any engine)
|
||||
HEYGEN_VOICE_GENERAL: list[tuple[str, str]] = [
|
||||
("Cassidy (English, female)", "16a09e4706f74997ba4ed05ea11470f6"),
|
||||
("Hope (English, female)", "42d00d4aac5441279d8536cd6b52c53c"),
|
||||
("Archer (English, male)", "453c20e1525a429080e2ad9e4b26f2cd"),
|
||||
("Brittney (English, female)", "4754e1ec667544b0bd18cdf4bec7d6a7"),
|
||||
("Mark (English, male)", "5d8c378ba8c3434586081a52ac368738"),
|
||||
("Andrew (English, male)", "6be73833ef9a4eb0aeee399b8fe9d62b"),
|
||||
("Spuds Oxley (English, male)", "76940a9adcd0490a9ce2cfe9a64a2664"),
|
||||
("Patrick (English, male)", "7e157ec62c9c45f1adca12faae72c86f"),
|
||||
("David Castlemore (English, male)", "828b59f834fd4c7188da322b6d9b6c75"),
|
||||
("Michael C (English, male)", "8661cd40d6c44c709e2d0031c0186ada"),
|
||||
("Adam Stone (English, male)", "88bb9ee1c81b466eb2a08fdde86d3619"),
|
||||
("Alex (English, male)", "897d6a9b2c844f56aa077238768fe10a"),
|
||||
("Monika Sogam (English, female)", "97dd67ab8ce242b6a9e7689cb00c6414"),
|
||||
("Jessica Anne Bogart (English, female)", "b966c31caf124c2a99f19ff1479c964f"),
|
||||
("John Doe (English, male)", "c4a8ceb7a2954500bc047fb092bcff3f"),
|
||||
("Ivy (English, female)", "cef3bc4e0a84424cafcde6f2cf466c97"),
|
||||
("Chill Brian (English, male)", "d2f4f24783d04e22ab49ee8fdc3715e0"),
|
||||
("Allison (English, female)", "f8c69e517f424cafaecde32dde57096b"),
|
||||
("Mia Starset (Norwegian, female)", "000466f8ac6d47a49f5743d50b3778de"),
|
||||
("William Shanks (Spanish, male)", "001248bb63f847888d37b766ee8b3a47"),
|
||||
("Zain (English, female)", "0047732240584155b1588455313e78ec"),
|
||||
("Jora Slobod (Romanian, male)", "00631519159a402ab5d8f719e51532bb"),
|
||||
("Narrator Mateo - Excited 🤩 (Spanish, male)", "0077225a877e457db4572ccaf245910b"),
|
||||
("Aria (English, female)", "007e1378fc454a9f976db570ba6164a7"),
|
||||
("Caryns (English, female)", "0082e70326864107823605db0d77f5e0"),
|
||||
("Klara (English, female)", "01209fdcd1c24a109c86dc24ee0f71c0"),
|
||||
("Son Tran (Vietnamese, male)", "0132f85950a94d11ba180f885101bf84"),
|
||||
("Bold Kasia - Excited 🤩 (Polish, female)", "015482a78b9a46ebae74bd0beb17765b"),
|
||||
("Marc Aurèle (French, male)", "018a94cf15574491a0bab7f6799ac15b"),
|
||||
("Shaun (English, male)", "01c42cddcfdc4665a57b8d89cba8ffc1"),
|
||||
("Senthil (English, male)", "01d674cfd32b4728a3fddd21b7e7d543"),
|
||||
("Cody (English, male)", "01f98ed43e6140349f47dbd37a416827"),
|
||||
("Saffron (English, female)", "0258bbc2cd8648cfa357adfb833f6d7b"),
|
||||
("Blanka - Lifelike (English, female)", "02880d1c6fd94b7799d91135581ed810"),
|
||||
("Rami (English, male)", "02d5366a90af4c7a87157808ff352e33"),
|
||||
("Rhodes (English, female)", "02dce0a169b3460084b6c914d18fb2c8"),
|
||||
("Michelle - Voice 1 (English, female)", "02X8sHnuxFpsq1caYWN0"),
|
||||
("Tuba (, female)", "034ca0c32b6542028748d6d365d90d6a"),
|
||||
("Autumn - UGC 3 (English, female)", "03dca9ebfca441dba55fb14afa0791b7"),
|
||||
("Reassuring Rupert (English, male)", "03fcf8ecb0a94b6b94e9007edb7c35f8"),
|
||||
]
|
||||
HEYGEN_VOICE_GENERAL_OPTIONS = [x[0] for x in HEYGEN_VOICE_GENERAL]
|
||||
HEYGEN_VOICE_GENERAL_MAP = dict(HEYGEN_VOICE_GENERAL)
|
||||
|
||||
HEYGEN_TRANSLATE_LANGUAGES = [
|
||||
"English",
|
||||
"Spanish",
|
||||
"Spanish (Spain)",
|
||||
"Spanish (Mexico)",
|
||||
"French",
|
||||
"French (France)",
|
||||
"German",
|
||||
"German (Germany)",
|
||||
"Portuguese",
|
||||
"Portuguese (Brazil)",
|
||||
"Italian",
|
||||
"Italian (Italy)",
|
||||
"Japanese",
|
||||
"Japanese (Japan)",
|
||||
"Korean",
|
||||
"Chinese (Mandarin, Simplified)",
|
||||
"Arabic",
|
||||
"Hindi",
|
||||
"Hindi (India)",
|
||||
"Russian",
|
||||
"Russian (Russia)",
|
||||
"Dutch",
|
||||
"Polish",
|
||||
"Turkish",
|
||||
"Indonesian",
|
||||
"Vietnamese",
|
||||
"Ukrainian",
|
||||
"Afrikaans (South Africa)",
|
||||
"Albanian (Albania)",
|
||||
"Amharic (Ethiopia)",
|
||||
"Arabic (Algeria)",
|
||||
"Arabic (Bahrain)",
|
||||
"Arabic (Egypt)",
|
||||
"Arabic (Iraq)",
|
||||
"Arabic (Jordan)",
|
||||
"Arabic (Kuwait)",
|
||||
"Arabic (Lebanon)",
|
||||
"Arabic (Libya)",
|
||||
"Arabic (Morocco)",
|
||||
"Arabic (Oman)",
|
||||
"Arabic (Qatar)",
|
||||
"Arabic (Saudi Arabia)",
|
||||
"Arabic (Syria)",
|
||||
"Arabic (Tunisia)",
|
||||
"Arabic (United Arab Emirates)",
|
||||
"Arabic (World)",
|
||||
"Arabic (Yemen)",
|
||||
"Armenian (Armenia)",
|
||||
"Azerbaijani (Latin, Azerbaijan)",
|
||||
"Bangla (Bangladesh)",
|
||||
"Basque",
|
||||
"Belarusian (Belarus)",
|
||||
"Bengali (India)",
|
||||
"Bosnian (Bosnia and Herzegovina)",
|
||||
"Bulgarian",
|
||||
"Bulgarian (Bulgaria)",
|
||||
"Burmese (Myanmar)",
|
||||
"Catalan",
|
||||
"Chinese (Cantonese, Traditional)",
|
||||
"Chinese (Jilu Mandarin, Simplified)",
|
||||
"Chinese (Northeastern Mandarin, Simplified)",
|
||||
"Chinese (Southwestern Mandarin, Simplified)",
|
||||
"Chinese (Taiwanese Mandarin, Traditional)",
|
||||
"Chinese (Wu, Simplified)",
|
||||
"Chinese (Zhongyuan Mandarin Henan, Simplified)",
|
||||
"Chinese (Zhongyuan Mandarin Shaanxi, Simplified)",
|
||||
"Croatian",
|
||||
"Croatian (Croatia)",
|
||||
"Czech",
|
||||
"Czech (Czechia)",
|
||||
"Danish",
|
||||
"Danish (Denmark)",
|
||||
"Dutch (Belgium)",
|
||||
"Dutch (Netherlands)",
|
||||
"English (Australia)",
|
||||
"English (Canada)",
|
||||
"English (Hong Kong SAR)",
|
||||
"English (India)",
|
||||
"English (Ireland)",
|
||||
"English (Kenya)",
|
||||
"English (New Zealand)",
|
||||
"English (Nigeria)",
|
||||
"English (Philippines)",
|
||||
"English (Singapore)",
|
||||
"English (South Africa)",
|
||||
"English (Tanzania)",
|
||||
"English (UK)",
|
||||
"English (United States)",
|
||||
"Estonian (Estonia)",
|
||||
"Filipino",
|
||||
"Filipino (Cebuano)",
|
||||
"Filipino (Philippines)",
|
||||
"Finnish",
|
||||
"Finnish (Finland)",
|
||||
"French (Belgium)",
|
||||
"French (Canada)",
|
||||
"French (Switzerland)",
|
||||
"Galician",
|
||||
"Georgian (Georgia)",
|
||||
"German (Austria)",
|
||||
"German (Switzerland)",
|
||||
"Greek",
|
||||
"Greek (Greece)",
|
||||
"Gujarati (India)",
|
||||
"Haitian Creole (Haiti)",
|
||||
"Hebrew (Israel)",
|
||||
"Hungarian (Hungary)",
|
||||
"Icelandic (Iceland)",
|
||||
"Indonesian (Indonesia)",
|
||||
"Irish (Ireland)",
|
||||
"Javanese (Latin, Indonesia)",
|
||||
"Kannada (India)",
|
||||
"Kazakh (Kazakhstan)",
|
||||
"Khmer (Cambodia)",
|
||||
"Konkani (India)",
|
||||
"Korean (Korea)",
|
||||
"Lao (Laos)",
|
||||
"Latin (Vatican City)",
|
||||
"Latvian (Latvia)",
|
||||
"Lithuanian (Lithuania)",
|
||||
"Luxembourgish (Luxembourg)",
|
||||
"Macedonian (North Macedonia)",
|
||||
"Maithili (India)",
|
||||
"Malagasy (Madagascar)",
|
||||
"Malay",
|
||||
"Malay (Malaysia)",
|
||||
"Malayalam (India)",
|
||||
"Maltese (Malta)",
|
||||
"Mandarin",
|
||||
"Marathi (India)",
|
||||
"Mongolian (Mongolia)",
|
||||
"Nepali (Nepal)",
|
||||
"Norwegian Bokmål (Norway)",
|
||||
"Norwegian Nynorsk (Norway)",
|
||||
"Odia (India)",
|
||||
"Pashto (Afghanistan)",
|
||||
"Persian (Iran)",
|
||||
"Polish (Poland)",
|
||||
"Portuguese (Portugal)",
|
||||
"Punjabi (India)",
|
||||
"Romanian",
|
||||
"Romanian (Romania)",
|
||||
"Serbian (Latin, Serbia)",
|
||||
"Sindhi (India)",
|
||||
"Sinhala (Sri Lanka)",
|
||||
"Slovak",
|
||||
"Slovak (Slovakia)",
|
||||
"Slovenian (Slovenia)",
|
||||
"Somali (Somalia)",
|
||||
"Spanish (Argentina)",
|
||||
"Spanish (Bolivia)",
|
||||
"Spanish (Chile)",
|
||||
"Spanish (Colombia)",
|
||||
"Spanish (Costa Rica)",
|
||||
"Spanish (Cuba)",
|
||||
"Spanish (Dominican Republic)",
|
||||
"Spanish (Ecuador)",
|
||||
"Spanish (El Salvador)",
|
||||
"Spanish (Equatorial Guinea)",
|
||||
"Spanish (Guatemala)",
|
||||
"Spanish (Honduras)",
|
||||
"Spanish (Latin America)",
|
||||
"Spanish (Nicaragua)",
|
||||
"Spanish (Panama)",
|
||||
"Spanish (Paraguay)",
|
||||
"Spanish (Peru)",
|
||||
"Spanish (Puerto Rico)",
|
||||
"Spanish (United States)",
|
||||
"Spanish (Uruguay)",
|
||||
"Spanish (Venezuela)",
|
||||
"Sundanese (Indonesia)",
|
||||
"Swahili (Kenya)",
|
||||
"Swahili (Tanzania)",
|
||||
"Swedish",
|
||||
"Swedish (Sweden)",
|
||||
"Tamil",
|
||||
"Tamil (India)",
|
||||
"Tamil (Malaysia)",
|
||||
"Tamil (Singapore)",
|
||||
"Tamil (Sri Lanka)",
|
||||
"Telugu (India)",
|
||||
"Thai (Thailand)",
|
||||
"Turkish (Türkiye)",
|
||||
"Ukrainian (Ukraine)",
|
||||
"Urdu (India)",
|
||||
"Urdu (Pakistan)",
|
||||
"Uzbek (Latin, Uzbekistan)",
|
||||
"Vietnamese (Vietnam)",
|
||||
"Welsh (United Kingdom)",
|
||||
"Zulu (South Africa)",
|
||||
]
|
||||
@ -128,7 +128,7 @@ class OpenAIResponse(ModelResponseProperties, ResponseProperties):
|
||||
parallel_tool_calls: bool | None = Field(True)
|
||||
status: str | None = Field(
|
||||
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)
|
||||
|
||||
|
||||
@ -274,6 +274,10 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
||||
input_tokens_price = 0.25
|
||||
output_text_tokens_price = 1.50
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion == "gemini-3.5-flash":
|
||||
input_tokens_price = 1.50
|
||||
output_text_tokens_price = 9.0
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion in ("gemini-3-pro-image-preview", "gemini-3-pro-image"):
|
||||
input_tokens_price = 2
|
||||
output_text_tokens_price = 12.0
|
||||
@ -619,11 +623,12 @@ class GeminiNode(IO.ComfyNode):
|
||||
|
||||
GEMINI_V2_MODELS: dict[str, str] = {
|
||||
"Gemini 3.1 Pro": "gemini-3.1-pro-preview",
|
||||
"Gemini 3.5 Flash": "gemini-3.5-flash",
|
||||
"Gemini 3.1 Flash-Lite": "gemini-3.1-flash-lite-preview",
|
||||
}
|
||||
|
||||
|
||||
def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
|
||||
def _gemini_text_model_inputs(thinking_default: str, thinking_options: list[str] | None = None) -> list[Input]:
|
||||
"""Per-model inputs revealed by the model DynamicCombo (shared media + sampling controls)."""
|
||||
return [
|
||||
IO.Autogrow.Input(
|
||||
@ -661,7 +666,7 @@ def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"thinking_level",
|
||||
options=["LOW", "HIGH"],
|
||||
options=thinking_options or ["LOW", "HIGH"],
|
||||
default=thinking_default,
|
||||
tooltip="How hard the model reasons internally before answering. "
|
||||
"HIGH improves quality on difficult tasks but costs more (thinking) tokens and is slower.",
|
||||
@ -719,6 +724,10 @@ class GeminiNodeV2(IO.ComfyNode):
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Gemini 3.5 Flash",
|
||||
_gemini_text_model_inputs("MEDIUM", ["MINIMAL", "LOW", "MEDIUM", "HIGH"]),
|
||||
),
|
||||
IO.DynamicCombo.Option("Gemini 3.1 Pro", _gemini_text_model_inputs("HIGH")),
|
||||
IO.DynamicCombo.Option("Gemini 3.1 Flash-Lite", _gemini_text_model_inputs("LOW")),
|
||||
],
|
||||
@ -759,7 +768,13 @@ class GeminiNodeV2(IO.ComfyNode):
|
||||
"type": "list_usd",
|
||||
"usd": [0.00025, 0.0015],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
} : {
|
||||
}
|
||||
: $contains($m, "3.5 flash") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.0015, 0.009],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
}
|
||||
: {
|
||||
"type": "list_usd",
|
||||
"usd": [0.002, 0.012],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
|
||||
799
comfy_api_nodes/nodes_heygen.py
Normal file
799
comfy_api_nodes/nodes_heygen.py
Normal file
@ -0,0 +1,799 @@
|
||||
import uuid
|
||||
|
||||
import torch
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.heygen import (
|
||||
HEYGEN_AVATAR_MAP,
|
||||
HEYGEN_AVATAR_OPTIONS,
|
||||
HEYGEN_TRANSLATE_LANGUAGES,
|
||||
HEYGEN_VOICE_GENERAL_MAP,
|
||||
HEYGEN_VOICE_GENERAL_OPTIONS,
|
||||
HEYGEN_VOICE_TTS_MAP,
|
||||
HEYGEN_VOICE_TTS_OPTIONS,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
audio_bytes_to_audio_input,
|
||||
download_url_as_bytesio,
|
||||
download_url_to_image_tensor,
|
||||
download_url_to_video_output,
|
||||
downscale_image_tensor_by_max_side,
|
||||
get_number_of_images,
|
||||
poll_op_raw,
|
||||
sync_op_raw,
|
||||
upload_audio_to_comfyapi,
|
||||
upload_image_to_comfyapi,
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_string,
|
||||
)
|
||||
from server import PromptServer
|
||||
|
||||
_VIDEOS_PATH = "/proxy/heygen/v3/videos"
|
||||
_TRANSLATIONS_PATH = "/proxy/heygen/v3/video-translations"
|
||||
_SPEECH_PATH = "/proxy/heygen/v3/voices/speech"
|
||||
_AVATARS_PATH = "/proxy/heygen/v3/avatars"
|
||||
_LOOKS_PATH = "/proxy/heygen/v3/avatars/looks"
|
||||
|
||||
_DEFAULT_VOICE_OPTION = "(avatar's default voice)"
|
||||
|
||||
_AVATARS_BY_ENGINE = {
|
||||
e: [label for label, (_aid, _atype, engines) in HEYGEN_AVATAR_MAP.items() if e in engines]
|
||||
for e in ("avatar_iv", "avatar_iii", "avatar_v")
|
||||
}
|
||||
|
||||
|
||||
async def _apply_speech_source(cls: type[IO.ComfyNode], payload: dict, speech: dict, require_voice: bool) -> None:
|
||||
"""Fill script/audio speech fields of a /v3/videos payload from the DynamicCombo dict."""
|
||||
if speech["speech"] == "audio":
|
||||
payload["audio_url"] = await upload_audio_to_comfyapi(
|
||||
cls, speech["audio"], container_format="mp3", codec_name="libmp3lame", mime_type="audio/mpeg"
|
||||
)
|
||||
elif speech["speech"] == "script":
|
||||
validate_string(speech["text"], strip_whitespace=True, min_length=1, max_length=5000)
|
||||
payload["script"] = speech["text"]
|
||||
voice_id = speech.get("custom_voice_id", "").strip()
|
||||
if not voice_id and speech["voice"] != _DEFAULT_VOICE_OPTION:
|
||||
voice_id = HEYGEN_VOICE_GENERAL_MAP[speech["voice"]]
|
||||
if voice_id:
|
||||
payload["voice_id"] = voice_id
|
||||
elif require_voice:
|
||||
raise ValueError("A voice is required when driving the video with a text script.")
|
||||
speed = speech.get("voice_speed", 1.0)
|
||||
if speed != 1.0:
|
||||
payload["voice_settings"] = {"speed": round(speed, 2)}
|
||||
|
||||
|
||||
async def _create_and_poll_video(cls: type[IO.ComfyNode], payload: dict) -> dict:
|
||||
"""POST a /v3/videos payload, poll until terminal, and return the final video data."""
|
||||
created = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_VIDEOS_PATH, method="POST", headers={"Idempotency-Key": uuid.uuid4().hex}),
|
||||
data=payload,
|
||||
)
|
||||
video_id = (created.get("data") or {}).get("video_id")
|
||||
if not video_id:
|
||||
raise ValueError(f"HeyGen did not return a video_id: {created}")
|
||||
final = await poll_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_VIDEOS_PATH}/{video_id}"),
|
||||
status_extractor=lambda r: (r.get("data") or {}).get("status"),
|
||||
queued_statuses=["pending", "waiting"],
|
||||
poll_interval=5.0,
|
||||
)
|
||||
data = final["data"]
|
||||
if not data.get("video_url"):
|
||||
raise ValueError(f"HeyGen returned no video_url for video {video_id}.")
|
||||
return data
|
||||
|
||||
|
||||
async def _resolve_avatar(
|
||||
cls: type[IO.ComfyNode], avatar_label: str, custom_avatar_id: str, engine_choice: str
|
||||
) -> tuple[str, str | None]:
|
||||
"""Resolve (avatar_id, engine_type) from the combo/override + engine widgets."""
|
||||
custom_avatar_id = custom_avatar_id.strip()
|
||||
if custom_avatar_id:
|
||||
look = (
|
||||
await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_LOOKS_PATH}/{custom_avatar_id}"),
|
||||
final_label_on_success=None,
|
||||
)
|
||||
).get("data") or {}
|
||||
avatar_id = custom_avatar_id
|
||||
avatar_label = look.get("name") or custom_avatar_id
|
||||
supported = look.get("supported_api_engines") or []
|
||||
else:
|
||||
avatar_id, avatar_type, supported = HEYGEN_AVATAR_MAP[avatar_label]
|
||||
|
||||
if engine_choice == "auto":
|
||||
engine = next((e for e in ("avatar_iv", "avatar_iii", "avatar_v") if e in supported), None)
|
||||
else:
|
||||
engine = engine_choice
|
||||
if supported and engine not in supported:
|
||||
raise ValueError(
|
||||
f"Avatar '{avatar_label}' does not support the {engine} engine "
|
||||
f"(supported: {', '.join(supported)}). Set engine to 'auto' to pick "
|
||||
"a compatible engine automatically."
|
||||
)
|
||||
return avatar_id, engine
|
||||
|
||||
|
||||
class HeyGenTalkingPhotoNode(IO.ComfyNode):
|
||||
"""Animate a still image of a person into a lip-synced talking video."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenTalkingPhotoNode",
|
||||
display_name="HeyGen Talking Photo",
|
||||
category="partner/video/HeyGen",
|
||||
description="Animate any image of a person into a lip-synced talking video "
|
||||
"(HeyGen Avatar IV). Drive it with a text script or your own audio.",
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
tooltip="Image of a person to animate. Downscaled automatically if larger than 2K.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"speech",
|
||||
display_name="speech source",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"script",
|
||||
[
|
||||
IO.String.Input(
|
||||
"text",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Text for the avatar to speak (up to 5000 characters). "
|
||||
"The generated speech must be at least 1 second long.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
options=HEYGEN_VOICE_GENERAL_OPTIONS,
|
||||
tooltip="Voice for the script (HeyGen's most popular voices).",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_voice_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
|
||||
"Any voice from HeyGen's library (2000+) can be used.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"voice_speed",
|
||||
default=1.0,
|
||||
min=0.5,
|
||||
max=1.5,
|
||||
step=0.05,
|
||||
optional=True,
|
||||
tooltip="Speech speed multiplier.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"audio",
|
||||
[
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Audio for the avatar to lip-sync, up to 10 minutes.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Drive the avatar with a text script (HeyGen text-to-speech) or your own audio.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720p", "1080p"],
|
||||
default="1080p",
|
||||
optional=True,
|
||||
tooltip="Output video resolution.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=["auto", "16:9", "9:16", "1:1", "4:5", "5:4"],
|
||||
default="auto",
|
||||
optional=True,
|
||||
tooltip="Output aspect ratio. 'auto' follows the input image.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"expressiveness",
|
||||
options=["low", "medium", "high"],
|
||||
default="low",
|
||||
optional=True,
|
||||
tooltip="How expressive the animated face and gestures are.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
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.0715,"format":{"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
speech: dict,
|
||||
resolution: str = "1080p",
|
||||
aspect_ratio: str = "auto",
|
||||
expressiveness: str = "low",
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
image = downscale_image_tensor_by_max_side(image, max_side=2000)
|
||||
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||
payload = {
|
||||
"type": "image",
|
||||
"image": {"type": "url", "url": image_url},
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"expressiveness": expressiveness,
|
||||
"title": "ComfyUI Talking Photo",
|
||||
}
|
||||
await _apply_speech_source(cls, payload, speech, require_voice=True)
|
||||
video = await _create_and_poll_video(cls, payload)
|
||||
return IO.NodeOutput(await download_url_to_video_output(video["video_url"]))
|
||||
|
||||
|
||||
class HeyGenAvatarVideoNode(IO.ComfyNode):
|
||||
"""Generate a presenter video from a HeyGen avatar look."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenAvatarVideoNode",
|
||||
display_name="HeyGen Avatar Video",
|
||||
category="partner/video/HeyGen",
|
||||
description="Generate a talking-presenter video from a HeyGen avatar. "
|
||||
"Includes HeyGen's most popular public avatars; any look ID can be supplied as an override.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"engine",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"auto",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=HEYGEN_AVATAR_OPTIONS,
|
||||
tooltip="Avatar look to present the video (curated from HeyGen's "
|
||||
"public library). The best engine the look supports is chosen "
|
||||
"automatically.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"avatar_iv",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=_AVATARS_BY_ENGINE["avatar_iv"],
|
||||
tooltip="Avatar looks that support the Avatar IV engine.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"avatar_iii",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=_AVATARS_BY_ENGINE["avatar_iii"],
|
||||
tooltip="Avatar looks that support the Avatar III engine.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"avatar_v",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=_AVATARS_BY_ENGINE["avatar_v"],
|
||||
tooltip="Avatar looks that support the Avatar V engine.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Rendering engine; each choice lists only the avatars that support it. "
|
||||
"'auto' offers every avatar and picks its best engine (Avatar IV preferred). "
|
||||
"Avatar V is highest fidelity, Avatar III is the most affordable.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_avatar_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen avatar look ID. When set, overrides the avatar selected above. "
|
||||
"Any of HeyGen's 3000+ public looks (or your private avatars) can be used.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"speech",
|
||||
display_name="speech source",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"script",
|
||||
[
|
||||
IO.String.Input(
|
||||
"text",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Text for the avatar to speak (up to 5000 characters). "
|
||||
"The generated speech must be at least 1 second long.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
options=[_DEFAULT_VOICE_OPTION] + HEYGEN_VOICE_GENERAL_OPTIONS,
|
||||
tooltip="Voice for the script. The default option uses the voice HeyGen assigned to the avatar.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_voice_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
|
||||
"Any voice from HeyGen's library (2000+) can be used.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"voice_speed",
|
||||
default=1.0,
|
||||
min=0.5,
|
||||
max=1.5,
|
||||
step=0.05,
|
||||
optional=True,
|
||||
tooltip="Speech speed multiplier.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"audio",
|
||||
[
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Audio for the avatar to lip-sync, up to 10 minutes.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Drive the avatar with a text script (HeyGen text-to-speech) or your own audio.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720p", "1080p"],
|
||||
default="1080p",
|
||||
optional=True,
|
||||
tooltip="Output video resolution.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=["auto", "16:9", "9:16", "1:1", "4:5", "5:4"],
|
||||
default="auto",
|
||||
optional=True,
|
||||
tooltip="Output aspect ratio. 'auto' follows the avatar's source footage.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"background_color",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional solid background color as a hex code (e.g. '#00ff00'). "
|
||||
"Leave empty for the avatar's own background.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
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(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["engine"]),
|
||||
expr="""
|
||||
widgets.engine = "avatar_iii"
|
||||
? {"type":"range_usd","min_usd":0.023881,"max_usd":0.061919,"format":{"suffix":"/second"}}
|
||||
: widgets.engine = "avatar_v"
|
||||
? {"type":"usd","usd":0.095381,"format":{"suffix":"/second"}}
|
||||
: widgets.engine = "avatar_iv"
|
||||
? {"type":"range_usd","min_usd":0.0715,"max_usd":0.095381,"format":{"suffix":"/second"}}
|
||||
: {"type":"range_usd","min_usd":0.023881,"max_usd":0.095381,"format":{"suffix":"/second"}}
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
engine: dict,
|
||||
speech: dict,
|
||||
custom_avatar_id: str = "",
|
||||
resolution: str = "1080p",
|
||||
aspect_ratio: str = "auto",
|
||||
background_color: str = "",
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
avatar_id, engine_type = await _resolve_avatar(cls, engine["avatar"], custom_avatar_id, engine["engine"])
|
||||
payload = {
|
||||
"type": "avatar",
|
||||
"avatar_id": avatar_id,
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"title": "ComfyUI Avatar Video",
|
||||
}
|
||||
if engine_type:
|
||||
payload["engine"] = {"type": engine_type}
|
||||
background_color = background_color.strip()
|
||||
if background_color:
|
||||
if not background_color.startswith("#"):
|
||||
raise ValueError("background_color must be a hex color code like '#00ff00'.")
|
||||
payload["background"] = {"type": "color", "value": background_color}
|
||||
await _apply_speech_source(cls, payload, speech, require_voice=False)
|
||||
video = await _create_and_poll_video(cls, payload)
|
||||
return IO.NodeOutput(await download_url_to_video_output(video["video_url"]))
|
||||
|
||||
|
||||
class HeyGenCreateAvatarNode(IO.ComfyNode):
|
||||
"""Create a reusable HeyGen avatar from a photo or a text prompt."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenCreateAvatarNode",
|
||||
display_name="HeyGen Create Avatar",
|
||||
category="partner/video/HeyGen",
|
||||
description="Create your own reusable HeyGen avatar from a photo of a person or "
|
||||
"from a text prompt (a generated character). Feed the resulting avatar_id into "
|
||||
"HeyGen Avatar Video's custom_avatar_id — and save the ID somewhere to reuse the "
|
||||
"avatar in future workflows.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"source",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"prompt",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Description of the avatar to generate (up to 1000 characters).",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_images",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Image.Input("ref_image"),
|
||||
names=[f"ref_image_{i}" for i in range(1, 4)],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Up to 3 reference images guiding the generated look.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"photo",
|
||||
[
|
||||
IO.Image.Input(
|
||||
"identity_photo",
|
||||
tooltip="Photo of the person to turn into an avatar. "
|
||||
"Downscaled automatically if larger than 2K.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Generate a new character from a text prompt, or create the avatar "
|
||||
"from a connected photo of a person.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.String.Output(
|
||||
display_name="avatar_id",
|
||||
tooltip="Avatar look ID. Pass it to HeyGen Avatar Video's custom_avatar_id; "
|
||||
"save it to reuse the avatar later.",
|
||||
),
|
||||
IO.Image.Output(display_name="preview"),
|
||||
],
|
||||
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":1.43}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
source: dict,
|
||||
) -> IO.NodeOutput:
|
||||
payload: dict = {"name": "ComfyUI Avatar"}
|
||||
if source["source"] == "photo":
|
||||
image = downscale_image_tensor_by_max_side(source["identity_photo"], max_side=2000)
|
||||
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||
payload["type"] = "photo"
|
||||
payload["file"] = {"type": "url", "url": image_url}
|
||||
else:
|
||||
validate_string(source["prompt"], strip_whitespace=True, min_length=1, max_length=1000)
|
||||
payload["type"] = "prompt"
|
||||
payload["prompt"] = source["prompt"]
|
||||
ref_tensors = [t for t in (source.get("reference_images") or {}).values() if t is not None]
|
||||
if ref_tensors:
|
||||
n_images = sum(get_number_of_images(t) for t in ref_tensors)
|
||||
if n_images > 3:
|
||||
raise ValueError(f"HeyGen accepts at most 3 reference images; got {n_images}.")
|
||||
scaled = [downscale_image_tensor_by_max_side(t, max_side=2000) for t in ref_tensors]
|
||||
ref_urls = await upload_images_to_comfyapi(
|
||||
cls, scaled, max_images=3, mime_type="image/png", total_pixels=None
|
||||
)
|
||||
payload["reference_images"] = [{"type": "url", "url": u} for u in ref_urls]
|
||||
created = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_AVATARS_PATH, method="POST"),
|
||||
data=payload,
|
||||
)
|
||||
look_id = ((created.get("data") or {}).get("avatar_item") or {}).get("id")
|
||||
if not look_id:
|
||||
raise ValueError(f"HeyGen did not return an avatar: {created}")
|
||||
final = await poll_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_LOOKS_PATH}/{look_id}"),
|
||||
# A missing status means the look needed no training and is ready.
|
||||
status_extractor=lambda r: (r.get("data") or {}).get("status") or "completed",
|
||||
failed_statuses=["failed", "pending_consent"],
|
||||
poll_interval=5.0,
|
||||
)
|
||||
data = final["data"]
|
||||
if data.get("preview_image_url"):
|
||||
preview = await download_url_to_image_tensor(data["preview_image_url"])
|
||||
else:
|
||||
preview = torch.zeros(1, 64, 64, 3)
|
||||
PromptServer.instance.send_progress_text(
|
||||
f"Please save the avatar_id for reuse.\n\navatar_id: {look_id}",
|
||||
cls.hidden.unique_id,
|
||||
)
|
||||
return IO.NodeOutput(look_id, preview)
|
||||
|
||||
|
||||
class HeyGenVideoTranslateNode(IO.ComfyNode):
|
||||
"""Translate a spoken video into another language with voice cloning and lip sync."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenVideoTranslateNode",
|
||||
display_name="HeyGen Video Translate",
|
||||
category="partner/video/HeyGen",
|
||||
description="Translate a spoken video into another language. Clones the original "
|
||||
"speaker's voice and re-animates the mouth to match the translated speech.",
|
||||
inputs=[
|
||||
IO.Video.Input(
|
||||
"video",
|
||||
tooltip="Video with speech to translate.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"output_language",
|
||||
options=HEYGEN_TRANSLATE_LANGUAGES,
|
||||
tooltip="Target language for the translated video.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"mode",
|
||||
options=["speed", "precision"],
|
||||
default="speed",
|
||||
tooltip="'speed' is faster; 'precision' produces higher-quality lip sync at twice the price.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"translate_audio_only",
|
||||
default=False,
|
||||
optional=True,
|
||||
tooltip="Only swap the audio track, keeping the original mouth movements (no lip sync).",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_count",
|
||||
default=0,
|
||||
min=0,
|
||||
max=10,
|
||||
optional=True,
|
||||
tooltip="Number of speakers in the video. 0 = detect automatically.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
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(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["mode"]),
|
||||
expr="""{"type":"usd","usd": widgets.mode = "precision" ? 0.095381 : 0.047619,"""
|
||||
""""format":{"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
video: Input.Video,
|
||||
output_language: str,
|
||||
mode: str,
|
||||
translate_audio_only: bool = False,
|
||||
speaker_count: int = 0,
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
video_url = await upload_video_to_comfyapi(cls, video)
|
||||
payload = {
|
||||
"video": {"type": "url", "url": video_url},
|
||||
"output_languages": [output_language],
|
||||
"mode": mode,
|
||||
"translate_audio_only": translate_audio_only,
|
||||
"title": "ComfyUI Video Translate",
|
||||
}
|
||||
if speaker_count > 0:
|
||||
payload["speaker_num"] = speaker_count
|
||||
created = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_TRANSLATIONS_PATH, method="POST"),
|
||||
data=payload,
|
||||
)
|
||||
translation_ids = (created.get("data") or {}).get("video_translation_ids") or []
|
||||
if not translation_ids:
|
||||
raise ValueError(f"HeyGen did not return a translation ID: {created}")
|
||||
final = await poll_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_TRANSLATIONS_PATH}/{translation_ids[0]}"),
|
||||
status_extractor=lambda r: (r.get("data") or {}).get("status"),
|
||||
queued_statuses=["pending"],
|
||||
poll_interval=5.0,
|
||||
)
|
||||
data = final["data"]
|
||||
if not data.get("video_url"):
|
||||
raise ValueError(f"HeyGen returned no video_url for translation {translation_ids[0]}.")
|
||||
return IO.NodeOutput(await download_url_to_video_output(data["video_url"]))
|
||||
|
||||
|
||||
class HeyGenTextToSpeechNode(IO.ComfyNode):
|
||||
"""Synthesize speech audio from text with HeyGen's Starfish TTS engine."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenTextToSpeechNode",
|
||||
display_name="HeyGen Text to Speech",
|
||||
category="partner/audio/HeyGen",
|
||||
description="Generate speech audio from text using HeyGen's Starfish TTS engine. "
|
||||
"Includes HeyGen's most popular voices across 17 languages.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"text",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Text to synthesize (up to 5000 characters). The generated speech "
|
||||
"must be at least 1 second long.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
options=HEYGEN_VOICE_TTS_OPTIONS,
|
||||
tooltip="Voice to use (curated from HeyGen's most popular Starfish-compatible voices).",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_voice_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
|
||||
"The voice must support the Starfish engine.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"speed",
|
||||
default=1.0,
|
||||
min=0.5,
|
||||
max=2.0,
|
||||
step=0.05,
|
||||
optional=True,
|
||||
tooltip="Speech speed multiplier.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"ssml",
|
||||
default=False,
|
||||
optional=True,
|
||||
tooltip="Treat the text as SSML markup (for pauses, emphasis, and pronunciation control).",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Audio.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.00095381,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
text: str,
|
||||
voice: str,
|
||||
custom_voice_id: str = "",
|
||||
speed: float = 1.0,
|
||||
ssml: bool = False,
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(text, strip_whitespace=True, min_length=1, max_length=5000)
|
||||
payload = {
|
||||
"text": text,
|
||||
"voice_id": custom_voice_id.strip() or HEYGEN_VOICE_TTS_MAP[voice],
|
||||
"speed": round(speed, 2),
|
||||
}
|
||||
if ssml:
|
||||
payload["input_type"] = "ssml"
|
||||
response = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_SPEECH_PATH, method="POST"),
|
||||
data=payload,
|
||||
)
|
||||
audio_url = (response.get("data") or {}).get("audio_url")
|
||||
if not audio_url:
|
||||
raise ValueError(f"HeyGen did not return an audio_url: {response}")
|
||||
audio_bytes = await download_url_as_bytesio(audio_url)
|
||||
return IO.NodeOutput(audio_bytes_to_audio_input(audio_bytes.getvalue()))
|
||||
|
||||
|
||||
class HeyGenExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
HeyGenTalkingPhotoNode,
|
||||
HeyGenAvatarVideoNode,
|
||||
HeyGenCreateAvatarNode,
|
||||
HeyGenVideoTranslateNode,
|
||||
HeyGenTextToSpeechNode,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> HeyGenExtension:
|
||||
return HeyGenExtension()
|
||||
@ -41,6 +41,9 @@ STARTING_POINT_ID_PATTERN = r"<starting_point_id:(.*)>"
|
||||
|
||||
|
||||
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 = "gpt-5.5"
|
||||
gpt_5 = "gpt-5"
|
||||
@ -1063,6 +1066,21 @@ class OpenAIChatNode(IO.ComfyNode):
|
||||
"usd": [0.002, 0.008],
|
||||
"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") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.03, 0.18],
|
||||
|
||||
@ -15,6 +15,7 @@ from comfy.comfy_api_env import normalize_comfy_api_base
|
||||
from comfy.deploy_environment import get_deploy_environment
|
||||
from comfy.model_management import processing_interrupted
|
||||
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
|
||||
@ -57,12 +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
|
||||
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
||||
"""
|
||||
return {
|
||||
headers = {
|
||||
**get_auth_header(node_cls),
|
||||
"Comfy-Env": get_deploy_environment(),
|
||||
"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:
|
||||
|
||||
102
comfy_extras/nodes_joyimage.py
Normal file
102
comfy_extras/nodes_joyimage.py
Normal file
@ -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()
|
||||
@ -8,6 +8,7 @@ import comfy.ldm.common_dit
|
||||
import comfy.latent_formats
|
||||
import comfy.ldm.lumina.controlnet
|
||||
import comfy.ldm.supir.supir_modules
|
||||
import comfy.ldm.anima.lllite
|
||||
from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
|
||||
from comfy_api.latest import io
|
||||
from comfy.ldm.supir.supir_patch import SUPIRPatch
|
||||
@ -236,10 +237,12 @@ class ModelPatchLoader:
|
||||
|
||||
def load_model_patch(self, name):
|
||||
model_patch_path = folder_paths.get_full_path_or_raise("model_patches", name)
|
||||
sd = comfy.utils.load_torch_file(model_patch_path, safe_load=True)
|
||||
sd, metadata = comfy.utils.load_torch_file(model_patch_path, safe_load=True, return_metadata=True)
|
||||
dtype = comfy.utils.weight_dtype(sd)
|
||||
|
||||
if 'controlnet_blocks.0.y_rms.weight' in sd:
|
||||
if 'lllite_conditioning1.conv1.weight' in sd:
|
||||
model = comfy.ldm.anima.lllite.AnimaLLLite(sd, metadata, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
|
||||
elif 'controlnet_blocks.0.y_rms.weight' in sd:
|
||||
additional_in_dim = sd["img_in.weight"].shape[1] - 64
|
||||
model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
|
||||
elif 'feature_embedder.mid_layer_norm.bias' in sd:
|
||||
@ -296,6 +299,50 @@ class ModelPatchLoader:
|
||||
return (model_patcher,)
|
||||
|
||||
|
||||
class AnimaLLLiteApply:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"model": ("MODEL",),
|
||||
"model_patch": ("MODEL_PATCH",),
|
||||
"image": ("IMAGE",),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {"mask": ("MASK",),
|
||||
}}
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "apply_patch"
|
||||
EXPERIMENTAL = True
|
||||
|
||||
CATEGORY = "model_patches/anima"
|
||||
|
||||
def apply_patch(self, model, model_patch, image, strength, start_percent, end_percent, mask=None):
|
||||
image = image[..., :3]
|
||||
|
||||
if model_patch.model.cond_in_channels == 4 and mask is None:
|
||||
mask = torch.zeros_like(image[..., 0])
|
||||
elif model_patch.model.cond_in_channels != 4:
|
||||
mask = None
|
||||
|
||||
model_sampling = model.get_model_object("model_sampling")
|
||||
sigma_start = float(model_sampling.percent_to_sigma(start_percent))
|
||||
sigma_end = float(model_sampling.percent_to_sigma(end_percent))
|
||||
patch = comfy.ldm.anima.lllite.AnimaLLLitePatch(model_patch, image, mask, strength, sigma_start, sigma_end)
|
||||
model_patched = model.clone()
|
||||
model_patched.set_model_post_input_patch(patch)
|
||||
model_patched.set_model_attn1_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(
|
||||
patch,
|
||||
{"q": "self_attn_q_proj", "k": "self_attn_k_proj", "v": "self_attn_v_proj"},
|
||||
))
|
||||
model_patched.set_model_attn2_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(
|
||||
patch,
|
||||
{"q": "cross_attn_q_proj"},
|
||||
))
|
||||
model_patched.set_model_patch(comfy.ldm.anima.lllite.AnimaLLLiteMLPPatch(patch), "mlp_patch")
|
||||
return (model_patched,)
|
||||
|
||||
|
||||
class DiffSynthCnetPatch:
|
||||
def __init__(self, model_patch, vae, image, strength, mask=None):
|
||||
self.model_patch = model_patch
|
||||
@ -674,6 +721,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"ZImageFunControlnet": ZImageFunControlnet,
|
||||
"USOStyleReference": USOStyleReference,
|
||||
"SUPIRApply": SUPIRApply,
|
||||
"AnimaLLLiteApply": AnimaLLLiteApply,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@ -682,4 +730,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ZImageFunControlnet": "Apply Z-Image Fun ControlNet",
|
||||
"USOStyleReference": "Apply USO Style Reference",
|
||||
"SUPIRApply": "Apply SUPIR Patch",
|
||||
"AnimaLLLiteApply": "Apply Anima LLLite",
|
||||
}
|
||||
|
||||
5
nodes.py
5
nodes.py
@ -992,7 +992,7 @@ class CLIPLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
||||
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2"], ),
|
||||
"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": {
|
||||
"device": (["default", "cpu"], {"advanced": True}),
|
||||
@ -1002,7 +1002,7 @@ class CLIPLoader:
|
||||
|
||||
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"):
|
||||
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_context_windows.py",
|
||||
"nodes_qwen.py",
|
||||
"nodes_joyimage.py",
|
||||
"nodes_boogu.py",
|
||||
"nodes_chroma_radiance.py",
|
||||
"nodes_pid.py",
|
||||
|
||||
51
openapi.yaml
51
openapi.yaml
@ -530,6 +530,10 @@ components:
|
||||
description: Job creation timestamp (Unix timestamp in milliseconds)
|
||||
format: int64
|
||||
type: integer
|
||||
execution_end_time:
|
||||
description: Workflow execution completion timestamp (Unix milliseconds, only present for terminal states)
|
||||
format: int64
|
||||
type: integer
|
||||
execution_error:
|
||||
allOf:
|
||||
- $ref: '#/components/schemas/ExecutionError'
|
||||
@ -538,6 +542,10 @@ components:
|
||||
additionalProperties: true
|
||||
description: Node-level execution metadata (only for terminal states)
|
||||
type: object
|
||||
execution_start_time:
|
||||
description: Workflow execution start timestamp (Unix milliseconds, only present once execution has started)
|
||||
format: int64
|
||||
type: integer
|
||||
execution_status:
|
||||
additionalProperties: true
|
||||
description: ComfyUI execution status and timeline (only for terminal states)
|
||||
@ -570,6 +578,12 @@ components:
|
||||
description: Last update timestamp (Unix timestamp in milliseconds)
|
||||
format: int64
|
||||
type: integer
|
||||
user_id:
|
||||
description: |
|
||||
ID of the user that owns this job (see the `workspace_id`
|
||||
description above for why this is always the caller's own id
|
||||
on a successful response).
|
||||
type: string
|
||||
workflow:
|
||||
additionalProperties: true
|
||||
description: |
|
||||
@ -583,6 +597,18 @@ components:
|
||||
workflow_id:
|
||||
description: UUID identifying the workflow graph definition
|
||||
type: string
|
||||
workspace_id:
|
||||
description: |
|
||||
ID of the workspace that owns this job. A successful (200)
|
||||
response from this operation is only ever returned for the
|
||||
caller's own job (see this operation's ownership-scoped
|
||||
query), so this is always the caller's own workspace —
|
||||
consumers that also need to correlate this job to its
|
||||
live-progress broadcast channel (workspace+user scoped; see
|
||||
the internal common/gateways/broadcast package) can use this
|
||||
value directly rather than resolving their own identity a
|
||||
second way.
|
||||
type: string
|
||||
required:
|
||||
- id
|
||||
- status
|
||||
@ -1565,7 +1591,13 @@ paths:
|
||||
schema:
|
||||
default: true
|
||||
type: boolean
|
||||
- description: Filter assets by exact content hash.
|
||||
- description: |
|
||||
Filter assets by content hash, in the canonical `blake3:<hex>`
|
||||
form. Matches regardless of which of this asset store's two
|
||||
internal hash storage formats the matching row was written
|
||||
under (the canonical form used by from-hash-created references,
|
||||
or the raw `<hex>.<ext>`/bare `<hex>` storage key used by direct
|
||||
uploads) — both represent the same content hash.
|
||||
in: query
|
||||
name: hash
|
||||
schema:
|
||||
@ -2464,6 +2496,23 @@ paths:
|
||||
schema:
|
||||
additionalProperties: true
|
||||
properties:
|
||||
free_tier_balance:
|
||||
description: Free-tier job allowance for an authenticated non-paid (FREE-tier) user in the rollout. Absent for paid users and unauthenticated requests. Synthesized from config before a grant row exists so a brand-new user still sees their full allowance.
|
||||
properties:
|
||||
allowance:
|
||||
description: Total free jobs granted for the current period
|
||||
type: integer
|
||||
remaining:
|
||||
description: Free jobs remaining (allowance - used, floored at 0)
|
||||
type: integer
|
||||
used:
|
||||
description: Free jobs consumed so far
|
||||
type: integer
|
||||
required:
|
||||
- allowance
|
||||
- used
|
||||
- remaining
|
||||
type: object
|
||||
max_upload_size:
|
||||
description: Maximum upload size in bytes
|
||||
type: integer
|
||||
|
||||
@ -1,5 +1,5 @@
|
||||
comfyui-frontend-package==1.45.21
|
||||
comfyui-workflow-templates==0.11.9
|
||||
comfyui-workflow-templates==0.11.11
|
||||
comfyui-embedded-docs==0.5.8
|
||||
torch
|
||||
torchsde
|
||||
@ -22,7 +22,7 @@ alembic
|
||||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.20
|
||||
comfy-kitchen==0.2.22
|
||||
comfy-aimdo==0.4.10
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
|
||||
@ -2,11 +2,12 @@ import pytest
|
||||
import torch
|
||||
import tempfile
|
||||
import os
|
||||
import sys
|
||||
import av
|
||||
import io
|
||||
from fractions import Fraction
|
||||
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 av.error import InvalidDataError
|
||||
|
||||
@ -237,3 +238,526 @@ def test_duration_consistency(video_components):
|
||||
manual_duration = float(components.images.shape[0] / components.frame_rate)
|
||||
|
||||
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)
|
||||
|
||||
@ -112,6 +112,17 @@ def _make_pid_v1_5_sd(latent_proj_channels=16):
|
||||
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):
|
||||
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
|
||||
|
||||
@ -258,6 +269,26 @@ class TestModelDetection:
|
||||
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):
|
||||
"""Each model in the registry must have a unique combination of
|
||||
``unet_config`` and ``required_keys``. If two models share the same
|
||||
|
||||
Loading…
Reference in New Issue
Block a user