Fredthearboro
8a3bec8391
VidiaVideoSaver: watermark + ffmpeg-mux + R2 upload + backend notify. V3 API (io.ComfyNode, ComfyExtension, comfy_entrypoint). not_idempotent=True forces re-execution (fixes upload/callback caching bug). Remote log shipping opt-in via VIDIA_REMOTE_LOGGING env. Env-only config (no hardcoded secrets).
401 lines
16 KiB
Python
401 lines
16 KiB
Python
"""
|
|
Video frame processing for the Vidia custom ComfyUI node.
|
|
Handles frame preparation, watermark overlays, dimension alignment,
|
|
frame resampling, ffmpeg process management, and format loading.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from string import Template
|
|
|
|
import folder_paths
|
|
from .utils import ffmpeg_path, tensor_to_bytes, logger
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Video format helpers
|
|
# ---------------------------------------------------------------------------
|
|
# Register our custom format path with ComfyUI
|
|
folder_paths.folder_names_and_paths["VIDIA_video_formats"] = (
|
|
[os.path.join(os.path.dirname(os.path.abspath(__file__)), "video_formats")],
|
|
[".json"],
|
|
)
|
|
|
|
|
|
def gen_format_widgets(video_format):
|
|
"""Generate format widgets from a video format dict."""
|
|
for k in video_format:
|
|
if k.endswith("_pass"):
|
|
for i in range(len(video_format[k])):
|
|
if isinstance(video_format[k][i], list):
|
|
item = [video_format[k][i]]
|
|
yield item
|
|
video_format[k][i] = item[0]
|
|
else:
|
|
if isinstance(video_format[k], list):
|
|
item = [video_format[k]]
|
|
yield item
|
|
video_format[k] = item[0]
|
|
|
|
|
|
def get_video_formats():
|
|
"""Get available video formats for the node's format dropdown."""
|
|
formats = []
|
|
for format_name in folder_paths.get_filename_list("VIDIA_video_formats"):
|
|
format_name = format_name[:-5] # strip .json
|
|
video_format_path = folder_paths.get_full_path("VIDIA_video_formats", format_name + ".json")
|
|
with open(video_format_path, 'r') as stream:
|
|
video_format = json.load(stream)
|
|
widgets = [w[0] for w in gen_format_widgets(video_format)]
|
|
if len(widgets) > 0:
|
|
formats.append(["video/" + format_name, widgets])
|
|
else:
|
|
formats.append("video/" + format_name)
|
|
return formats
|
|
|
|
|
|
def apply_format_widgets(format_name, kwargs):
|
|
"""Apply user-provided widget values to a video format dict."""
|
|
video_format_path = folder_paths.get_full_path("VIDIA_video_formats", format_name + ".json")
|
|
with open(video_format_path, 'r') as stream:
|
|
video_format = json.load(stream)
|
|
for w in gen_format_widgets(video_format):
|
|
if w[0][0] not in kwargs:
|
|
if len(w[0]) > 2 and 'default' in w[0][2]:
|
|
default = w[0][2]['default']
|
|
else:
|
|
if type(w[0][1]) is list:
|
|
default = w[0][1][0]
|
|
else:
|
|
default = {"BOOLEAN": False, "INT": 0, "FLOAT": 0, "STRING": ""}[w[0][1]]
|
|
kwargs[w[0][0]] = default
|
|
logger.warn(f"Missing input for {w[0][0]} has been set to {default}")
|
|
if len(w[0]) > 3:
|
|
w[0] = Template(w[0][3]).substitute(val=kwargs[w[0][0]])
|
|
else:
|
|
w[0] = str(kwargs[w[0][0]])
|
|
return video_format
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ffmpeg process (generator-based frame piping)
|
|
# ---------------------------------------------------------------------------
|
|
def ffmpeg_process(args, video_format, video_metadata, file_path, env, custom_logger=None):
|
|
"""Run an ffmpeg encode process via a generator that accepts raw frame data.
|
|
|
|
Usage:
|
|
proc = ffmpeg_process(args, fmt, meta, path, env)
|
|
next(proc) # advance to first yield
|
|
for frame in frames:
|
|
proc.send(frame_bytes)
|
|
total = proc.send(None) # finalize
|
|
"""
|
|
log = custom_logger or logger
|
|
res = None
|
|
frame_data = yield
|
|
total_frames_output = 0
|
|
|
|
if video_format.get('save_metadata', 'False') != 'False':
|
|
os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)
|
|
metadata = json.dumps(video_metadata)
|
|
metadata_path = os.path.join(folder_paths.get_temp_directory(), "metadata.txt")
|
|
metadata = metadata.replace("\\", "\\\\")
|
|
metadata = metadata.replace(";", "\\;")
|
|
metadata = metadata.replace("#", "\\#")
|
|
metadata = metadata.replace("=", "\\=")
|
|
metadata = metadata.replace("\n", "\\\n")
|
|
metadata = "comment=" + metadata
|
|
|
|
with open(metadata_path, "w") as f:
|
|
f.write(";FFMETADATA1\n")
|
|
f.write(metadata)
|
|
|
|
m_args = args[:1] + ["-i", metadata_path] + args[1:] + ["-metadata", "creation_time=now"]
|
|
with subprocess.Popen(m_args + [file_path], stderr=subprocess.PIPE,
|
|
stdin=subprocess.PIPE, env=env) as proc:
|
|
try:
|
|
while frame_data is not None:
|
|
proc.stdin.write(frame_data)
|
|
frame_data = yield
|
|
total_frames_output += 1
|
|
proc.stdin.flush()
|
|
proc.stdin.close()
|
|
res = proc.stderr.read()
|
|
except BrokenPipeError:
|
|
err = proc.stderr.read()
|
|
log.error(f"FFMPEG BROKEN PIPE (metadata pass). STDERR: {err.decode('utf-8', errors='ignore')}")
|
|
if os.path.exists(file_path):
|
|
raise Exception("ffmpeg subprocess error:\n" + err.decode('utf-8'))
|
|
print(err.decode('utf-8'), end="", file=sys.stderr)
|
|
log.warn("Error occurred when saving with metadata")
|
|
|
|
# The main-pass fallback runs only when the metadata pass emitted stderr.
|
|
# This is safe because callers build ffmpeg with "-v error", so a successful
|
|
# metadata encode yields empty stderr and the fallback is skipped (preserving
|
|
# metadata). A non-empty stderr means a real error; the fallback re-encodes.
|
|
if res != b'':
|
|
with subprocess.Popen(args + [file_path], stderr=subprocess.PIPE,
|
|
stdin=subprocess.PIPE, env=env) as proc:
|
|
try:
|
|
while frame_data is not None:
|
|
proc.stdin.write(frame_data)
|
|
frame_data = yield
|
|
total_frames_output += 1
|
|
proc.stdin.flush()
|
|
proc.stdin.close()
|
|
res = proc.stderr.read()
|
|
except BrokenPipeError:
|
|
res = proc.stderr.read()
|
|
log.error(f"FFMPEG BROKEN PIPE (main pass). STDERR: {res.decode('utf-8', errors='ignore')}")
|
|
raise Exception("ffmpeg subprocess error:\n" + res.decode('utf-8'))
|
|
|
|
if res is not None and len(res) > 0:
|
|
log.warn(f"FFMPEG STDERR: {res.decode('utf-8', errors='ignore')}")
|
|
|
|
yield total_frames_output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Frame resampling
|
|
# ---------------------------------------------------------------------------
|
|
def resample_frames_ffmpeg(frames, source_fps, target_fps, log):
|
|
"""Resample frames from source_fps to target_fps using ffmpeg."""
|
|
if not frames or source_fps <= 0 or target_fps <= 0 or source_fps == target_fps:
|
|
return frames
|
|
|
|
log.info(f"Resampling frames from {source_fps} to {target_fps} using ffmpeg.")
|
|
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
uid = uuid.uuid4().hex[:8]
|
|
temp_input_path = os.path.join(temp_dir, f"temp_resample_in_{uid}.mkv")
|
|
|
|
h, w, c = frames[0].shape
|
|
pix_fmt_in = 'rgba' if c == 4 else 'rgb24'
|
|
|
|
encode_cmd = [
|
|
ffmpeg_path, '-y', '-f', 'rawvideo',
|
|
'-s', f'{w}x{h}', '-pix_fmt', pix_fmt_in, '-r', str(source_fps),
|
|
'-i', '-', '-an', '-vcodec', 'ffv1', temp_input_path
|
|
]
|
|
|
|
input_bytes = b''.join([tensor_to_bytes(frame).tobytes() for frame in frames])
|
|
proc_encode = subprocess.Popen(encode_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
_, stderr = proc_encode.communicate(input=input_bytes)
|
|
|
|
if proc_encode.returncode != 0:
|
|
log.error(f"ffmpeg encode failed: {stderr.decode('utf-8', errors='ignore')}")
|
|
if os.path.exists(temp_input_path):
|
|
os.unlink(temp_input_path)
|
|
raise Exception("ffmpeg frame resampling (encode step) failed.")
|
|
|
|
pix_fmt_out = 'rgb24'
|
|
decode_cmd = [
|
|
ffmpeg_path, '-i', temp_input_path, '-vf', f'fps={target_fps}',
|
|
'-f', 'rawvideo', '-pix_fmt', pix_fmt_out, '-'
|
|
]
|
|
|
|
proc_decode = subprocess.Popen(decode_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
output_frames = []
|
|
while True:
|
|
frame_bytes = proc_decode.stdout.read(w * h * 3)
|
|
if not frame_bytes:
|
|
break
|
|
frame = np.frombuffer(frame_bytes, dtype=np.uint8).reshape((h, w, 3))
|
|
output_frames.append(torch.from_numpy(frame).float() / 255.0)
|
|
|
|
_, stderr = proc_decode.communicate()
|
|
if proc_decode.returncode != 0:
|
|
log.error(f"ffmpeg decode failed: {stderr.decode('utf-8', errors='ignore')}")
|
|
|
|
if os.path.exists(temp_input_path):
|
|
os.unlink(temp_input_path)
|
|
|
|
log.info(f"Resampling complete. Original: {len(frames)} frames, Resampled: {len(output_frames)} frames.")
|
|
return output_frames
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Watermark processing
|
|
# ---------------------------------------------------------------------------
|
|
def process_watermark_frames(watermark_frames, main_shape):
|
|
"""Process watermark frames to match main video dimensions (for append watermark)."""
|
|
if isinstance(watermark_frames, torch.Tensor):
|
|
if watermark_frames.ndim == 4:
|
|
watermark_frames = [watermark_frames[i] for i in range(watermark_frames.size(0))]
|
|
elif watermark_frames.ndim == 3:
|
|
watermark_frames = [watermark_frames]
|
|
else:
|
|
print(f"[process_watermark_frames] ERROR: Unexpected tensor dimension: {watermark_frames.ndim}")
|
|
return []
|
|
|
|
if not watermark_frames:
|
|
return []
|
|
|
|
h, w = main_shape[:2]
|
|
wm_h, wm_w = watermark_frames[0].shape[0], watermark_frames[0].shape[1]
|
|
black_bg = torch.zeros((h, w, main_shape[2]), dtype=torch.float32)
|
|
|
|
if h >= wm_h and w >= wm_w:
|
|
x_offset = (w - wm_w) // 2
|
|
y_offset = (h - wm_h) // 2
|
|
for i, frame in enumerate(watermark_frames):
|
|
new_bg = black_bg.clone()
|
|
new_bg[y_offset:y_offset + wm_h, x_offset:x_offset + wm_w] = frame
|
|
watermark_frames[i] = new_bg
|
|
else:
|
|
scale_h = h / wm_h
|
|
scale_w = w / wm_w
|
|
scale = min(scale_h, scale_w)
|
|
new_h = int(wm_h * scale)
|
|
new_w = int(wm_w * scale)
|
|
new_size = (new_h, new_w)
|
|
|
|
watermark_tensor = torch.stack(watermark_frames)
|
|
resized = F.interpolate(watermark_tensor.permute(0, 3, 1, 2), size=new_size)
|
|
resized_frames = list(resized.permute(0, 2, 3, 1))
|
|
|
|
x_offset = (w - new_w) // 2
|
|
y_offset = (h - new_h) // 2
|
|
for i, frame in enumerate(resized_frames):
|
|
new_bg = black_bg.clone()
|
|
new_bg[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = frame
|
|
watermark_frames[i] = new_bg
|
|
|
|
return watermark_frames
|
|
|
|
|
|
def process_corner_watermark(watermark_frames, mask_frames, main_shape, opacity=0.4, logger=None):
|
|
"""Process corner watermark with mask for transparency."""
|
|
import logging as _logging
|
|
log = logger or _logging.getLogger("VidiaOpenStudioNode")
|
|
|
|
if isinstance(watermark_frames, torch.Tensor) and watermark_frames.ndim == 4:
|
|
watermark_frames = [frame for frame in watermark_frames]
|
|
elif isinstance(watermark_frames, torch.Tensor) and watermark_frames.ndim == 3:
|
|
watermark_frames = [watermark_frames]
|
|
|
|
if isinstance(mask_frames, torch.Tensor) and mask_frames.ndim == 3:
|
|
mask_frames = [mask for mask in mask_frames]
|
|
elif isinstance(mask_frames, torch.Tensor) and mask_frames.ndim == 2:
|
|
mask_frames = [mask_frames]
|
|
|
|
if not watermark_frames or watermark_frames[0].numel() == 0 or watermark_frames[0].shape[0] == 0 or watermark_frames[0].shape[1] == 0:
|
|
log.warn("Corner watermark input is empty. Skipping.")
|
|
return [], [], (0, 0)
|
|
|
|
if not mask_frames or mask_frames[0].numel() == 0 or mask_frames[0].shape[0] == 0 or mask_frames[0].shape[1] == 0:
|
|
log.warn("Corner watermark mask is empty. Skipping.")
|
|
return [], [], (0, 0)
|
|
|
|
h, w = main_shape[:2]
|
|
watermark_h, watermark_w = watermark_frames[0].shape[:2]
|
|
|
|
if watermark_h == 0 or watermark_w == 0:
|
|
log.warn(f"Corner watermark has zero dimensions ({watermark_h}x{watermark_w}). Skipping.")
|
|
return [], [], (0, 0)
|
|
|
|
target_width = int(w * 0.2)
|
|
if target_width == 0:
|
|
log.warn(f"Main video width too small ({w}) for corner watermark. Skipping.")
|
|
return [], [], (0, 0)
|
|
|
|
scale = target_width / watermark_w
|
|
new_h = int(watermark_h * scale)
|
|
new_w = target_width
|
|
|
|
if new_h == 0 or new_w == 0:
|
|
log.warn(f"Calculated corner watermark size invalid ({new_h}x{new_w}). Skipping.")
|
|
return [], [], (0, 0)
|
|
|
|
new_size = (new_h, new_w)
|
|
x_offset = w - new_size[1] - 20
|
|
y_offset = h - new_size[0] - 20
|
|
|
|
try:
|
|
watermark_tensor = torch.stack(watermark_frames)
|
|
mask_tensor = torch.stack(mask_frames)
|
|
|
|
resized_watermark = F.interpolate(watermark_tensor.permute(0, 3, 1, 2), size=new_size)
|
|
resized_mask = F.interpolate(mask_tensor.unsqueeze(1), size=new_size).squeeze(1)
|
|
|
|
processed_watermark = list(resized_watermark.permute(0, 2, 3, 1))
|
|
processed_mask = list(resized_mask)
|
|
|
|
return processed_watermark, processed_mask, (x_offset, y_offset)
|
|
except Exception as e:
|
|
log.error(f"Error processing corner watermark: {e}")
|
|
return [], [], (0, 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Path management
|
|
# ---------------------------------------------------------------------------
|
|
class VideoPathManager:
|
|
"""Manages all file paths for video processing to prevent overwrites."""
|
|
|
|
def __init__(self, filename_prefix, output_dir, save_output=True):
|
|
self.filename_prefix = filename_prefix
|
|
|
|
if not filename_prefix or filename_prefix.strip() == "":
|
|
self.filename_prefix = f"vidia_{uuid.uuid4().hex[:8]}"
|
|
print(f"[VideoPathManager] WARNING: Empty filename_prefix, using: {self.filename_prefix}")
|
|
|
|
self.output_dir = output_dir if save_output else folder_paths.get_temp_directory()
|
|
self.full_output_folder, _, _, self.subfolder, _ = folder_paths.get_save_image_path(
|
|
self.filename_prefix, self.output_dir
|
|
)
|
|
self.output_files = []
|
|
|
|
def get_first_frame_path(self):
|
|
return os.path.join(self.full_output_folder, f"{self.filename_prefix}_00001.png")
|
|
|
|
def get_animation_path(self, format_ext):
|
|
return os.path.join(self.full_output_folder, f"{self.filename_prefix}.{format_ext}")
|
|
|
|
def get_video_path(self, extension):
|
|
return os.path.join(self.full_output_folder, f"{self.filename_prefix}.{extension}")
|
|
|
|
def get_temp_video_path(self, extension):
|
|
uid = uuid.uuid4().hex[:8]
|
|
return os.path.join(folder_paths.get_temp_directory(), f"temp_{self.filename_prefix}_{uid}.{extension}")
|
|
|
|
def add_output_file(self, path):
|
|
self.output_files.append(path)
|
|
return path
|
|
|
|
def get_output_files(self):
|
|
return self.output_files
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Processing state tracker
|
|
# ---------------------------------------------------------------------------
|
|
class VideoProcessingState:
|
|
"""Tracks the state of video processing."""
|
|
|
|
def __init__(self):
|
|
self.total_frames = 0
|
|
self.processed_frames = 0
|
|
self.has_corner_watermark = False
|
|
self.has_append_watermark = False
|
|
self.has_audio = False
|
|
self.audio_processed = False
|
|
self.saved_first_frame = False
|
|
self.video_created = False
|
|
self.upload_attempted = False
|
|
self.upload_successful = False
|
|
self.errors = []
|
|
|
|
def add_error(self, error_msg):
|
|
self.errors.append(error_msg)
|
|
|
|
def has_errors(self):
|
|
return len(self.errors) > 0
|