mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 12:28:17 +08:00
Merge branch 'master' into feat/torchao-int4
This commit is contained in:
commit
879833ee37
@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None):
|
||||
The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
|
||||
"""
|
||||
|
||||
def two_pass_attention(q, k, v, heads, **kwargs):
|
||||
def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs):
|
||||
B, H, T, D = q.shape
|
||||
|
||||
if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
elif ar_len >= T:
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
|
||||
elif ar_len <= 0:
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
else:
|
||||
out_ar = comfy.ops.scaled_dot_product_attention(
|
||||
q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True,
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
|
||||
)
|
||||
out_gen = optimized_attention(
|
||||
q[:, :, ar_len:], k, v, heads,
|
||||
mask=None, skip_reshape=True, skip_output_reshape=True,
|
||||
transformer_options=transformer_options,
|
||||
transformer_options=transformer_options, enable_gqa=enable_gqa,
|
||||
)
|
||||
out = torch.cat([out_ar, out_gen], dim=2)
|
||||
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from av.container import InputContainer
|
||||
from av.subtitles.stream import SubtitleStream
|
||||
from av.video.reformatter import ColorRange
|
||||
from fractions import Fraction
|
||||
from typing import Optional
|
||||
from .._input import AudioInput, VideoInput
|
||||
@ -9,6 +10,7 @@ import itertools
|
||||
import json
|
||||
import numpy as np
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||
import logging
|
||||
@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
|
||||
return max(component.bits for component in stream.format.components)
|
||||
|
||||
|
||||
def last_decodable_audio_stream(container: InputContainer):
|
||||
"""Streams FFmpeg has no decoder for have no codec context, and decoding their
|
||||
packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
|
||||
stream = next(
|
||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||
None,
|
||||
)
|
||||
if stream is None and len(container.streams.audio):
|
||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||
return stream
|
||||
|
||||
|
||||
def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
|
||||
"""Containers probed only up to a window (mpegts) leave audio codec parameters unset when
|
||||
audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
|
||||
Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
|
||||
for i, packet in enumerate(container.demux(audio_stream)):
|
||||
try:
|
||||
frames = packet.decode()
|
||||
except av.error.FFmpegError:
|
||||
frames = ()
|
||||
if frames:
|
||||
return frames[0].sample_rate, frames[0].layout.nb_channels
|
||||
if i >= max_packets:
|
||||
break
|
||||
return 0, 0
|
||||
|
||||
|
||||
def write_output_metadata(container: InputContainer, output, metadata: dict | None):
|
||||
"""Copy the source container's metadata, then overlay the caller's tags."""
|
||||
for key, value in container.metadata.items():
|
||||
if metadata is None or key not in metadata:
|
||||
output.metadata[key] = value
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
|
||||
|
||||
|
||||
def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
|
||||
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
||||
raise ValueError("Only MP4 format is supported for now")
|
||||
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
||||
raise ValueError("Only H264 codec is supported for now")
|
||||
open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
|
||||
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
||||
open_kwargs["format"] = format.value
|
||||
elif isinstance(path, io.BytesIO):
|
||||
open_kwargs["format"] = "mp4" # no file extension to infer the format from
|
||||
return open_kwargs
|
||||
|
||||
|
||||
class VideoFromFile(VideoInput):
|
||||
"""
|
||||
Class representing video input from a file.
|
||||
@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
|
||||
return estimated_frames
|
||||
|
||||
# 3. Last resort: decode frames and count them (streaming)
|
||||
if self.__start_time < 0:
|
||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
||||
else:
|
||||
start_time = self.__start_time
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
frame_count = 1
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
frame_iterator = (
|
||||
container.decode(video_stream)
|
||||
@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
|
||||
|
||||
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
if self.__start_time < 0:
|
||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
||||
else:
|
||||
start_time = self.__start_time
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
|
||||
# Get video frames
|
||||
frames = []
|
||||
audio_frames = []
|
||||
alphas = None
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||
|
||||
if start_pts != 0:
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
@ -281,18 +328,11 @@ class VideoFromFile(VideoInput):
|
||||
video_done = False
|
||||
audio_done = True
|
||||
|
||||
# Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
|
||||
# and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
|
||||
audio_stream = next(
|
||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||
None,
|
||||
)
|
||||
audio_stream = last_decodable_audio_stream(container)
|
||||
if audio_stream is not None:
|
||||
streams += [audio_stream]
|
||||
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
||||
audio_done = False
|
||||
elif len(container.streams.audio):
|
||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||
|
||||
for packet in container.demux(*streams):
|
||||
if video_done and audio_done:
|
||||
@ -305,7 +345,7 @@ class VideoFromFile(VideoInput):
|
||||
for frame in packet.decode():
|
||||
if frame.pts < start_pts:
|
||||
continue
|
||||
if self.__duration and frame.pts >= end_pts:
|
||||
if duration and frame.pts >= end_pts:
|
||||
video_done = True
|
||||
break
|
||||
|
||||
@ -372,7 +412,7 @@ class VideoFromFile(VideoInput):
|
||||
map(resampler.resample, packet.decode())
|
||||
)
|
||||
for frame in aframes:
|
||||
if self.__duration and frame.time > start_time + self.__duration:
|
||||
if duration and frame.time > start_time + duration:
|
||||
audio_done = True
|
||||
break
|
||||
|
||||
@ -394,8 +434,8 @@ class VideoFromFile(VideoInput):
|
||||
|
||||
if len(audio_frames) > 0:
|
||||
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
||||
if self.__duration:
|
||||
audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
|
||||
if duration:
|
||||
audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
|
||||
|
||||
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
|
||||
audio = AudioInput({
|
||||
@ -441,28 +481,14 @@ class VideoFromFile(VideoInput):
|
||||
if not reuse_streams:
|
||||
if bit_depth is None:
|
||||
bit_depth = source_bit_depth
|
||||
components = self.get_components_internal(container)
|
||||
video = VideoFromComponents(components)
|
||||
return video.save_to(
|
||||
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
|
||||
)
|
||||
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
|
||||
|
||||
streams = container.streams
|
||||
|
||||
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
||||
with av.open(path, **open_kwargs) as output_container:
|
||||
# Copy over the original metadata
|
||||
for key, value in container.metadata.items():
|
||||
if metadata is None or key not in metadata:
|
||||
output_container.metadata[key] = value
|
||||
|
||||
# Add our new metadata
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
if isinstance(value, str):
|
||||
output_container.metadata[key] = value
|
||||
else:
|
||||
output_container.metadata[key] = json.dumps(value)
|
||||
# Add metadata before writing any streams
|
||||
write_output_metadata(container, output_container, metadata)
|
||||
|
||||
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
||||
stream_map = {}
|
||||
@ -480,6 +506,282 @@ class VideoFromFile(VideoInput):
|
||||
packet.stream = stream_map[packet.stream]
|
||||
output_container.mux(packet)
|
||||
|
||||
def _save_transcoded(
|
||||
self,
|
||||
container: InputContainer,
|
||||
path: str | io.BytesIO,
|
||||
format: VideoContainer,
|
||||
codec: VideoCodec,
|
||||
metadata: dict | None,
|
||||
bit_depth: int,
|
||||
):
|
||||
"""Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
|
||||
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
|
||||
stream_end_pts = None
|
||||
if video_stream.duration is not None:
|
||||
stream_end_pts = (video_stream.start_time or 0) + video_stream.duration
|
||||
output_end_pts = end_pts
|
||||
if stream_end_pts is not None and (output_end_pts is None or stream_end_pts < output_end_pts):
|
||||
output_end_pts = stream_end_pts
|
||||
if start_pts != 0:
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
|
||||
audio_stream = last_decodable_audio_stream(container)
|
||||
pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p"
|
||||
rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||
|
||||
resampler = None
|
||||
sample_rate = 0
|
||||
audio_time_base = None
|
||||
duration_cap = None
|
||||
if audio_stream is not None:
|
||||
sample_rate = audio_stream.codec_context.sample_rate
|
||||
channels = audio_stream.codec_context.channels
|
||||
if not sample_rate:
|
||||
sample_rate, channels = probe_audio_params(container, audio_stream)
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
if sample_rate:
|
||||
audio_stream.codec_context.flush_buffers()
|
||||
else:
|
||||
logging.warning("Audio stream parameters could not be determined; ignoring audio.")
|
||||
audio_stream = None
|
||||
if audio_stream is not None:
|
||||
audio_time_base = Fraction(1, sample_rate)
|
||||
layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
|
||||
resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
|
||||
if duration:
|
||||
duration_cap = math.ceil(duration * sample_rate)
|
||||
|
||||
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
|
||||
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
|
||||
video_done = False
|
||||
audio_done = audio_stream is None
|
||||
video_pts_offset = None
|
||||
last_video_pts = None
|
||||
last_video_end = None
|
||||
# rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise
|
||||
video_frame_durations = {}
|
||||
source_size = None
|
||||
rotation_k = 0
|
||||
rotation_filter = None
|
||||
audio_started = False
|
||||
samples_written = 0
|
||||
pending_audio = []
|
||||
# The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims),
|
||||
# and never seeking back keeps webm/mkv leading audio intact.
|
||||
output = None
|
||||
out_video = None
|
||||
out_audio = None
|
||||
|
||||
def audio_frame_from_ndarray(nd_planar):
|
||||
frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout)
|
||||
frame.sample_rate = sample_rate
|
||||
return frame
|
||||
|
||||
def drain_audio(final=False):
|
||||
# Audio may cover the pts span of the video written so far, capped by the requested duration
|
||||
nonlocal samples_written, audio_done
|
||||
if last_video_end is None:
|
||||
cap = 0
|
||||
else:
|
||||
cap = math.ceil(last_video_end * video_stream.time_base * sample_rate)
|
||||
if duration_cap is not None:
|
||||
cap = min(cap, duration_cap)
|
||||
while pending_audio and not audio_done:
|
||||
frame = pending_audio[0]
|
||||
if samples_written + frame.samples <= cap:
|
||||
frame.pts = samples_written
|
||||
frame.time_base = audio_time_base
|
||||
output.mux(out_audio.encode(frame))
|
||||
samples_written += frame.samples
|
||||
pending_audio.pop(0)
|
||||
continue
|
||||
if final:
|
||||
keep = frame.to_ndarray()[..., :cap - samples_written]
|
||||
if keep.shape[-1] > 0:
|
||||
tail = audio_frame_from_ndarray(keep)
|
||||
tail.pts = samples_written
|
||||
tail.time_base = audio_time_base
|
||||
output.mux(out_audio.encode(tail))
|
||||
samples_written += keep.shape[-1]
|
||||
pending_audio.clear()
|
||||
break
|
||||
if duration_cap is not None and samples_written >= duration_cap:
|
||||
audio_done = True
|
||||
return cap
|
||||
|
||||
try:
|
||||
for packet in container.demux(*streams):
|
||||
if video_done and audio_done:
|
||||
break
|
||||
|
||||
if packet.stream == video_stream and not video_done:
|
||||
try:
|
||||
frames = packet.decode()
|
||||
except av.error.InvalidDataError:
|
||||
logging.info("pyav decode error")
|
||||
continue
|
||||
for frame in frames:
|
||||
if frame.pts is not None and frame.pts < start_pts:
|
||||
continue
|
||||
if end_pts is not None and frame.pts is not None and frame.pts >= end_pts:
|
||||
video_done = True
|
||||
if last_video_pts is not None:
|
||||
# the source continues past the window: hold the last kept frame to the window end
|
||||
end_offset = video_pts_offset if video_pts_offset is not None else start_pts
|
||||
last_video_end = max(last_video_end, end_pts - end_offset)
|
||||
break
|
||||
# the source's true display duration of this frame; average_rate is not a
|
||||
# frame duration (sparse/VFR sources), so it is only the fallback
|
||||
frame_duration = frame.duration if frame.duration else pts_step
|
||||
if end_pts is not None and frame.pts is not None:
|
||||
frame_duration = min(frame_duration, end_pts - frame.pts)
|
||||
if output is None:
|
||||
rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
|
||||
if rotation_k % 2:
|
||||
out_width, out_height = frame.height, frame.width
|
||||
else:
|
||||
out_width, out_height = frame.width, frame.height
|
||||
if out_width % 2 or out_height % 2:
|
||||
raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
|
||||
source_size = (frame.width, frame.height)
|
||||
output = av.open(path, **open_kwargs)
|
||||
# Add metadata before writing any streams
|
||||
write_output_metadata(container, output, metadata)
|
||||
out_video = output.add_stream("h264", rate=rate)
|
||||
# no B-frames: reordering makes mp4 sample durations follow decode order,
|
||||
# so irregular-VFR spans and trim windows land wrong
|
||||
out_video.codec_context.max_b_frames = 0
|
||||
out_video.width = out_width
|
||||
out_video.height = out_height
|
||||
out_video.pix_fmt = pix_fmt
|
||||
# source pts pass through (rebased to 0), so variable frame rate survives
|
||||
out_video.codec_context.time_base = video_stream.time_base
|
||||
if audio_stream is not None:
|
||||
out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
|
||||
if (frame.width, frame.height) != source_size:
|
||||
# encoding would silently rescale the new geometry into the old one
|
||||
raise ValueError(
|
||||
f"Video resolution changes mid-stream "
|
||||
f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode"
|
||||
)
|
||||
if rotation_k:
|
||||
if rotation_filter is None:
|
||||
g = av.filter.Graph()
|
||||
g_src = g.add_buffer(width=frame.width, height=frame.height,
|
||||
format=frame.format.name, time_base=video_stream.time_base)
|
||||
tail = g_src
|
||||
for filter_name, filter_args in {1: [("transpose", "cclock")],
|
||||
2: [("hflip", None), ("vflip", None)],
|
||||
3: [("transpose", "clock")]}[rotation_k]:
|
||||
step = g.add(filter_name, filter_args)
|
||||
tail.link_to(step)
|
||||
tail = step
|
||||
g_sink = g.add("buffersink")
|
||||
tail.link_to(g_sink)
|
||||
g.configure()
|
||||
rotation_filter = (g_src, g_sink)
|
||||
rotation_filter[0].push(frame)
|
||||
frame = rotation_filter[1].pull()
|
||||
if frame.color_range == ColorRange.JPEG:
|
||||
# compress full-range sources (yuvj/MJPEG) to limited range
|
||||
frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
|
||||
else:
|
||||
frame = frame.reformat(format=pix_fmt)
|
||||
frame_output_end = None
|
||||
if frame.pts is not None:
|
||||
if video_pts_offset is None:
|
||||
video_pts_offset = frame.pts
|
||||
frame.pts -= video_pts_offset
|
||||
if output_end_pts is not None:
|
||||
frame_output_end = output_end_pts - video_pts_offset
|
||||
if frame.pts + frame_duration > frame_output_end:
|
||||
clamped_pts = frame_output_end - frame_duration
|
||||
if clamped_pts >= 0 and (last_video_pts is None or clamped_pts > last_video_pts):
|
||||
frame.pts = min(frame.pts, clamped_pts)
|
||||
elif frame.pts < frame_output_end:
|
||||
frame_duration = frame_output_end - frame.pts
|
||||
else:
|
||||
continue
|
||||
if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts):
|
||||
# broken sources emit missing/backward timestamps mid-stream, which the
|
||||
# muxer rejects; nudge them forward by one nominal frame interval
|
||||
frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step
|
||||
if frame_output_end is not None and frame.pts + frame_duration > frame_output_end:
|
||||
if frame.pts >= frame_output_end:
|
||||
continue
|
||||
frame_duration = frame_output_end - frame.pts
|
||||
last_video_pts = frame.pts
|
||||
last_video_end = frame.pts + frame_duration
|
||||
video_frame_durations[frame.pts] = frame_duration
|
||||
# the decoded pict_type would force x264's frame types (intra-only
|
||||
# sources like MJPEG/ProRes would come out all-keyframe)
|
||||
frame.pict_type = 0
|
||||
for out_packet in out_video.encode(frame):
|
||||
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||
output.mux(out_packet)
|
||||
drain_audio()
|
||||
|
||||
elif packet.stream == audio_stream and not audio_done:
|
||||
for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
|
||||
frame_start = None
|
||||
if resampled.pts is not None:
|
||||
# passthrough frames keep the source stream's time base
|
||||
tb = resampled.time_base if resampled.time_base else audio_time_base
|
||||
frame_start = float(resampled.pts * tb)
|
||||
if duration and not audio_started and frame_start >= start_time + duration:
|
||||
audio_done = True
|
||||
break
|
||||
if not audio_started:
|
||||
if frame_start is None:
|
||||
frame_start = 0.0
|
||||
to_skip = max(0, int((start_time - frame_start) * sample_rate))
|
||||
if to_skip >= resampled.samples:
|
||||
continue
|
||||
audio_started = True
|
||||
if duration and frame_start > start_time:
|
||||
duration_cap = min(duration_cap, math.ceil((start_time + duration - frame_start) * sample_rate))
|
||||
if to_skip:
|
||||
pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
|
||||
continue
|
||||
pending_audio.append(resampled)
|
||||
if video_done:
|
||||
# the video window is complete so the cap is final, but containers
|
||||
# that interleave audio behind video (fragmented mp4) still owe most
|
||||
# of it: stop only once the demuxed audio covers the cap
|
||||
cap = drain_audio()
|
||||
if pending_audio or samples_written >= cap:
|
||||
drain_audio(final=True)
|
||||
audio_done = True
|
||||
break
|
||||
|
||||
if output is None:
|
||||
raise ValueError(f"No decodable video frames found in file '{self.__file}'")
|
||||
if out_audio is not None and not audio_done:
|
||||
drain_audio(final=True)
|
||||
window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
|
||||
for out_packet in out_video.encode(None):
|
||||
duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||
if out_packet.pts == last_video_pts:
|
||||
duration = max(duration, window_fill)
|
||||
out_packet.duration = duration
|
||||
output.mux(out_packet)
|
||||
if out_audio is not None:
|
||||
output.mux(out_audio.encode(None))
|
||||
except BaseException:
|
||||
if output is not None:
|
||||
output.close()
|
||||
if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
|
||||
os.remove(path)
|
||||
raise
|
||||
else:
|
||||
if output is not None:
|
||||
output.close()
|
||||
|
||||
def _get_first_video_stream(self, container: InputContainer):
|
||||
if len(container.streams.video):
|
||||
return container.streams.video[0]
|
||||
@ -527,22 +829,12 @@ class VideoFromComponents(VideoInput):
|
||||
bit_depth: int | None = None,
|
||||
):
|
||||
"""Save the video to a file path or BytesIO buffer."""
|
||||
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
||||
raise ValueError("Only MP4 format is supported for now")
|
||||
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
||||
raise ValueError("Only H264 codec is supported for now")
|
||||
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||
# None means "use the depth this video was created with" (CreateVideo's choice).
|
||||
if bit_depth is None:
|
||||
bit_depth = self.__bit_depth
|
||||
is_10bit = bit_depth >= 10
|
||||
extra_kwargs = {}
|
||||
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
||||
extra_kwargs["format"] = format.value
|
||||
elif isinstance(path, io.BytesIO):
|
||||
# BytesIO has no file extension, so av.open can't infer the format.
|
||||
# Default to mp4 since that's the only supported format anyway.
|
||||
extra_kwargs["format"] = "mp4"
|
||||
with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
|
||||
with av.open(path, **open_kwargs) as output:
|
||||
# Add metadata before writing any streams
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
|
||||
49
comfy_api_nodes/apis/sync_so.py
Normal file
49
comfy_api_nodes/apis/sync_so.py
Normal file
@ -0,0 +1,49 @@
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SyncInputItem(BaseModel):
|
||||
type: str = Field(..., description="Input kind: 'video', 'image' or 'audio'.")
|
||||
url: str = Field(...)
|
||||
|
||||
|
||||
class SyncActiveSpeakerDetection(BaseModel):
|
||||
auto_detect: bool | None = Field(
|
||||
None, description="Detect the active speaker automatically. Video input only; rejected for images."
|
||||
)
|
||||
frame_number: int | None = Field(
|
||||
None, description="Frame used for manual speaker selection. Must be 0 for image inputs."
|
||||
)
|
||||
coordinates: list[int] | None = Field(
|
||||
None, description="Pixel [x, y] of the speaker's face in the frame selected by frame_number."
|
||||
)
|
||||
|
||||
|
||||
class SyncGenerationOptions(BaseModel):
|
||||
sync_mode: str | None = Field(
|
||||
None,
|
||||
description="How to resolve an audio/video duration mismatch: "
|
||||
"cut_off, bounce, loop, silence or remap. Ignored for image inputs.",
|
||||
)
|
||||
i2v_prompt: str | None = Field(
|
||||
None, description="Motion prompt for image-to-video generation. Image input only."
|
||||
)
|
||||
active_speaker_detection: SyncActiveSpeakerDetection | None = Field(None)
|
||||
|
||||
|
||||
class SyncGenerationRequest(BaseModel):
|
||||
model: str = Field(..., description="Generation model, e.g. 'sync-3'.")
|
||||
input: list[SyncInputItem] = Field(
|
||||
..., description="Exactly one visual input (video or image) plus one audio input."
|
||||
)
|
||||
options: SyncGenerationOptions | None = Field(None)
|
||||
|
||||
|
||||
class SyncGeneration(BaseModel):
|
||||
"""Subset of the Generation object returned by POST /v2/generate and GET /v2/generate/{id}."""
|
||||
|
||||
id: str = Field(...)
|
||||
status: str = Field(..., description="PENDING | PROCESSING | COMPLETED | FAILED | REJECTED")
|
||||
outputUrl: str | None = Field(None)
|
||||
outputDuration: float | None = Field(None)
|
||||
error: str | None = Field(None, description="Human-readable failure message.")
|
||||
errorCode: str | None = Field(None, description="Stable machine-readable code from the GET /v2/errors catalog.")
|
||||
391
comfy_api_nodes/nodes_sync_so.py
Normal file
391
comfy_api_nodes/nodes_sync_so.py
Normal file
@ -0,0 +1,391 @@
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.sync_so import (
|
||||
SyncActiveSpeakerDetection,
|
||||
SyncGeneration,
|
||||
SyncGenerationOptions,
|
||||
SyncGenerationRequest,
|
||||
SyncInputItem,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
download_url_to_video_output,
|
||||
downscale_image_tensor,
|
||||
downscale_image_tensor_by_max_side,
|
||||
get_image_dimensions,
|
||||
get_number_of_images,
|
||||
poll_op,
|
||||
sync_op,
|
||||
upload_audio_to_comfyapi,
|
||||
upload_image_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_audio_duration,
|
||||
)
|
||||
|
||||
|
||||
class SyncLipSyncNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="SyncLipSyncNode",
|
||||
display_name="sync.so Lip Sync",
|
||||
category="partner/video/sync.so",
|
||||
description=(
|
||||
"Re-sync mouth movement in a video to new speech audio using sync.so. "
|
||||
"Handles close-ups, profiles and obstructions automatically while preserving "
|
||||
"the speaker's expression. Cost scales with output duration."
|
||||
),
|
||||
inputs=[
|
||||
IO.Video.Input(
|
||||
"video",
|
||||
tooltip="Footage of the speaker to re-sync. Up to 4K (4096x2160); "
|
||||
"a constant frame rate of 24/25/30 fps works best.",
|
||||
),
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Speech audio to sync the mouth to.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed controls whether the node should re-run; "
|
||||
"results are non-deterministic regardless of seed.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"sync-3",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"sync_mode",
|
||||
options=["bounce", "cut_off", "loop", "silence", "remap"],
|
||||
default="bounce",
|
||||
tooltip=(
|
||||
"How to handle a duration mismatch between video and audio; "
|
||||
"this also sets the output length. "
|
||||
"bounce: video plays forward then backward until the audio ends "
|
||||
"(output = audio length). "
|
||||
"loop: video restarts until the audio ends (output = audio length). "
|
||||
"remap: video is time-stretched to match the audio (output = audio length). "
|
||||
"cut_off: the longer track is trimmed (output = shorter length). "
|
||||
"silence: nothing is trimmed; the shorter track is padded "
|
||||
"(output = longer length)."
|
||||
),
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"speaker_selection",
|
||||
options=["default", "auto-detect", "coordinates"],
|
||||
default="default",
|
||||
tooltip=(
|
||||
"Which face to lipsync when several people are visible. "
|
||||
"default: let the model decide. "
|
||||
"auto-detect: detect and follow the active speaker. "
|
||||
"coordinates: target the face at pixel (speaker_x, speaker_y) "
|
||||
"in the frame chosen by speaker_frame."
|
||||
),
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_frame",
|
||||
default=0,
|
||||
min=0,
|
||||
max=1_000_000,
|
||||
advanced=True,
|
||||
tooltip="Video frame used to locate the speaker. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_x",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="X pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_y",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="Y pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
tooltip="sync.so generation model.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
video: Input.Video,
|
||||
audio: Input.Audio,
|
||||
seed: int,
|
||||
model: dict,
|
||||
) -> IO.NodeOutput:
|
||||
try:
|
||||
width, height = video.get_dimensions()
|
||||
except Exception:
|
||||
width = height = None
|
||||
if width and height and (max(width, height) > 4096 or width * height > 4096 * 2160):
|
||||
raise ValueError(
|
||||
f"sync.so rejects videos above 4K (4096x2160); got {width}x{height}. Downscale the video first."
|
||||
)
|
||||
validate_audio_duration(audio, max_duration=600)
|
||||
|
||||
if model["speaker_selection"] == "auto-detect":
|
||||
speaker_detection = SyncActiveSpeakerDetection(auto_detect=True)
|
||||
elif model["speaker_selection"] == "coordinates":
|
||||
speaker_detection = SyncActiveSpeakerDetection(
|
||||
frame_number=model["speaker_frame"],
|
||||
coordinates=[model["speaker_x"], model["speaker_y"]],
|
||||
)
|
||||
else:
|
||||
speaker_detection = None
|
||||
|
||||
video_url = await upload_video_to_comfyapi(cls, video, max_duration=600)
|
||||
audio_url = await upload_audio_to_comfyapi(cls, audio)
|
||||
|
||||
generation = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
|
||||
response_model=SyncGeneration,
|
||||
data=SyncGenerationRequest(
|
||||
model=model["model"],
|
||||
input=[
|
||||
SyncInputItem(type="video", url=video_url),
|
||||
SyncInputItem(type="audio", url=audio_url),
|
||||
],
|
||||
options=SyncGenerationOptions(
|
||||
sync_mode=model["sync_mode"],
|
||||
active_speaker_detection=speaker_detection,
|
||||
),
|
||||
),
|
||||
)
|
||||
generation = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
|
||||
response_model=SyncGeneration,
|
||||
status_extractor=lambda g: g.status,
|
||||
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
|
||||
failed_statuses=[],
|
||||
queued_statuses=["PENDING"],
|
||||
poll_interval=10.0,
|
||||
)
|
||||
if generation.status != "COMPLETED":
|
||||
code = f" [{generation.errorCode}]" if generation.errorCode else ""
|
||||
raise ValueError(
|
||||
f"sync.so generation {generation.status.lower()}{code}: "
|
||||
f"{generation.error or 'no error details provided'}"
|
||||
)
|
||||
if not generation.outputUrl:
|
||||
raise ValueError("sync.so generation completed but no output URL was returned.")
|
||||
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
|
||||
|
||||
|
||||
class SyncTalkingImageNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="SyncTalkingImageNode",
|
||||
display_name="sync.so Talking Image",
|
||||
category="partner/video/sync.so",
|
||||
description=(
|
||||
"Animate a still portrait into a talking video driven by speech audio, "
|
||||
"using sync.so's sync-3 model. The output duration matches the audio. "
|
||||
"Cost scales with output duration."
|
||||
),
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
tooltip="A single image with a clearly visible face, up to 4K (4096x2160).",
|
||||
),
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Speech audio driving the talking video; the output duration matches it. "
|
||||
"Chain any TTS node here to drive the animation from text.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Optional guidance for how the portrait comes to life, e.g. "
|
||||
"'make the subject smile and look at the camera'. "
|
||||
"Leave empty for natural talking motion.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed controls whether the node should re-run; "
|
||||
"results are non-deterministic regardless of seed.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"sync-3",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"speaker_selection",
|
||||
options=["default", "coordinates"],
|
||||
default="default",
|
||||
tooltip=(
|
||||
"Which face to animate when several people are visible. "
|
||||
"default: let the model decide. "
|
||||
"coordinates: target the face at pixel (speaker_x, speaker_y) "
|
||||
"in the image. Auto-detection is not supported for images."
|
||||
),
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_x",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="X pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_y",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="Y pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"auto_downscale",
|
||||
default=True,
|
||||
advanced=True,
|
||||
tooltip="Automatically downscale the image if it exceeds the 4K "
|
||||
"(4096x2160) input limit; speaker coordinates are scaled to match. "
|
||||
"When disabled, an oversized image raises an error instead.",
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
tooltip="sync.so generation model. Image input is exclusive to sync-3.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
audio: Input.Audio,
|
||||
prompt: str,
|
||||
seed: int,
|
||||
model: dict,
|
||||
) -> IO.NodeOutput:
|
||||
if get_number_of_images(image) != 1:
|
||||
raise ValueError("Exactly one image is required; got a batch. Pick one frame first.")
|
||||
validate_audio_duration(audio, max_duration=600)
|
||||
|
||||
height, width = get_image_dimensions(image)
|
||||
speaker_x, speaker_y = model["speaker_x"], model["speaker_y"]
|
||||
if max(width, height) > 4096 or width * height > 4096 * 2160:
|
||||
if not model["auto_downscale"]:
|
||||
raise ValueError(
|
||||
f"sync.so rejects images above 4K (4096x2160); got {width}x{height}. "
|
||||
"Downscale the image first or enable auto_downscale."
|
||||
)
|
||||
image = downscale_image_tensor(image, total_pixels=4096 * 2160)
|
||||
image = downscale_image_tensor_by_max_side(image, max_side=4096)
|
||||
new_height, new_width = get_image_dimensions(image)
|
||||
# speaker coordinates are given in the original image's pixel space
|
||||
speaker_x = min(new_width - 1, round(speaker_x * new_width / width))
|
||||
speaker_y = min(new_height - 1, round(speaker_y * new_height / height))
|
||||
|
||||
if model["speaker_selection"] == "coordinates":
|
||||
speaker_detection = SyncActiveSpeakerDetection(
|
||||
frame_number=0, # images have a single frame; auto_detect is rejected by the API
|
||||
coordinates=[speaker_x, speaker_y],
|
||||
)
|
||||
else:
|
||||
speaker_detection = None
|
||||
|
||||
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||
audio_url = await upload_audio_to_comfyapi(cls, audio)
|
||||
|
||||
generation = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
|
||||
response_model=SyncGeneration,
|
||||
data=SyncGenerationRequest(
|
||||
model=model["model"],
|
||||
input=[
|
||||
SyncInputItem(type="image", url=image_url),
|
||||
SyncInputItem(type="audio", url=audio_url),
|
||||
],
|
||||
options=SyncGenerationOptions(
|
||||
i2v_prompt=prompt.strip() or None,
|
||||
active_speaker_detection=speaker_detection,
|
||||
),
|
||||
),
|
||||
)
|
||||
generation = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
|
||||
response_model=SyncGeneration,
|
||||
status_extractor=lambda g: g.status,
|
||||
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
|
||||
failed_statuses=[],
|
||||
queued_statuses=["PENDING"],
|
||||
poll_interval=10.0,
|
||||
)
|
||||
if generation.status != "COMPLETED":
|
||||
code = f" [{generation.errorCode}]" if generation.errorCode else ""
|
||||
raise ValueError(
|
||||
f"sync.so generation {generation.status.lower()}{code}: "
|
||||
f"{generation.error or 'no error details provided'}"
|
||||
)
|
||||
if not generation.outputUrl:
|
||||
raise ValueError("sync.so generation completed but no output URL was returned.")
|
||||
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
|
||||
|
||||
|
||||
class SyncExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
SyncLipSyncNode,
|
||||
SyncTalkingImageNode,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> SyncExtension:
|
||||
return SyncExtension()
|
||||
@ -15,6 +15,7 @@ from comfy.comfy_api_env import normalize_comfy_api_base
|
||||
from comfy.deploy_environment import get_deploy_environment
|
||||
from comfy.model_management import processing_interrupted
|
||||
from comfy_api.latest import IO
|
||||
from comfy_execution.utils import get_executing_context
|
||||
from comfyui_version import __version__ as comfyui_version
|
||||
|
||||
from .common_exceptions import ProcessingInterrupted
|
||||
@ -57,12 +58,16 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
|
||||
relative/cloud URLs resolved against ``default_base_url()``; because the result
|
||||
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
||||
"""
|
||||
return {
|
||||
headers = {
|
||||
**get_auth_header(node_cls),
|
||||
"Comfy-Env": get_deploy_environment(),
|
||||
"Comfy-Usage-Source": get_usage_source(node_cls),
|
||||
"Comfy-Core-Version": comfyui_version,
|
||||
}
|
||||
ctx = get_executing_context()
|
||||
if ctx is not None:
|
||||
headers["Comfy-Job-Id"] = ctx.prompt_id
|
||||
return headers
|
||||
|
||||
|
||||
def default_base_url() -> str:
|
||||
|
||||
@ -174,8 +174,9 @@ class Preview3DAdvanced(IO.ComfyNode):
|
||||
filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
@ -243,8 +244,9 @@ class PreviewGaussianSplat(IO.ComfyNode):
|
||||
filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
@ -303,8 +305,9 @@ class PreviewPointCloud(IO.ComfyNode):
|
||||
filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
@ -375,8 +378,9 @@ class Load3DAdvanced(IO.ComfyNode):
|
||||
file_3d = None
|
||||
if model_file and model_file != "none":
|
||||
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
model_3d_info = viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height)
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height)
|
||||
|
||||
|
||||
class Load3DExtension(ComfyExtension):
|
||||
|
||||
@ -418,8 +418,9 @@ def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str:
|
||||
|
||||
def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
|
||||
model_file = _save_file3d_to_output(model_3d, filename_prefix)
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
|
||||
@ -1,3 +1,3 @@
|
||||
# This file is automatically generated by the build process when version is
|
||||
# updated in pyproject.toml.
|
||||
__version__ = "0.27.0"
|
||||
__version__ = "0.28.0"
|
||||
|
||||
@ -426,12 +426,12 @@ def _is_intermediate_output(dynprompt, node_id):
|
||||
|
||||
|
||||
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
|
||||
if cached.ui is not None:
|
||||
ui_outputs[node_id] = cached.ui
|
||||
if server.client_id is None:
|
||||
return
|
||||
cached_ui = cached.ui or {}
|
||||
server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
|
||||
if cached.ui is not None:
|
||||
ui_outputs[node_id] = cached.ui
|
||||
|
||||
async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
|
||||
unique_id = current_item
|
||||
|
||||
45
openapi.yaml
45
openapi.yaml
@ -7,18 +7,18 @@ components:
|
||||
description: Timestamp when the asset was created
|
||||
format: date-time
|
||||
type: string
|
||||
display_name:
|
||||
description: Display name of the asset. Mirrors name for backwards compatibility.
|
||||
nullable: true
|
||||
type: string
|
||||
file_path:
|
||||
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
|
||||
nullable: true
|
||||
type: string
|
||||
hash:
|
||||
description: Blake3 hash of the asset content.
|
||||
pattern: ^blake3:[a-f0-9]{64}$
|
||||
type: string
|
||||
loader_path:
|
||||
description: The value a loader consumes to load this asset. Null when no loader can resolve the file.
|
||||
nullable: true
|
||||
type: string
|
||||
display_name:
|
||||
description: Human-facing label for the asset. Not unique.
|
||||
nullable: true
|
||||
type: string
|
||||
id:
|
||||
description: Unique identifier for the asset
|
||||
format: uuid
|
||||
@ -144,6 +144,14 @@ components:
|
||||
AssetUpdated:
|
||||
description: Response returned when an existing asset is successfully updated.
|
||||
properties:
|
||||
display_name:
|
||||
description: Display name of the asset. Mirrors name for backwards compatibility.
|
||||
nullable: true
|
||||
type: string
|
||||
file_path:
|
||||
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
|
||||
nullable: true
|
||||
type: string
|
||||
hash:
|
||||
description: Blake3 hash of the asset content.
|
||||
pattern: ^blake3:[a-f0-9]{64}$
|
||||
@ -775,14 +783,6 @@ components:
|
||||
ModelFolder:
|
||||
description: Represents a folder containing models
|
||||
properties:
|
||||
extensions:
|
||||
description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all).
|
||||
example:
|
||||
- .ckpt
|
||||
- .safetensors
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
folders:
|
||||
description: List of paths where models of this type are stored
|
||||
example:
|
||||
@ -1644,7 +1644,7 @@ paths:
|
||||
format: uuid
|
||||
type: string
|
||||
tags:
|
||||
description: JSON-encoded array of tag strings. For new byte uploads, include exactly one destination role (`input`, `output`, or `models`); `models` uploads also require exactly one `model_type:<folder_name>` tag. Extra tags are stored as labels and do not create path components.
|
||||
description: JSON-encoded array of freeform tag strings, e.g. '["models","checkpoint"]'. Common types include "models", "input", "output", and "temp", but any tag can be used in any order.
|
||||
type: string
|
||||
user_metadata:
|
||||
description: Custom JSON metadata as a string
|
||||
@ -1829,7 +1829,7 @@ paths:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/Asset'
|
||||
$ref: '#/components/schemas/AssetUpdated'
|
||||
description: Asset updated successfully
|
||||
"400":
|
||||
content:
|
||||
@ -2470,9 +2470,6 @@ paths:
|
||||
supports_preview_metadata:
|
||||
description: Whether the server supports preview metadata
|
||||
type: boolean
|
||||
supports_model_type_tags:
|
||||
description: Whether the server supports namespaced model type asset tags
|
||||
type: boolean
|
||||
type: object
|
||||
description: Success
|
||||
headers:
|
||||
@ -3300,6 +3297,12 @@ paths:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Invalid request parameters
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unauthorized - Authentication required
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "ComfyUI"
|
||||
version = "0.27.0"
|
||||
version = "0.28.0"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
comfyui-frontend-package==1.45.20
|
||||
comfyui-workflow-templates==0.11.6
|
||||
comfyui-embedded-docs==0.5.7
|
||||
comfyui-frontend-package==1.45.21
|
||||
comfyui-workflow-templates==0.11.9
|
||||
comfyui-embedded-docs==0.5.8
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
@ -22,7 +22,7 @@ alembic
|
||||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.18
|
||||
comfy-kitchen==0.2.20
|
||||
comfy-aimdo==0.4.10
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
|
||||
@ -2,11 +2,12 @@ import pytest
|
||||
import torch
|
||||
import tempfile
|
||||
import os
|
||||
import sys
|
||||
import av
|
||||
import io
|
||||
from fractions import Fraction
|
||||
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
|
||||
from comfy_api.util.video_types import VideoComponents
|
||||
from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
|
||||
from comfy_api.input.basic_types import AudioInput
|
||||
from av.error import InvalidDataError
|
||||
|
||||
@ -237,3 +238,526 @@ def test_duration_consistency(video_components):
|
||||
manual_duration = float(components.images.shape[0] / components.frame_rate)
|
||||
|
||||
assert duration == pytest.approx(manual_duration)
|
||||
|
||||
|
||||
def create_transcode_source(
|
||||
width=64, height=64, frames=30, fps=30, audio_streams=1, undecodable_audio=0, rotation=False,
|
||||
container_format="mov", audio_codec="pcm_s16le",
|
||||
):
|
||||
"""Create a temp video that save_to must transcode (mpeg4 video, so codec != h264).
|
||||
|
||||
``undecodable_audio`` trailing PCM streams get their fourcc corrupted so no decoder exists
|
||||
(``codec_context is None``), like the APAC track in iPhone spatial-audio recordings.
|
||||
``rotation`` patches a 90-degree display matrix into the video track header.
|
||||
"""
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format=container_format) as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||
video_stream.width = width
|
||||
video_stream.height = height
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
audio = []
|
||||
for _ in range(audio_streams + undecodable_audio):
|
||||
stream = container.add_stream(audio_codec, rate=44100)
|
||||
stream.sample_rate = 44100
|
||||
audio.append(stream)
|
||||
|
||||
for i in range(frames):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((height, width, 3), (i * 7) % 256, dtype=torch.uint8).numpy(),
|
||||
format="rgb24",
|
||||
)
|
||||
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||
# write audio in 1024-sample frames, like real decoders produce, so the
|
||||
# per-frame skip/cap logic in the transcode path actually runs
|
||||
for stream in audio:
|
||||
for offset in range(0, 44100 * frames // fps, 1024):
|
||||
n = min(1024, 44100 * frames // fps - offset)
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(1, n, dtype=torch.int16).numpy(), format="s16", layout="mono"
|
||||
)
|
||||
audio_frame.sample_rate = 44100
|
||||
audio_frame.pts = offset
|
||||
container.mux(stream.encode(audio_frame))
|
||||
for stream in [video_stream, *audio]:
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
data = bytearray(buffer.getvalue())
|
||||
end = len(data)
|
||||
for _ in range(undecodable_audio):
|
||||
end = data.rindex(b"sowt", 0, end)
|
||||
data[end:end + 4] = b"Xpac"
|
||||
if rotation:
|
||||
# the 3x3 display matrix sits 40 bytes into the version-0 tkhd payload; first tkhd
|
||||
# inside moov = video track (search from moov so mdat bytes can't false-match)
|
||||
matrix_offset = data.index(b"tkhd", data.rindex(b"moov")) + 4 + 40
|
||||
values = [0, 1 << 16, 0, -(1 << 16), 0, 0, 0, 0, 1 << 30]
|
||||
data[matrix_offset:matrix_offset + 36] = b"".join(v.to_bytes(4, "big", signed=True) for v in values)
|
||||
|
||||
tmp = tempfile.NamedTemporaryFile(suffix=f".{container_format}", delete=False)
|
||||
tmp.write(bytes(data))
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
def transcode_and_probe(video):
|
||||
buffer = io.BytesIO()
|
||||
video.save_to(buffer, format=VideoContainer.MP4, codec=VideoCodec.H264)
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
video_stream = container.streams.video[0]
|
||||
audio_stream = container.streams.audio[0] if container.streams.audio else None
|
||||
frames = 0
|
||||
first_pts = None
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
if first_pts is None:
|
||||
first_pts = frame.pts
|
||||
frames += 1
|
||||
return {
|
||||
"codec": video_stream.codec_context.name,
|
||||
"width": video_stream.codec_context.width,
|
||||
"height": video_stream.codec_context.height,
|
||||
"frames": frames,
|
||||
"first_pts": first_pts,
|
||||
"video_seconds": float(video_stream.duration * video_stream.time_base) if video_stream.duration else None,
|
||||
"audio_seconds": float(audio_stream.duration * audio_stream.time_base)
|
||||
if audio_stream and audio_stream.duration else None,
|
||||
"audio_codecs": [s.codec_context.name for s in container.streams.audio],
|
||||
}
|
||||
|
||||
|
||||
def test_save_to_transcode_streams_without_buffering_frames():
|
||||
"""Transcoding must not decode the whole video into memory first (~2 GiB for this source)"""
|
||||
resource = pytest.importorskip("resource") # no getrusage on Windows
|
||||
rss_scale = 1 if sys.platform == "darwin" else 1024 # ru_maxrss: bytes on macOS, KiB elsewhere
|
||||
# ru_maxrss is a lifetime peak: a heavier test running earlier would shrink the measured
|
||||
# delta and quietly defang this canary, so keep this source the biggest thing in the suite
|
||||
file_path = create_transcode_source(width=640, height=480, frames=300)
|
||||
try:
|
||||
rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale
|
||||
result = transcode_and_probe(VideoFromFile(file_path))
|
||||
rss_delta = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale - rss_before
|
||||
|
||||
assert result["codec"] == "h264"
|
||||
assert result["frames"] == 300
|
||||
assert rss_delta < 500 * 2**20, f"transcode buffered frames in RAM (peak grew {rss_delta / 2**20:.0f} MiB)"
|
||||
finally:
|
||||
os.unlink(file_path)
|
||||
|
||||
|
||||
def test_save_to_transcode_honors_trim_window():
|
||||
"""start_time/duration trim applies to both video and audio on the streaming path"""
|
||||
file_path = create_transcode_source(frames=90) # 3s @ 30fps
|
||||
try:
|
||||
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
|
||||
assert result["frames"] == pytest.approx(30, abs=2)
|
||||
assert result["first_pts"] == 0 # trimmed output is rebased to start at zero
|
||||
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||
finally:
|
||||
os.unlink(file_path)
|
||||
|
||||
|
||||
def test_save_to_transcode_keeps_audio_of_sparse_video():
|
||||
"""Audio that runs ahead of a sparse video track (slideshows, timelapses) must be
|
||||
kept in full — it is only clamped to the video's end, never to the video cursor."""
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mp4") as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=30)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
|
||||
for t in (0, 30, 60): # 3 frames spread over 60 seconds
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), t * 4, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
).reformat(format="yuv420p")
|
||||
frame.pts = t * 15360
|
||||
frame.time_base = Fraction(1, 15360)
|
||||
container.mux(video_stream.encode(frame))
|
||||
container.mux(video_stream.encode(None))
|
||||
for offset in range(0, 48000 * 60, 1024):
|
||||
n = min(1024, 48000 * 60 - offset)
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
|
||||
)
|
||||
audio_frame.sample_rate = 48000
|
||||
audio_frame.pts = offset
|
||||
audio_frame.time_base = Fraction(1, 48000)
|
||||
container.mux(audio_stream.encode(audio_frame))
|
||||
container.mux(audio_stream.encode(None))
|
||||
|
||||
buffer.seek(0)
|
||||
result = transcode_and_probe(VideoFromFile(buffer))
|
||||
assert result["audio_seconds"] == pytest.approx(60.0, abs=1.0)
|
||||
|
||||
|
||||
def test_save_to_transcode_vfr_audio_covers_video_span():
|
||||
"""A trim window in the sparse region of a VFR file keeps audio for the true pts span
|
||||
of the kept frames. Deriving the span as frames/average_rate undercuts it badly: the
|
||||
average is dominated by the dense region (and can be plain wrong on MediaRecorder files)."""
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mp4") as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=30)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
audio_stream = container.add_stream("aac", rate=48000, layout="stereo")
|
||||
# 10 frames inside the first second, then one every 1.25 s
|
||||
for i, t in enumerate([x / 10 for x in range(10)] + [1.0, 2.25, 3.5, 4.75]):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), (i * 16) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
).reformat(format="yuv420p")
|
||||
frame.pts = int(t * 15360)
|
||||
frame.time_base = Fraction(1, 15360)
|
||||
container.mux(video_stream.encode(frame))
|
||||
container.mux(video_stream.encode(None))
|
||||
for offset in range(0, 48000 * 6, 1024):
|
||||
n = min(1024, 48000 * 6 - offset)
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo"
|
||||
)
|
||||
audio_frame.sample_rate = 48000
|
||||
audio_frame.pts = offset
|
||||
audio_frame.time_base = Fraction(1, 48000)
|
||||
container.mux(audio_stream.encode(audio_frame))
|
||||
container.mux(audio_stream.encode(None))
|
||||
|
||||
buffer.seek(0)
|
||||
result = transcode_and_probe(VideoFromFile(buffer, start_time=1, duration=5))
|
||||
# kept frames: 1.0/2.25/3.5/4.75 s -> rebased span 3.75 s + one nominal interval
|
||||
assert result["frames"] == 4
|
||||
assert result["audio_seconds"] == pytest.approx(4.0, abs=0.45)
|
||||
|
||||
|
||||
def test_save_to_transcode_trims_audio_in_stream_time_base_units():
|
||||
"""Matroska audio timestamps tick in 1/1000, not 1/sample_rate; trim and audio timing
|
||||
must convert through the frame's time base instead of assuming sample units. AAC audio,
|
||||
because it decodes straight to the encoder's format and hits the resampler passthrough
|
||||
that keeps the source time base on the frames."""
|
||||
file_path = create_transcode_source(frames=90, container_format="matroska", audio_codec="aac")
|
||||
try:
|
||||
result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1))
|
||||
assert result["audio_codecs"] == ["aac"]
|
||||
assert result["video_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||
finally:
|
||||
os.unlink(file_path)
|
||||
|
||||
|
||||
def test_save_to_transcode_learns_unprobed_audio_params():
|
||||
"""mpegts is only probed a few seconds deep at open, so an audio stream whose first
|
||||
packet comes later (live captures where audio kicks in late) still has sample_rate 0
|
||||
when the transcode starts; the parameters must be learned from the stream itself."""
|
||||
sample_rate, fps, video_seconds, audio_start = 48000, 30, 13, 12
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mpegts") as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||
for i in range(video_seconds * fps):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
)
|
||||
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||
for offset in range(0, (video_seconds - audio_start) * sample_rate, 1024):
|
||||
n = min(1024, (video_seconds - audio_start) * sample_rate - offset)
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||
)
|
||||
audio_frame.sample_rate = sample_rate
|
||||
audio_frame.pts = audio_start * sample_rate + offset
|
||||
container.mux(audio_stream.encode(audio_frame))
|
||||
for stream in (video_stream, audio_stream):
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
# the scenario requires unprobed parameters; if a future FFmpeg probes deeper,
|
||||
# push audio_start/video_seconds further out to restore it
|
||||
assert container.streams.audio[0].codec_context.sample_rate == 0
|
||||
result = transcode_and_probe(VideoFromFile(buffer))
|
||||
assert result["frames"] == video_seconds * fps
|
||||
assert result["audio_codecs"] == ["aac"]
|
||||
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||
|
||||
buffer.seek(0)
|
||||
trimmed_before_audio = transcode_and_probe(VideoFromFile(buffer, duration=1))
|
||||
assert trimmed_before_audio["frames"] == fps
|
||||
assert trimmed_before_audio["audio_codecs"] == []
|
||||
assert trimmed_before_audio["audio_seconds"] is None
|
||||
|
||||
buffer.seek(0)
|
||||
trimmed_crossing_audio = transcode_and_probe(VideoFromFile(buffer, start_time=11.5, duration=1))
|
||||
assert trimmed_crossing_audio["frames"] == fps
|
||||
assert trimmed_crossing_audio["audio_codecs"] == ["aac"]
|
||||
assert trimmed_crossing_audio["video_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||
assert trimmed_crossing_audio["audio_seconds"] == pytest.approx(0.5, abs=0.1)
|
||||
|
||||
|
||||
def test_save_to_transcode_trimmed_fragmented_mp4_keeps_audio():
|
||||
"""Fragmented mp4 (MediaRecorder, DASH/HLS-derived files) delivers audio well behind
|
||||
video, so when the trim window's last video frame arrives the audio demuxed so far
|
||||
does not cover the window yet; the transcode must keep demuxing audio until it does
|
||||
instead of finalizing on the first audio frame it sees afterwards."""
|
||||
sample_rate, fps, seconds = 48000, 30, 6
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mp4", options={"movflags": "frag_keyframe+empty_moov"}) as container:
|
||||
video_stream = container.add_stream("h264", rate=fps)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||
next_audio_pts = 0
|
||||
for i in range(seconds * fps):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
)
|
||||
container.mux(video_stream.encode(frame.reformat(format="yuv420p")))
|
||||
while next_audio_pts / sample_rate <= i / fps: # feed audio alongside, like a live pipeline
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(1, 1024, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||
)
|
||||
audio_frame.sample_rate = sample_rate
|
||||
audio_frame.pts = next_audio_pts
|
||||
container.mux(audio_stream.encode(audio_frame))
|
||||
next_audio_pts += 1024
|
||||
for stream in (video_stream, audio_stream):
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
result = transcode_and_probe(VideoFromFile(buffer, start_time=0.5, duration=1.0))
|
||||
assert result["video_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.05)
|
||||
|
||||
|
||||
def test_save_to_transcode_sparse_video_keeps_true_duration():
|
||||
"""average_rate is not a frame duration: a 3-frame video spanning 60 s averages
|
||||
0.05 fps, and padding the last frame with 1/average_rate used to extend the
|
||||
output — and the audio kept with it — about 20 s past the source span."""
|
||||
sample_rate = 48000
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mp4") as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=30)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono")
|
||||
for i, second in enumerate((0, 30, 60)):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), i * 80, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
).reformat(format="yuv420p")
|
||||
frame.pts = second * 30
|
||||
frame.time_base = Fraction(1, 30)
|
||||
container.mux(video_stream.encode(frame))
|
||||
for offset in range(0, 90 * sample_rate, 1024):
|
||||
n = min(1024, 90 * sample_rate - offset)
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono"
|
||||
)
|
||||
audio_frame.sample_rate = sample_rate
|
||||
audio_frame.pts = offset
|
||||
container.mux(audio_stream.encode(audio_frame))
|
||||
for stream in (video_stream, audio_stream):
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
result = transcode_and_probe(VideoFromFile(buffer))
|
||||
assert result["frames"] == 3
|
||||
# the last frame keeps its true stts duration (1/30 s), not 1/average_rate (~20 s)
|
||||
assert result["video_seconds"] == pytest.approx(60.03, abs=0.05)
|
||||
assert result["audio_seconds"] == pytest.approx(60.03, abs=0.1)
|
||||
|
||||
trimmed = transcode_and_probe(VideoFromFile(buffer, duration=45))
|
||||
assert trimmed["frames"] == 2
|
||||
# a kept frame whose source duration crosses the window end is clamped to it
|
||||
assert trimmed["video_seconds"] == pytest.approx(45.0, abs=0.05)
|
||||
assert trimmed["audio_seconds"] == pytest.approx(45.0, abs=0.1)
|
||||
|
||||
|
||||
def test_save_to_transcode_clamps_final_pts_to_declared_stream_duration():
|
||||
"""Some iPhone MOVs report a video stream duration that ends before the final
|
||||
decoded frame's nominal duration. A transcode must not turn that trailing
|
||||
timestamp quirk into an extra frame interval compared to the source/remux path."""
|
||||
fps = 30
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mp4") as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=fps)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
for i, pts in enumerate([*range(31), 32]):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
).reformat(format="yuv420p")
|
||||
frame.pts = pts
|
||||
frame.time_base = Fraction(1, fps)
|
||||
container.mux(video_stream.encode(frame))
|
||||
container.mux(video_stream.encode(None))
|
||||
|
||||
class _StreamProxy:
|
||||
def __init__(self, stream, duration):
|
||||
self._stream = stream
|
||||
self.duration = duration
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._stream, name)
|
||||
|
||||
class _StreamsProxy:
|
||||
def __init__(self, video_stream):
|
||||
self.video = [video_stream]
|
||||
self.audio = []
|
||||
|
||||
class _PacketProxy:
|
||||
def __init__(self, packet, stream):
|
||||
self._packet = packet
|
||||
self.stream = stream
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._packet, name)
|
||||
|
||||
class _ContainerProxy:
|
||||
def __init__(self, container, stream):
|
||||
self._container = container
|
||||
self._stream = stream
|
||||
self.streams = _StreamsProxy(stream)
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._container, name)
|
||||
|
||||
def demux(self, *streams):
|
||||
for packet in self._container.demux(self._stream._stream):
|
||||
yield _PacketProxy(packet, self._stream)
|
||||
|
||||
buffer.seek(0)
|
||||
output = io.BytesIO()
|
||||
with av.open(buffer) as container:
|
||||
real_stream = container.streams.video[0]
|
||||
declared_duration = 32 * int(round((1 / fps) / real_stream.time_base))
|
||||
stream = _StreamProxy(real_stream, declared_duration)
|
||||
VideoFromFile(buffer)._save_transcoded(
|
||||
_ContainerProxy(container, stream), output, VideoContainer.MP4, VideoCodec.H264, None, 8
|
||||
)
|
||||
|
||||
output.seek(0)
|
||||
with av.open(output) as container:
|
||||
video_stream = container.streams.video[0]
|
||||
frames = [f for p in container.demux(video_stream) for f in p.decode()]
|
||||
assert len(frames) == 32
|
||||
assert float(video_stream.duration * video_stream.time_base) == pytest.approx(32 / fps, abs=0.01)
|
||||
assert float(frames[-1].pts * frames[-1].time_base) == pytest.approx(31 / fps, abs=0.01)
|
||||
|
||||
|
||||
def test_save_to_transcode_irregular_vfr_keeps_span():
|
||||
"""B-frames reorder packets, and mp4 sample durations follow decode order: the dts
|
||||
timeline ends before the pts timeline, so an irregular-VFR source's tail holds fell
|
||||
out of the container (this 20.23 s span used to come out as 15.27 s, and the 10 s
|
||||
trim as 6.03 s). The transcode encodes without B-frames so every sample keeps its
|
||||
true display duration."""
|
||||
durations = [1, 1, 60, 1, 1, 120, 1, 180, 1, 1, 150, 90] # 1/30 s ticks, span 20.2333 s
|
||||
generator = torch.Generator().manual_seed(7)
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mp4") as container:
|
||||
video_stream = container.add_stream("mpeg4", rate=30)
|
||||
video_stream.width = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
pts = 0
|
||||
for duration in durations:
|
||||
# textured frames, so an encoder with default settings has B-frames to gain from
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.randint(0, 255, (64, 64, 3), generator=generator, dtype=torch.uint8).numpy(),
|
||||
format="rgb24",
|
||||
).reformat(format="yuv420p")
|
||||
frame.pts = pts
|
||||
frame.time_base = Fraction(1, 30)
|
||||
pts += duration
|
||||
for packet in video_stream.encode(frame):
|
||||
packet.duration = duration # exact stts in the source
|
||||
container.mux(packet)
|
||||
container.mux(video_stream.encode(None))
|
||||
|
||||
result = transcode_and_probe(VideoFromFile(buffer))
|
||||
assert result["frames"] == len(durations)
|
||||
assert result["video_seconds"] == pytest.approx(sum(durations) / 30, abs=0.05)
|
||||
|
||||
trimmed = transcode_and_probe(VideoFromFile(buffer, duration=10))
|
||||
assert trimmed["frames"] == 8 # frames at 12.167 s+ fall outside the window
|
||||
assert trimmed["video_seconds"] == pytest.approx(10.0, abs=0.05)
|
||||
|
||||
|
||||
def test_save_to_transcode_trim_survives_missing_leading_pts():
|
||||
"""A trim should survive pts-less kept frames followed by a real-pts frame past the window."""
|
||||
nulled_frames = 0
|
||||
|
||||
class _PacketProxy:
|
||||
def __init__(self, packet):
|
||||
self._packet = packet
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._packet, name)
|
||||
|
||||
@property
|
||||
def stream(self):
|
||||
return self._packet.stream
|
||||
|
||||
def decode(self):
|
||||
nonlocal nulled_frames
|
||||
frames = self._packet.decode()
|
||||
for frame in frames:
|
||||
if nulled_frames < 2:
|
||||
frame.pts = None
|
||||
nulled_frames += 1
|
||||
return frames
|
||||
|
||||
class _ContainerProxy:
|
||||
def __init__(self, real):
|
||||
self._real = real
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._real, name)
|
||||
|
||||
def demux(self, *streams):
|
||||
for packet in self._real.demux(*streams):
|
||||
yield _PacketProxy(packet)
|
||||
|
||||
file_path = create_transcode_source(frames=10, audio_streams=0)
|
||||
try:
|
||||
buffer = io.BytesIO()
|
||||
with av.open(file_path) as container:
|
||||
# 0.05 s window: both pts-less frames are kept (synthesized pts 0 and 512),
|
||||
# and the first real-pts frame (1024 ticks) already lies past end_pts (768)
|
||||
VideoFromFile(file_path, duration=0.05)._save_transcoded(
|
||||
_ContainerProxy(container), buffer, VideoContainer.MP4, VideoCodec.H264, None, 8
|
||||
)
|
||||
assert nulled_frames == 2
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
video_stream = container.streams.video[0]
|
||||
frames = [f for p in container.demux(video_stream) for f in p.decode()]
|
||||
assert len(frames) == 2
|
||||
assert float(video_stream.duration * video_stream.time_base) == pytest.approx(2 / 30, abs=0.01)
|
||||
finally:
|
||||
os.unlink(file_path)
|
||||
|
||||
|
||||
def test_save_to_transcode_bakes_rotation():
|
||||
"""A 90-degree display-matrix rotation swaps the output dimensions (portrait video)"""
|
||||
file_path = create_transcode_source(width=64, height=32, rotation=True)
|
||||
try:
|
||||
result = transcode_and_probe(VideoFromFile(file_path))
|
||||
assert (result["width"], result["height"]) == (32, 64)
|
||||
assert result["frames"] == 30
|
||||
finally:
|
||||
os.unlink(file_path)
|
||||
|
||||
|
||||
def test_save_to_transcode_skips_undecodable_audio():
|
||||
"""Streaming transcode keeps the decodable audio track and drops undecodable ones;
|
||||
with no decodable audio at all the output is video-only instead of crashing."""
|
||||
mixed = all_bad = None
|
||||
try:
|
||||
mixed = create_transcode_source(audio_streams=1, undecodable_audio=1)
|
||||
all_bad = create_transcode_source(audio_streams=0, undecodable_audio=2)
|
||||
result = transcode_and_probe(VideoFromFile(mixed))
|
||||
assert result["audio_codecs"] == ["aac"]
|
||||
assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1)
|
||||
assert transcode_and_probe(VideoFromFile(all_bad))["audio_codecs"] == []
|
||||
finally:
|
||||
for path in (mixed, all_bad):
|
||||
if path:
|
||||
os.unlink(path)
|
||||
|
||||
@ -818,6 +818,30 @@ class TestExecution:
|
||||
except urllib.error.HTTPError:
|
||||
pass # Expected behavior
|
||||
|
||||
def test_cached_outputs_in_job_without_client_id(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
output = g.node("SaveImage", images=image.out(0))
|
||||
|
||||
# Prime the cache with a normal run.
|
||||
client.run(g)
|
||||
|
||||
# Resubmit anonymously (no client_id) so output nodes are cache hits with no websocket client.
|
||||
data = json.dumps({"prompt": g.finalize()}).encode('utf-8')
|
||||
req = urllib.request.Request(f"http://{client.server_address}/prompt", data=data)
|
||||
prompt_id = json.loads(urllib.request.urlopen(req).read())['prompt_id']
|
||||
|
||||
for _ in range(100):
|
||||
job = client.get_job(prompt_id)
|
||||
if job is not None and job['status'] not in ('pending', 'in_progress'):
|
||||
break
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
raise AssertionError("Prompt did not complete in time")
|
||||
|
||||
assert job['status'] == 'completed'
|
||||
assert output.id in job['outputs'], "Cached outputs must appear in job outputs without a client_id"
|
||||
|
||||
def _create_history_item(self, client, builder):
|
||||
g = GraphBuilder(prefix="offset_test")
|
||||
input_node = g.node(
|
||||
|
||||
Loading…
Reference in New Issue
Block a user