241 lines
8.3 KiB
Python
241 lines
8.3 KiB
Python
import io
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
import torch
|
|
import folder_paths
|
|
|
|
from .ffmpeg_utils import FFmpegError, run_ffmpeg
|
|
from .image_utils import load_frames_from_dir, save_frames_to_dir
|
|
|
|
try:
|
|
from comfy_api.input_impl import VideoFromFile
|
|
except Exception:
|
|
try:
|
|
from comfy_api.latest import InputImpl
|
|
|
|
VideoFromFile = InputImpl.VideoFromFile
|
|
except Exception:
|
|
VideoFromFile = None
|
|
|
|
|
|
def _is_video_input(value: Any) -> bool:
|
|
return hasattr(value, "get_stream_source") and hasattr(value, "get_components")
|
|
|
|
|
|
def _interleave_tensors(batch_a: torch.Tensor, batch_b: torch.Tensor) -> torch.Tensor:
|
|
if batch_a.dim() < 1 or batch_b.dim() < 1:
|
|
raise ValueError("Inputs must include a batch dimension.")
|
|
if batch_a.shape[0] != batch_b.shape[0]:
|
|
logging.warning(
|
|
"Interleave Frames: batch sizes differ (%s vs %s); truncating.",
|
|
batch_a.shape[0],
|
|
batch_b.shape[0],
|
|
)
|
|
count = min(batch_a.shape[0], batch_b.shape[0])
|
|
if count == 0:
|
|
raise ValueError("Inputs must contain at least one frame.")
|
|
batch_a = batch_a[:count]
|
|
batch_b = batch_b[:count]
|
|
return torch.stack([batch_a, batch_b], dim=1).reshape(2 * count, *batch_a.shape[1:])
|
|
|
|
|
|
def _detect_batch_type(value: Any) -> str:
|
|
if isinstance(value, dict):
|
|
if "samples" not in value:
|
|
raise ValueError("LATENT input is missing 'samples'.")
|
|
return "LATENT"
|
|
if isinstance(value, torch.Tensor):
|
|
return "IMAGE"
|
|
raise ValueError("Input must be IMAGE or LATENT.")
|
|
|
|
|
|
def _interleave_latents(batch_a: dict, batch_b: dict) -> dict:
|
|
samples_a = batch_a["samples"]
|
|
samples_b = batch_b["samples"]
|
|
if not isinstance(samples_a, torch.Tensor) or not isinstance(samples_b, torch.Tensor):
|
|
raise ValueError("LATENT samples must be torch tensors.")
|
|
|
|
output = {"samples": _interleave_tensors(samples_a, samples_b)}
|
|
for key, value_a in batch_a.items():
|
|
if key == "samples" or key not in batch_b:
|
|
continue
|
|
value_b = batch_b[key]
|
|
if not isinstance(value_a, torch.Tensor) or not isinstance(value_b, torch.Tensor):
|
|
continue
|
|
if value_a.shape[0] != samples_a.shape[0] or value_b.shape[0] != samples_b.shape[0]:
|
|
continue
|
|
output[key] = _interleave_tensors(value_a, value_b)
|
|
return output
|
|
|
|
|
|
def _ensure_video_path(video: Any) -> tuple[str, bool]:
|
|
if isinstance(video, str):
|
|
return video, False
|
|
if not _is_video_input(video):
|
|
raise ValueError("Input is not a supported video type.")
|
|
|
|
source = video.get_stream_source()
|
|
if isinstance(source, str):
|
|
return source, False
|
|
if isinstance(source, io.BytesIO):
|
|
source.seek(0)
|
|
extension = "mp4"
|
|
if hasattr(video, "get_container_format"):
|
|
try:
|
|
extension = video.get_container_format().split(",")[0] or extension
|
|
except Exception:
|
|
pass
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
handle, temp_path = tempfile.mkstemp(
|
|
prefix="ffmpeg_input_", suffix=f".{extension}", dir=temp_dir
|
|
)
|
|
os.close(handle)
|
|
with open(temp_path, "wb") as target:
|
|
target.write(source.read())
|
|
return temp_path, True
|
|
raise ValueError("Unsupported video stream source.")
|
|
|
|
|
|
class FFmpegSplitVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video": ("IMAGE,VIDEO", {"tooltip": "Video input or frame batch."}),
|
|
"fps": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 120.0, "step": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("frames",)
|
|
FUNCTION = "split_video"
|
|
CATEGORY = "ffmpeg"
|
|
|
|
def split_video(self, video, fps=0.0):
|
|
if isinstance(video, torch.Tensor):
|
|
if video.dim() != 4:
|
|
raise ValueError("IMAGE input must be a 4D tensor [N, H, W, C].")
|
|
return (video,)
|
|
|
|
input_path = None
|
|
cleanup_input = False
|
|
frames_dir = None
|
|
try:
|
|
input_path, cleanup_input = _ensure_video_path(video)
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
frames_dir = tempfile.mkdtemp(prefix="ffmpeg_split_", dir=temp_dir)
|
|
output_pattern = os.path.join(frames_dir, "frame_%06d.png")
|
|
|
|
args = ["-y", "-i", input_path]
|
|
if fps and fps > 0:
|
|
args += ["-vf", f"fps={fps}"]
|
|
args += ["-vsync", "0", output_pattern]
|
|
run_ffmpeg(args)
|
|
|
|
frames = load_frames_from_dir(frames_dir)
|
|
return (frames,)
|
|
except FFmpegError as exc:
|
|
raise RuntimeError(str(exc)) from exc
|
|
finally:
|
|
if cleanup_input and input_path and os.path.exists(input_path):
|
|
os.remove(input_path)
|
|
if frames_dir and os.path.isdir(frames_dir):
|
|
shutil.rmtree(frames_dir, ignore_errors=True)
|
|
|
|
|
|
class FFmpegMergeFrames:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"frames": ("IMAGE,VIDEO", {"tooltip": "Frame batch to merge into a video."}),
|
|
"fps": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 120.0, "step": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO", "IMAGE")
|
|
RETURN_NAMES = ("video", "frames")
|
|
FUNCTION = "merge_frames"
|
|
CATEGORY = "ffmpeg"
|
|
|
|
def merge_frames(self, frames, fps=30.0):
|
|
if _is_video_input(frames):
|
|
components = frames.get_components()
|
|
return (frames, components.images)
|
|
|
|
if not isinstance(frames, torch.Tensor) or frames.dim() != 4:
|
|
raise ValueError("Frames input must be a 4D IMAGE tensor [N, H, W, C].")
|
|
if frames.shape[0] == 0:
|
|
raise ValueError("Frames input is empty.")
|
|
if VideoFromFile is None:
|
|
raise RuntimeError("Video output is not available in this ComfyUI build.")
|
|
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
temp_root = tempfile.mkdtemp(prefix="ffmpeg_merge_", dir=temp_dir)
|
|
frames_dir = os.path.join(temp_root, "frames")
|
|
os.makedirs(frames_dir, exist_ok=True)
|
|
output_path = os.path.join(temp_root, "output.mp4")
|
|
|
|
try:
|
|
save_frames_to_dir(frames, frames_dir, "frame_%06d.png")
|
|
input_pattern = os.path.join(frames_dir, "frame_%06d.png")
|
|
args = [
|
|
"-y",
|
|
"-framerate",
|
|
str(float(fps)),
|
|
"-i",
|
|
input_pattern,
|
|
"-c:v",
|
|
"libx264",
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
output_path,
|
|
]
|
|
run_ffmpeg(args)
|
|
if not os.path.exists(output_path):
|
|
raise RuntimeError("FFmpeg did not produce an output video.")
|
|
return (VideoFromFile(output_path), frames)
|
|
except FFmpegError as exc:
|
|
raise RuntimeError(str(exc)) from exc
|
|
finally:
|
|
if os.path.isdir(frames_dir):
|
|
shutil.rmtree(frames_dir, ignore_errors=True)
|
|
|
|
|
|
class InterleaveFrames:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"batch_a": ("IMAGE,LATENT", {"tooltip": "First batch to interleave."}),
|
|
"batch_b": ("IMAGE,LATENT", {"tooltip": "Second batch to interleave."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE,LATENT",)
|
|
RETURN_NAMES = ("batch",)
|
|
FUNCTION = "interleave"
|
|
CATEGORY = "ffmpeg"
|
|
|
|
def interleave(self, batch_a, batch_b):
|
|
type_a = _detect_batch_type(batch_a)
|
|
type_b = _detect_batch_type(batch_b)
|
|
if type_a != type_b:
|
|
raise ValueError("Input types must match (IMAGE+IMAGE or LATENT+LATENT).")
|
|
|
|
if type_a == "IMAGE":
|
|
if not isinstance(batch_a, torch.Tensor) or not isinstance(batch_b, torch.Tensor):
|
|
raise ValueError("IMAGE inputs must be torch tensors.")
|
|
return (_interleave_tensors(batch_a, batch_b),)
|
|
|
|
if not isinstance(batch_a, dict) or not isinstance(batch_b, dict):
|
|
raise ValueError("LATENT inputs must be dicts.")
|
|
return (_interleave_latents(batch_a, batch_b),)
|