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.
|
better to remove a broken feature path than keep a complicated partial fix.
|
||||||
- Preserve existing APIs, node names, model-loading behavior, file layout, and
|
- Preserve existing APIs, node names, model-loading behavior, file layout, and
|
||||||
workflow compatibility unless the change is explicitly about replacing them.
|
workflow compatibility unless the change is explicitly about replacing them.
|
||||||
|
- When compatibility is explicitly out of scope, remove compatibility-only
|
||||||
|
aliases, duplicate nodes, legacy entry points, and preset wrappers instead of
|
||||||
|
retaining parallel ways to perform the same operation.
|
||||||
- Code must look hand-written for this repository. Changes that read like
|
- Code must look hand-written for this repository. Changes that read like
|
||||||
generic AI-generated code will be rejected automatically: unnecessary helper
|
generic AI-generated code will be rejected automatically: unnecessary helper
|
||||||
layers, vague names, boilerplate comments, defensive branches without a real
|
layers, vague names, boilerplate comments, defensive branches without a real
|
||||||
@ -96,6 +99,13 @@
|
|||||||
unless they are read by current code and change current behavior. Remove
|
unless they are read by current code and change current behavior. Remove
|
||||||
pass-through or stored-but-unused values instead of preserving upstream or
|
pass-through or stored-but-unused values instead of preserving upstream or
|
||||||
deprecated API baggage.
|
deprecated API baggage.
|
||||||
|
- Do not add a model-specific option to a shared helper when only one caller
|
||||||
|
needs it. Keep one-off behavior at the model integration boundary, or extend
|
||||||
|
the shared helper only when the option is a coherent reusable capability.
|
||||||
|
- Implementations of shared model interfaces should accept the standard caller
|
||||||
|
contract without model-specific rejection branches for optional capabilities
|
||||||
|
they do not consume. Let supported behavior be determined by implementation
|
||||||
|
paths that actually use those inputs.
|
||||||
- If an implementation needs auxiliary values for its own workflow, expose them
|
- If an implementation needs auxiliary values for its own workflow, expose them
|
||||||
through a private helper or a clearly named implementation-specific method
|
through a private helper or a clearly named implementation-specific method
|
||||||
instead of overloading the public method's return contract.
|
instead of overloading the public method's return contract.
|
||||||
@ -154,6 +164,10 @@
|
|||||||
`comfy-kitchen` helpers where they already solve the problem.
|
`comfy-kitchen` helpers where they already solve the problem.
|
||||||
- Use optimized comfy-kitchen ops in places where they improve performance
|
- Use optimized comfy-kitchen ops in places where they improve performance
|
||||||
without changing the expected dtype, device, memory, or interface behavior.
|
without changing the expected dtype, device, memory, or interface behavior.
|
||||||
|
- Prefer ComfyUI's shared optimized kernels and backend dispatchers over
|
||||||
|
handwritten implementations of the same operation. Remove duplicate local
|
||||||
|
kernels and adapt inputs to the shared operation's documented layout while
|
||||||
|
preserving the model's original math and output contract.
|
||||||
- All models should use the optimized attention function selected by ComfyUI.
|
- All models should use the optimized attention function selected by ComfyUI.
|
||||||
Treat optimized backend functions, dispatch helpers, and capability-selected
|
Treat optimized backend functions, dispatch helpers, and capability-selected
|
||||||
callables as opaque. Higher-level code must not inspect function identity,
|
callables as opaque. Higher-level code must not inspect function identity,
|
||||||
@ -176,6 +190,12 @@
|
|||||||
- Model detection code that inspects linear weight shapes should only use the
|
- Model detection code that inspects linear weight shapes should only use the
|
||||||
first dimension. The second dimension may be half the original size for
|
first dimension. The second dimension may be half the original size for
|
||||||
NVFP4 or other 4-bit quantized models.
|
NVFP4 or other 4-bit quantized models.
|
||||||
|
- A model-detection signature must guard every state-dict key it dereferences.
|
||||||
|
Do not partially match a format and then raise an incidental `KeyError` while
|
||||||
|
extracting its configuration.
|
||||||
|
- Order model-detection checks from established or more-specific signatures to
|
||||||
|
newer or broader signatures. Put a broad new detector near the generic
|
||||||
|
fallback when giving it higher precedence could steal another model family.
|
||||||
- Avoid adding `einops` usage in core inference code. Use native torch tensor
|
- Avoid adding `einops` usage in core inference code. Use native torch tensor
|
||||||
ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
|
ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
|
||||||
`unsqueeze`, and `squeeze` instead.
|
`unsqueeze`, and `squeeze` instead.
|
||||||
@ -192,11 +212,23 @@
|
|||||||
methods for scalar or structural calculations.
|
methods for scalar or structural calculations.
|
||||||
- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
|
- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
|
||||||
storage dtype, bias dtype, and original tensor shape metadata.
|
storage dtype, bias dtype, and original tensor shape metadata.
|
||||||
|
- Do not cast the result of an optimized backend operation back to its input
|
||||||
|
dtype unless that backend's documented result contract requires normalization.
|
||||||
|
In particular, trust the selected optimized-attention implementation to honor
|
||||||
|
its dtype contract.
|
||||||
- Keep model-native latent layout handling inside the model or latent-format
|
- Keep model-native latent layout handling inside the model or latent-format
|
||||||
owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
|
owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
|
||||||
dimensions in nodes or other caller-side adapters just to satisfy a model
|
dimensions in nodes or other caller-side adapters just to satisfy a model
|
||||||
forward; the model path should consume and return the native latent shape for
|
forward; the model path should consume and return the native latent shape for
|
||||||
that model family.
|
that model family.
|
||||||
|
- DiT models should accept latent dimensions that are not exact patch-size
|
||||||
|
multiples. Use `comfy.ldm.common_dit.pad_to_patch_size` on every patchified
|
||||||
|
target or reference input, then crop only the target output back to its
|
||||||
|
original dimensions.
|
||||||
|
- Avoid defensive shape and configuration checks that merely replace the clear
|
||||||
|
failure from the tensor operation immediately below them. Add explicit
|
||||||
|
validation only when it provides materially better context at a real boundary
|
||||||
|
or prevents silent incorrect output.
|
||||||
- Assume inputs to the main model forward are already in the compute dtype by
|
- Assume inputs to the main model forward are already in the compute dtype by
|
||||||
default, except integer inputs such as some model timestep tensors. Do not add
|
default, except integer inputs such as some model timestep tensors. Do not add
|
||||||
defensive or convenience casts in model code; it is better for invalid dtype
|
defensive or convenience casts in model code; it is better for invalid dtype
|
||||||
@ -260,6 +292,15 @@
|
|||||||
- Model implementations should add the minimal number of ComfyUI nodes required
|
- Model implementations should add the minimal number of ComfyUI nodes required
|
||||||
to run the model. Reuse existing nodes as much as possible; adapting the model
|
to run the model. Reuse existing nodes as much as possible; adapting the model
|
||||||
to work with existing nodes is strongly preferred over creating new nodes.
|
to work with existing nodes is strongly preferred over creating new nodes.
|
||||||
|
- Use `io.Autogrow` for a variable number of repeated inputs instead of a fixed
|
||||||
|
series of numbered optional sockets. Set its minimum to zero when the model
|
||||||
|
has a valid no-item path, and cap it only when the model has a real limit.
|
||||||
|
- Mark inputs optional when execution has a valid path that does not read them.
|
||||||
|
If one optional input is needed only to process another optional input, do not
|
||||||
|
force users on the path that supplies neither to connect it.
|
||||||
|
- Conditioning nodes should normally output conditioning only. Do not expose
|
||||||
|
input or intermediate images as convenience outputs for downstream sizing or
|
||||||
|
routing; use the existing image path or a dedicated image operation instead.
|
||||||
- Nodes should output only values they own. Do not add pass-through outputs for
|
- Nodes should output only values they own. Do not add pass-through outputs for
|
||||||
workflow convenience unless the node is explicitly an output node. Existing
|
workflow convenience unless the node is explicitly an output node. Existing
|
||||||
models, latents, conditioning, or other inputs should flow directly to the
|
models, latents, conditioning, or other inputs should flow directly to the
|
||||||
|
|||||||
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
|
import comfy.patcher_extension
|
||||||
from comfy.ldm.modules.attention import optimized_attention
|
from comfy.ldm.modules.attention import optimized_attention
|
||||||
import comfy.ldm.common_dit
|
import comfy.ldm.common_dit
|
||||||
|
import comfy.ops
|
||||||
import comfy.quant_ops
|
import comfy.quant_ops
|
||||||
|
|
||||||
|
|
||||||
@ -148,11 +149,29 @@ class Attention(nn.Module):
|
|||||||
x: torch.Tensor,
|
x: torch.Tensor,
|
||||||
context: Optional[torch.Tensor] = None,
|
context: Optional[torch.Tensor] = None,
|
||||||
rope_emb: Optional[torch.Tensor] = None,
|
rope_emb: Optional[torch.Tensor] = None,
|
||||||
|
transformer_options: Optional[dict] = {},
|
||||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||||
q = self.q_proj(x)
|
|
||||||
context = x if context is None else context
|
context = x if context is None else context
|
||||||
k = self.k_proj(context)
|
q_input = x
|
||||||
v = self.v_proj(context)
|
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(
|
q, k, v = map(
|
||||||
lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
|
lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
|
||||||
(q, k, v),
|
(q, k, v),
|
||||||
@ -161,11 +180,16 @@ class Attention(nn.Module):
|
|||||||
def apply_norm_and_rotary_pos_emb(
|
def apply_norm_and_rotary_pos_emb(
|
||||||
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
|
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||||
q = self.q_norm(q)
|
|
||||||
k = self.k_norm(k)
|
|
||||||
v = self.v_norm(v)
|
v = self.v_norm(v)
|
||||||
if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
|
if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
|
||||||
q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rope_emb)
|
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, q, offloadable=True)
|
||||||
|
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, k, offloadable=True)
|
||||||
|
q, k = comfy.quant_ops.ck.rms_rope_split_half(q, k, rope_emb, q_scale, k_scale, self.q_norm.eps)
|
||||||
|
comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
|
||||||
|
comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
|
||||||
|
else:
|
||||||
|
q = self.q_norm(q)
|
||||||
|
k = self.k_norm(k)
|
||||||
return q, k, v
|
return q, k, v
|
||||||
|
|
||||||
q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
|
q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
|
||||||
@ -188,7 +212,7 @@ class Attention(nn.Module):
|
|||||||
x (Tensor): The query tensor of shape [B, Mq, K]
|
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
|
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)
|
return self.compute_attention(q, k, v, transformer_options=transformer_options)
|
||||||
|
|
||||||
|
|
||||||
@ -555,8 +579,14 @@ class Block(nn.Module):
|
|||||||
self.layer_norm_mlp,
|
self.layer_norm_mlp,
|
||||||
scale_mlp_B_T_1_1_D,
|
scale_mlp_B_T_1_1_D,
|
||||||
shift_mlp_B_T_1_1_D,
|
shift_mlp_B_T_1_1_D,
|
||||||
)
|
).to(compute_dtype)
|
||||||
result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.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))
|
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
|
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
|
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}"
|
), 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 = {
|
block_kwargs = {
|
||||||
"rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0),
|
"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,
|
"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,
|
"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
|
# 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:
|
if x_B_T_H_W_D.dtype == torch.float16:
|
||||||
x_B_T_H_W_D = x_B_T_H_W_D.float()
|
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 = block(
|
||||||
x_B_T_H_W_D,
|
x_B_T_H_W_D,
|
||||||
t_embedding_B_T_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.seedvr.model
|
||||||
import comfy.ldm.boogu.model
|
import comfy.ldm.boogu.model
|
||||||
import comfy.ldm.qwen_image.model
|
import comfy.ldm.qwen_image.model
|
||||||
|
import comfy.ldm.joyimage.model
|
||||||
import comfy.ldm.ideogram4.model
|
import comfy.ldm.ideogram4.model
|
||||||
import comfy.ldm.krea2.model
|
import comfy.ldm.krea2.model
|
||||||
import comfy.ldm.kandinsky5.model
|
import comfy.ldm.kandinsky5.model
|
||||||
@ -2276,6 +2277,28 @@ class QwenImage(BaseModel):
|
|||||||
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
class JoyImage(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel)
|
||||||
|
self.memory_usage_factor_conds = ("ref_latents",)
|
||||||
|
|
||||||
|
def extra_conds(self, **kwargs):
|
||||||
|
out = super().extra_conds(**kwargs)
|
||||||
|
cross_attn = kwargs.get("cross_attn", None)
|
||||||
|
if cross_attn is not None:
|
||||||
|
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
|
||||||
|
ref_latents = kwargs.get("reference_latents", None)
|
||||||
|
if ref_latents is not None:
|
||||||
|
out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
|
||||||
|
return out
|
||||||
|
|
||||||
|
def extra_conds_shapes(self, **kwargs):
|
||||||
|
out = {}
|
||||||
|
ref_latents = kwargs.get("reference_latents", None)
|
||||||
|
if ref_latents is not None:
|
||||||
|
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||||
|
return out
|
||||||
|
|
||||||
class Ideogram4(BaseModel):
|
class Ideogram4(BaseModel):
|
||||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
|
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
|
||||||
|
|||||||
@ -1058,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
|||||||
dit_config["image_model"] = "SAM31"
|
dit_config["image_model"] = "SAM31"
|
||||||
return dit_config
|
return dit_config
|
||||||
|
|
||||||
|
if (
|
||||||
|
'{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and '{}img_in.weight'.format(key_prefix) in state_dict_keys
|
||||||
|
and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5
|
||||||
|
):
|
||||||
|
img_in = state_dict['{}img_in.weight'.format(key_prefix)]
|
||||||
|
head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
|
||||||
|
return {
|
||||||
|
"image_model": "joyimage",
|
||||||
|
"in_channels": img_in.shape[1],
|
||||||
|
"hidden_size": img_in.shape[0],
|
||||||
|
"patch_size": list(img_in.shape[2:]),
|
||||||
|
"num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'),
|
||||||
|
"num_attention_heads": img_in.shape[0] // head_dim,
|
||||||
|
"text_dim": 4096,
|
||||||
|
}
|
||||||
|
|
||||||
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|||||||
@ -76,6 +76,7 @@ import comfy.text_encoders.gemma4
|
|||||||
import comfy.text_encoders.cogvideo
|
import comfy.text_encoders.cogvideo
|
||||||
import comfy.text_encoders.sa3
|
import comfy.text_encoders.sa3
|
||||||
import comfy.text_encoders.gpt_oss
|
import comfy.text_encoders.gpt_oss
|
||||||
|
import comfy.text_encoders.joyimage
|
||||||
|
|
||||||
import comfy.model_patcher
|
import comfy.model_patcher
|
||||||
import comfy.lora
|
import comfy.lora
|
||||||
@ -1377,6 +1378,7 @@ class CLIPType(Enum):
|
|||||||
IDEOGRAM4 = 30
|
IDEOGRAM4 = 30
|
||||||
BOOGU = 31
|
BOOGU = 31
|
||||||
KREA2 = 32
|
KREA2 = 32
|
||||||
|
JOYIMAGE = 33
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@ -1706,6 +1708,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
|
|||||||
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
|
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
|
||||||
clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
|
clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
|
||||||
clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
|
clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
|
||||||
|
elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B: # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx.
|
||||||
|
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
|
||||||
|
clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data))
|
||||||
|
clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer
|
||||||
elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
|
elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
|
||||||
klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
|
klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
|
||||||
clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
|
clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
|
||||||
|
|||||||
@ -27,6 +27,7 @@ import comfy.text_encoders.z_image
|
|||||||
import comfy.text_encoders.ideogram4
|
import comfy.text_encoders.ideogram4
|
||||||
import comfy.text_encoders.boogu
|
import comfy.text_encoders.boogu
|
||||||
import comfy.text_encoders.krea2
|
import comfy.text_encoders.krea2
|
||||||
|
import comfy.text_encoders.joyimage
|
||||||
import comfy.text_encoders.anima
|
import comfy.text_encoders.anima
|
||||||
import comfy.text_encoders.ace15
|
import comfy.text_encoders.ace15
|
||||||
import comfy.text_encoders.longcat_image
|
import comfy.text_encoders.longcat_image
|
||||||
@ -1911,6 +1912,38 @@ class QwenImage(supported_models_base.BASE):
|
|||||||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
||||||
return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
|
return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
|
||||||
|
|
||||||
|
class JoyImage(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"image_model": "joyimage",
|
||||||
|
}
|
||||||
|
|
||||||
|
sampling_settings = {
|
||||||
|
"multiplier": 1000,
|
||||||
|
"shift": 1.5,
|
||||||
|
}
|
||||||
|
|
||||||
|
memory_usage_factor = 1.8
|
||||||
|
|
||||||
|
unet_extra_config = {
|
||||||
|
"theta": 10000,
|
||||||
|
"rope_dim_list": [16, 56, 56],
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.Wan21
|
||||||
|
|
||||||
|
supported_inference_dtypes = [torch.bfloat16, torch.float32]
|
||||||
|
|
||||||
|
vae_key_prefix = ["vae."]
|
||||||
|
text_encoder_key_prefix = ["text_encoders."]
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
return model_base.JoyImage(self, device=device)
|
||||||
|
|
||||||
|
def clip_target(self, state_dict={}):
|
||||||
|
pref = self.text_encoder_key_prefix[0]
|
||||||
|
qwen3vl_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl.transformer.".format(pref))
|
||||||
|
return supported_models_base.ClipTarget(comfy.text_encoders.joyimage.JoyImageTokenizer, comfy.text_encoders.joyimage.te(**qwen3vl_detect))
|
||||||
|
|
||||||
class HunyuanImage21(HunyuanVideo):
|
class HunyuanImage21(HunyuanVideo):
|
||||||
unet_config = {
|
unet_config = {
|
||||||
"image_model": "hunyuan_video",
|
"image_model": "hunyuan_video",
|
||||||
@ -2389,6 +2422,7 @@ models = [
|
|||||||
Omnigen2,
|
Omnigen2,
|
||||||
Boogu,
|
Boogu,
|
||||||
QwenImage,
|
QwenImage,
|
||||||
|
JoyImage,
|
||||||
Ideogram4,
|
Ideogram4,
|
||||||
Krea2,
|
Krea2,
|
||||||
Flux2,
|
Flux2,
|
||||||
|
|||||||
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,
|
merge_size: int = 2,
|
||||||
image_mean: list = None,
|
image_mean: list = None,
|
||||||
image_std: list = None,
|
image_std: list = None,
|
||||||
|
interpolation: str = "bilinear",
|
||||||
):
|
):
|
||||||
if image_mean is None:
|
if image_mean is None:
|
||||||
image_mean = [0.48145466, 0.4578275, 0.40821073]
|
image_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||||
@ -47,10 +48,9 @@ def process_qwen2vl_images(
|
|||||||
img_resized = F.interpolate(
|
img_resized = F.interpolate(
|
||||||
img.unsqueeze(0),
|
img.unsqueeze(0),
|
||||||
size=(h_bar, w_bar),
|
size=(h_bar, w_bar),
|
||||||
mode='bilinear',
|
mode=interpolation,
|
||||||
align_corners=False
|
align_corners=False
|
||||||
).squeeze(0)
|
).squeeze(0)
|
||||||
|
|
||||||
normalized = img_resized.clone()
|
normalized = img_resized.clone()
|
||||||
for c in range(3):
|
for c in range(3):
|
||||||
normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]
|
normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]
|
||||||
|
|||||||
@ -1,5 +1,6 @@
|
|||||||
from av.container import InputContainer
|
from av.container import InputContainer
|
||||||
from av.subtitles.stream import SubtitleStream
|
from av.subtitles.stream import SubtitleStream
|
||||||
|
from av.video.reformatter import ColorRange
|
||||||
from fractions import Fraction
|
from fractions import Fraction
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from .._input import AudioInput, VideoInput
|
from .._input import AudioInput, VideoInput
|
||||||
@ -9,6 +10,7 @@ import itertools
|
|||||||
import json
|
import json
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import math
|
import math
|
||||||
|
import os
|
||||||
import torch
|
import torch
|
||||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||||
import logging
|
import logging
|
||||||
@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
|
|||||||
return max(component.bits for component in stream.format.components)
|
return max(component.bits for component in stream.format.components)
|
||||||
|
|
||||||
|
|
||||||
|
def last_decodable_audio_stream(container: InputContainer):
|
||||||
|
"""Streams FFmpeg has no decoder for have no codec context, and decoding their
|
||||||
|
packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
|
||||||
|
stream = next(
|
||||||
|
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
if stream is None and len(container.streams.audio):
|
||||||
|
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||||
|
return stream
|
||||||
|
|
||||||
|
|
||||||
|
def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
|
||||||
|
"""Containers probed only up to a window (mpegts) leave audio codec parameters unset when
|
||||||
|
audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
|
||||||
|
Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
|
||||||
|
for i, packet in enumerate(container.demux(audio_stream)):
|
||||||
|
try:
|
||||||
|
frames = packet.decode()
|
||||||
|
except av.error.FFmpegError:
|
||||||
|
frames = ()
|
||||||
|
if frames:
|
||||||
|
return frames[0].sample_rate, frames[0].layout.nb_channels
|
||||||
|
if i >= max_packets:
|
||||||
|
break
|
||||||
|
return 0, 0
|
||||||
|
|
||||||
|
|
||||||
|
def write_output_metadata(container: InputContainer, output, metadata: dict | None):
|
||||||
|
"""Copy the source container's metadata, then overlay the caller's tags."""
|
||||||
|
for key, value in container.metadata.items():
|
||||||
|
if metadata is None or key not in metadata:
|
||||||
|
output.metadata[key] = value
|
||||||
|
if metadata is not None:
|
||||||
|
for key, value in metadata.items():
|
||||||
|
output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
|
||||||
|
|
||||||
|
|
||||||
|
def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
|
||||||
|
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
||||||
|
raise ValueError("Only MP4 format is supported for now")
|
||||||
|
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
||||||
|
raise ValueError("Only H264 codec is supported for now")
|
||||||
|
open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
|
||||||
|
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
||||||
|
open_kwargs["format"] = format.value
|
||||||
|
elif isinstance(path, io.BytesIO):
|
||||||
|
open_kwargs["format"] = "mp4" # no file extension to infer the format from
|
||||||
|
return open_kwargs
|
||||||
|
|
||||||
|
|
||||||
class VideoFromFile(VideoInput):
|
class VideoFromFile(VideoInput):
|
||||||
"""
|
"""
|
||||||
Class representing video input from a file.
|
Class representing video input from a file.
|
||||||
@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
|
|||||||
return estimated_frames
|
return estimated_frames
|
||||||
|
|
||||||
# 3. Last resort: decode frames and count them (streaming)
|
# 3. Last resort: decode frames and count them (streaming)
|
||||||
if self.__start_time < 0:
|
start_time, duration = self.get_active_trim_window()
|
||||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
|
||||||
else:
|
|
||||||
start_time = self.__start_time
|
|
||||||
frame_count = 1
|
frame_count = 1
|
||||||
start_pts = int(start_time / video_stream.time_base)
|
start_pts = int(start_time / video_stream.time_base)
|
||||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||||
container.seek(start_pts, stream=video_stream)
|
container.seek(start_pts, stream=video_stream)
|
||||||
frame_iterator = (
|
frame_iterator = (
|
||||||
container.decode(video_stream)
|
container.decode(video_stream)
|
||||||
@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
|
|||||||
|
|
||||||
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
||||||
video_stream = self._get_first_video_stream(container)
|
video_stream = self._get_first_video_stream(container)
|
||||||
if self.__start_time < 0:
|
start_time, duration = self.get_active_trim_window()
|
||||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
|
||||||
else:
|
|
||||||
start_time = self.__start_time
|
|
||||||
|
|
||||||
# Get video frames
|
# Get video frames
|
||||||
frames = []
|
frames = []
|
||||||
audio_frames = []
|
audio_frames = []
|
||||||
alphas = None
|
alphas = None
|
||||||
start_pts = int(start_time / video_stream.time_base)
|
start_pts = int(start_time / video_stream.time_base)
|
||||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||||
|
|
||||||
if start_pts != 0:
|
if start_pts != 0:
|
||||||
container.seek(start_pts, stream=video_stream)
|
container.seek(start_pts, stream=video_stream)
|
||||||
@ -281,18 +328,11 @@ class VideoFromFile(VideoInput):
|
|||||||
video_done = False
|
video_done = False
|
||||||
audio_done = True
|
audio_done = True
|
||||||
|
|
||||||
# Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
|
audio_stream = last_decodable_audio_stream(container)
|
||||||
# and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
|
|
||||||
audio_stream = next(
|
|
||||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
if audio_stream is not None:
|
if audio_stream is not None:
|
||||||
streams += [audio_stream]
|
streams += [audio_stream]
|
||||||
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
||||||
audio_done = False
|
audio_done = False
|
||||||
elif len(container.streams.audio):
|
|
||||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
|
||||||
|
|
||||||
for packet in container.demux(*streams):
|
for packet in container.demux(*streams):
|
||||||
if video_done and audio_done:
|
if video_done and audio_done:
|
||||||
@ -305,7 +345,7 @@ class VideoFromFile(VideoInput):
|
|||||||
for frame in packet.decode():
|
for frame in packet.decode():
|
||||||
if frame.pts < start_pts:
|
if frame.pts < start_pts:
|
||||||
continue
|
continue
|
||||||
if self.__duration and frame.pts >= end_pts:
|
if duration and frame.pts >= end_pts:
|
||||||
video_done = True
|
video_done = True
|
||||||
break
|
break
|
||||||
|
|
||||||
@ -372,7 +412,7 @@ class VideoFromFile(VideoInput):
|
|||||||
map(resampler.resample, packet.decode())
|
map(resampler.resample, packet.decode())
|
||||||
)
|
)
|
||||||
for frame in aframes:
|
for frame in aframes:
|
||||||
if self.__duration and frame.time > start_time + self.__duration:
|
if duration and frame.time > start_time + duration:
|
||||||
audio_done = True
|
audio_done = True
|
||||||
break
|
break
|
||||||
|
|
||||||
@ -394,8 +434,8 @@ class VideoFromFile(VideoInput):
|
|||||||
|
|
||||||
if len(audio_frames) > 0:
|
if len(audio_frames) > 0:
|
||||||
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
||||||
if self.__duration:
|
if duration:
|
||||||
audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
|
audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
|
||||||
|
|
||||||
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
|
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
|
||||||
audio = AudioInput({
|
audio = AudioInput({
|
||||||
@ -441,28 +481,14 @@ class VideoFromFile(VideoInput):
|
|||||||
if not reuse_streams:
|
if not reuse_streams:
|
||||||
if bit_depth is None:
|
if bit_depth is None:
|
||||||
bit_depth = source_bit_depth
|
bit_depth = source_bit_depth
|
||||||
components = self.get_components_internal(container)
|
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
|
||||||
video = VideoFromComponents(components)
|
|
||||||
return video.save_to(
|
|
||||||
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
|
|
||||||
)
|
|
||||||
|
|
||||||
streams = container.streams
|
streams = container.streams
|
||||||
|
|
||||||
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
||||||
with av.open(path, **open_kwargs) as output_container:
|
with av.open(path, **open_kwargs) as output_container:
|
||||||
# Copy over the original metadata
|
# Add metadata before writing any streams
|
||||||
for key, value in container.metadata.items():
|
write_output_metadata(container, output_container, metadata)
|
||||||
if metadata is None or key not in metadata:
|
|
||||||
output_container.metadata[key] = value
|
|
||||||
|
|
||||||
# Add our new metadata
|
|
||||||
if metadata is not None:
|
|
||||||
for key, value in metadata.items():
|
|
||||||
if isinstance(value, str):
|
|
||||||
output_container.metadata[key] = value
|
|
||||||
else:
|
|
||||||
output_container.metadata[key] = json.dumps(value)
|
|
||||||
|
|
||||||
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
||||||
stream_map = {}
|
stream_map = {}
|
||||||
@ -480,6 +506,282 @@ class VideoFromFile(VideoInput):
|
|||||||
packet.stream = stream_map[packet.stream]
|
packet.stream = stream_map[packet.stream]
|
||||||
output_container.mux(packet)
|
output_container.mux(packet)
|
||||||
|
|
||||||
|
def _save_transcoded(
|
||||||
|
self,
|
||||||
|
container: InputContainer,
|
||||||
|
path: str | io.BytesIO,
|
||||||
|
format: VideoContainer,
|
||||||
|
codec: VideoCodec,
|
||||||
|
metadata: dict | None,
|
||||||
|
bit_depth: int,
|
||||||
|
):
|
||||||
|
"""Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
|
||||||
|
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||||
|
video_stream = self._get_first_video_stream(container)
|
||||||
|
start_time, duration = self.get_active_trim_window()
|
||||||
|
start_pts = int(start_time / video_stream.time_base)
|
||||||
|
end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
|
||||||
|
stream_end_pts = None
|
||||||
|
if video_stream.duration is not None:
|
||||||
|
stream_end_pts = (video_stream.start_time or 0) + video_stream.duration
|
||||||
|
output_end_pts = end_pts
|
||||||
|
if stream_end_pts is not None and (output_end_pts is None or stream_end_pts < output_end_pts):
|
||||||
|
output_end_pts = stream_end_pts
|
||||||
|
if start_pts != 0:
|
||||||
|
container.seek(start_pts, stream=video_stream)
|
||||||
|
|
||||||
|
audio_stream = last_decodable_audio_stream(container)
|
||||||
|
pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p"
|
||||||
|
rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||||
|
|
||||||
|
resampler = None
|
||||||
|
sample_rate = 0
|
||||||
|
audio_time_base = None
|
||||||
|
duration_cap = None
|
||||||
|
if audio_stream is not None:
|
||||||
|
sample_rate = audio_stream.codec_context.sample_rate
|
||||||
|
channels = audio_stream.codec_context.channels
|
||||||
|
if not sample_rate:
|
||||||
|
sample_rate, channels = probe_audio_params(container, audio_stream)
|
||||||
|
container.seek(start_pts, stream=video_stream)
|
||||||
|
if sample_rate:
|
||||||
|
audio_stream.codec_context.flush_buffers()
|
||||||
|
else:
|
||||||
|
logging.warning("Audio stream parameters could not be determined; ignoring audio.")
|
||||||
|
audio_stream = None
|
||||||
|
if audio_stream is not None:
|
||||||
|
audio_time_base = Fraction(1, sample_rate)
|
||||||
|
layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
|
||||||
|
resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
|
||||||
|
if duration:
|
||||||
|
duration_cap = math.ceil(duration * sample_rate)
|
||||||
|
|
||||||
|
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
|
||||||
|
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
|
||||||
|
video_done = False
|
||||||
|
audio_done = audio_stream is None
|
||||||
|
video_pts_offset = None
|
||||||
|
last_video_pts = None
|
||||||
|
last_video_end = None
|
||||||
|
# rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise
|
||||||
|
video_frame_durations = {}
|
||||||
|
source_size = None
|
||||||
|
rotation_k = 0
|
||||||
|
rotation_filter = None
|
||||||
|
audio_started = False
|
||||||
|
samples_written = 0
|
||||||
|
pending_audio = []
|
||||||
|
# The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims),
|
||||||
|
# and never seeking back keeps webm/mkv leading audio intact.
|
||||||
|
output = None
|
||||||
|
out_video = None
|
||||||
|
out_audio = None
|
||||||
|
|
||||||
|
def audio_frame_from_ndarray(nd_planar):
|
||||||
|
frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout)
|
||||||
|
frame.sample_rate = sample_rate
|
||||||
|
return frame
|
||||||
|
|
||||||
|
def drain_audio(final=False):
|
||||||
|
# Audio may cover the pts span of the video written so far, capped by the requested duration
|
||||||
|
nonlocal samples_written, audio_done
|
||||||
|
if last_video_end is None:
|
||||||
|
cap = 0
|
||||||
|
else:
|
||||||
|
cap = math.ceil(last_video_end * video_stream.time_base * sample_rate)
|
||||||
|
if duration_cap is not None:
|
||||||
|
cap = min(cap, duration_cap)
|
||||||
|
while pending_audio and not audio_done:
|
||||||
|
frame = pending_audio[0]
|
||||||
|
if samples_written + frame.samples <= cap:
|
||||||
|
frame.pts = samples_written
|
||||||
|
frame.time_base = audio_time_base
|
||||||
|
output.mux(out_audio.encode(frame))
|
||||||
|
samples_written += frame.samples
|
||||||
|
pending_audio.pop(0)
|
||||||
|
continue
|
||||||
|
if final:
|
||||||
|
keep = frame.to_ndarray()[..., :cap - samples_written]
|
||||||
|
if keep.shape[-1] > 0:
|
||||||
|
tail = audio_frame_from_ndarray(keep)
|
||||||
|
tail.pts = samples_written
|
||||||
|
tail.time_base = audio_time_base
|
||||||
|
output.mux(out_audio.encode(tail))
|
||||||
|
samples_written += keep.shape[-1]
|
||||||
|
pending_audio.clear()
|
||||||
|
break
|
||||||
|
if duration_cap is not None and samples_written >= duration_cap:
|
||||||
|
audio_done = True
|
||||||
|
return cap
|
||||||
|
|
||||||
|
try:
|
||||||
|
for packet in container.demux(*streams):
|
||||||
|
if video_done and audio_done:
|
||||||
|
break
|
||||||
|
|
||||||
|
if packet.stream == video_stream and not video_done:
|
||||||
|
try:
|
||||||
|
frames = packet.decode()
|
||||||
|
except av.error.InvalidDataError:
|
||||||
|
logging.info("pyav decode error")
|
||||||
|
continue
|
||||||
|
for frame in frames:
|
||||||
|
if frame.pts is not None and frame.pts < start_pts:
|
||||||
|
continue
|
||||||
|
if end_pts is not None and frame.pts is not None and frame.pts >= end_pts:
|
||||||
|
video_done = True
|
||||||
|
if last_video_pts is not None:
|
||||||
|
# the source continues past the window: hold the last kept frame to the window end
|
||||||
|
end_offset = video_pts_offset if video_pts_offset is not None else start_pts
|
||||||
|
last_video_end = max(last_video_end, end_pts - end_offset)
|
||||||
|
break
|
||||||
|
# the source's true display duration of this frame; average_rate is not a
|
||||||
|
# frame duration (sparse/VFR sources), so it is only the fallback
|
||||||
|
frame_duration = frame.duration if frame.duration else pts_step
|
||||||
|
if end_pts is not None and frame.pts is not None:
|
||||||
|
frame_duration = min(frame_duration, end_pts - frame.pts)
|
||||||
|
if output is None:
|
||||||
|
rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
|
||||||
|
if rotation_k % 2:
|
||||||
|
out_width, out_height = frame.height, frame.width
|
||||||
|
else:
|
||||||
|
out_width, out_height = frame.width, frame.height
|
||||||
|
if out_width % 2 or out_height % 2:
|
||||||
|
raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
|
||||||
|
source_size = (frame.width, frame.height)
|
||||||
|
output = av.open(path, **open_kwargs)
|
||||||
|
# Add metadata before writing any streams
|
||||||
|
write_output_metadata(container, output, metadata)
|
||||||
|
out_video = output.add_stream("h264", rate=rate)
|
||||||
|
# no B-frames: reordering makes mp4 sample durations follow decode order,
|
||||||
|
# so irregular-VFR spans and trim windows land wrong
|
||||||
|
out_video.codec_context.max_b_frames = 0
|
||||||
|
out_video.width = out_width
|
||||||
|
out_video.height = out_height
|
||||||
|
out_video.pix_fmt = pix_fmt
|
||||||
|
# source pts pass through (rebased to 0), so variable frame rate survives
|
||||||
|
out_video.codec_context.time_base = video_stream.time_base
|
||||||
|
if audio_stream is not None:
|
||||||
|
out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
|
||||||
|
if (frame.width, frame.height) != source_size:
|
||||||
|
# encoding would silently rescale the new geometry into the old one
|
||||||
|
raise ValueError(
|
||||||
|
f"Video resolution changes mid-stream "
|
||||||
|
f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode"
|
||||||
|
)
|
||||||
|
if rotation_k:
|
||||||
|
if rotation_filter is None:
|
||||||
|
g = av.filter.Graph()
|
||||||
|
g_src = g.add_buffer(width=frame.width, height=frame.height,
|
||||||
|
format=frame.format.name, time_base=video_stream.time_base)
|
||||||
|
tail = g_src
|
||||||
|
for filter_name, filter_args in {1: [("transpose", "cclock")],
|
||||||
|
2: [("hflip", None), ("vflip", None)],
|
||||||
|
3: [("transpose", "clock")]}[rotation_k]:
|
||||||
|
step = g.add(filter_name, filter_args)
|
||||||
|
tail.link_to(step)
|
||||||
|
tail = step
|
||||||
|
g_sink = g.add("buffersink")
|
||||||
|
tail.link_to(g_sink)
|
||||||
|
g.configure()
|
||||||
|
rotation_filter = (g_src, g_sink)
|
||||||
|
rotation_filter[0].push(frame)
|
||||||
|
frame = rotation_filter[1].pull()
|
||||||
|
if frame.color_range == ColorRange.JPEG:
|
||||||
|
# compress full-range sources (yuvj/MJPEG) to limited range
|
||||||
|
frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
|
||||||
|
else:
|
||||||
|
frame = frame.reformat(format=pix_fmt)
|
||||||
|
frame_output_end = None
|
||||||
|
if frame.pts is not None:
|
||||||
|
if video_pts_offset is None:
|
||||||
|
video_pts_offset = frame.pts
|
||||||
|
frame.pts -= video_pts_offset
|
||||||
|
if output_end_pts is not None:
|
||||||
|
frame_output_end = output_end_pts - video_pts_offset
|
||||||
|
if frame.pts + frame_duration > frame_output_end:
|
||||||
|
clamped_pts = frame_output_end - frame_duration
|
||||||
|
if clamped_pts >= 0 and (last_video_pts is None or clamped_pts > last_video_pts):
|
||||||
|
frame.pts = min(frame.pts, clamped_pts)
|
||||||
|
elif frame.pts < frame_output_end:
|
||||||
|
frame_duration = frame_output_end - frame.pts
|
||||||
|
else:
|
||||||
|
continue
|
||||||
|
if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts):
|
||||||
|
# broken sources emit missing/backward timestamps mid-stream, which the
|
||||||
|
# muxer rejects; nudge them forward by one nominal frame interval
|
||||||
|
frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step
|
||||||
|
if frame_output_end is not None and frame.pts + frame_duration > frame_output_end:
|
||||||
|
if frame.pts >= frame_output_end:
|
||||||
|
continue
|
||||||
|
frame_duration = frame_output_end - frame.pts
|
||||||
|
last_video_pts = frame.pts
|
||||||
|
last_video_end = frame.pts + frame_duration
|
||||||
|
video_frame_durations[frame.pts] = frame_duration
|
||||||
|
# the decoded pict_type would force x264's frame types (intra-only
|
||||||
|
# sources like MJPEG/ProRes would come out all-keyframe)
|
||||||
|
frame.pict_type = 0
|
||||||
|
for out_packet in out_video.encode(frame):
|
||||||
|
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||||
|
output.mux(out_packet)
|
||||||
|
drain_audio()
|
||||||
|
|
||||||
|
elif packet.stream == audio_stream and not audio_done:
|
||||||
|
for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
|
||||||
|
frame_start = None
|
||||||
|
if resampled.pts is not None:
|
||||||
|
# passthrough frames keep the source stream's time base
|
||||||
|
tb = resampled.time_base if resampled.time_base else audio_time_base
|
||||||
|
frame_start = float(resampled.pts * tb)
|
||||||
|
if duration and not audio_started and frame_start >= start_time + duration:
|
||||||
|
audio_done = True
|
||||||
|
break
|
||||||
|
if not audio_started:
|
||||||
|
if frame_start is None:
|
||||||
|
frame_start = 0.0
|
||||||
|
to_skip = max(0, int((start_time - frame_start) * sample_rate))
|
||||||
|
if to_skip >= resampled.samples:
|
||||||
|
continue
|
||||||
|
audio_started = True
|
||||||
|
if duration and frame_start > start_time:
|
||||||
|
duration_cap = min(duration_cap, math.ceil((start_time + duration - frame_start) * sample_rate))
|
||||||
|
if to_skip:
|
||||||
|
pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
|
||||||
|
continue
|
||||||
|
pending_audio.append(resampled)
|
||||||
|
if video_done:
|
||||||
|
# the video window is complete so the cap is final, but containers
|
||||||
|
# that interleave audio behind video (fragmented mp4) still owe most
|
||||||
|
# of it: stop only once the demuxed audio covers the cap
|
||||||
|
cap = drain_audio()
|
||||||
|
if pending_audio or samples_written >= cap:
|
||||||
|
drain_audio(final=True)
|
||||||
|
audio_done = True
|
||||||
|
break
|
||||||
|
|
||||||
|
if output is None:
|
||||||
|
raise ValueError(f"No decodable video frames found in file '{self.__file}'")
|
||||||
|
if out_audio is not None and not audio_done:
|
||||||
|
drain_audio(final=True)
|
||||||
|
window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
|
||||||
|
for out_packet in out_video.encode(None):
|
||||||
|
duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||||
|
if out_packet.pts == last_video_pts:
|
||||||
|
duration = max(duration, window_fill)
|
||||||
|
out_packet.duration = duration
|
||||||
|
output.mux(out_packet)
|
||||||
|
if out_audio is not None:
|
||||||
|
output.mux(out_audio.encode(None))
|
||||||
|
except BaseException:
|
||||||
|
if output is not None:
|
||||||
|
output.close()
|
||||||
|
if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
|
||||||
|
os.remove(path)
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
if output is not None:
|
||||||
|
output.close()
|
||||||
|
|
||||||
def _get_first_video_stream(self, container: InputContainer):
|
def _get_first_video_stream(self, container: InputContainer):
|
||||||
if len(container.streams.video):
|
if len(container.streams.video):
|
||||||
return container.streams.video[0]
|
return container.streams.video[0]
|
||||||
@ -527,22 +829,12 @@ class VideoFromComponents(VideoInput):
|
|||||||
bit_depth: int | None = None,
|
bit_depth: int | None = None,
|
||||||
):
|
):
|
||||||
"""Save the video to a file path or BytesIO buffer."""
|
"""Save the video to a file path or BytesIO buffer."""
|
||||||
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||||
raise ValueError("Only MP4 format is supported for now")
|
|
||||||
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
|
||||||
raise ValueError("Only H264 codec is supported for now")
|
|
||||||
# None means "use the depth this video was created with" (CreateVideo's choice).
|
# None means "use the depth this video was created with" (CreateVideo's choice).
|
||||||
if bit_depth is None:
|
if bit_depth is None:
|
||||||
bit_depth = self.__bit_depth
|
bit_depth = self.__bit_depth
|
||||||
is_10bit = bit_depth >= 10
|
is_10bit = bit_depth >= 10
|
||||||
extra_kwargs = {}
|
with av.open(path, **open_kwargs) as output:
|
||||||
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
|
||||||
extra_kwargs["format"] = format.value
|
|
||||||
elif isinstance(path, io.BytesIO):
|
|
||||||
# BytesIO has no file extension, so av.open can't infer the format.
|
|
||||||
# Default to mp4 since that's the only supported format anyway.
|
|
||||||
extra_kwargs["format"] = "mp4"
|
|
||||||
with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
|
|
||||||
# Add metadata before writing any streams
|
# Add metadata before writing any streams
|
||||||
if metadata is not None:
|
if metadata is not None:
|
||||||
for key, value in metadata.items():
|
for key, value in metadata.items():
|
||||||
|
|||||||
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)
|
parallel_tool_calls: bool | None = Field(True)
|
||||||
status: str | None = Field(
|
status: str | None = Field(
|
||||||
None,
|
None,
|
||||||
description="One of `completed`, `failed`, `in_progress`, or `incomplete`.",
|
description="One of `completed`, `failed`, `in_progress`, `incomplete`, `queued`, or `cancelled`.",
|
||||||
)
|
)
|
||||||
usage: ResponseUsage | None = Field(None)
|
usage: ResponseUsage | None = Field(None)
|
||||||
|
|
||||||
|
|||||||
@ -274,6 +274,10 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||||||
input_tokens_price = 0.25
|
input_tokens_price = 0.25
|
||||||
output_text_tokens_price = 1.50
|
output_text_tokens_price = 1.50
|
||||||
output_image_tokens_price = 0.0
|
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"):
|
elif response.modelVersion in ("gemini-3-pro-image-preview", "gemini-3-pro-image"):
|
||||||
input_tokens_price = 2
|
input_tokens_price = 2
|
||||||
output_text_tokens_price = 12.0
|
output_text_tokens_price = 12.0
|
||||||
@ -619,11 +623,12 @@ class GeminiNode(IO.ComfyNode):
|
|||||||
|
|
||||||
GEMINI_V2_MODELS: dict[str, str] = {
|
GEMINI_V2_MODELS: dict[str, str] = {
|
||||||
"Gemini 3.1 Pro": "gemini-3.1-pro-preview",
|
"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",
|
"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)."""
|
"""Per-model inputs revealed by the model DynamicCombo (shared media + sampling controls)."""
|
||||||
return [
|
return [
|
||||||
IO.Autogrow.Input(
|
IO.Autogrow.Input(
|
||||||
@ -661,7 +666,7 @@ def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
|
|||||||
),
|
),
|
||||||
IO.Combo.Input(
|
IO.Combo.Input(
|
||||||
"thinking_level",
|
"thinking_level",
|
||||||
options=["LOW", "HIGH"],
|
options=thinking_options or ["LOW", "HIGH"],
|
||||||
default=thinking_default,
|
default=thinking_default,
|
||||||
tooltip="How hard the model reasons internally before answering. "
|
tooltip="How hard the model reasons internally before answering. "
|
||||||
"HIGH improves quality on difficult tasks but costs more (thinking) tokens and is slower.",
|
"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(
|
IO.DynamicCombo.Input(
|
||||||
"model",
|
"model",
|
||||||
options=[
|
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 Pro", _gemini_text_model_inputs("HIGH")),
|
||||||
IO.DynamicCombo.Option("Gemini 3.1 Flash-Lite", _gemini_text_model_inputs("LOW")),
|
IO.DynamicCombo.Option("Gemini 3.1 Flash-Lite", _gemini_text_model_inputs("LOW")),
|
||||||
],
|
],
|
||||||
@ -759,7 +768,13 @@ class GeminiNodeV2(IO.ComfyNode):
|
|||||||
"type": "list_usd",
|
"type": "list_usd",
|
||||||
"usd": [0.00025, 0.0015],
|
"usd": [0.00025, 0.0015],
|
||||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
"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",
|
"type": "list_usd",
|
||||||
"usd": [0.002, 0.012],
|
"usd": [0.002, 0.012],
|
||||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
"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):
|
class SupportedOpenAIModel(str, Enum):
|
||||||
|
gpt_5_6_sol = "gpt-5.6-sol"
|
||||||
|
gpt_5_6_terra = "gpt-5.6-terra"
|
||||||
|
gpt_5_6_luna = "gpt-5.6-luna"
|
||||||
gpt_5_5_pro = "gpt-5.5-pro"
|
gpt_5_5_pro = "gpt-5.5-pro"
|
||||||
gpt_5_5 = "gpt-5.5"
|
gpt_5_5 = "gpt-5.5"
|
||||||
gpt_5 = "gpt-5"
|
gpt_5 = "gpt-5"
|
||||||
@ -1063,6 +1066,21 @@ class OpenAIChatNode(IO.ComfyNode):
|
|||||||
"usd": [0.002, 0.008],
|
"usd": [0.002, 0.008],
|
||||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
}
|
}
|
||||||
|
: $contains($m, "gpt-5.6-terra") ? {
|
||||||
|
"type": "list_usd",
|
||||||
|
"usd": [0.0025, 0.015],
|
||||||
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
|
}
|
||||||
|
: $contains($m, "gpt-5.6-luna") ? {
|
||||||
|
"type": "list_usd",
|
||||||
|
"usd": [0.001, 0.006],
|
||||||
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
|
}
|
||||||
|
: $contains($m, "gpt-5.6") ? {
|
||||||
|
"type": "list_usd",
|
||||||
|
"usd": [0.005, 0.03],
|
||||||
|
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||||
|
}
|
||||||
: $contains($m, "gpt-5.5-pro") ? {
|
: $contains($m, "gpt-5.5-pro") ? {
|
||||||
"type": "list_usd",
|
"type": "list_usd",
|
||||||
"usd": [0.03, 0.18],
|
"usd": [0.03, 0.18],
|
||||||
|
|||||||
@ -15,6 +15,7 @@ from comfy.comfy_api_env import normalize_comfy_api_base
|
|||||||
from comfy.deploy_environment import get_deploy_environment
|
from comfy.deploy_environment import get_deploy_environment
|
||||||
from comfy.model_management import processing_interrupted
|
from comfy.model_management import processing_interrupted
|
||||||
from comfy_api.latest import IO
|
from comfy_api.latest import IO
|
||||||
|
from comfy_execution.utils import get_executing_context
|
||||||
from comfyui_version import __version__ as comfyui_version
|
from comfyui_version import __version__ as comfyui_version
|
||||||
|
|
||||||
from .common_exceptions import ProcessingInterrupted
|
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
|
relative/cloud URLs resolved against ``default_base_url()``; because the result
|
||||||
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
||||||
"""
|
"""
|
||||||
return {
|
headers = {
|
||||||
**get_auth_header(node_cls),
|
**get_auth_header(node_cls),
|
||||||
"Comfy-Env": get_deploy_environment(),
|
"Comfy-Env": get_deploy_environment(),
|
||||||
"Comfy-Usage-Source": get_usage_source(node_cls),
|
"Comfy-Usage-Source": get_usage_source(node_cls),
|
||||||
"Comfy-Core-Version": comfyui_version,
|
"Comfy-Core-Version": comfyui_version,
|
||||||
}
|
}
|
||||||
|
ctx = get_executing_context()
|
||||||
|
if ctx is not None:
|
||||||
|
headers["Comfy-Job-Id"] = ctx.prompt_id
|
||||||
|
return headers
|
||||||
|
|
||||||
|
|
||||||
def default_base_url() -> str:
|
def default_base_url() -> str:
|
||||||
|
|||||||
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.latent_formats
|
||||||
import comfy.ldm.lumina.controlnet
|
import comfy.ldm.lumina.controlnet
|
||||||
import comfy.ldm.supir.supir_modules
|
import comfy.ldm.supir.supir_modules
|
||||||
|
import comfy.ldm.anima.lllite
|
||||||
from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
|
from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
|
||||||
from comfy_api.latest import io
|
from comfy_api.latest import io
|
||||||
from comfy.ldm.supir.supir_patch import SUPIRPatch
|
from comfy.ldm.supir.supir_patch import SUPIRPatch
|
||||||
@ -236,10 +237,12 @@ class ModelPatchLoader:
|
|||||||
|
|
||||||
def load_model_patch(self, name):
|
def load_model_patch(self, name):
|
||||||
model_patch_path = folder_paths.get_full_path_or_raise("model_patches", 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)
|
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
|
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)
|
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:
|
elif 'feature_embedder.mid_layer_norm.bias' in sd:
|
||||||
@ -296,6 +299,50 @@ class ModelPatchLoader:
|
|||||||
return (model_patcher,)
|
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:
|
class DiffSynthCnetPatch:
|
||||||
def __init__(self, model_patch, vae, image, strength, mask=None):
|
def __init__(self, model_patch, vae, image, strength, mask=None):
|
||||||
self.model_patch = model_patch
|
self.model_patch = model_patch
|
||||||
@ -674,6 +721,7 @@ NODE_CLASS_MAPPINGS = {
|
|||||||
"ZImageFunControlnet": ZImageFunControlnet,
|
"ZImageFunControlnet": ZImageFunControlnet,
|
||||||
"USOStyleReference": USOStyleReference,
|
"USOStyleReference": USOStyleReference,
|
||||||
"SUPIRApply": SUPIRApply,
|
"SUPIRApply": SUPIRApply,
|
||||||
|
"AnimaLLLiteApply": AnimaLLLiteApply,
|
||||||
}
|
}
|
||||||
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
@ -682,4 +730,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
|||||||
"ZImageFunControlnet": "Apply Z-Image Fun ControlNet",
|
"ZImageFunControlnet": "Apply Z-Image Fun ControlNet",
|
||||||
"USOStyleReference": "Apply USO Style Reference",
|
"USOStyleReference": "Apply USO Style Reference",
|
||||||
"SUPIRApply": "Apply SUPIR Patch",
|
"SUPIRApply": "Apply SUPIR Patch",
|
||||||
|
"AnimaLLLiteApply": "Apply Anima LLLite",
|
||||||
}
|
}
|
||||||
|
|||||||
5
nodes.py
5
nodes.py
@ -992,7 +992,7 @@ class CLIPLoader:
|
|||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def INPUT_TYPES(s):
|
||||||
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
||||||
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2"], ),
|
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage"], ),
|
||||||
},
|
},
|
||||||
"optional": {
|
"optional": {
|
||||||
"device": (["default", "cpu"], {"advanced": True}),
|
"device": (["default", "cpu"], {"advanced": True}),
|
||||||
@ -1002,7 +1002,7 @@ class CLIPLoader:
|
|||||||
|
|
||||||
CATEGORY = "model/loaders"
|
CATEGORY = "model/loaders"
|
||||||
|
|
||||||
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
||||||
|
|
||||||
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
|
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
|
||||||
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||||
@ -2462,6 +2462,7 @@ async def init_builtin_extra_nodes():
|
|||||||
"nodes_seedvr.py",
|
"nodes_seedvr.py",
|
||||||
"nodes_context_windows.py",
|
"nodes_context_windows.py",
|
||||||
"nodes_qwen.py",
|
"nodes_qwen.py",
|
||||||
|
"nodes_joyimage.py",
|
||||||
"nodes_boogu.py",
|
"nodes_boogu.py",
|
||||||
"nodes_chroma_radiance.py",
|
"nodes_chroma_radiance.py",
|
||||||
"nodes_pid.py",
|
"nodes_pid.py",
|
||||||
|
|||||||
51
openapi.yaml
51
openapi.yaml
@ -530,6 +530,10 @@ components:
|
|||||||
description: Job creation timestamp (Unix timestamp in milliseconds)
|
description: Job creation timestamp (Unix timestamp in milliseconds)
|
||||||
format: int64
|
format: int64
|
||||||
type: integer
|
type: integer
|
||||||
|
execution_end_time:
|
||||||
|
description: Workflow execution completion timestamp (Unix milliseconds, only present for terminal states)
|
||||||
|
format: int64
|
||||||
|
type: integer
|
||||||
execution_error:
|
execution_error:
|
||||||
allOf:
|
allOf:
|
||||||
- $ref: '#/components/schemas/ExecutionError'
|
- $ref: '#/components/schemas/ExecutionError'
|
||||||
@ -538,6 +542,10 @@ components:
|
|||||||
additionalProperties: true
|
additionalProperties: true
|
||||||
description: Node-level execution metadata (only for terminal states)
|
description: Node-level execution metadata (only for terminal states)
|
||||||
type: object
|
type: object
|
||||||
|
execution_start_time:
|
||||||
|
description: Workflow execution start timestamp (Unix milliseconds, only present once execution has started)
|
||||||
|
format: int64
|
||||||
|
type: integer
|
||||||
execution_status:
|
execution_status:
|
||||||
additionalProperties: true
|
additionalProperties: true
|
||||||
description: ComfyUI execution status and timeline (only for terminal states)
|
description: ComfyUI execution status and timeline (only for terminal states)
|
||||||
@ -570,6 +578,12 @@ components:
|
|||||||
description: Last update timestamp (Unix timestamp in milliseconds)
|
description: Last update timestamp (Unix timestamp in milliseconds)
|
||||||
format: int64
|
format: int64
|
||||||
type: integer
|
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:
|
workflow:
|
||||||
additionalProperties: true
|
additionalProperties: true
|
||||||
description: |
|
description: |
|
||||||
@ -583,6 +597,18 @@ components:
|
|||||||
workflow_id:
|
workflow_id:
|
||||||
description: UUID identifying the workflow graph definition
|
description: UUID identifying the workflow graph definition
|
||||||
type: string
|
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:
|
required:
|
||||||
- id
|
- id
|
||||||
- status
|
- status
|
||||||
@ -1565,7 +1591,13 @@ paths:
|
|||||||
schema:
|
schema:
|
||||||
default: true
|
default: true
|
||||||
type: boolean
|
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
|
in: query
|
||||||
name: hash
|
name: hash
|
||||||
schema:
|
schema:
|
||||||
@ -2464,6 +2496,23 @@ paths:
|
|||||||
schema:
|
schema:
|
||||||
additionalProperties: true
|
additionalProperties: true
|
||||||
properties:
|
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:
|
max_upload_size:
|
||||||
description: Maximum upload size in bytes
|
description: Maximum upload size in bytes
|
||||||
type: integer
|
type: integer
|
||||||
|
|||||||
@ -1,5 +1,5 @@
|
|||||||
comfyui-frontend-package==1.45.21
|
comfyui-frontend-package==1.45.21
|
||||||
comfyui-workflow-templates==0.11.9
|
comfyui-workflow-templates==0.11.11
|
||||||
comfyui-embedded-docs==0.5.8
|
comfyui-embedded-docs==0.5.8
|
||||||
torch
|
torch
|
||||||
torchsde
|
torchsde
|
||||||
@ -22,7 +22,7 @@ alembic
|
|||||||
SQLAlchemy>=2.0.0
|
SQLAlchemy>=2.0.0
|
||||||
filelock
|
filelock
|
||||||
av>=16.0.0
|
av>=16.0.0
|
||||||
comfy-kitchen==0.2.20
|
comfy-kitchen==0.2.22
|
||||||
comfy-aimdo==0.4.10
|
comfy-aimdo==0.4.10
|
||||||
requests
|
requests
|
||||||
simpleeval>=1.0.0
|
simpleeval>=1.0.0
|
||||||
|
|||||||
@ -2,11 +2,12 @@ import pytest
|
|||||||
import torch
|
import torch
|
||||||
import tempfile
|
import tempfile
|
||||||
import os
|
import os
|
||||||
|
import sys
|
||||||
import av
|
import av
|
||||||
import io
|
import io
|
||||||
from fractions import Fraction
|
from fractions import Fraction
|
||||||
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
|
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
|
||||||
from comfy_api.util.video_types import VideoComponents
|
from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
|
||||||
from comfy_api.input.basic_types import AudioInput
|
from comfy_api.input.basic_types import AudioInput
|
||||||
from av.error import InvalidDataError
|
from av.error import InvalidDataError
|
||||||
|
|
||||||
@ -237,3 +238,526 @@ def test_duration_consistency(video_components):
|
|||||||
manual_duration = float(components.images.shape[0] / components.frame_rate)
|
manual_duration = float(components.images.shape[0] / components.frame_rate)
|
||||||
|
|
||||||
assert duration == pytest.approx(manual_duration)
|
assert duration == pytest.approx(manual_duration)
|
||||||
|
|
||||||
|
|
||||||
|
def create_transcode_source(
|
||||||
|
width=64, height=64, frames=30, fps=30, audio_streams=1, undecodable_audio=0, rotation=False,
|
||||||
|
container_format="mov", audio_codec="pcm_s16le",
|
||||||
|
):
|
||||||
|
"""Create a temp video that save_to must transcode (mpeg4 video, so codec != h264).
|
||||||
|
|
||||||
|
``undecodable_audio`` trailing PCM streams get their fourcc corrupted so no decoder exists
|
||||||
|
(``codec_context is None``), like the APAC track in iPhone spatial-audio recordings.
|
||||||
|
``rotation`` patches a 90-degree display matrix into the video track header.
|
||||||
|
"""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format=container_format) as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||||
|
video_stream.width = width
|
||||||
|
video_stream.height = height
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio = []
|
||||||
|
for _ in range(audio_streams + undecodable_audio):
|
||||||
|
stream = container.add_stream(audio_codec, rate=44100)
|
||||||
|
stream.sample_rate = 44100
|
||||||
|
audio.append(stream)
|
||||||
|
|
||||||
|
for i in range(frames):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((height, width, 3), (i * 7) % 256, dtype=torch.uint8).numpy(),
|
||||||
|
format="rgb24",
|
||||||
|
)
|
||||||
|
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||||
|
# write audio in 1024-sample frames, like real decoders produce, so the
|
||||||
|
# per-frame skip/cap logic in the transcode path actually runs
|
||||||
|
for stream in audio:
|
||||||
|
for offset in range(0, 44100 * frames // fps, 1024):
|
||||||
|
n = min(1024, 44100 * frames // fps - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, n, dtype=torch.int16).numpy(), format="s16", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = 44100
|
||||||
|
audio_frame.pts = offset
|
||||||
|
container.mux(stream.encode(audio_frame))
|
||||||
|
for stream in [video_stream, *audio]:
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
data = bytearray(buffer.getvalue())
|
||||||
|
end = len(data)
|
||||||
|
for _ in range(undecodable_audio):
|
||||||
|
end = data.rindex(b"sowt", 0, end)
|
||||||
|
data[end:end + 4] = b"Xpac"
|
||||||
|
if rotation:
|
||||||
|
# the 3x3 display matrix sits 40 bytes into the version-0 tkhd payload; first tkhd
|
||||||
|
# inside moov = video track (search from moov so mdat bytes can't false-match)
|
||||||
|
matrix_offset = data.index(b"tkhd", data.rindex(b"moov")) + 4 + 40
|
||||||
|
values = [0, 1 << 16, 0, -(1 << 16), 0, 0, 0, 0, 1 << 30]
|
||||||
|
data[matrix_offset:matrix_offset + 36] = b"".join(v.to_bytes(4, "big", signed=True) for v in values)
|
||||||
|
|
||||||
|
tmp = tempfile.NamedTemporaryFile(suffix=f".{container_format}", delete=False)
|
||||||
|
tmp.write(bytes(data))
|
||||||
|
tmp.close()
|
||||||
|
return tmp.name
|
||||||
|
|
||||||
|
|
||||||
|
def transcode_and_probe(video):
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
video.save_to(buffer, format=VideoContainer.MP4, codec=VideoCodec.H264)
|
||||||
|
buffer.seek(0)
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
video_stream = container.streams.video[0]
|
||||||
|
audio_stream = container.streams.audio[0] if container.streams.audio else None
|
||||||
|
frames = 0
|
||||||
|
first_pts = None
|
||||||
|
for packet in container.demux(video_stream):
|
||||||
|
for frame in packet.decode():
|
||||||
|
if first_pts is None:
|
||||||
|
first_pts = frame.pts
|
||||||
|
frames += 1
|
||||||
|
return {
|
||||||
|
"codec": video_stream.codec_context.name,
|
||||||
|
"width": video_stream.codec_context.width,
|
||||||
|
"height": video_stream.codec_context.height,
|
||||||
|
"frames": frames,
|
||||||
|
"first_pts": first_pts,
|
||||||
|
"video_seconds": float(video_stream.duration * video_stream.time_base) if video_stream.duration else None,
|
||||||
|
"audio_seconds": float(audio_stream.duration * audio_stream.time_base)
|
||||||
|
if audio_stream and audio_stream.duration else None,
|
||||||
|
"audio_codecs": [s.codec_context.name for s in container.streams.audio],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_streams_without_buffering_frames():
|
||||||
|
"""Transcoding must not decode the whole video into memory first (~2 GiB for this source)"""
|
||||||
|
resource = pytest.importorskip("resource") # no getrusage on Windows
|
||||||
|
rss_scale = 1 if sys.platform == "darwin" else 1024 # ru_maxrss: bytes on macOS, KiB elsewhere
|
||||||
|
# ru_maxrss is a lifetime peak: a heavier test running earlier would shrink the measured
|
||||||
|
# delta and quietly defang this canary, so keep this source the biggest thing in the suite
|
||||||
|
file_path = create_transcode_source(width=640, height=480, frames=300)
|
||||||
|
try:
|
||||||
|
rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path))
|
||||||
|
rss_delta = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale - rss_before
|
||||||
|
|
||||||
|
assert result["codec"] == "h264"
|
||||||
|
assert result["frames"] == 300
|
||||||
|
assert rss_delta < 500 * 2**20, f"transcode buffered frames in RAM (peak grew {rss_delta / 2**20:.0f} MiB)"
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_honors_trim_window():
|
||||||
|
"""start_time/duration trim applies to both video and audio on the streaming path"""
|
||||||
|
file_path = create_transcode_source(frames=90) # 3s @ 30fps
|
||||||
|
try:
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
|
||||||
|
assert result["frames"] == pytest.approx(30, abs=2)
|
||||||
|
assert result["first_pts"] == 0 # trimmed output is rebased to start at zero
|
||||||
|
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_keeps_audio_of_sparse_video():
|
||||||
|
"""Audio that runs ahead of a sparse video track (slideshows, timelapses) must be
|
||||||
|
kept in full — it is only clamped to the video's end, never to the video cursor."""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
|
||||||
|
for t in (0, 30, 60): # 3 frames spread over 60 seconds
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), t * 4, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = t * 15360
|
||||||
|
frame.time_base = Fraction(1, 15360)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
for offset in range(0, 48000 * 60, 1024):
|
||||||
|
n = min(1024, 48000 * 60 - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = 48000
|
||||||
|
audio_frame.pts = offset
|
||||||
|
audio_frame.time_base = Fraction(1, 48000)
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
container.mux(audio_stream.encode(None))
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["audio_seconds"] == pytest.approx(60.0, abs=1.0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_vfr_audio_covers_video_span():
|
||||||
|
"""A trim window in the sparse region of a VFR file keeps audio for the true pts span
|
||||||
|
of the kept frames. Deriving the span as frames/average_rate undercuts it badly: the
|
||||||
|
average is dominated by the dense region (and can be plain wrong on MediaRecorder files)."""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
|
||||||
|
# 10 frames inside the first second, then one every 1.25 s
|
||||||
|
for i, t in enumerate([x / 10 for x in range(10)] + [1.0, 2.25, 3.5, 4.75]):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 16) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = int(t * 15360)
|
||||||
|
frame.time_base = Fraction(1, 15360)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
for offset in range(0, 48000 * 6, 1024):
|
||||||
|
n = min(1024, 48000 * 6 - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = 48000
|
||||||
|
audio_frame.pts = offset
|
||||||
|
audio_frame.time_base = Fraction(1, 48000)
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
container.mux(audio_stream.encode(None))
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer, start_time=1, duration=5))
|
||||||
|
# kept frames: 1.0/2.25/3.5/4.75 s -> rebased span 3.75 s + one nominal interval
|
||||||
|
assert result["frames"] == 4
|
||||||
|
assert result["audio_seconds"] == pytest.approx(4.0, abs=0.45)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_trims_audio_in_stream_time_base_units():
|
||||||
|
"""Matroska audio timestamps tick in 1/1000, not 1/sample_rate; trim and audio timing
|
||||||
|
must convert through the frame's time base instead of assuming sample units. AAC audio,
|
||||||
|
because it decodes straight to the encoder's format and hits the resampler passthrough
|
||||||
|
that keeps the source time base on the frames."""
|
||||||
|
file_path = create_transcode_source(frames=90, container_format="matroska", audio_codec="aac")
|
||||||
|
try:
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
|
||||||
|
assert result["audio_codecs"] == ["aac"]
|
||||||
|
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_learns_unprobed_audio_params():
|
||||||
|
"""mpegts is only probed a few seconds deep at open, so an audio stream whose first
|
||||||
|
packet comes later (live captures where audio kicks in late) still has sample_rate 0
|
||||||
|
when the transcode starts; the parameters must be learned from the stream itself."""
|
||||||
|
sample_rate, fps, video_seconds, audio_start = 48000, 30, 13, 12
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mpegts") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||||
|
for i in range(video_seconds * fps):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
)
|
||||||
|
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||||
|
for offset in range(0, (video_seconds - audio_start) * sample_rate, 1024):
|
||||||
|
n = min(1024, (video_seconds - audio_start) * sample_rate - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = sample_rate
|
||||||
|
audio_frame.pts = audio_start * sample_rate + offset
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
for stream in (video_stream, audio_stream):
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
# the scenario requires unprobed parameters; if a future FFmpeg probes deeper,
|
||||||
|
# push audio_start/video_seconds further out to restore it
|
||||||
|
assert container.streams.audio[0].codec_context.sample_rate == 0
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["frames"] == video_seconds * fps
|
||||||
|
assert result["audio_codecs"] == ["aac"]
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
trimmed_before_audio = transcode_and_probe(VideoFromFile(buffer, duration=1))
|
||||||
|
assert trimmed_before_audio["frames"] == fps
|
||||||
|
assert trimmed_before_audio["audio_codecs"] == []
|
||||||
|
assert trimmed_before_audio["audio_seconds"] is None
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
trimmed_crossing_audio = transcode_and_probe(VideoFromFile(buffer, start_time=11.5, duration=1))
|
||||||
|
assert trimmed_crossing_audio["frames"] == fps
|
||||||
|
assert trimmed_crossing_audio["audio_codecs"] == ["aac"]
|
||||||
|
assert trimmed_crossing_audio["video_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||||
|
assert trimmed_crossing_audio["audio_seconds"] == pytest.approx(0.5, abs=0.1)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_trimmed_fragmented_mp4_keeps_audio():
|
||||||
|
"""Fragmented mp4 (MediaRecorder, DASH/HLS-derived files) delivers audio well behind
|
||||||
|
video, so when the trim window's last video frame arrives the audio demuxed so far
|
||||||
|
does not cover the window yet; the transcode must keep demuxing audio until it does
|
||||||
|
instead of finalizing on the first audio frame it sees afterwards."""
|
||||||
|
sample_rate, fps, seconds = 48000, 30, 6
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4", options={"movflags": "frag_keyframe+empty_moov"}) as container:
|
||||||
|
video_stream = container.add_stream("h264", rate=fps)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||||
|
next_audio_pts = 0
|
||||||
|
for i in range(seconds * fps):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
)
|
||||||
|
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||||
|
while next_audio_pts / sample_rate <= i / fps: # feed audio alongside, like a live pipeline
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, 1024, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = sample_rate
|
||||||
|
audio_frame.pts = next_audio_pts
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
next_audio_pts += 1024
|
||||||
|
for stream in (video_stream, audio_stream):
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer, start_time=0.5, duration=1.0))
|
||||||
|
assert result["video_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_sparse_video_keeps_true_duration():
|
||||||
|
"""average_rate is not a frame duration: a 3-frame video spanning 60 s averages
|
||||||
|
0.05 fps, and padding the last frame with 1/average_rate used to extend the
|
||||||
|
output — and the audio kept with it — about 20 s past the source span."""
|
||||||
|
sample_rate = 48000
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||||
|
for i, second in enumerate((0, 30, 60)):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), i * 80, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = second * 30
|
||||||
|
frame.time_base = Fraction(1, 30)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
for offset in range(0, 90 * sample_rate, 1024):
|
||||||
|
n = min(1024, 90 * sample_rate - offset)
|
||||||
|
audio_frame = av.AudioFrame.from_ndarray(
|
||||||
|
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||||
|
)
|
||||||
|
audio_frame.sample_rate = sample_rate
|
||||||
|
audio_frame.pts = offset
|
||||||
|
container.mux(audio_stream.encode(audio_frame))
|
||||||
|
for stream in (video_stream, audio_stream):
|
||||||
|
container.mux(stream.encode(None))
|
||||||
|
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["frames"] == 3
|
||||||
|
# the last frame keeps its true stts duration (1/30 s), not 1/average_rate (~20 s)
|
||||||
|
assert result["video_seconds"] == pytest.approx(60.03, abs=0.05)
|
||||||
|
assert result["audio_seconds"] == pytest.approx(60.03, abs=0.1)
|
||||||
|
|
||||||
|
trimmed = transcode_and_probe(VideoFromFile(buffer, duration=45))
|
||||||
|
assert trimmed["frames"] == 2
|
||||||
|
# a kept frame whose source duration crosses the window end is clamped to it
|
||||||
|
assert trimmed["video_seconds"] == pytest.approx(45.0, abs=0.05)
|
||||||
|
assert trimmed["audio_seconds"] == pytest.approx(45.0, abs=0.1)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_clamps_final_pts_to_declared_stream_duration():
|
||||||
|
"""Some iPhone MOVs report a video stream duration that ends before the final
|
||||||
|
decoded frame's nominal duration. A transcode must not turn that trailing
|
||||||
|
timestamp quirk into an extra frame interval compared to the source/remux path."""
|
||||||
|
fps = 30
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
for i, pts in enumerate([*range(31), 32]):
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = pts
|
||||||
|
frame.time_base = Fraction(1, fps)
|
||||||
|
container.mux(video_stream.encode(frame))
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
|
||||||
|
class _StreamProxy:
|
||||||
|
def __init__(self, stream, duration):
|
||||||
|
self._stream = stream
|
||||||
|
self.duration = duration
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._stream, name)
|
||||||
|
|
||||||
|
class _StreamsProxy:
|
||||||
|
def __init__(self, video_stream):
|
||||||
|
self.video = [video_stream]
|
||||||
|
self.audio = []
|
||||||
|
|
||||||
|
class _PacketProxy:
|
||||||
|
def __init__(self, packet, stream):
|
||||||
|
self._packet = packet
|
||||||
|
self.stream = stream
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._packet, name)
|
||||||
|
|
||||||
|
class _ContainerProxy:
|
||||||
|
def __init__(self, container, stream):
|
||||||
|
self._container = container
|
||||||
|
self._stream = stream
|
||||||
|
self.streams = _StreamsProxy(stream)
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._container, name)
|
||||||
|
|
||||||
|
def demux(self, *streams):
|
||||||
|
for packet in self._container.demux(self._stream._stream):
|
||||||
|
yield _PacketProxy(packet, self._stream)
|
||||||
|
|
||||||
|
buffer.seek(0)
|
||||||
|
output = io.BytesIO()
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
real_stream = container.streams.video[0]
|
||||||
|
declared_duration = 32 * int(round((1 / fps) / real_stream.time_base))
|
||||||
|
stream = _StreamProxy(real_stream, declared_duration)
|
||||||
|
VideoFromFile(buffer)._save_transcoded(
|
||||||
|
_ContainerProxy(container, stream), output, VideoContainer.MP4, VideoCodec.H264, None, 8
|
||||||
|
)
|
||||||
|
|
||||||
|
output.seek(0)
|
||||||
|
with av.open(output) as container:
|
||||||
|
video_stream = container.streams.video[0]
|
||||||
|
frames = [f for p in container.demux(video_stream) for f in p.decode()]
|
||||||
|
assert len(frames) == 32
|
||||||
|
assert float(video_stream.duration * video_stream.time_base) == pytest.approx(32 / fps, abs=0.01)
|
||||||
|
assert float(frames[-1].pts * frames[-1].time_base) == pytest.approx(31 / fps, abs=0.01)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_irregular_vfr_keeps_span():
|
||||||
|
"""B-frames reorder packets, and mp4 sample durations follow decode order: the dts
|
||||||
|
timeline ends before the pts timeline, so an irregular-VFR source's tail holds fell
|
||||||
|
out of the container (this 20.23 s span used to come out as 15.27 s, and the 10 s
|
||||||
|
trim as 6.03 s). The transcode encodes without B-frames so every sample keeps its
|
||||||
|
true display duration."""
|
||||||
|
durations = [1, 1, 60, 1, 1, 120, 1, 180, 1, 1, 150, 90] # 1/30 s ticks, span 20.2333 s
|
||||||
|
generator = torch.Generator().manual_seed(7)
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(buffer, mode="w", format="mp4") as container:
|
||||||
|
video_stream = container.add_stream("mpeg4", rate=30)
|
||||||
|
video_stream.width = video_stream.height = 64
|
||||||
|
video_stream.pix_fmt = "yuv420p"
|
||||||
|
pts = 0
|
||||||
|
for duration in durations:
|
||||||
|
# textured frames, so an encoder with default settings has B-frames to gain from
|
||||||
|
frame = av.VideoFrame.from_ndarray(
|
||||||
|
torch.randint(0, 255, (64, 64, 3), generator=generator, dtype=torch.uint8).numpy(),
|
||||||
|
format="rgb24",
|
||||||
|
).reformat(format="yuv420p")
|
||||||
|
frame.pts = pts
|
||||||
|
frame.time_base = Fraction(1, 30)
|
||||||
|
pts += duration
|
||||||
|
for packet in video_stream.encode(frame):
|
||||||
|
packet.duration = duration # exact stts in the source
|
||||||
|
container.mux(packet)
|
||||||
|
container.mux(video_stream.encode(None))
|
||||||
|
|
||||||
|
result = transcode_and_probe(VideoFromFile(buffer))
|
||||||
|
assert result["frames"] == len(durations)
|
||||||
|
assert result["video_seconds"] == pytest.approx(sum(durations) / 30, abs=0.05)
|
||||||
|
|
||||||
|
trimmed = transcode_and_probe(VideoFromFile(buffer, duration=10))
|
||||||
|
assert trimmed["frames"] == 8 # frames at 12.167 s+ fall outside the window
|
||||||
|
assert trimmed["video_seconds"] == pytest.approx(10.0, abs=0.05)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_trim_survives_missing_leading_pts():
|
||||||
|
"""A trim should survive pts-less kept frames followed by a real-pts frame past the window."""
|
||||||
|
nulled_frames = 0
|
||||||
|
|
||||||
|
class _PacketProxy:
|
||||||
|
def __init__(self, packet):
|
||||||
|
self._packet = packet
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._packet, name)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def stream(self):
|
||||||
|
return self._packet.stream
|
||||||
|
|
||||||
|
def decode(self):
|
||||||
|
nonlocal nulled_frames
|
||||||
|
frames = self._packet.decode()
|
||||||
|
for frame in frames:
|
||||||
|
if nulled_frames < 2:
|
||||||
|
frame.pts = None
|
||||||
|
nulled_frames += 1
|
||||||
|
return frames
|
||||||
|
|
||||||
|
class _ContainerProxy:
|
||||||
|
def __init__(self, real):
|
||||||
|
self._real = real
|
||||||
|
|
||||||
|
def __getattr__(self, name):
|
||||||
|
return getattr(self._real, name)
|
||||||
|
|
||||||
|
def demux(self, *streams):
|
||||||
|
for packet in self._real.demux(*streams):
|
||||||
|
yield _PacketProxy(packet)
|
||||||
|
|
||||||
|
file_path = create_transcode_source(frames=10, audio_streams=0)
|
||||||
|
try:
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with av.open(file_path) as container:
|
||||||
|
# 0.05 s window: both pts-less frames are kept (synthesized pts 0 and 512),
|
||||||
|
# and the first real-pts frame (1024 ticks) already lies past end_pts (768)
|
||||||
|
VideoFromFile(file_path, duration=0.05)._save_transcoded(
|
||||||
|
_ContainerProxy(container), buffer, VideoContainer.MP4, VideoCodec.H264, None, 8
|
||||||
|
)
|
||||||
|
assert nulled_frames == 2
|
||||||
|
buffer.seek(0)
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
video_stream = container.streams.video[0]
|
||||||
|
frames = [f for p in container.demux(video_stream) for f in p.decode()]
|
||||||
|
assert len(frames) == 2
|
||||||
|
assert float(video_stream.duration * video_stream.time_base) == pytest.approx(2 / 30, abs=0.01)
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_bakes_rotation():
|
||||||
|
"""A 90-degree display-matrix rotation swaps the output dimensions (portrait video)"""
|
||||||
|
file_path = create_transcode_source(width=64, height=32, rotation=True)
|
||||||
|
try:
|
||||||
|
result = transcode_and_probe(VideoFromFile(file_path))
|
||||||
|
assert (result["width"], result["height"]) == (32, 64)
|
||||||
|
assert result["frames"] == 30
|
||||||
|
finally:
|
||||||
|
os.unlink(file_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_to_transcode_skips_undecodable_audio():
|
||||||
|
"""Streaming transcode keeps the decodable audio track and drops undecodable ones;
|
||||||
|
with no decodable audio at all the output is video-only instead of crashing."""
|
||||||
|
mixed = all_bad = None
|
||||||
|
try:
|
||||||
|
mixed = create_transcode_source(audio_streams=1, undecodable_audio=1)
|
||||||
|
all_bad = create_transcode_source(audio_streams=0, undecodable_audio=2)
|
||||||
|
result = transcode_and_probe(VideoFromFile(mixed))
|
||||||
|
assert result["audio_codecs"] == ["aac"]
|
||||||
|
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||||
|
assert transcode_and_probe(VideoFromFile(all_bad))["audio_codecs"] == []
|
||||||
|
finally:
|
||||||
|
for path in (mixed, all_bad):
|
||||||
|
if path:
|
||||||
|
os.unlink(path)
|
||||||
|
|||||||
@ -112,6 +112,17 @@ def _make_pid_v1_5_sd(latent_proj_channels=16):
|
|||||||
return sd
|
return sd
|
||||||
|
|
||||||
|
|
||||||
|
def _make_joyimage_edit_plus_sd():
|
||||||
|
sd = {
|
||||||
|
"img_in.weight": torch.empty(4096, 16, 1, 2, 2, device="meta"),
|
||||||
|
"condition_embedder.time_embedder.linear_1.weight": torch.empty(1, device="meta"),
|
||||||
|
"double_blocks.0.attn.img_attn_q_norm.weight": torch.empty(128, device="meta"),
|
||||||
|
}
|
||||||
|
for i in range(40):
|
||||||
|
sd[f"double_blocks.{i}.attn.img_attn_qkv.weight"] = torch.empty(1, device="meta")
|
||||||
|
return sd
|
||||||
|
|
||||||
|
|
||||||
def _add_model_diffusion_prefix(sd):
|
def _add_model_diffusion_prefix(sd):
|
||||||
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
|
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
|
||||||
|
|
||||||
@ -258,6 +269,26 @@ class TestModelDetection:
|
|||||||
assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,)
|
assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,)
|
||||||
assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,)
|
assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,)
|
||||||
|
|
||||||
|
def test_joyimage_edit_plus_detection(self):
|
||||||
|
sd = _make_joyimage_edit_plus_sd()
|
||||||
|
unet_config = detect_unet_config(sd, "")
|
||||||
|
|
||||||
|
assert unet_config == {
|
||||||
|
"image_model": "joyimage",
|
||||||
|
"in_channels": 16,
|
||||||
|
"hidden_size": 4096,
|
||||||
|
"patch_size": [1, 2, 2],
|
||||||
|
"num_layers": 40,
|
||||||
|
"num_attention_heads": 32,
|
||||||
|
"text_dim": 4096,
|
||||||
|
}
|
||||||
|
assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "JoyImage"
|
||||||
|
|
||||||
|
def test_incomplete_joyimage_signature_is_not_detected(self):
|
||||||
|
sd = _make_joyimage_edit_plus_sd()
|
||||||
|
del sd["double_blocks.0.attn.img_attn_q_norm.weight"]
|
||||||
|
assert detect_unet_config(sd, "") is None
|
||||||
|
|
||||||
def test_unet_config_and_required_keys_combination_is_unique(self):
|
def test_unet_config_and_required_keys_combination_is_unique(self):
|
||||||
"""Each model in the registry must have a unique combination of
|
"""Each model in the registry must have a unique combination of
|
||||||
``unet_config`` and ``required_keys``. If two models share the same
|
``unet_config`` and ``required_keys``. If two models share the same
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user