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).
486 lines
20 KiB
Python
486 lines
20 KiB
Python
"""
|
|
VidiaVideoSaver -- ComfyUI V3 custom node for Vidia Open Studio
|
|
(node pack: comfyui-vidia-open-studio-node).
|
|
Orchestrates video creation, audio processing, and cloud upload.
|
|
|
|
Requires a ComfyUI build that ships the V3 node API (comfy_api.latest).
|
|
Target: ComfyUI >= v0.24. The V1 NODE_CLASS_MAPPINGS registration was removed.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import datetime
|
|
import torch
|
|
from PIL import Image, ExifTags
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import folder_paths
|
|
from comfy_api.latest import io
|
|
|
|
from .config import HAS_BOTO3
|
|
from .utils import ffmpeg_path, tensor_to_bytes, VideoProcessingLogger
|
|
from .video_processing import (
|
|
get_video_formats,
|
|
apply_format_widgets,
|
|
ffmpeg_process,
|
|
resample_frames_ffmpeg,
|
|
process_watermark_frames,
|
|
process_corner_watermark,
|
|
VideoPathManager,
|
|
VideoProcessingState,
|
|
)
|
|
from .audio_processing import process_audio, mux_audio_into_video
|
|
from .upload import upload_to_cloudflare_s3
|
|
|
|
|
|
# Custom output type. The string is kept identical to the V1 node so existing
|
|
# workflows that wire VidiaVideoSaver as the [out] terminal still resolve.
|
|
VIDIA_FILENAMES = io.Custom("VIDIA_FILENAMES")
|
|
|
|
|
|
class VidiaVideoSaver(io.ComfyNode):
|
|
"""Saves generated frames to a video file with optional watermarks and audio,
|
|
uploads to Cloudflare R2, and notifies the Vidia backend.
|
|
|
|
Pipeline:
|
|
1. Preprocess and validate input frames
|
|
2. Generate metadata from prompt information
|
|
3. Save the first frame with metadata
|
|
4. Process frames (watermarks, dimension alignment)
|
|
5. Create the video file via ffmpeg using format JSON settings
|
|
6. Process and concatenate audio to match video segments
|
|
7. Upload to Cloudflare R2
|
|
8. Notify backend of completion
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
ffmpeg_formats = get_video_formats()
|
|
# VHS-style format entries are either "video/<name>" or
|
|
# ["video/<name>", [widgets]]. V3 combos take flat option strings, so
|
|
# flatten to the name. The shipped h264-mp4 format declares no widgets;
|
|
# a widget-bearing format would need DynamicCombo to inject sub-inputs
|
|
# under V3 (no shipped format does today).
|
|
format_options = [f[0] if isinstance(f, list) else f for f in ffmpeg_formats]
|
|
return io.Schema(
|
|
node_id="VidiaVideoSaver",
|
|
display_name="Vidia Video Saver",
|
|
category="Vidia",
|
|
description=(
|
|
"Saves video frames to a file with optional watermarks and audio, "
|
|
"uploads to Cloudflare R2, and notifies the backend."
|
|
),
|
|
is_output_node=True,
|
|
# Side effects (R2 upload + backend callback) must run every time;
|
|
# never serve this node from cache.
|
|
not_idempotent=True,
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.Float.Input("frame_rate", default=8, min=1, step=1),
|
|
io.String.Input("filename_prefix", default="Vidia"),
|
|
io.Combo.Input("format", options=["image/gif", "image/webp"] + format_options),
|
|
io.Boolean.Input("save_output", default=True),
|
|
io.Audio.Input("audio", optional=True, lazy=True),
|
|
io.Image.Input("append_watermark", optional=True),
|
|
io.Audio.Input("watermark_audio", optional=True, lazy=True),
|
|
io.Image.Input("corner_watermark", optional=True),
|
|
io.Mask.Input("corner_watermark_mask", optional=True),
|
|
io.Float.Input("corner_opacity", default=0.4, min=0.0, max=1.0, step=0.1),
|
|
io.Float.Input("append_watermark_source_fps", default=0, min=0, max=120, step=1),
|
|
],
|
|
outputs=[
|
|
VIDIA_FILENAMES.Output("Filenames"),
|
|
],
|
|
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
|
|
)
|
|
|
|
@classmethod
|
|
def validate_inputs(cls, images, format, **kwargs):
|
|
if "/" not in format:
|
|
return "format must be of the form <type>/<ext> (e.g. video/h264-mp4)"
|
|
return True
|
|
|
|
@classmethod
|
|
def check_lazy_status(cls, images, audio=None, watermark_audio=None, **kwargs):
|
|
"""Tell the executor which lazy audio inputs are still needed."""
|
|
needed = []
|
|
if audio is None:
|
|
needed.append("audio")
|
|
if watermark_audio is None:
|
|
needed.append("watermark_audio")
|
|
return needed
|
|
|
|
# ------------------------------------------------------------------
|
|
# Main entry point
|
|
# ------------------------------------------------------------------
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images,
|
|
frame_rate,
|
|
filename_prefix,
|
|
format,
|
|
save_output,
|
|
audio=None,
|
|
append_watermark=None,
|
|
watermark_audio=None,
|
|
corner_watermark=None,
|
|
corner_watermark_mask=None,
|
|
corner_opacity=0.4,
|
|
append_watermark_source_fps=0,
|
|
**kwargs,
|
|
):
|
|
logger = VideoProcessingLogger("VidiaVideoSaver")
|
|
path_manager = VideoPathManager(
|
|
filename_prefix,
|
|
folder_paths.get_output_directory() if save_output else folder_paths.get_temp_directory(),
|
|
save_output,
|
|
)
|
|
state = VideoProcessingState()
|
|
|
|
logger.info(
|
|
f"Starting video processing with "
|
|
f"{len(images) if isinstance(images, list) else images.size(0)} frames"
|
|
)
|
|
|
|
images = cls._preprocess_inputs(images)
|
|
if images is None:
|
|
return io.NodeOutput((save_output, []), ui={"error": "No input frames"})
|
|
|
|
video_metadata = cls._prepare_metadata(cls.hidden.prompt, cls.hidden.extra_pnginfo)
|
|
|
|
cls._save_first_frame(images[0], video_metadata, path_manager, state, logger)
|
|
|
|
format_type, format_ext = format.split("/")
|
|
logger.info(f"Using format: {format_type}/{format_ext}")
|
|
|
|
if format_type == "image":
|
|
cls._process_image_format(images, format_ext, frame_rate, path_manager, state, logger)
|
|
else:
|
|
# Video format via ffmpeg
|
|
if ffmpeg_path is None:
|
|
raise RuntimeError(
|
|
"ffmpeg is required for video outputs and could not be found."
|
|
)
|
|
|
|
video_format = apply_format_widgets(format_ext, kwargs)
|
|
user_frames_count = len(images)
|
|
state.processed_frames = user_frames_count
|
|
|
|
processed_frames = cls._prepare_frames(
|
|
images, video_format, corner_watermark,
|
|
corner_watermark_mask, corner_opacity, append_watermark,
|
|
frame_rate, append_watermark_source_fps, logger, state,
|
|
)
|
|
|
|
watermark_frames_count = (
|
|
len(processed_frames) - user_frames_count if state.has_append_watermark else 0
|
|
)
|
|
logger.info(
|
|
f"Frame counts - User: {user_frames_count}, "
|
|
f"Append Watermark: {watermark_frames_count}, Total: {len(processed_frames)}"
|
|
)
|
|
|
|
success = cls._create_video_file(
|
|
processed_frames, video_format, video_metadata, frame_rate,
|
|
path_manager, state, logger,
|
|
)
|
|
|
|
if success:
|
|
audio_result = process_audio(
|
|
audio, watermark_audio, video_format, frame_rate,
|
|
user_frames_count=user_frames_count,
|
|
watermark_frames_count=watermark_frames_count,
|
|
logger=logger,
|
|
)
|
|
# audio_result is False/True, or (combined_audio_path, temp_files)
|
|
if isinstance(audio_result, tuple):
|
|
combined_audio_path, temp_files = audio_result
|
|
video_path = path_manager.get_output_files()[-1]
|
|
temp_video_path = path_manager.get_temp_video_path(
|
|
video_format.get("extension", "mp4")
|
|
)
|
|
mux_success = mux_audio_into_video(
|
|
video_path, combined_audio_path, temp_video_path, logger
|
|
)
|
|
if mux_success:
|
|
state.audio_processed = True
|
|
for tf in temp_files:
|
|
if os.path.exists(tf):
|
|
try:
|
|
os.unlink(tf)
|
|
except OSError:
|
|
pass
|
|
|
|
# Upload to Cloudflare
|
|
try:
|
|
if HAS_BOTO3:
|
|
upload_result = upload_to_cloudflare_s3(
|
|
format, path_manager.get_output_files()[-1],
|
|
filename_prefix, logger, state,
|
|
)
|
|
else:
|
|
logger.error("boto3 is not installed, cannot upload to S3.")
|
|
return io.NodeOutput(
|
|
(save_output, path_manager.get_output_files()),
|
|
ui={"error": "boto3 is not installed. Video saved locally."},
|
|
)
|
|
|
|
if upload_result.get("success", False):
|
|
return io.NodeOutput(
|
|
(save_output, path_manager.get_output_files()),
|
|
ui={"response": upload_result},
|
|
)
|
|
else:
|
|
error_msg = upload_result.get('error', 'Unknown error')
|
|
if isinstance(error_msg, list):
|
|
error_msg = ''.join(error_msg)
|
|
return io.NodeOutput(
|
|
(save_output, path_manager.get_output_files()),
|
|
ui={"error": f"Upload failed but video saved locally: {error_msg}"},
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Upload failed: {str(e)}")
|
|
return io.NodeOutput(
|
|
(save_output, path_manager.get_output_files()),
|
|
ui={"error": f"Upload failed but video saved locally: {str(e)}"},
|
|
)
|
|
|
|
return io.NodeOutput(
|
|
(save_output, path_manager.get_output_files()),
|
|
ui={"success": "Video processing completed"},
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers (classmethods; V3 nodes hold no instance state)
|
|
# ------------------------------------------------------------------
|
|
@classmethod
|
|
def _preprocess_inputs(cls, images):
|
|
"""Validate and preprocess input images. Returns None if empty."""
|
|
if isinstance(images, dict):
|
|
images = images['samples']
|
|
if isinstance(images, torch.Tensor) and images.size(0) == 0:
|
|
return None
|
|
return images
|
|
|
|
@classmethod
|
|
def _prepare_metadata(cls, prompt, extra_pnginfo):
|
|
video_metadata = {}
|
|
if prompt is not None:
|
|
video_metadata["prompt"] = json.dumps(prompt)
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
video_metadata[x] = extra_pnginfo[x]
|
|
video_metadata["CreationTime"] = datetime.datetime.now().isoformat(" ")[:19]
|
|
return video_metadata
|
|
|
|
@classmethod
|
|
def _save_first_frame(cls, first_frame, video_metadata, path_manager, state, logger):
|
|
metadata = PngInfo()
|
|
for key, value in video_metadata.items():
|
|
if isinstance(value, str):
|
|
metadata.add_text(key, value)
|
|
else:
|
|
metadata.add_text(key, json.dumps(value))
|
|
|
|
first_image_path = path_manager.get_first_frame_path()
|
|
logger.info(f"Saving first frame to {first_image_path}")
|
|
|
|
try:
|
|
Image.fromarray(tensor_to_bytes(first_frame)).save(
|
|
first_image_path, pnginfo=metadata, compress_level=4,
|
|
)
|
|
path_manager.add_output_file(first_image_path)
|
|
state.saved_first_frame = True
|
|
except Exception as e:
|
|
logger.error(f"Failed to save first frame: {str(e)}")
|
|
state.add_error(f"First frame save failed: {str(e)}")
|
|
|
|
@classmethod
|
|
def _process_image_format(cls, images, format_ext, frame_rate, path_manager, state, logger):
|
|
image_kwargs = {}
|
|
if format_ext == "gif":
|
|
image_kwargs['disposal'] = 2
|
|
if format_ext == "webp":
|
|
exif = Image.Exif()
|
|
exif[ExifTags.IFD.Exif] = {36867: datetime.datetime.now().isoformat(" ")[:19]}
|
|
image_kwargs['exif'] = exif
|
|
|
|
file_path = path_manager.get_animation_path(format_ext)
|
|
logger.info(f"Saving animation to {file_path}")
|
|
|
|
try:
|
|
frames = map(lambda x: Image.fromarray(tensor_to_bytes(x)), images)
|
|
next_frame = next(frames)
|
|
next_frame.save(
|
|
file_path, format=format_ext.upper(), save_all=True,
|
|
append_images=frames, duration=round(1000 / frame_rate),
|
|
loop=0, compress_level=4, **image_kwargs,
|
|
)
|
|
path_manager.add_output_file(file_path)
|
|
state.video_created = True
|
|
except Exception as e:
|
|
logger.error(f"Failed to create {format_ext} animation: {str(e)}")
|
|
state.add_error(f"Animation creation failed: {str(e)}")
|
|
|
|
@classmethod
|
|
def _prepare_frames(cls, images, video_format, corner_watermark=None,
|
|
corner_watermark_mask=None, corner_opacity=0.4,
|
|
append_watermark=None, frame_rate=None,
|
|
append_watermark_source_fps=0, logger=None, state=None):
|
|
"""Prepare frames: dimension alignment, watermarks, resampling."""
|
|
if isinstance(images, torch.Tensor):
|
|
images = list(images)
|
|
first_frame = images[0]
|
|
|
|
# Dimension alignment
|
|
dim_alignment = video_format.get("dim_alignment", 2)
|
|
if (first_frame.shape[1] % dim_alignment) or (first_frame.shape[0] % dim_alignment):
|
|
to_pad = (-first_frame.shape[1] % dim_alignment,
|
|
-first_frame.shape[0] % dim_alignment)
|
|
padding = (to_pad[0] // 2, to_pad[0] - to_pad[0] // 2,
|
|
to_pad[1] // 2, to_pad[1] - to_pad[1] // 2)
|
|
padfunc = torch.nn.ReplicationPad2d(padding)
|
|
|
|
def pad(image):
|
|
return padfunc(image.permute(2, 0, 1).to(dtype=torch.float32)).permute(1, 2, 0)
|
|
|
|
logger.warn("Frames not aligned to required dimensions; padding applied")
|
|
images = list(map(pad, images))
|
|
|
|
# Corner watermark (only on user frames)
|
|
if corner_watermark is not None and corner_watermark_mask is not None:
|
|
logger.info("Applying corner watermark")
|
|
|
|
if isinstance(corner_watermark, torch.Tensor) and corner_watermark.ndim == 3:
|
|
corner_watermark = [corner_watermark] * len(images)
|
|
if isinstance(corner_watermark_mask, torch.Tensor) and corner_watermark_mask.ndim == 2:
|
|
corner_watermark_mask = [corner_watermark_mask] * len(images)
|
|
|
|
processed_wm, processed_mask, (x_offset, y_offset) = process_corner_watermark(
|
|
corner_watermark, corner_watermark_mask, first_frame.shape, corner_opacity, logger=logger,
|
|
)
|
|
|
|
if processed_wm and processed_mask:
|
|
for i in range(len(images)):
|
|
try:
|
|
frame = images[i]
|
|
wm = processed_wm[i % len(processed_wm)]
|
|
mask = processed_mask[i % len(processed_mask)]
|
|
mask_3d = mask.unsqueeze(-1).repeat(1, 1, 3)
|
|
blended = frame.clone()
|
|
|
|
h, w = wm.shape[:2]
|
|
y_start, x_start = max(0, y_offset), max(0, x_offset)
|
|
y_end = min(frame.shape[0], y_offset + h)
|
|
x_end = min(frame.shape[1], x_offset + w)
|
|
h_slice, w_slice = y_end - y_start, x_end - x_start
|
|
|
|
if h_slice <= 0 or w_slice <= 0:
|
|
continue
|
|
|
|
wm_y_start = -y_offset if y_offset < 0 else 0
|
|
wm_x_start = -x_offset if x_offset < 0 else 0
|
|
|
|
frame_area = frame[y_start:y_end, x_start:x_end]
|
|
mask_area = mask_3d[wm_y_start:wm_y_start + h_slice, wm_x_start:wm_x_start + w_slice]
|
|
wm_area = wm[wm_y_start:wm_y_start + h_slice, wm_x_start:wm_x_start + w_slice] * corner_opacity
|
|
|
|
if frame_area.shape != mask_area.shape or frame_area.shape != wm_area.shape:
|
|
continue
|
|
|
|
blended[y_start:y_end, x_start:x_end] = frame_area * (1.0 - mask_area) + wm_area * mask_area
|
|
images[i] = blended
|
|
except Exception as e:
|
|
logger.error(f"Error applying watermark to frame {i}: {e}")
|
|
|
|
state.has_corner_watermark = True
|
|
|
|
# Append watermark
|
|
watermark_frames = []
|
|
if append_watermark is not None:
|
|
logger.info("Processing append watermark")
|
|
if isinstance(append_watermark, torch.Tensor) and append_watermark.ndim == 3:
|
|
append_watermark = [append_watermark] * len(images)
|
|
|
|
watermark_frames = process_watermark_frames(append_watermark, first_frame.shape)
|
|
|
|
if append_watermark_source_fps > 0 and frame_rate > 0 and abs(append_watermark_source_fps - frame_rate) > 0.01:
|
|
watermark_frames = resample_frames_ffmpeg(watermark_frames, append_watermark_source_fps, frame_rate, logger)
|
|
|
|
state.has_append_watermark = True
|
|
|
|
if watermark_frames:
|
|
images.extend(watermark_frames)
|
|
|
|
return images
|
|
|
|
@classmethod
|
|
def _create_video_file(cls, processed_frames, video_format, video_metadata, frame_rate,
|
|
path_manager, state, logger):
|
|
"""Create the video file with ffmpeg, using main_pass settings from the format JSON."""
|
|
format_ext = video_format.get("extension", "mp4")
|
|
file_path = path_manager.get_video_path(format_ext)
|
|
logger.info(f"Creating video at {file_path}")
|
|
|
|
try:
|
|
if isinstance(processed_frames, torch.Tensor):
|
|
processed_frames = list(processed_frames)
|
|
|
|
env = os.environ.copy()
|
|
|
|
# Build input args
|
|
pix_fmt_in = "rgb24" if processed_frames[0].shape[2] == 3 else "rgba"
|
|
args = [
|
|
ffmpeg_path,
|
|
"-y",
|
|
"-v", "error",
|
|
"-f", "rawvideo",
|
|
"-pix_fmt", pix_fmt_in,
|
|
"-s", f"{processed_frames[0].shape[1]}x{processed_frames[0].shape[0]}",
|
|
"-r", str(frame_rate),
|
|
"-i", "-",
|
|
]
|
|
|
|
# Apply main_pass settings from the format JSON
|
|
main_pass = video_format.get("main_pass", [])
|
|
if main_pass:
|
|
args.extend(main_pass)
|
|
logger.info(f"Applying format main_pass args: {main_pass}")
|
|
else:
|
|
args.extend(["-c:v", "libx264", "-preset", "medium", "-crf", "23", "-pix_fmt", "yuv420p"])
|
|
logger.warn("No main_pass in format JSON, using default h264 settings")
|
|
|
|
logger.info(f"ffmpeg command: {' '.join(args)} {file_path}")
|
|
|
|
proc = ffmpeg_process(args, video_format, video_metadata, file_path, env, logger)
|
|
next(proc)
|
|
|
|
for frame in processed_frames:
|
|
proc.send(tensor_to_bytes(frame).tobytes())
|
|
|
|
total_frames_output = proc.send(None)
|
|
logger.info(f"Finished encoding with {total_frames_output} frames")
|
|
|
|
if os.path.exists(file_path):
|
|
file_size = os.path.getsize(file_path)
|
|
if file_size < 1000:
|
|
logger.error(f"Video file is too small, likely corrupt: {file_path} ({file_size} bytes)")
|
|
state.add_error(f"Video file too small: {file_size} bytes")
|
|
return False
|
|
else:
|
|
logger.info(f"Created video: {file_path} ({file_size} bytes)")
|
|
path_manager.add_output_file(file_path)
|
|
state.video_created = True
|
|
return True
|
|
else:
|
|
logger.error(f"Video file not created: {file_path}")
|
|
state.add_error("Video file not created")
|
|
return False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to create video: {str(e)}")
|
|
state.add_error(f"Video creation failed: {str(e)}")
|
|
return False
|