Compare commits

..

20 Commits

Author SHA1 Message Date
524414b117 fix: trimmed transcode before delayed audio starts
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-07-08 18:57:41 +03:00
891258882d fix: continue probing audio parameters after a packet decode error instead of dropping stream immediately
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-07-08 18:56:36 +03:00
e7c39d0c20 fix: trim transcode with missing leading PTS
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-07-08 18:56:36 +03:00
6a889609bd fix(Video): stream the video transcode instead of buffering every frame in RAM
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-07-08 18:56:35 +03:00
091b70edda add models-directory launch argument (#9113) 2026-07-08 22:20:47 +08:00
ffbecfffb9 Fix crash when using UNetSelfAttentionMultiply (#14823) 2026-07-07 21:17:31 -07:00
b481bc15af Support gqa on all attention backends, drop support for pytorch 2.4 (#14772) 2026-07-07 22:57:52 -04:00
6880614319 Update AGENTS.md (#14819) 2026-07-07 18:36:13 -07:00
51bf508a0b feat: Implement basic text overlay node (CORE-137) (#14610) 2026-07-07 21:26:52 +08:00
a3020f107e fix(Video): don't crash on videos with undecodable audio streams (#14746)
* fix(Video): don't crash on videos with undecodable audio streams

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* Update comfy_api_nodes/util/upload_helpers.py

---------

Signed-off-by: bigcat88 <bigcat88@icloud.com>
Co-authored-by: Alexis Rolland <alexisrolland@hotmail.com>
2026-07-07 15:59:49 +03:00
7cf4e78335 Delete symlink that breaks our updates. (#14803) 2026-07-06 22:24:05 -04:00
7747c342d4 ci: add CLA Assistant workflow (#14582) 2026-07-07 06:44:19 +08:00
439bd807f8 Skip unloading dynamic model patchers in current workflow. (#14799) 2026-07-06 14:35:12 -07:00
b08debceca chore: update embedded docs to v0.5.7 (#14783) 2026-07-06 09:56:09 +08:00
000c6b784e Small speedup for text model sampling. (#14773) 2026-07-05 18:39:24 -07:00
985fb9d6ad [Partner Nodes] fix(logs-auth): mask authorization headers in logs (#14774)
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-07-05 13:55:29 +03:00
7f287b705e fix: Bug when setting transparency in color picker (#14764) 2026-07-04 19:13:38 -04:00
b7ba504e06 Try to make coderabbit enforce AGENTS.md (#14759) 2026-07-04 14:25:24 -04:00
6c62ca0b6b fix: error when embedding is loaded with models using llama_template (#14744) 2026-07-04 17:06:09 +08:00
3fe9f5fecb Add CLAUDE.md as symlink to AGENTS.md (#14757) 2026-07-04 13:12:47 +08:00
28 changed files with 1279 additions and 222 deletions

View File

@ -4,12 +4,12 @@ early_access: false
tone_instructions: "Only comment on issues introduced by this PR's changes. Do not flag pre-existing problems in moved, re-indented, or reformatted code."
reviews:
profile: "chill"
request_changes_workflow: false
profile: "assertive"
request_changes_workflow: true
high_level_summary: false
poem: false
review_status: false
review_details: false
review_details: true
commit_status: true
collapse_walkthrough: true
changed_files_summary: false
@ -39,6 +39,14 @@ reviews:
- path: "**"
instructions: |
IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing a `with:` block),
@ -123,5 +131,10 @@ chat:
knowledge_base:
opt_out: false
code_guidelines:
enabled: true
filePatterns:
- files: "AGENTS.md"
applyTo: "**"
learnings:
scope: "auto"

91
.github/workflows/cla.yml vendored Normal file
View File

@ -0,0 +1,91 @@
name: CLA Assistant
on:
issue_comment:
types: [created]
pull_request_target:
types: [opened, synchronize, closed]
permissions:
actions: write
contents: read # 'read' is enough because signatures live in a REMOTE repo
pull-requests: write
statuses: write
jobs:
cla-assistant:
runs-on: ubuntu-latest
steps:
# The CLA action normally requires every commit author in a PR to sign.
# We only want the PR author to sign, so we allowlist all other committers
# by computing them from the PR's commits and excluding the PR author.
- name: Build author-only allowlist
id: allowlist
if: >
github.event_name == 'pull_request_target' ||
(github.event_name == 'issue_comment' && github.event.issue.pull_request && (
github.event.comment.body == 'recheck' ||
github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
))
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
run: |
others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
--jq '.[] | (.author.login // empty), (.committer.login // empty)' \
| sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
if [ -n "$others" ]; then
echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
else
echo "allowlist=${BASE_ALLOWLIST}" >> "$GITHUB_OUTPUT"
fi
- name: CLA Assistant
# Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase.
# IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
if: >
github.event_name == 'pull_request_target' ||
(github.event_name == 'issue_comment' && github.event.issue.pull_request && (
github.event.comment.body == 'recheck' ||
github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
))
uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# PAT required to write to the centralized signatures repo.
PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
with:
# Where the CLA document lives (shown to contributors)
path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
# Centralized signature storage
remote-organization-name: comfy-org
remote-repository-name: comfy-cla
path-to-signatures: signatures/cla.json
branch: main
# Only the PR author must sign: bots plus every non-author committer
# are allowlisted via the "Build author-only allowlist" step above.
# *[bot] is a catch-all for any GitHub App bot account.
allowlist: ${{ steps.allowlist.outputs.allowlist }}
# Custom PR comment messages
custom-notsigned-prcomment: |
🎉 Thank you for your contribution, we really appreciate it! 🎉
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
- Confirm that you own your contribution.
- Keep the right to reuse your own code.
- Grant us a copyright license to include and share it within our projects.
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.
✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍
custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
custom-allsigned-prcomment: |
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.

View File

@ -127,6 +127,8 @@
- Do not add unnecessary `try`/`except` blocks. Use them for optional dependency,
platform, or backend capability detection only when the program has a useful
fallback. Prefer specific exception types when changing new code.
- If a library version is pinned in `requirements.txt`, do not add code to
ComfyUI to handle older versions of that library.
- Remove any workarounds for PyTorch versions that ComfyUI no longer officially
supports. Deprecated workarounds include catching an exception and rerunning
the same op with the input cast to float. If a workaround does not have a

View File

@ -229,7 +229,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia
Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
torch 2.4 and above is supported but some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old.
torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
### Instructions:

View File

@ -225,6 +225,7 @@ parser.add_argument(
)
parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path. Overrides --base-directory.")
parser.add_argument("--models-directory", type=is_valid_directory, default=None, help="Set the ComfyUI models directory. Overrides the models folder in --base-directory.")
parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.")

View File

@ -217,10 +217,7 @@ class AceStepAttention(nn.Module):
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
n_rep = self.num_heads // self.num_kv_heads
if n_rep > 1:
key_states = key_states.repeat_interleave(n_rep, dim=1)
value_states = value_states.repeat_interleave(n_rep, dim=1)
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
attn_bias = None
if self.sliding_window is not None and not self.is_cross_attention:
@ -244,7 +241,7 @@ class AceStepAttention(nn.Module):
else:
attn_bias = window_bias
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False)
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False, **gqa_kwargs)
attn_output = self.o_proj(attn_output)
return attn_output

View File

@ -425,19 +425,16 @@ class Attention(nn.Module):
if n == 1 and causal:
causal = False
if h != kv_h:
# Repeat interleave kv_heads to match q_heads
heads_per_kv_head = h // kv_h
k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
gqa_kwargs = {"enable_gqa": True} if h != kv_h else {}
if self.differential:
q, q_diff = q.unbind(dim=1)
k, k_diff = k.unbind(dim=1)
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
out = out - out_diff
else:
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options)
out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs)
out = self.to_out(out)

View File

@ -74,11 +74,8 @@ class BooguDoubleStreamProcessor(nn.Module):
key = key.transpose(1, 2)
value = value.transpose(1, 2)
if attn.kv_heads < attn.heads:
key = key.repeat_interleave(attn.heads // attn.kv_heads, dim=1)
value = value.repeat_interleave(attn.heads // attn.kv_heads, dim=1)
hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
gqa_kwargs = {"enable_gqa": True} if attn.kv_heads < attn.heads else {}
hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs)
# Split back to instruction/image, apply per-stream output projections, recombine.
instruct_hidden_states = self.instruct_out(hidden_states[:, :L_instruct])

View File

@ -1,5 +1,6 @@
import math
import sys
import inspect
import torch
import torch.nn.functional as F
@ -14,16 +15,16 @@ from .sub_quadratic_attention import efficient_dot_product_attention
from comfy import model_management
TORCH_HAS_GQA = model_management.torch_version_numeric >= (2, 5)
if model_management.xformers_enabled():
import xformers
import xformers.ops
SAGE_ATTENTION_IS_AVAILABLE = False
SAGE_ATTENTION_SUPPORTS_MASK = False
try:
from sageattention import sageattn
SAGE_ATTENTION_IS_AVAILABLE = True
SAGE_ATTENTION_SUPPORTS_MASK = "attn_mask" in inspect.signature(sageattn).parameters
except ImportError as e:
if model_management.sage_attention_enabled():
if e.name == "sageattention":
@ -89,6 +90,44 @@ def default(val, d):
return val
return d
def _gqa_repeat_factor(query_heads, key_heads, value_heads):
if key_heads != value_heads:
raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
if query_heads == key_heads:
return 1
if query_heads % key_heads != 0:
raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
return query_heads // key_heads
def _repeat_kv_for_gqa(k, v, query_heads, head_dim):
n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
if n_rep > 1:
k = k.repeat_interleave(n_rep, dim=head_dim)
v = v.repeat_interleave(n_rep, dim=head_dim)
return k, v
def _heads_from_dim(tensor, dim_head, name):
inner_dim = tensor.shape[-1]
if inner_dim % dim_head != 0:
raise ValueError(f"{name} inner dimension {inner_dim} is not divisible by head dimension {dim_head}")
return inner_dim // dim_head
def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_kv=True):
q = q.unsqueeze(3).reshape(b, -1, heads, dim_head)
if enable_gqa:
key_heads = _heads_from_dim(k, dim_head, "Key")
value_heads = _heads_from_dim(v, dim_head, "Value")
else:
key_heads = heads
value_heads = heads
k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head)
v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head)
if enable_gqa:
_gqa_repeat_factor(heads, key_heads, value_heads)
if expand_kv:
k, v = _repeat_kv_for_gqa(k, v, heads, -2)
return q, k, v
# feedforward
class GEGLU(nn.Module):
@ -152,28 +191,19 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
b, _, dim_head = q.shape
dim_head //= heads
if kwargs.get("enable_gqa", False) and q.shape[-3] != k.shape[-3]:
n_rep = q.shape[-3] // k.shape[-3]
k = k.repeat_interleave(n_rep, dim=-3)
v = v.repeat_interleave(n_rep, dim=-3)
scale = kwargs.get("scale", dim_head ** -0.5)
h = heads
if skip_reshape:
q, k, v = map(
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head),
(q, k, v),
)
else:
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, -1, heads, dim_head)
.permute(0, 2, 1, 3)
.reshape(b * heads, -1, dim_head)
.contiguous(),
(q, k, v),
)
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v))
# force cast to fp32 to avoid overflowing
if attn_precision == torch.float32:
@ -231,13 +261,16 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
query = query * (kwargs["scale"] * dim_head ** 0.5)
if skip_reshape:
if kwargs.get("enable_gqa", False):
key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3)
query = query.reshape(b * heads, -1, dim_head)
value = value.reshape(b * heads, -1, dim_head)
key = key.reshape(b * heads, -1, dim_head).movedim(1, 2)
else:
query = query.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
value = value.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
query, key, value = _reshape_qkv_to_heads(query, key, value, b, heads, dim_head, kwargs.get("enable_gqa", False))
query = query.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
value = value.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
key = key.permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
dtype = query.dtype
@ -304,19 +337,15 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
scale = kwargs.get("scale", dim_head ** -0.5)
if skip_reshape:
q, k, v = map(
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head),
(q, k, v),
)
else:
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, -1, heads, dim_head)
.permute(0, 2, 1, 3)
.reshape(b * heads, -1, dim_head)
.contiguous(),
(q, k, v),
)
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v))
r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
@ -438,7 +467,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
disabled_xformers = True
if disabled_xformers:
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs)
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
if skip_reshape:
# b h k d -> b k h d
@ -446,13 +475,12 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
lambda t: t.permute(0, 2, 1, 3),
(q, k, v),
)
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2)
# actually do the reshaping
else:
dim_head //= heads
q, k, v = map(
lambda t: t.reshape(b, -1, heads, dim_head),
(q, k, v),
)
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
if mask is not None:
# add a singleton batch dimension
@ -474,7 +502,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
mask = mask_out[..., :mask.shape[-1]]
mask = mask.expand(b, heads, -1, -1)
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, scale=kwargs.get("scale", None))
if skip_output_reshape:
out = out.permute(0, 2, 1, 3)
@ -498,10 +526,8 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
else:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = map(
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
(q, k, v),
)
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False)
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
if mask is not None:
# add a batch dimension if there isn't already one
@ -511,9 +537,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
if mask.ndim == 3:
mask = mask.unsqueeze(1)
# Pass through extra SDPA kwargs (scale, enable_gqa) if provided
# enable_gqa requires PyTorch 2.5+; older versions use manual KV expansion above
sdpa_keys = ("scale", "enable_gqa") if TORCH_HAS_GQA else ("scale",)
sdpa_keys = ("scale", "enable_gqa")
sdpa_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys}
if SDP_BATCH_LIMIT >= b:
@ -541,20 +565,19 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
@wrap_attn
def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
if kwargs.get("low_precision_attention", True) is False:
if kwargs.get("low_precision_attention", True) is False or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK):
return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
exception_fallback = False
if skip_reshape:
b, _, _, dim_head = q.shape
tensor_layout = "HND"
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
else:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = map(
lambda t: t.view(b, -1, heads, dim_head),
(q, k, v),
)
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False))
tensor_layout = "NHD"
if mask is not None:
@ -565,8 +588,12 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
if mask.ndim == 3:
mask = mask.unsqueeze(1)
sage_kwargs = {"is_causal": False, "tensor_layout": tensor_layout, "sm_scale": kwargs.get("scale", None), "smooth_k": False}
if mask is not None:
sage_kwargs["attn_mask"] = mask
try:
out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout)
out = sageattn(q, k, v, **sage_kwargs)
except Exception as e:
logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e))
exception_fallback = True
@ -616,7 +643,6 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
skip_output_reshape=skip_output_reshape,
**kwargs
)
q_s, k_s, v_s = q, k, v
N = q.shape[2]
dim_head = D
else:
@ -642,11 +668,15 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
**kwargs
)
if not skip_reshape:
q_s, k_s, v_s = map(
lambda t: t.view(B, -1, heads, dim_head).permute(0, 2, 1, 3).contiguous(),
(q, k, v),
)
if skip_reshape:
q_s = q
if kwargs.get("enable_gqa", False):
k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3)
else:
k_s, v_s = k, v
else:
q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False))
q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s))
B, H, L, D = q_s.shape
try:
@ -662,7 +692,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
q, k, v, heads,
mask=mask,
attn_precision=attn_precision,
skip_reshape=False,
skip_reshape=skip_reshape,
skip_output_reshape=skip_output_reshape,
**kwargs
)
@ -681,19 +711,20 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
try:
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal, softmax_scale=softmax_scale_arg)
@flash_attn_wrapper.register_fake
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False):
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False, softmax_scale=-1.0):
# Output shape is the same as q
return q.new_empty(q.shape)
except AttributeError as error:
FLASH_ATTN_ERROR = error
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}"
@wrap_attn
@ -703,10 +734,8 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
else:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = map(
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
(q, k, v),
)
q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False)
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
if mask is not None:
# add a batch dimension if there isn't already one
@ -725,10 +754,16 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
v.transpose(1, 2),
dropout_p=0.0,
causal=False,
softmax_scale=kwargs.get("scale", -1.0),
).transpose(1, 2)
except Exception as e:
logging.warning(f"Flash Attention failed, using default SDPA: {e}")
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
sdpa_extra = {}
if kwargs.get("enable_gqa", False):
sdpa_extra["enable_gqa"] = True
if "scale" in kwargs:
sdpa_extra["scale"] = kwargs["scale"]
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra)
if not skip_output_reshape:
out = (
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
@ -1209,5 +1244,3 @@ class SpatialVideoTransformer(SpatialTransformer):
x = self.proj_out(x)
out = x + x_in
return out

View File

@ -141,11 +141,8 @@ class Attention(nn.Module):
key = key.transpose(1, 2)
value = value.transpose(1, 2)
if self.kv_heads < self.heads:
key = key.repeat_interleave(self.heads // self.kv_heads, dim=1)
value = value.repeat_interleave(self.heads // self.kv_heads, dim=1)
hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
gqa_kwargs = {"enable_gqa": True} if self.kv_heads < self.heads else {}
hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs)
hidden_states = self.to_out[0](hidden_states)
return hidden_states

View File

@ -174,6 +174,8 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
elif xfer_dest2 is not None:
xfer_source.prepare(xfer_dest2, stream, copy=True, commit=False)
return
else:
return
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream, r2=xfer_dest2)
def handle_pin(m, pin, source, dest, subset="weights", size=None):

View File

@ -468,6 +468,9 @@ class CLIP:
def decode(self, token_ids, skip_special_tokens=True):
return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
def is_dynamic(self):
return self.patcher.is_dynamic()
class VAE:
def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
@ -1251,6 +1254,8 @@ class VAE:
except:
return None
def is_dynamic(self):
return self.patcher.is_dynamic()
class StyleModel:
def __init__(self, model, device="cpu"):

View File

@ -543,18 +543,24 @@ class SDTokenizer:
def _try_get_embedding(self, embedding_name:str):
'''
Takes a potential embedding name and tries to retrieve it.
Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
Returns a Tuple consisting of the embedding, the cleaned embedding name, and any leftover string, embedding can be None.
'''
split_embed = embedding_name.split()
embedding_name = split_embed[0]
leftover = ' '.join(split_embed[1:])
match = re.search(r'[<\[]', embedding_name)
if match is not None:
leftover = embedding_name[match.start():] + (" " + leftover if leftover else "")
embedding_name = embedding_name[:match.start()]
embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
if embed is None:
stripped = embedding_name.strip(',')
if len(stripped) < len(embedding_name):
embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
return (embed, "{} {}".format(embedding_name[len(stripped):], leftover))
return (embed, leftover)
return (embed, embedding_name, "{} {}".format(embedding_name[len(stripped):], leftover))
return (embed, embedding_name, leftover)
def pad_tokens(self, tokens, amount):
if self.pad_left:
@ -585,7 +591,7 @@ class SDTokenizer:
tokens = []
for weighted_segment, weight in parsed_weights:
to_tokenize = unescape_important(weighted_segment)
split = re.split(' {0}|\n{0}'.format(self.embedding_identifier), to_tokenize)
split = re.split(r'(?<=\s){}'.format(re.escape(self.embedding_identifier)), to_tokenize)
to_tokenize = [split[0]]
for i in range(1, len(split)):
to_tokenize.append("{}{}".format(self.embedding_identifier, split[i]))
@ -595,7 +601,7 @@ class SDTokenizer:
# if we find an embedding, deal with the embedding
if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
embedding_name = word[len(self.embedding_identifier):].strip('\n')
embed, leftover = self._try_get_embedding(embedding_name)
embed, embedding_name, leftover = self._try_get_embedding(embedding_name)
if embed is None:
logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
else:

View File

@ -12,7 +12,7 @@ import torch.nn.functional as F
import comfy.ops
from comfy import sd1_clip
from comfy.ldm.modules.attention import TORCH_HAS_GQA, optimized_attention_for_device
from comfy.ldm.modules.attention import optimized_attention_for_device
from comfy.text_encoders.llama import RMSNorm, apply_rope
@ -110,10 +110,6 @@ def _attention_with_sinks(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sin
putting the sink logit in the mask at that column.
"""
if num_kv_groups > 1 and not TORCH_HAS_GQA:
k = k.repeat_interleave(num_kv_groups, dim=1)
v = v.repeat_interleave(num_kv_groups, dim=1)
B, _, S_q, D = q.shape
H_kv = k.shape[1]
S_kv = k.shape[-2]

View File

@ -550,10 +550,8 @@ class Attention(nn.Module):
xv = xv[:, :, -sliding_window:]
attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
return self.o_proj(output), present_key_value
class MLP(nn.Module):
@ -937,22 +935,41 @@ class BaseGenerate:
return torch.argmax(logits, dim=-1, keepdim=True)
# Sampling mode
if repetition_penalty != 1.0:
for i in range(logits.shape[0]):
for token_id in set(token_history):
logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
if presence_penalty is not None and presence_penalty != 0.0:
for i in range(logits.shape[0]):
for token_id in set(token_history):
logits[i, token_id] -= presence_penalty
if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
token_ids = torch.tensor(list(set(token_history)), device=logits.device)
token_logits = logits[:, token_ids]
if repetition_penalty != 1.0:
token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty)
if presence_penalty is not None and presence_penalty != 0.0:
token_logits = token_logits - presence_penalty
logits[:, token_ids] = token_logits
if temperature != 1.0:
logits = logits / temperature
if top_k > 0:
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
logits[indices_to_remove] = torch.finfo(logits.dtype).min
top_k = min(top_k, logits.shape[-1])
logits, top_indices = torch.topk(logits, top_k)
if min_p > 0.0:
probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
min_threshold = min_p * top_probs
indices_to_remove = probs_before_filter < min_threshold
logits[indices_to_remove] = torch.finfo(logits.dtype).min
if top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 0] = False
indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
logits[indices_to_remove] = torch.finfo(logits.dtype).min
probs = torch.nn.functional.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1, generator=generator)
return top_indices.gather(1, next_token)
if min_p > 0.0:
probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)

View File

@ -366,12 +366,8 @@ class GatedAttention(nn.Module):
xv = torch.cat((past_value[:, :, :index], xv), dim=2)
present_key_value = (xk, xv, index + num_tokens)
# Expand KV heads for GQA
if self.num_heads != self.num_kv_heads:
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True)
gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs)
output = output * gate.sigmoid()
return self.o_proj(output), present_key_value

View File

@ -1,5 +1,6 @@
from av.container import InputContainer
from av.subtitles.stream import SubtitleStream
from av.video.reformatter import ColorRange
from fractions import Fraction
from typing import Optional
from .._input import AudioInput, VideoInput
@ -9,6 +10,7 @@ import itertools
import json
import numpy as np
import math
import os
import torch
from .._util import VideoContainer, VideoCodec, VideoComponents
import logging
@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
return max(component.bits for component in stream.format.components)
def last_decodable_audio_stream(container: InputContainer):
"""Streams FFmpeg has no decoder for have no codec context, and decoding their
packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
stream = next(
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
None,
)
if stream is None and len(container.streams.audio):
logging.warning("No decodable audio stream found in video; ignoring audio.")
return stream
def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
"""Containers probed only up to a window (mpegts) leave audio codec parameters unset when
audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
for i, packet in enumerate(container.demux(audio_stream)):
try:
frames = packet.decode()
except av.error.FFmpegError:
frames = ()
if frames:
return frames[0].sample_rate, frames[0].layout.nb_channels
if i >= max_packets:
break
return 0, 0
def write_output_metadata(container: InputContainer, output, metadata: dict | None):
"""Copy the source container's metadata, then overlay the caller's tags."""
for key, value in container.metadata.items():
if metadata is None or key not in metadata:
output.metadata[key] = value
if metadata is not None:
for key, value in metadata.items():
output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
raise ValueError("Only MP4 format is supported for now")
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
raise ValueError("Only H264 codec is supported for now")
open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
open_kwargs["format"] = format.value
elif isinstance(path, io.BytesIO):
open_kwargs["format"] = "mp4" # no file extension to infer the format from
return open_kwargs
class VideoFromFile(VideoInput):
"""
Class representing video input from a file.
@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
return estimated_frames
# 3. Last resort: decode frames and count them (streaming)
if self.__start_time < 0:
start_time = max(self._get_raw_duration() + self.__start_time, 0)
else:
start_time = self.__start_time
start_time, duration = self.get_active_trim_window()
frame_count = 1
start_pts = int(start_time / video_stream.time_base)
end_pts = int((start_time + self.__duration) / video_stream.time_base)
end_pts = int((start_time + duration) / video_stream.time_base)
container.seek(start_pts, stream=video_stream)
frame_iterator = (
container.decode(video_stream)
@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
def get_components_internal(self, container: InputContainer) -> VideoComponents:
video_stream = self._get_first_video_stream(container)
if self.__start_time < 0:
start_time = max(self._get_raw_duration() + self.__start_time, 0)
else:
start_time = self.__start_time
start_time, duration = self.get_active_trim_window()
# Get video frames
frames = []
audio_frames = []
alphas = None
start_pts = int(start_time / video_stream.time_base)
end_pts = int((start_time + self.__duration) / video_stream.time_base)
end_pts = int((start_time + duration) / video_stream.time_base)
if start_pts != 0:
container.seek(start_pts, stream=video_stream)
@ -281,8 +328,8 @@ class VideoFromFile(VideoInput):
video_done = False
audio_done = True
if len(container.streams.audio):
audio_stream = container.streams.audio[-1]
audio_stream = last_decodable_audio_stream(container)
if audio_stream is not None:
streams += [audio_stream]
resampler = av.audio.resampler.AudioResampler(format='fltp')
audio_done = False
@ -298,7 +345,7 @@ class VideoFromFile(VideoInput):
for frame in packet.decode():
if frame.pts < start_pts:
continue
if self.__duration and frame.pts >= end_pts:
if duration and frame.pts >= end_pts:
video_done = True
break
@ -365,7 +412,7 @@ class VideoFromFile(VideoInput):
map(resampler.resample, packet.decode())
)
for frame in aframes:
if self.__duration and frame.time > start_time + self.__duration:
if duration and frame.time > start_time + duration:
audio_done = True
break
@ -387,8 +434,8 @@ class VideoFromFile(VideoInput):
if len(audio_frames) > 0:
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
if self.__duration:
audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
if duration:
audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
audio = AudioInput({
@ -434,33 +481,22 @@ class VideoFromFile(VideoInput):
if not reuse_streams:
if bit_depth is None:
bit_depth = source_bit_depth
components = self.get_components_internal(container)
video = VideoFromComponents(components)
return video.save_to(
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
)
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
streams = container.streams
open_kwargs = get_open_write_kwargs(path, container_format, format)
with av.open(path, **open_kwargs) as output_container:
# Copy over the original metadata
for key, value in container.metadata.items():
if metadata is None or key not in metadata:
output_container.metadata[key] = value
# Add metadata before writing any streams
write_output_metadata(container, output_container, metadata)
# 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
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
stream_map = {}
for stream in streams:
if isinstance(stream, (av.VideoStream, av.AudioStream, SubtitleStream)):
if stream.codec_context is None:
logging.warning("Skipping %s stream %d with unsupported codec", stream.type, stream.index)
continue
out_stream = output_container.add_stream_from_template(template=stream, opaque=True)
stream_map[stream] = out_stream
@ -470,6 +506,259 @@ class VideoFromFile(VideoInput):
packet.stream = stream_map[packet.stream]
output_container.mux(packet)
def _save_transcoded(
self,
container: InputContainer,
path: str | io.BytesIO,
format: VideoContainer,
codec: VideoCodec,
metadata: dict | None,
bit_depth: int,
):
"""Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
open_kwargs = mp4_output_open_kwargs(path, format, codec)
video_stream = self._get_first_video_stream(container)
start_time, duration = self.get_active_trim_window()
start_pts = int(start_time / video_stream.time_base)
end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
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)
if frame.pts is not None:
if video_pts_offset is None:
video_pts_offset = frame.pts
frame.pts -= video_pts_offset
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
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 to_skip:
pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
continue
pending_audio.append(resampled)
if video_done:
# the video window is complete so the cap is final, but containers
# that interleave audio behind video (fragmented mp4) still owe most
# of it: stop only once the demuxed audio covers the cap
cap = drain_audio()
if pending_audio or samples_written >= cap:
drain_audio(final=True)
audio_done = True
break
if output is None:
raise ValueError(f"No decodable video frames found in file '{self.__file}'")
if out_audio is not None and not audio_done:
drain_audio(final=True)
window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
for out_packet in out_video.encode(None):
duration = video_frame_durations.pop(out_packet.pts, 0)
if out_packet.pts == last_video_pts:
duration = max(duration, window_fill)
out_packet.duration = duration
output.mux(out_packet)
if out_audio is not None:
output.mux(out_audio.encode(None))
except BaseException:
if output is not None:
output.close()
if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
os.remove(path)
raise
else:
if output is not None:
output.close()
def _get_first_video_stream(self, container: InputContainer):
if len(container.streams.video):
return container.streams.video[0]
@ -517,22 +806,12 @@ class VideoFromComponents(VideoInput):
bit_depth: int | None = None,
):
"""Save the video to a file path or BytesIO buffer."""
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
raise ValueError("Only MP4 format is supported for now")
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
raise ValueError("Only H264 codec is supported for now")
open_kwargs = mp4_output_open_kwargs(path, format, codec)
# None means "use the depth this video was created with" (CreateVideo's choice).
if bit_depth is None:
bit_depth = self.__bit_depth
is_10bit = bit_depth >= 10
extra_kwargs = {}
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
extra_kwargs["format"] = format.value
elif isinstance(path, io.BytesIO):
# BytesIO has no file extension, so av.open can't infer the format.
# Default to mp4 since that's the only supported format anyway.
extra_kwargs["format"] = "mp4"
with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
with av.open(path, **open_kwargs) as output:
# Add metadata before writing any streams
if metadata is not None:
for key, value in metadata.items():

View File

@ -9,6 +9,7 @@ from typing import Any
import folder_paths
logger = logging.getLogger(__name__)
_SENSITIVE_HEADERS = {"authorization", "x-api-key"}
def get_log_directory():
@ -73,6 +74,10 @@ def _format_data_for_logging(data: Any) -> str:
return str(data)
def _redact_headers(headers: dict) -> dict:
return {k: ("***" if k.lower() in _SENSITIVE_HEADERS else v) for k, v in headers.items()}
def log_request_response(
operation_id: str,
request_method: str,
@ -101,7 +106,7 @@ def log_request_response(
log_content.append(f"Method: {request_method}")
log_content.append(f"URL: {request_url}")
if request_headers:
log_content.append(f"Headers:\n{_format_data_for_logging(request_headers)}")
log_content.append(f"Headers:\n{_format_data_for_logging(_redact_headers(request_headers))}")
if request_params:
log_content.append(f"Params:\n{_format_data_for_logging(request_params)}")
if request_data is not None:

View File

@ -158,7 +158,14 @@ async def upload_video_to_comfyapi(
# Convert VideoInput to BytesIO using specified container/codec
video_bytes_io = BytesIO()
video.save_to(video_bytes_io, format=container, codec=codec)
try:
video.save_to(video_bytes_io, format=container, codec=codec)
except Exception as e:
raise ValueError(
f"Could not convert the input video to {container.value.upper()} for upload; "
f"the file may be corrupted or use an unsupported codec. "
f"Try re-exporting it as MP4 (H.264). Original error: {e}"
) from e
video_bytes_io.seek(0)
return await upload_file_to_comfyapi(cls, video_bytes_io, filename, upload_mime_type, wait_label)

View File

@ -1,6 +1,5 @@
import asyncio
import bisect
import gc
import itertools
import psutil
import time
@ -504,6 +503,21 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
def all_outputs_dynamic(outputs):
if outputs is None:
return False
for output in outputs:
if isinstance(output, (list, tuple)):
if not all_outputs_dynamic(output):
return False
elif not hasattr(output, "is_dynamic") or not output.is_dynamic():
return False
return True
class RAMPressureCache(LRUCache):
def __init__(self, key_class, enable_providers=False):
@ -529,44 +543,16 @@ class RAMPressureCache(LRUCache):
if psutil.virtual_memory().available >= target:
return
def remove_cache_key(key):
del self.cache[key]
self.used_generation.pop(key, None)
self.timestamps.pop(key, None)
self.children.pop(key, None)
def has_old_model_patcher(outputs):
if outputs is None:
return False
for output in outputs:
if isinstance(output, (list, tuple)):
if has_old_model_patcher(output):
return True
elif isinstance(output, ModelPatcher):
return True
return False
old_modelpatcher_keys = []
for key, cache_entry in self.cache.items():
if self.used_generation[key] == self.generation:
continue
if has_old_model_patcher(cache_entry.outputs):
old_modelpatcher_keys.append(key)
for key in old_modelpatcher_keys:
remove_cache_key(key)
if old_modelpatcher_keys:
gc.collect()
if psutil.virtual_memory().available >= target:
return
clean_list = []
for key, cache_entry in self.cache.items():
if not free_active and self.used_generation[key] == self.generation:
continue
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
if all_outputs_dynamic(cache_entry.outputs) and self.used_generation[key] == self.generation:
continue
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
def scan_list_for_ram_usage(outputs):
@ -578,17 +564,19 @@ class RAMPressureCache(LRUCache):
scan_list_for_ram_usage(output)
elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
ram_usage += output.numel() * output.element_size()
elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
#old ModelPatchers are the first to go
ram_usage = 1e30
scan_list_for_ram_usage(cache_entry.outputs)
oom_score *= ram_usage
#In the case where we have no information on the node ram usage at all,
#break OOM score ties on the last touch timestamp (pure LRU)
bisect.insort(clean_list, (oom_score, self.timestamps[key], ram_usage, key))
bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
to_free = target - psutil.virtual_memory().available
while to_free > 0 and clean_list:
_, _, ram_usage, key = clean_list.pop()
remove_cache_key(key)
to_free -= ram_usage
gc.collect()
while psutil.virtual_memory().available < target and clean_list:
_, _, key = clean_list.pop()
del self.cache[key]
self.used_generation.pop(key, None)
self.timestamps.pop(key, None)
self.children.pop(key, None)

View File

@ -16,23 +16,30 @@ class ColorToRGBInt(io.ComfyNode):
],
outputs=[
io.Int.Output(display_name="rgb_int"),
io.Color.Output(display_name="hex")
io.Color.Output(display_name="hex"),
io.Float.Output(display_name="alpha"),
],
)
@classmethod
def execute(cls, color: str) -> io.NodeOutput:
# expect format #RRGGBB
if len(color) != 7 or color[0] != "#":
raise ValueError("Color must be in format #RRGGBB")
# expect format #RRGGBB or #RRGGBBAA
if len(color) not in (7, 9) or color[0] != "#":
raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA")
try:
int(color[1:], 16)
except ValueError:
raise ValueError("Color must be in format #RRGGBB") from None
raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA") from None
alpha = 1.0
if len(color) == 9:
alpha = int(color[7:9], 16) / 255.0
color = color[:7]
r, g, b = hex_to_rgb(color)
rgb_int = r * 256 * 256 + g * 256 + b
return io.NodeOutput(rgb_int, color)
return io.NodeOutput(rgb_int, color, alpha)
class ColorExtension(ComfyExtension):

View File

@ -0,0 +1,150 @@
import numpy as np
import torch
from PIL import Image as PILImage, ImageColor, ImageDraw, ImageFont
from typing_extensions import override
from comfy_api.latest import ComfyExtension, IO
class TextOverlay(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="TextOverlay",
display_name="Draw Text Overlay",
category="text",
description="Draw text overlay on an image or batch of images.",
search_aliases=["text", "label", "caption", "subtitle", "watermark", "title", "addlabel", "overlay"],
inputs=[
IO.Image.Input("images"),
IO.String.Input("text", multiline=True, default=""),
IO.Float.Input("font_size", default=5.0, min=0.5, max=50.0, step=0.5, tooltip="Font size as a percentage of the image height."),
IO.Color.Input("color", default="#ffffff", tooltip="Color of the text."),
IO.Combo.Input("position", options=["top", "bottom"], default="top"),
IO.Combo.Input("align", options=["left", "center", "right"], default="left"),
IO.Boolean.Input("outline", default=True, tooltip="Draw a black outline around the text."),
],
outputs=[IO.Image.Output(display_name="images")],
)
@classmethod
def execute(cls, images, text, font_size, color, position, align, outline) -> IO.NodeOutput:
if text.strip() == "":
return IO.NodeOutput(images)
text = text.replace("\\n", "\n").replace("\\t", "\t")
text_rgba = cls.parse_color_to_rgba(color)
outline_rgba = (0, 0, 0, 255) if outline else (0, 0, 0, 0)
# Render the overlay once and composite it across all frames in the batch
height = images.shape[1]
width = images.shape[2]
overlay_rgb, overlay_alpha = cls.render_overlay_text(width, height, text, position, align, font_size, text_rgba, outline_rgba)
overlay_rgb = overlay_rgb.to(device=images.device, dtype=images.dtype)
overlay_alpha = overlay_alpha.to(device=images.device, dtype=images.dtype)
result = images * (1.0 - overlay_alpha) + overlay_rgb * overlay_alpha
return IO.NodeOutput(result)
@staticmethod
def parse_color_to_rgba(color_string):
parsed = ImageColor.getrgb(color_string)
if len(parsed) == 3:
return (*parsed, 255)
return parsed
@classmethod
def render_overlay_text(cls, width, height, text, position, align, font_size, text_rgba, outline_rgba):
line_spacing = 1.2
margin_percent = 1.0
min_font_percent = 2.0
min_font_pixels = 10
outline_thickness_factor = 0.04
# Draw onto a transparent layer so the result can be alpha-composited over any frame.
layer = PILImage.new("RGBA", (width, height), (0, 0, 0, 0))
draw = ImageDraw.Draw(layer)
margin = int(round(margin_percent / 100.0 * min(width, height)))
max_width = max(1, width - 2 * margin)
max_height = max(1, height - 2 * margin)
# Font scales with resolution, then shrinks to fit the height.
size = max(1, int(round(font_size / 100.0 * height)))
floor = min(size, max(min_font_pixels, int(round(min_font_percent / 100.0 * height))))
while True:
font = ImageFont.load_default(size=size)
stroke = max(1, int(round(size * outline_thickness_factor))) if outline_rgba[3] > 0 else 0
block = "\n".join(cls.wrap_text(text, font, max_width))
# convert line spacing to pixel spacing
single = draw.textbbox((0, 0), "Ay", font=font, stroke_width=stroke)
double = draw.multiline_textbbox((0, 0), "Ay\nAy", font=font, spacing=0, stroke_width=stroke)
natural_advance = (double[3] - double[1]) - (single[3] - single[1])
pixel_spacing = int(round(size * line_spacing - natural_advance))
box = draw.multiline_textbbox((0, 0), block, font=font, spacing=pixel_spacing, stroke_width=stroke)
block_height = box[3] - box[1]
if block_height <= max_height or size <= floor:
break
size = max(floor, int(size * 0.9))
anchor_h, x = {"left": ("l", margin), "center": ("m", width / 2), "right": ("r", width - margin)}[align]
# Offset y so the rendered text sits flush against the margin
if position == "bottom":
y = height - margin - box[3]
else:
y = margin - box[1]
draw.multiline_text((x, y), block, font=font, fill=text_rgba, anchor=anchor_h + "a",
align=align, spacing=pixel_spacing, stroke_width=stroke, stroke_fill=outline_rgba)
overlay = np.array(layer).astype(np.float32) / 255.0
overlay_rgb = torch.from_numpy(overlay[:, :, :3])
overlay_alpha = torch.from_numpy(overlay[:, :, 3:4])
return overlay_rgb, overlay_alpha
@staticmethod
def wrap_text(text, font, max_width):
lines = []
for raw_line in text.split("\n"):
words = raw_line.split()
if not words:
lines.append("")
continue
current = ""
# Break the line into words and split words that are too long
for word in words:
while font.getlength(word) > max_width and len(word) > 1:
cut = 1
while cut < len(word) and font.getlength(word[:cut + 1]) <= max_width:
cut += 1
if current:
lines.append(current)
current = ""
lines.append(word[:cut])
word = word[cut:]
candidate = word if not current else current + " " + word
if not current or font.getlength(candidate) <= max_width:
current = candidate
else:
lines.append(current)
current = word
if current:
lines.append(current)
return lines
class TextOverlayExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
return [TextOverlay]
async def comfy_entrypoint() -> TextOverlayExtension:
return TextOverlayExtension()

View File

@ -17,7 +17,11 @@ if args.base_directory:
else:
base_path = os.path.dirname(os.path.realpath(__file__))
models_dir = os.path.join(base_path, "models")
if args.models_directory:
models_dir = os.path.abspath(args.models_directory)
else:
models_dir = os.path.join(base_path, "models")
folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])

View File

@ -131,6 +131,10 @@ def apply_custom_paths():
if args.base_directory:
logging.info(f"Setting base directory to: {folder_paths.base_path}")
# --models-directory
if args.models_directory:
logging.info(f"Setting models directory to: {folder_paths.models_dir}")
# --output-directory, --input-directory, --user-directory
if args.output_directory:
output_dir = os.path.abspath(args.output_directory)
@ -314,7 +318,7 @@ def prompt_worker(q, server_instance):
cache_ram = 0
cache_ram_inactive = 0
if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
cache_ram = min(10.0, max(1.5, comfy.model_management.total_ram * 0.05 / 1024.0))
cache_ram = min(10.0, max(2.0, comfy.model_management.total_ram * 0.10 / 1024.0))
cache_ram_inactive = min(96.0, comfy.model_management.total_ram / 1024.0)
if len(args.cache_ram) > 0:
cache_ram = args.cache_ram[0]

View File

@ -2478,6 +2478,7 @@ async def init_builtin_extra_nodes():
"nodes_glsl.py",
"nodes_lora_debug.py",
"nodes_textgen.py",
"nodes_text_overlay.py",
"nodes_color.py",
"nodes_toolkit.py",
"nodes_replacements.py",

View File

@ -1,6 +1,6 @@
comfyui-frontend-package==1.45.20
comfyui-workflow-templates==0.11.2
comfyui-embedded-docs==0.5.6
comfyui-embedded-docs==0.5.7
torch
torchsde
torchvision

View File

@ -2,11 +2,12 @@ import pytest
import torch
import tempfile
import os
import sys
import av
import io
from fractions import Fraction
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
from comfy_api.util.video_types import VideoComponents
from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
from comfy_api.input.basic_types import AudioInput
from av.error import InvalidDataError
@ -237,3 +238,447 @@ def test_duration_consistency(video_components):
manual_duration = float(components.images.shape[0] / components.frame_rate)
assert duration == pytest.approx(manual_duration)
def create_transcode_source(
width=64, height=64, frames=30, fps=30, audio_streams=1, undecodable_audio=0, rotation=False,
container_format="mov", audio_codec="pcm_s16le",
):
"""Create a temp video that save_to must transcode (mpeg4 video, so codec != h264).
``undecodable_audio`` trailing PCM streams get their fourcc corrupted so no decoder exists
(``codec_context is None``), like the APAC track in iPhone spatial-audio recordings.
``rotation`` patches a 90-degree display matrix into the video track header.
"""
buffer = io.BytesIO()
with av.open(buffer, mode="w", format=container_format) as container:
video_stream = container.add_stream("mpeg4", rate=fps)
video_stream.width = width
video_stream.height = height
video_stream.pix_fmt = "yuv420p"
audio = []
for _ in range(audio_streams + undecodable_audio):
stream = container.add_stream(audio_codec, rate=44100)
stream.sample_rate = 44100
audio.append(stream)
for i in range(frames):
frame = av.VideoFrame.from_ndarray(
torch.full((height, width, 3), (i * 7) % 256, dtype=torch.uint8).numpy(),
format="rgb24",
)
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
# write audio in 1024-sample frames, like real decoders produce, so the
# per-frame skip/cap logic in the transcode path actually runs
for stream in audio:
for offset in range(0, 44100 * frames // fps, 1024):
n = min(1024, 44100 * frames // fps - offset)
audio_frame = av.AudioFrame.from_ndarray(
torch.zeros(1, n, dtype=torch.int16).numpy(), format="s16", layout="mono"
)
audio_frame.sample_rate = 44100
audio_frame.pts = offset
container.mux(stream.encode(audio_frame))
for stream in [video_stream, *audio]:
container.mux(stream.encode(None))
data = bytearray(buffer.getvalue())
end = len(data)
for _ in range(undecodable_audio):
end = data.rindex(b"sowt", 0, end)
data[end:end + 4] = b"Xpac"
if rotation:
# the 3x3 display matrix sits 40 bytes into the version-0 tkhd payload; first tkhd
# inside moov = video track (search from moov so mdat bytes can't false-match)
matrix_offset = data.index(b"tkhd", data.rindex(b"moov")) + 4 + 40
values = [0, 1 << 16, 0, -(1 << 16), 0, 0, 0, 0, 1 << 30]
data[matrix_offset:matrix_offset + 36] = b"".join(v.to_bytes(4, "big", signed=True) for v in values)
tmp = tempfile.NamedTemporaryFile(suffix=f".{container_format}", delete=False)
tmp.write(bytes(data))
tmp.close()
return tmp.name
def transcode_and_probe(video):
buffer = io.BytesIO()
video.save_to(buffer, format=VideoContainer.MP4, codec=VideoCodec.H264)
buffer.seek(0)
with av.open(buffer) as container:
video_stream = container.streams.video[0]
audio_stream = container.streams.audio[0] if container.streams.audio else None
frames = 0
first_pts = None
for packet in container.demux(video_stream):
for frame in packet.decode():
if first_pts is None:
first_pts = frame.pts
frames += 1
return {
"codec": video_stream.codec_context.name,
"width": video_stream.codec_context.width,
"height": video_stream.codec_context.height,
"frames": frames,
"first_pts": first_pts,
"video_seconds": float(video_stream.duration * video_stream.time_base) if video_stream.duration else None,
"audio_seconds": float(audio_stream.duration * audio_stream.time_base)
if audio_stream and audio_stream.duration else None,
"audio_codecs": [s.codec_context.name for s in container.streams.audio],
}
def test_save_to_transcode_streams_without_buffering_frames():
"""Transcoding must not decode the whole video into memory first (~2 GiB for this source)"""
resource = pytest.importorskip("resource") # no getrusage on Windows
rss_scale = 1 if sys.platform == "darwin" else 1024 # ru_maxrss: bytes on macOS, KiB elsewhere
# ru_maxrss is a lifetime peak: a heavier test running earlier would shrink the measured
# delta and quietly defang this canary, so keep this source the biggest thing in the suite
file_path = create_transcode_source(width=640, height=480, frames=300)
try:
rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale
result = transcode_and_probe(VideoFromFile(file_path))
rss_delta = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale - rss_before
assert result["codec"] == "h264"
assert result["frames"] == 300
assert rss_delta < 500 * 2**20, f"transcode buffered frames in RAM (peak grew {rss_delta / 2**20:.0f} MiB)"
finally:
os.unlink(file_path)
def test_save_to_transcode_honors_trim_window():
"""start_time/duration trim applies to both video and audio on the streaming path"""
file_path = create_transcode_source(frames=90) # 3s @ 30fps
try:
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
assert result["frames"] == pytest.approx(30, abs=2)
assert result["first_pts"] == 0 # trimmed output is rebased to start at zero
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
finally:
os.unlink(file_path)
def test_save_to_transcode_keeps_audio_of_sparse_video():
"""Audio that runs ahead of a sparse video track (slideshows, timelapses) must be
kept in full — it is only clamped to the video's end, never to the video cursor."""
buffer = io.BytesIO()
with av.open(buffer, mode="w", format="mp4") as container:
video_stream = container.add_stream("mpeg4", rate=30)
video_stream.width = video_stream.height = 64
video_stream.pix_fmt = "yuv420p"
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
for t in (0, 30, 60): # 3 frames spread over 60 seconds
frame = av.VideoFrame.from_ndarray(
torch.full((64, 64, 3), t * 4, dtype=torch.uint8).numpy(), format="rgb24"
).reformat(format="yuv420p")
frame.pts = t * 15360
frame.time_base = Fraction(1, 15360)
container.mux(video_stream.encode(frame))
container.mux(video_stream.encode(None))
for offset in range(0, 48000 * 60, 1024):
n = min(1024, 48000 * 60 - offset)
audio_frame = av.AudioFrame.from_ndarray(
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
)
audio_frame.sample_rate = 48000
audio_frame.pts = offset
audio_frame.time_base = Fraction(1, 48000)
container.mux(audio_stream.encode(audio_frame))
container.mux(audio_stream.encode(None))
buffer.seek(0)
result = transcode_and_probe(VideoFromFile(buffer))
assert result["audio_seconds"] == pytest.approx(60.0, abs=1.0)
def test_save_to_transcode_vfr_audio_covers_video_span():
"""A trim window in the sparse region of a VFR file keeps audio for the true pts span
of the kept frames. Deriving the span as frames/average_rate undercuts it badly: the
average is dominated by the dense region (and can be plain wrong on MediaRecorder files)."""
buffer = io.BytesIO()
with av.open(buffer, mode="w", format="mp4") as container:
video_stream = container.add_stream("mpeg4", rate=30)
video_stream.width = video_stream.height = 64
video_stream.pix_fmt = "yuv420p"
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
# 10 frames inside the first second, then one every 1.25 s
for i, t in enumerate([x / 10 for x in range(10)] + [1.0, 2.25, 3.5, 4.75]):
frame = av.VideoFrame.from_ndarray(
torch.full((64, 64, 3), (i * 16) % 256, dtype=torch.uint8).numpy(), format="rgb24"
).reformat(format="yuv420p")
frame.pts = int(t * 15360)
frame.time_base = Fraction(1, 15360)
container.mux(video_stream.encode(frame))
container.mux(video_stream.encode(None))
for offset in range(0, 48000 * 6, 1024):
n = min(1024, 48000 * 6 - offset)
audio_frame = av.AudioFrame.from_ndarray(
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
)
audio_frame.sample_rate = 48000
audio_frame.pts = offset
audio_frame.time_base = Fraction(1, 48000)
container.mux(audio_stream.encode(audio_frame))
container.mux(audio_stream.encode(None))
buffer.seek(0)
result = transcode_and_probe(VideoFromFile(buffer, start_time=1, duration=5))
# kept frames: 1.0/2.25/3.5/4.75 s -> rebased span 3.75 s + one nominal interval
assert result["frames"] == 4
assert result["audio_seconds"] == pytest.approx(4.0, abs=0.45)
def test_save_to_transcode_trims_audio_in_stream_time_base_units():
"""Matroska audio timestamps tick in 1/1000, not 1/sample_rate; trim and audio timing
must convert through the frame's time base instead of assuming sample units. AAC audio,
because it decodes straight to the encoder's format and hits the resampler passthrough
that keeps the source time base on the frames."""
file_path = create_transcode_source(frames=90, container_format="matroska", audio_codec="aac")
try:
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
assert result["audio_codecs"] == ["aac"]
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
finally:
os.unlink(file_path)
def test_save_to_transcode_learns_unprobed_audio_params():
"""mpegts is only probed a few seconds deep at open, so an audio stream whose first
packet comes later (live captures where audio kicks in late) still has sample_rate 0
when the transcode starts; the parameters must be learned from the stream itself."""
sample_rate, fps, video_seconds, audio_start = 48000, 30, 13, 12
buffer = io.BytesIO()
with av.open(buffer, mode="w", format="mpegts") as container:
video_stream = container.add_stream("mpeg4", rate=fps)
video_stream.width = video_stream.height = 64
video_stream.pix_fmt = "yuv420p"
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
for i in range(video_seconds * fps):
frame = av.VideoFrame.from_ndarray(
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
)
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
for offset in range(0, (video_seconds - audio_start) * sample_rate, 1024):
n = min(1024, (video_seconds - audio_start) * sample_rate - offset)
audio_frame = av.AudioFrame.from_ndarray(
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
)
audio_frame.sample_rate = sample_rate
audio_frame.pts = audio_start * sample_rate + offset
container.mux(audio_stream.encode(audio_frame))
for stream in (video_stream, audio_stream):
container.mux(stream.encode(None))
buffer.seek(0)
with av.open(buffer) as container:
# the scenario requires unprobed parameters; if a future FFmpeg probes deeper,
# push audio_start/video_seconds further out to restore it
assert container.streams.audio[0].codec_context.sample_rate == 0
result = transcode_and_probe(VideoFromFile(buffer))
assert result["frames"] == video_seconds * fps
assert result["audio_codecs"] == ["aac"]
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
buffer.seek(0)
trimmed_before_audio = transcode_and_probe(VideoFromFile(buffer, duration=1))
assert trimmed_before_audio["frames"] == fps
assert trimmed_before_audio["audio_codecs"] == []
assert trimmed_before_audio["audio_seconds"] is None
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_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)

View File

@ -163,3 +163,20 @@ def test_base_path_change_clears_old(set_base_dir):
for name in ["controlnet", "diffusion_models", "text_encoders"]:
assert len(folder_paths.get_folder_paths(name)) == 2
def test_models_directory_cli_and_getters(temp_dir):
try:
with patch.object(sys, 'argv', ["main.py", "--models-directory", temp_dir]):
reload(comfy.cli_args)
reload(folder_paths)
assert folder_paths.models_dir == os.path.abspath(temp_dir)
with pytest.raises(Exception):
comfy.cli_args.is_valid_directory(os.path.join(temp_dir, "non_existent_folder_path"))
finally:
with patch.object(sys, 'argv', ["main.py"]):
reload(comfy.cli_args)
reload(folder_paths)