277 lines
9.5 KiB
Python
277 lines
9.5 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_many_tensors(batches: list[torch.Tensor]) -> torch.Tensor:
|
|
if len(batches) < 2:
|
|
raise ValueError("At least two batches are required.")
|
|
for batch in batches:
|
|
if batch.dim() < 1:
|
|
raise ValueError("Inputs must include a batch dimension.")
|
|
base_shape = batches[0].shape[1:]
|
|
for batch in batches[1:]:
|
|
if batch.shape[1:] != base_shape:
|
|
raise ValueError("All inputs must have the same shape (except batch size).")
|
|
|
|
lengths = [batch.shape[0] for batch in batches]
|
|
if len(set(lengths)) > 1:
|
|
logging.warning(
|
|
"Interleave Frames: batch sizes differ %s; truncating to shortest.",
|
|
lengths,
|
|
)
|
|
count = min(lengths)
|
|
if count == 0:
|
|
raise ValueError("Inputs must contain at least one frame.")
|
|
trimmed = [batch[:count] for batch in batches]
|
|
return torch.stack(trimmed, dim=1).reshape(count * len(trimmed), *base_shape)
|
|
|
|
|
|
def _interleave_tensors(batch_a: torch.Tensor, batch_b: torch.Tensor) -> torch.Tensor:
|
|
return _interleave_many_tensors([batch_a, batch_b])
|
|
|
|
|
|
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_latent_batches(batches: list[dict]) -> dict:
|
|
samples = [batch["samples"] for batch in batches]
|
|
for sample in samples:
|
|
if not isinstance(sample, torch.Tensor):
|
|
raise ValueError("LATENT samples must be torch tensors.")
|
|
|
|
output = {"samples": _interleave_many_tensors(samples)}
|
|
common_keys = set(batches[0].keys())
|
|
for batch in batches[1:]:
|
|
common_keys &= set(batch.keys())
|
|
common_keys.discard("samples")
|
|
|
|
for key in common_keys:
|
|
values = [batch[key] for batch in batches]
|
|
if not all(isinstance(value, torch.Tensor) for value in values):
|
|
continue
|
|
sample_lengths = [sample.shape[0] for sample in samples]
|
|
if any(value.shape[0] != length for value, length in zip(values, sample_lengths)):
|
|
continue
|
|
try:
|
|
output[key] = _interleave_many_tensors(values)
|
|
except ValueError:
|
|
continue
|
|
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": {
|
|
"batch1": ("IMAGE,LATENT", {"tooltip": "First batch to interleave."}),
|
|
"batch2": ("IMAGE,LATENT", {"tooltip": "Second batch to interleave."}),
|
|
},
|
|
"optional": {
|
|
"batch3": ("IMAGE,LATENT", {"tooltip": "Third batch to interleave."}),
|
|
"batch4": ("IMAGE,LATENT", {"tooltip": "Fourth batch to interleave."}),
|
|
"batch5": ("IMAGE,LATENT", {"tooltip": "Fifth batch to interleave."}),
|
|
"batch6": ("IMAGE,LATENT", {"tooltip": "Sixth batch to interleave."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE,LATENT",)
|
|
RETURN_NAMES = ("batch",)
|
|
FUNCTION = "interleave"
|
|
CATEGORY = "ffmpeg"
|
|
|
|
def interleave(
|
|
self,
|
|
batch1,
|
|
batch2,
|
|
batch3=None,
|
|
batch4=None,
|
|
batch5=None,
|
|
batch6=None,
|
|
):
|
|
batches = [batch for batch in [batch1, batch2, batch3, batch4, batch5, batch6] if batch is not None]
|
|
if len(batches) < 2:
|
|
raise ValueError("At least two batches must be provided.")
|
|
|
|
types = {_detect_batch_type(batch) for batch in batches}
|
|
if len(types) != 1:
|
|
raise ValueError("All inputs must have the same type (IMAGE or LATENT).")
|
|
batch_type = types.pop()
|
|
|
|
if batch_type == "IMAGE":
|
|
if not all(isinstance(batch, torch.Tensor) for batch in batches):
|
|
raise ValueError("IMAGE inputs must be torch tensors.")
|
|
return (_interleave_many_tensors(batches),)
|
|
|
|
if not all(isinstance(batch, dict) for batch in batches):
|
|
raise ValueError("LATENT inputs must be dicts.")
|
|
return (_interleave_latent_batches(batches),)
|