1153 lines
50 KiB
Python
1153 lines
50 KiB
Python
import os
|
||
import itertools
|
||
import numpy as np
|
||
import torch
|
||
import cv2
|
||
import hashlib
|
||
import re
|
||
import time
|
||
import subprocess
|
||
import shutil as _shutil
|
||
|
||
import folder_paths
|
||
from comfy.utils import common_upscale, ProgressBar
|
||
import nodes
|
||
|
||
from .xb_av import open_av
|
||
from .xb_paths import resolve_filename_prefix
|
||
|
||
# ============================================================
|
||
# 画幅比例 / 倍数 公共逻辑
|
||
# 画幅比例 = 自定义 + 官方 Resolution Selector 的比例；
|
||
# 倍数（multiple）约束宽高：既作为调节步长，也保证宽高能整除。
|
||
# ============================================================
|
||
XB_ASPECT_RATIOS = [
|
||
("自定义", None),
|
||
("1:1 (Square)", (1, 1)),
|
||
("2:3 (Portrait Photo)", (2, 3)),
|
||
("3:2 (Photo)", (3, 2)),
|
||
("3:4 (Portrait Standard)", (3, 4)),
|
||
("4:3 (Standard)", (4, 3)),
|
||
("9:16 (Portrait Widescreen)", (9, 16)),
|
||
("16:9 (Widescreen)", (16, 9)),
|
||
("21:9 (Ultrawide)", (21, 9)),
|
||
]
|
||
XB_ASPECT_VALUES = [name for name, _ in XB_ASPECT_RATIOS]
|
||
XB_ASPECT_MAP = {name: r for name, r in XB_ASPECT_RATIOS if r}
|
||
|
||
# 旧工作流里的值 → 新值（保证老图照常开）
|
||
XB_ASPECT_LEGACY = {
|
||
"Free": "自定义",
|
||
"1:1": "1:1 (Square)",
|
||
"2:3": "2:3 (Portrait Photo)",
|
||
"3:2": "3:2 (Photo)",
|
||
"3:4": "3:4 (Portrait Standard)",
|
||
"4:3": "4:3 (Standard)",
|
||
"9:16": "9:16 (Portrait Widescreen)",
|
||
"16:9": "16:9 (Widescreen)",
|
||
"21:9": "21:9 (Ultrawide)",
|
||
}
|
||
|
||
XB_MULTIPLE_DEFAULT = 16
|
||
|
||
|
||
def xb_multiple(value, default=XB_MULTIPLE_DEFAULT):
|
||
"""把「倍数」规整成 >= 8 的整数。"""
|
||
try:
|
||
m = int(round(float(value)))
|
||
except Exception:
|
||
m = default
|
||
return max(8, m)
|
||
|
||
|
||
def xb_snap(value, m):
|
||
"""吸附到 m 的整数倍（并保证 >= m）。"""
|
||
try:
|
||
v = int(round(float(value)))
|
||
except Exception:
|
||
v = m
|
||
return max(m, int(round(v / m)) * m)
|
||
|
||
|
||
def xb_aspect_ratio(name):
|
||
"""画幅比例 → 宽高比（宽/高）；自定义 / 未知返回 None。"""
|
||
if name in XB_ASPECT_MAP:
|
||
w, h = XB_ASPECT_MAP[name]
|
||
return w / h
|
||
name = XB_ASPECT_LEGACY.get(name, name)
|
||
if name in XB_ASPECT_MAP:
|
||
w, h = XB_ASPECT_MAP[name]
|
||
return w / h
|
||
# 兼容 "16:9 (LTX)" 这类旧写法
|
||
for key, ratio in XB_ASPECT_MAP.items():
|
||
if name and name.startswith(key.split(" ")[0]):
|
||
w, h = ratio
|
||
return w / h
|
||
return None
|
||
|
||
|
||
def xb_fit_size(name, width, height, m):
|
||
"""按比例 + 倍数算出最终宽高（两边都能被 m 整除）。"""
|
||
w = xb_snap(width, m)
|
||
h = xb_snap(height, m)
|
||
ratio = xb_aspect_ratio(name)
|
||
if ratio:
|
||
if w >= h:
|
||
h = xb_snap(w / ratio, m)
|
||
else:
|
||
w = xb_snap(h * ratio, m)
|
||
return max(m, w), max(m, h)
|
||
|
||
|
||
def xb_multiple_input():
|
||
"""「倍数」控件定义（放在最下面）。"""
|
||
return ("INT", {
|
||
"default": XB_MULTIPLE_DEFAULT, "min": 8, "max": 128, "step": 4,
|
||
"tooltip": "宽高必须是它的整数倍；调节分辨率时也按这个步长走（默认 16）。",
|
||
})
|
||
|
||
# ============================================================
|
||
# XB_ImageParamsMaster — 图像参数主控
|
||
# ============================================================
|
||
class XB_ImageParamsMaster:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"aspect_ratio": (XB_ASPECT_VALUES, {"default": "自定义"}),
|
||
"width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
|
||
"height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
|
||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
|
||
"strength_float": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 10.00, "step": 0.01}),
|
||
"strength_int": ("INT", {"default": 1, "min": 0, "max": 1000, "step": 1}),
|
||
"multiple": xb_multiple_input(),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("INT", "INT", "INT", "FLOAT", "INT", "INT")
|
||
RETURN_NAMES = ("Image Width", "Image Height", "Batch Size", "Float Control", "Int Control", "Scale Size")
|
||
FUNCTION = "process"
|
||
CATEGORY = "XB_ToolBox/Image_Params"
|
||
|
||
def process(self, aspect_ratio, width, height, batch_size, strength_float, strength_int, multiple=XB_MULTIPLE_DEFAULT):
|
||
m = xb_multiple(multiple)
|
||
safe_w, safe_h = xb_fit_size(aspect_ratio, width, height, m)
|
||
return (safe_w, safe_h, batch_size, float(strength_float), int(strength_int), max(safe_w, safe_h))
|
||
|
||
# ============================================================
|
||
# XB_ImageParamsMasterMini — 图像参数主控（精简版）
|
||
# 去掉「强度(小数)」「强度(整数)」两个控件与对应输出
|
||
# 保留画幅比例联动 + 16 步长联动（JS 端负责）
|
||
# ============================================================
|
||
class XB_ImageParamsMasterMini:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"aspect_ratio": (XB_ASPECT_VALUES, {"default": "自定义"}),
|
||
"width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
|
||
"height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
|
||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
|
||
"multiple": xb_multiple_input(),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("INT", "INT", "INT", "INT")
|
||
RETURN_NAMES = ("Image Width", "Image Height", "Batch Size", "Scale Size")
|
||
FUNCTION = "process"
|
||
CATEGORY = "XB_ToolBox/Image_Params"
|
||
|
||
def process(self, aspect_ratio, width, height, batch_size, multiple=XB_MULTIPLE_DEFAULT):
|
||
m = xb_multiple(multiple)
|
||
safe_w, safe_h = xb_fit_size(aspect_ratio, width, height, m)
|
||
return (safe_w, safe_h, batch_size, max(safe_w, safe_h))
|
||
|
||
# ============================================================
|
||
# XB_VideoParamsMaster — 视频参数主控
|
||
# ============================================================
|
||
class XB_VideoParamsMaster:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"aspect_ratio": (XB_ASPECT_VALUES, {"default": "自定义"}),
|
||
"duration_display": ("STRING", {"default": "Video Duration: 0.00 s", "multiline": False}),
|
||
"width": ("INT", {"default": 480, "min": 64, "max": 8192, "step": 16}),
|
||
"height": ("INT", {"default": 832, "min": 64, "max": 8192, "step": 16}),
|
||
"length": ("INT", {"default": 81, "min": 1, "max": 9999, "step": 1}),
|
||
"fps": ("INT", {"default": 16, "min": 1, "max": 120, "step": 1}),
|
||
"fps_float": ("FLOAT", {"default": 16.0, "min": 1.0, "max": 120.0, "step": 0.01}),
|
||
"multiple": xb_multiple_input(),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("INT", "INT", "INT", "INT", "FLOAT", "INT")
|
||
RETURN_NAMES = ("Width", "Height", "Frames", "FPS", "FPS_Float", "Scale Size")
|
||
FUNCTION = "process"
|
||
CATEGORY = "XB_ToolBox/Image_Params"
|
||
|
||
def process(self, aspect_ratio, duration_display, width, height, length, fps, fps_float,
|
||
multiple=XB_MULTIPLE_DEFAULT):
|
||
final_fps = int(round(fps))
|
||
m = xb_multiple(multiple)
|
||
safe_w, safe_h = xb_fit_size(aspect_ratio, width, height, m)
|
||
|
||
is_free = xb_aspect_ratio(aspect_ratio) is None
|
||
if is_free:
|
||
safe_len = max(1, int(length))
|
||
else:
|
||
safe_len = max(1, round((length - 1) / 8) * 8 + 1)
|
||
|
||
return (safe_w, safe_h, safe_len, final_fps, float(final_fps), max(safe_w, safe_h))
|
||
|
||
# ============================================================
|
||
# XB_MasterParameter — 万能参数控制器
|
||
# ============================================================
|
||
class XB_MasterParameter:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"mode": (["Free Mode", "Model Mode", "Encode Mode", "Decode Mode", "Ratio Mode", "Other Mode"], {"default": "Free Mode"}),
|
||
"value": ("FLOAT", {"default": 1024.0, "min": 0.0, "max": 9999.0, "step": 0.01}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("INT", "FLOAT")
|
||
RETURN_NAMES = ("Integer (INT)", "Float (FLOAT)")
|
||
FUNCTION = "get_value"
|
||
CATEGORY = "XB_ToolBox/Params_Control"
|
||
|
||
def get_value(self, mode, value):
|
||
if mode == "Free Mode":
|
||
val = max(0.0, min(9999.0, value))
|
||
elif mode == "Model Mode":
|
||
val = max(0.0, min(50.0, value))
|
||
elif mode in ["Encode Mode", "Decode Mode"]:
|
||
val = max(64.0, min(3840.0, value))
|
||
val = round(val / 32) * 32
|
||
elif mode == "Ratio Mode":
|
||
val = max(0.0, min(1.0, value))
|
||
elif mode == "Other Mode":
|
||
val = max(0.0, min(9999.0, value))
|
||
else:
|
||
val = value
|
||
|
||
return (int(round(val)), float(val))
|
||
|
||
|
||
# ============================================================
|
||
# 以下代码一字不差搬运自 VHS load_video_nodes.py 和 utils.py
|
||
# ============================================================
|
||
|
||
DIMMAX = 8192
|
||
BIGMAX = (2**53-1)
|
||
|
||
video_extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov']
|
||
|
||
VHSLoadFormats = {
|
||
'None': {},
|
||
'AnimateDiff': {'target_rate': 8, 'dim': (8,0,512,512)},
|
||
'Mochi': {'target_rate': 24, 'dim': (16,0,848,480), 'frames':(6,1)},
|
||
'LTXV': {'target_rate': 24, 'dim': (32,0,768,512), 'frames':(8,1)},
|
||
'Hunyuan': {'target_rate': 24, 'dim': (16,0,848,480), 'frames':(4,1)},
|
||
'Cosmos': {'target_rate': 24, 'dim': (16,0,1280,704), 'frames':(8,1)},
|
||
'Wan': {'target_rate': 16, 'dim': (8,0,832,480), 'frames':(4,1)},
|
||
}
|
||
|
||
if not hasattr(nodes, 'VHSLoadFormats'):
|
||
nodes.VHSLoadFormats = {}
|
||
|
||
def get_load_formats():
|
||
formats = {}
|
||
formats.update(nodes.VHSLoadFormats)
|
||
formats.update(VHSLoadFormats)
|
||
return (list(formats.keys()),
|
||
{'default': 'AnimateDiff', 'formats': formats})
|
||
|
||
def get_format(format):
|
||
if format in VHSLoadFormats:
|
||
return VHSLoadFormats[format]
|
||
return nodes.VHSLoadFormats.get(format, {})
|
||
|
||
def strip_path(path):
|
||
path = path.strip()
|
||
if path.startswith("\""):
|
||
path = path[1:]
|
||
if path.endswith("\""):
|
||
path = path[:-1]
|
||
return path
|
||
|
||
def calculate_file_hash(filename: str, hash_every_n: int = 1):
|
||
h = hashlib.sha256()
|
||
h.update(filename.encode())
|
||
h.update(str(os.path.getmtime(filename)).encode())
|
||
return h.hexdigest()
|
||
|
||
def target_size(width, height, custom_width, custom_height, downscale_ratio=8) -> tuple[int, int]:
|
||
if downscale_ratio is None:
|
||
downscale_ratio = 8
|
||
if custom_width == 0 and custom_height == 0:
|
||
pass
|
||
elif custom_height == 0:
|
||
height *= custom_width/width
|
||
width = custom_width
|
||
elif custom_width == 0:
|
||
width *= custom_height/height
|
||
height = custom_height
|
||
else:
|
||
width = custom_width
|
||
height = custom_height
|
||
width = int(width/downscale_ratio + 0.5) * downscale_ratio
|
||
height = int(height/downscale_ratio + 0.5) * downscale_ratio
|
||
return (width, height)
|
||
|
||
def cv_frame_generator(video, force_rate, frame_load_cap, skip_first_frames,
|
||
select_every_nth, meta_batch=None, unique_id=None):
|
||
video_cap = cv2.VideoCapture(video)
|
||
if not video_cap.isOpened() or not video_cap.grab():
|
||
raise ValueError(f"{video} could not be loaded with cv.")
|
||
|
||
# extract video metadata
|
||
fps = video_cap.get(cv2.CAP_PROP_FPS)
|
||
width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
total_frames = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
duration = total_frames / fps
|
||
|
||
width = 0
|
||
|
||
if width <=0 or height <=0:
|
||
_, frame = video_cap.retrieve()
|
||
height, width, _ = frame.shape
|
||
|
||
# set video_cap to look at start_index frame
|
||
total_frame_count = 0
|
||
total_frames_evaluated = -1
|
||
frames_added = 0
|
||
base_frame_time = 1 / fps
|
||
prev_frame = None
|
||
|
||
if force_rate == 0:
|
||
target_frame_time = base_frame_time
|
||
else:
|
||
target_frame_time = 1/force_rate
|
||
|
||
if total_frames > 0:
|
||
if force_rate != 0:
|
||
yieldable_frames = int(total_frames / fps * force_rate)
|
||
else:
|
||
yieldable_frames = total_frames
|
||
if select_every_nth:
|
||
yieldable_frames //= select_every_nth
|
||
if frame_load_cap != 0:
|
||
yieldable_frames = min(frame_load_cap, yieldable_frames)
|
||
else:
|
||
yieldable_frames = 0
|
||
yield (width, height, fps, duration, total_frames, target_frame_time, yieldable_frames)
|
||
pbar = ProgressBar(yieldable_frames)
|
||
time_offset=target_frame_time
|
||
while video_cap.isOpened():
|
||
if time_offset < target_frame_time:
|
||
is_returned = video_cap.grab()
|
||
# if didn't return frame, video has ended
|
||
if not is_returned:
|
||
break
|
||
time_offset += base_frame_time
|
||
if time_offset < target_frame_time:
|
||
continue
|
||
time_offset -= target_frame_time
|
||
# if not at start_index, skip doing anything with frame
|
||
total_frame_count += 1
|
||
if total_frame_count <= skip_first_frames:
|
||
continue
|
||
else:
|
||
total_frames_evaluated += 1
|
||
|
||
# if should not be selected, skip doing anything with frame
|
||
if total_frames_evaluated%select_every_nth != 0:
|
||
continue
|
||
|
||
# opencv loads images in BGR format (yuck), so need to convert to RGB for ComfyUI use
|
||
# follow up: can videos ever have an alpha channel?
|
||
# To my testing: No. opencv has no support for alpha
|
||
unused, frame = video_cap.retrieve()
|
||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
# convert frame to comfyui's expected format
|
||
# TODO: frame contains no exif information. Check if opencv2 has already applied
|
||
frame = np.array(frame, dtype=np.float32)
|
||
torch.from_numpy(frame).div_(255)
|
||
if prev_frame is not None:
|
||
inp = yield prev_frame
|
||
if inp is not None:
|
||
#ensure the finally block is called
|
||
return
|
||
prev_frame = frame
|
||
frames_added += 1
|
||
if pbar is not None:
|
||
pbar.update_absolute(frames_added, yieldable_frames)
|
||
# if cap exists and we've reached it, stop processing frames
|
||
if frame_load_cap > 0 and frames_added >= frame_load_cap:
|
||
break
|
||
if meta_batch is not None:
|
||
meta_batch.inputs.pop(unique_id)
|
||
meta_batch.has_closed_inputs = True
|
||
if prev_frame is not None:
|
||
yield prev_frame
|
||
|
||
#Python 3.12 adds an itertools.batched, but it's easily replicated for legacy support
|
||
def batched(it, n):
|
||
while batch := tuple(itertools.islice(it, n)):
|
||
yield batch
|
||
|
||
def batched_vae_encode(images, vae, frames_per_batch):
|
||
for batch in batched(images, frames_per_batch):
|
||
image_batch = torch.from_numpy(np.array(batch))
|
||
yield from vae.encode(image_batch).numpy()
|
||
|
||
def resized_cv_frame_gen(custom_width, custom_height, downscale_ratio, **kwargs):
|
||
gen = cv_frame_generator(**kwargs)
|
||
info = next(gen)
|
||
width, height = info[0], info[1]
|
||
frames_per_batch = (1920 * 1080 * 16) // (width * height) or 1
|
||
if kwargs.get('meta_batch', None) is not None:
|
||
frames_per_batch = min(frames_per_batch, kwargs['meta_batch'].frames_per_batch)
|
||
if custom_width != 0 or custom_height != 0 or downscale_ratio is not None:
|
||
new_size = target_size(width, height, custom_width, custom_height, downscale_ratio)
|
||
yield (*info, new_size[0], new_size[1], False)
|
||
if new_size[0] != width or new_size[1] != height:
|
||
def rescale(frame):
|
||
s = torch.from_numpy(np.fromiter(frame, np.dtype((np.float32, (height, width, 3)))))
|
||
s = s.movedim(-1,1)
|
||
s = common_upscale(s, new_size[0], new_size[1], "lanczos", "center")
|
||
return s.movedim(1,-1).numpy()
|
||
yield from itertools.chain.from_iterable(map(rescale, batched(gen, frames_per_batch)))
|
||
return
|
||
else:
|
||
yield (*info, info[0], info[1], False)
|
||
yield from gen
|
||
|
||
def _lazy_get_audio(file, start_time=0, duration=0):
|
||
"""替代 VHS 的 lazy_get_audio，使用 PyAV。视频无音轨时返回 None，不伪造静音。"""
|
||
try:
|
||
import av
|
||
except ImportError as exc:
|
||
raise RuntimeError("Video audio loading requires PyAV (av)") from exc
|
||
# 只探测流元数据，不解码；这样可选 AUDIO 消费端能拿到 None 而不是非法的单采样点张量
|
||
try:
|
||
with open_av(file) as container:
|
||
if not any(stream.type == 'audio' for stream in container.streams):
|
||
return None
|
||
except Exception as exc:
|
||
raise RuntimeError("Unable to inspect video audio streams") from exc
|
||
class _LazyAudioMap:
|
||
def __init__(s, file, st, dur):
|
||
s._d = None; s.file = file; s.st = st; s.dur = dur
|
||
def _load(s):
|
||
if s._d is not None: return
|
||
try:
|
||
with open_av(s.file) as c:
|
||
a = next((x for x in c.streams if x.type == 'audio'), None)
|
||
if a is None:
|
||
raise ValueError("Audio stream disappeared after probing")
|
||
sr = a.codec_context.sample_rate or 44100
|
||
frames = []
|
||
sp = int(s.st*sr) if s.st>0 else 0
|
||
ep = int((s.st+s.dur)*sr) if s.dur>0 else None
|
||
for f in c.decode(a):
|
||
if f.pts is not None:
|
||
if sp>0 and f.pts<sp: continue
|
||
if ep is not None and f.pts>=ep: break
|
||
frames.append(f.to_ndarray())
|
||
if not frames:
|
||
raise ValueError("Audio stream has no decoded frames in the requested interval")
|
||
wav = np.concatenate(frames, axis=-1)
|
||
if wav.size == 0:
|
||
raise ValueError("Audio stream contains no decoded samples")
|
||
if wav.ndim==1: wav=wav[np.newaxis,:]
|
||
elif wav.shape[0]>1: wav=wav.mean(axis=0, keepdims=True)
|
||
wf = torch.from_numpy(wav.astype(np.float32))
|
||
if wf.dim()==2: wf=wf.unsqueeze(0)
|
||
s._d = {'waveform': wf, 'sample_rate': sr}
|
||
except Exception as exc:
|
||
raise RuntimeError("Unable to decode video audio") from exc
|
||
def __getitem__(s,k): s._load(); return s._d[k]
|
||
def __iter__(s): s._load(); return iter(s._d)
|
||
def __len__(s): s._load(); return len(s._d)
|
||
return _LazyAudioMap(file, start_time, duration)
|
||
|
||
def load_video(meta_batch=None, unique_id=None, memory_limit_mb=None, vae=None,
|
||
generator=resized_cv_frame_gen, format='None', **kwargs):
|
||
if 'force_size' in kwargs:
|
||
kwargs.pop('force_size')
|
||
format = get_format(format)
|
||
kwargs['video'] = strip_path(kwargs['video'])
|
||
if vae is not None:
|
||
downscale_ratio = getattr(vae, "downscale_ratio", 8)
|
||
else:
|
||
downscale_ratio = format.get('dim', (1,))[0]
|
||
if meta_batch is None or unique_id not in meta_batch.inputs:
|
||
gen = generator(meta_batch=meta_batch, unique_id=unique_id, downscale_ratio=downscale_ratio, **kwargs)
|
||
(width, height, fps, duration, total_frames, target_frame_time, yieldable_frames, new_width, new_height, alpha) = next(gen)
|
||
|
||
if meta_batch is not None:
|
||
meta_batch.inputs[unique_id] = (gen, width, height, fps, duration, total_frames, target_frame_time, yieldable_frames, new_width, new_height, alpha)
|
||
if yieldable_frames:
|
||
meta_batch.total_frames = min(meta_batch.total_frames, yieldable_frames)
|
||
else:
|
||
(gen, width, height, fps, duration, total_frames, target_frame_time, yieldable_frames, new_width, new_height, alpha) = meta_batch.inputs[unique_id]
|
||
|
||
memory_limit = None
|
||
if memory_limit_mb is not None:
|
||
memory_limit *= 2 ** 20
|
||
else:
|
||
try:
|
||
import psutil
|
||
memory_limit = (psutil.virtual_memory().available + psutil.swap_memory().free) - 2 ** 27
|
||
except:
|
||
memory_limit = BIGMAX
|
||
if vae is not None:
|
||
max_loadable_frames = int(memory_limit//(width*height*3*(4+4+1/10)))
|
||
else:
|
||
max_loadable_frames = int(memory_limit//(width*height*3*(.1)))
|
||
if meta_batch is not None:
|
||
if 'frames' in format:
|
||
if meta_batch.frames_per_batch % format['frames'][0] != format['frames'][1]:
|
||
error = (meta_batch.frames_per_batch - format['frames'][1]) % format['frames'][0]
|
||
suggested = meta_batch.frames_per_batch - error
|
||
if error > format['frames'][0] / 2:
|
||
suggested += format['frames'][0]
|
||
raise RuntimeError(f"The chosen frames per batch is incompatible with the selected format. Try {suggested}")
|
||
if meta_batch.frames_per_batch > max_loadable_frames:
|
||
raise RuntimeError(f"Meta Batch set to {meta_batch.frames_per_batch} frames but only {max_loadable_frames} can fit in memory")
|
||
gen = itertools.islice(gen, meta_batch.frames_per_batch)
|
||
else:
|
||
original_gen = gen
|
||
gen = itertools.islice(gen, max_loadable_frames)
|
||
frames_per_batch = (1920 * 1080 * 16) // (width * height) or 1
|
||
if vae is not None:
|
||
gen = batched_vae_encode(gen, vae, frames_per_batch)
|
||
vw,vh = new_width//downscale_ratio, new_height//downscale_ratio
|
||
channels = getattr(vae, 'latent_channels', 4)
|
||
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (channels,vh,vw)))))
|
||
else:
|
||
#Some minor wizardry to eliminate a copy and reduce max memory by a factor of ~2
|
||
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (new_height, new_width, 4 if alpha else 3)))))
|
||
if meta_batch is None and memory_limit is not None:
|
||
try:
|
||
next(original_gen)
|
||
raise RuntimeError(f"Memory limit hit after loading {len(images)} frames. Stopping execution.")
|
||
except StopIteration:
|
||
pass
|
||
if len(images) == 0:
|
||
raise RuntimeError("No frames generated")
|
||
if 'frames' in format and len(images) % format['frames'][0] != format['frames'][1]:
|
||
err_msg = f"The number of frames loaded {len(images)}, does not match the requirements of the currently selected format."
|
||
if len(format['frames']) > 2 and format['frames'][2]:
|
||
raise RuntimeError(err_msg)
|
||
div, mod = format['frames'][:2]
|
||
frames = (len(images) - mod) // div * div + mod
|
||
images = images[:frames]
|
||
if 'start_time' in kwargs:
|
||
start_time = kwargs['start_time']
|
||
else:
|
||
start_time = kwargs['skip_first_frames'] * target_frame_time
|
||
target_frame_time *= kwargs.get('select_every_nth', 1)
|
||
#Setup lambda for lazy audio capture
|
||
audio = _lazy_get_audio(kwargs['video'], start_time, kwargs['frame_load_cap']*target_frame_time)
|
||
#Adjust target_frame_time for select_every_nth
|
||
video_info = {
|
||
"source_fps": fps,
|
||
"source_frame_count": total_frames,
|
||
"source_duration": duration,
|
||
"source_width": width,
|
||
"source_height": height,
|
||
"loaded_fps": 1/target_frame_time,
|
||
"loaded_frame_count": len(images),
|
||
"loaded_duration": len(images) * target_frame_time,
|
||
"loaded_width": new_width,
|
||
"loaded_height": new_height,
|
||
}
|
||
if vae is None:
|
||
return (images, len(images), audio, video_info)
|
||
else:
|
||
return ({"samples": images}, len(images), audio, video_info)
|
||
|
||
|
||
# ============================================================
|
||
# XB_VideoLoader
|
||
# ============================================================
|
||
|
||
class XB_VideoLoader:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
input_dir = folder_paths.get_input_directory()
|
||
files = []
|
||
for f in os.listdir(input_dir):
|
||
if os.path.isfile(os.path.join(input_dir, f)):
|
||
file_parts = f.split('.')
|
||
if len(file_parts) > 1 and (file_parts[-1].lower() in video_extensions):
|
||
files.append(f)
|
||
return {"required": {
|
||
"video": (sorted(files),),
|
||
"force_rate": ("FLOAT", {"default": 0, "min": 0, "max": 60, "step": 1}),
|
||
"custom_width": ("INT", {"default": 0, "min": 0, "max": DIMMAX}),
|
||
"custom_height": ("INT", {"default": 0, "min": 0, "max": DIMMAX}),
|
||
"frame_load_cap": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
|
||
"skip_first_frames": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
|
||
"select_every_nth": ("INT", {"default": 1, "min": 1, "max": BIGMAX, "step": 1}),
|
||
},
|
||
"optional": {
|
||
"meta_batch": ("VHS_BatchManager",),
|
||
"vae": ("VAE",),
|
||
"format": get_load_formats(),
|
||
},
|
||
"hidden": {
|
||
"unique_id": "UNIQUE_ID"
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "INT", "AUDIO", "VHS_VIDEOINFO")
|
||
RETURN_NAMES = ("IMAGE", "frame_count", "audio", "video_info")
|
||
FUNCTION = "load_video"
|
||
CATEGORY = "XB_ToolBox/Video"
|
||
|
||
def load_video(self, **kwargs):
|
||
kwargs['video'] = folder_paths.get_annotated_filepath(strip_path(kwargs['video']))
|
||
return load_video(**kwargs)
|
||
|
||
@classmethod
|
||
def IS_CHANGED(s, video, **kwargs):
|
||
image_path = folder_paths.get_annotated_filepath(video)
|
||
return calculate_file_hash(image_path)
|
||
|
||
@classmethod
|
||
def VALIDATE_INPUTS(s, video):
|
||
if not folder_paths.exists_annotated_filepath(video):
|
||
return "Invalid video file: {}".format(video)
|
||
return True
|
||
|
||
|
||
# ============================================================
|
||
# 以下代码搬运自 VHS utils.py 和 nodes.py
|
||
# ============================================================
|
||
|
||
# ffmpeg_path
|
||
if "VHS_FORCE_FFMPEG_PATH" in os.environ:
|
||
_ffmpeg = os.environ.get("VHS_FORCE_FFMPEG_PATH")
|
||
else:
|
||
_ffmpeg_paths = []
|
||
try:
|
||
from imageio_ffmpeg import get_ffmpeg_exe
|
||
_ffmpeg_paths.append(get_ffmpeg_exe())
|
||
except:
|
||
pass
|
||
system_ffmpeg = _shutil.which("ffmpeg")
|
||
if system_ffmpeg:
|
||
_ffmpeg_paths.append(system_ffmpeg)
|
||
if os.path.isfile("ffmpeg"): _ffmpeg_paths.append(os.path.abspath("ffmpeg"))
|
||
if os.path.isfile("ffmpeg.exe"): _ffmpeg_paths.append(os.path.abspath("ffmpeg.exe"))
|
||
_ffmpeg = _ffmpeg_paths[0] if _ffmpeg_paths else None
|
||
|
||
import json as _json
|
||
import datetime as _dt
|
||
import sys as _sys
|
||
import functools as _functools
|
||
from PIL import Image as _PILImage, ExifTags as _ExifTags
|
||
from PIL.PngImagePlugin import PngInfo as _PngInfo
|
||
from string import Template as _Template
|
||
|
||
# gifski_path
|
||
_gifski_path = os.environ.get("VHS_GIFSKI", None)
|
||
if _gifski_path is None:
|
||
_gifski_path = os.environ.get("JOV_GIFSKI", None)
|
||
if _gifski_path is None:
|
||
_gifski_path = _shutil.which("gifski")
|
||
|
||
def _merge_filter_args(args, ftype="-vf"):
|
||
try:
|
||
start_index = args.index(ftype)+1
|
||
index = start_index
|
||
while True:
|
||
index = args.index(ftype, index)
|
||
args[start_index] += ',' + args[index+1]
|
||
args.pop(index)
|
||
args.pop(index)
|
||
except ValueError:
|
||
pass
|
||
|
||
def _cached(duration):
|
||
def dec(f):
|
||
cached_ret = None
|
||
cache_time = 0
|
||
def cached_func():
|
||
nonlocal cache_time, cached_ret
|
||
if time.time() > cache_time + duration or cached_ret is None:
|
||
cache_time = time.time()
|
||
cached_ret = f()
|
||
return cached_ret
|
||
return cached_func
|
||
return dec
|
||
|
||
class _MultiInput(str):
|
||
def __new__(cls, string, allowed_types="*"):
|
||
res = super().__new__(cls, string)
|
||
res.allowed_types=allowed_types
|
||
return res
|
||
def __ne__(self, other):
|
||
if self.allowed_types == "*" or other == "*":
|
||
return False
|
||
return other not in self.allowed_types
|
||
|
||
_imageOrLatent = _MultiInput("IMAGE", ["IMAGE", "LATENT"])
|
||
_floatOrInt = _MultiInput("FLOAT", ["FLOAT", "INT"])
|
||
|
||
class _ContainsAll(dict):
|
||
def __contains__(self, other):
|
||
return True
|
||
def __getitem__(self, key):
|
||
return super().get(key, (None, {}))
|
||
|
||
def _flatten_list(l):
|
||
ret = []
|
||
for e in l:
|
||
if isinstance(e, list):
|
||
ret.extend(e)
|
||
else:
|
||
ret.append(e)
|
||
return ret
|
||
|
||
def _iterate_format(video_format, for_widgets=True):
|
||
def indirector(cont, index):
|
||
if isinstance(cont[index], list) and (not for_widgets
|
||
or len(cont[index])> 1 and not isinstance(cont[index][1], dict)):
|
||
inp = yield cont[index]
|
||
if inp is not None:
|
||
cont[index] = inp
|
||
yield
|
||
for k in video_format:
|
||
if k == "extra_widgets":
|
||
if for_widgets:
|
||
yield from video_format["extra_widgets"]
|
||
elif k.endswith("_pass"):
|
||
for i in range(len(video_format[k])):
|
||
yield from indirector(video_format[k], i)
|
||
if not for_widgets:
|
||
video_format[k] = _flatten_list(video_format[k])
|
||
else:
|
||
yield from indirector(video_format, k)
|
||
|
||
if 'VHS_video_formats' not in folder_paths.folder_names_and_paths:
|
||
folder_paths.folder_names_and_paths["VHS_video_formats"] = ((),{".json"})
|
||
if len(folder_paths.folder_names_and_paths['VHS_video_formats'][1]) == 0:
|
||
folder_paths.folder_names_and_paths["VHS_video_formats"][1].add(".json")
|
||
|
||
_base_formats_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "video_formats")
|
||
|
||
@_cached(5)
|
||
def _get_video_formats():
|
||
format_files = {}
|
||
for format_name in folder_paths.get_filename_list("VHS_video_formats"):
|
||
format_files[format_name] = folder_paths.get_full_path("VHS_video_formats", format_name)
|
||
os.makedirs(_base_formats_dir, exist_ok=True)
|
||
for item in os.scandir(_base_formats_dir):
|
||
if not item.is_file() or not item.name.endswith('.json'):
|
||
continue
|
||
format_files[item.name[:-5]] = item.path
|
||
formats = []
|
||
format_widgets = {}
|
||
for format_name, path in format_files.items():
|
||
with open(path, 'r') as stream:
|
||
video_format = _json.load(stream)
|
||
if "gifski_pass" in video_format and _gifski_path is None:
|
||
continue
|
||
widgets = list(_iterate_format(video_format))
|
||
formats.append("video/" + format_name)
|
||
if (len(widgets) > 0):
|
||
format_widgets["video/"+ format_name] = widgets
|
||
return formats, format_widgets
|
||
|
||
def _apply_format_widgets(format_name, kwargs):
|
||
if os.path.exists(os.path.join(_base_formats_dir, format_name + ".json")):
|
||
video_format_path = os.path.join(_base_formats_dir, format_name + ".json")
|
||
else:
|
||
video_format_path = folder_paths.get_full_path("VHS_video_formats", format_name)
|
||
with open(video_format_path, 'r') as stream:
|
||
video_format = _json.load(stream)
|
||
for w in _iterate_format(video_format):
|
||
if w[0] not in kwargs:
|
||
if len(w) > 2 and 'default' in w[2]:
|
||
default = w[2]['default']
|
||
else:
|
||
if type(w[1]) is list:
|
||
default = w[1][0]
|
||
else:
|
||
default = {"BOOLEAN": False, "INT": 0, "FLOAT": 0, "STRING": ""}[w[1]]
|
||
kwargs[w[0]] = default
|
||
wit = _iterate_format(video_format, False)
|
||
for w in wit:
|
||
while isinstance(w, list):
|
||
if len(w) == 1:
|
||
w = [_Template(x).substitute(**kwargs) for x in w[0]]
|
||
break
|
||
elif isinstance(w[1], dict):
|
||
w = w[1][str(kwargs[w[0]])]
|
||
elif len(w) > 3:
|
||
w = _Template(w[3]).substitute(val=kwargs[w[0]])
|
||
else:
|
||
w = str(kwargs[w[0]])
|
||
wit.send(w)
|
||
return video_format
|
||
|
||
def _tensor_to_int(tensor, bits):
|
||
tensor = tensor.cpu().numpy() * (2**bits-1) + 0.5
|
||
return np.clip(tensor, 0, (2**bits-1))
|
||
def _tensor_to_shorts(tensor):
|
||
return _tensor_to_int(tensor, 16).astype(np.uint16)
|
||
def _tensor_to_bytes(tensor):
|
||
return _tensor_to_int(tensor, 8).astype(np.uint8)
|
||
|
||
def _ffmpeg_process(args, video_format, video_metadata, file_path, env):
|
||
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_path = os.path.join(folder_paths.get_temp_directory(), "metadata.txt")
|
||
def escape_ffmpeg_metadata(key, value):
|
||
value = str(value)
|
||
value = value.replace("\\","\\\\")
|
||
value = value.replace(";","\\;")
|
||
value = value.replace("#","\\#")
|
||
value = value.replace("=","\\=")
|
||
value = value.replace("\n","\\\n")
|
||
return f"{key}={value}"
|
||
with open(metadata_path, "w") as f:
|
||
f.write(";FFMETADATA1\n")
|
||
if "prompt" in video_metadata:
|
||
f.write(escape_ffmpeg_metadata("prompt", _json.dumps(video_metadata["prompt"])) + "\n")
|
||
if "workflow" in video_metadata:
|
||
f.write(escape_ffmpeg_metadata("workflow", _json.dumps(video_metadata["workflow"])) + "\n")
|
||
for k, v in video_metadata.items():
|
||
if k not in ["prompt", "workflow"]:
|
||
f.write(escape_ffmpeg_metadata(k, _json.dumps(v)) + "\n")
|
||
m_args = args[:1] + ["-i", metadata_path] + args[1:] + ["-metadata", "creation_time=now", "-movflags", "use_metadata_tags"]
|
||
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 as e:
|
||
err = proc.stderr.read()
|
||
if os.path.exists(file_path):
|
||
raise Exception("An error occurred in the ffmpeg subprocess:\n" + err.decode(*ENCODE_ARGS))
|
||
print(err.decode(*ENCODE_ARGS), end="", file=_sys.stderr)
|
||
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 as e:
|
||
res = proc.stderr.read()
|
||
raise Exception("An error occurred in the ffmpeg subprocess:\n" + res.decode(*ENCODE_ARGS))
|
||
yield total_frames_output
|
||
if len(res) > 0:
|
||
print(res.decode(*ENCODE_ARGS), end="", file=_sys.stderr)
|
||
|
||
def _gifski_process(args, dimensions, frame_rate, video_format, file_path, env):
|
||
frame_data = yield
|
||
with subprocess.Popen(args + video_format['main_pass'] + ['-f', 'yuv4mpegpipe', '-'],
|
||
stderr=subprocess.PIPE, stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE, env=env) as procff:
|
||
with subprocess.Popen([_gifski_path] + video_format['gifski_pass']
|
||
+ ['-W', f'{dimensions[0]}', '-H', f'{dimensions[1]}']
|
||
+ ['-r', f'{frame_rate}']
|
||
+ ['-q', '-o', file_path, '-'], stderr=subprocess.PIPE,
|
||
stdin=procff.stdout, stdout=subprocess.PIPE,
|
||
env=env) as procgs:
|
||
try:
|
||
while frame_data is not None:
|
||
procff.stdin.write(frame_data)
|
||
frame_data = yield
|
||
procff.stdin.flush()
|
||
procff.stdin.close()
|
||
resff = procff.stderr.read()
|
||
resgs = procgs.stderr.read()
|
||
outgs = procgs.stdout.read()
|
||
except BrokenPipeError as e:
|
||
procff.stdin.close()
|
||
resff = procff.stderr.read()
|
||
resgs = procgs.stderr.read()
|
||
raise Exception("An error occurred while creating gifski output\n"
|
||
+ "Make sure you are using gifski --version >=1.32.0\nffmpeg: "
|
||
+ resff.decode(*ENCODE_ARGS) + '\ngifski: ' + resgs.decode(*ENCODE_ARGS))
|
||
if len(resff) > 0:
|
||
print(resff.decode(*ENCODE_ARGS), end="", file=_sys.stderr)
|
||
if len(resgs) > 0:
|
||
print(resgs.decode(*ENCODE_ARGS), end="", file=_sys.stderr)
|
||
if len(outgs) > 0:
|
||
print(outgs.decode(*ENCODE_ARGS))
|
||
|
||
def _to_pingpong(inp):
|
||
if not hasattr(inp, "__getitem__"):
|
||
inp = list(inp)
|
||
yield from inp
|
||
for i in range(len(inp)-2,0,-1):
|
||
yield inp[i]
|
||
|
||
|
||
# ============================================================
|
||
# XB_VideoCombine — 完整搬运 VHS VideoCombine，完全独立
|
||
# ============================================================
|
||
|
||
class XB_VideoCombine:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
ffmpeg_formats, format_widgets = _get_video_formats()
|
||
format_widgets["image/webp"] = [['lossless', "BOOLEAN", {'default': True}]]
|
||
return {
|
||
"required": {
|
||
"images": (_imageOrLatent,),
|
||
"frame_rate": (_floatOrInt, {"default": 8, "min": 1, "step": 1}),
|
||
"loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
||
"filename_prefix": ("STRING", {"default": "视频"}),
|
||
"format": (["image/gif", "image/webp"] + ffmpeg_formats, {'formats': format_widgets}),
|
||
"pingpong": ("BOOLEAN", {"default": False}),
|
||
"save_output": ("BOOLEAN", {"default": True}),
|
||
},
|
||
"optional": {
|
||
"audio": ("AUDIO",),
|
||
},
|
||
"hidden": _ContainsAll({
|
||
"prompt": "PROMPT",
|
||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||
"unique_id": "UNIQUE_ID"
|
||
}),
|
||
}
|
||
|
||
RETURN_TYPES = ("VHS_FILENAMES",)
|
||
RETURN_NAMES = ("Filenames",)
|
||
OUTPUT_NODE = True
|
||
FUNCTION = "combine_video"
|
||
CATEGORY = "XB_ToolBox/Video"
|
||
|
||
def combine_video(self, frame_rate, loop_count, images=None, latents=None,
|
||
filename_prefix="视频", format="image/gif", pingpong=False,
|
||
save_output=True, prompt=None, extra_pnginfo=None,
|
||
audio=None, unique_id=None, manual_format_widgets=None, **kwargs):
|
||
if latents is not None:
|
||
images = latents
|
||
if images is None:
|
||
return ((save_output, []),)
|
||
if isinstance(images, torch.Tensor) and images.size(0) == 0:
|
||
return ((save_output, []),)
|
||
num_frames = len(images)
|
||
pbar = ProgressBar(num_frames)
|
||
first_image = images[0]
|
||
while len(first_image.shape) > 3:
|
||
first_image = first_image[0]
|
||
images = iter(images)
|
||
|
||
output_dir = folder_paths.get_output_directory() if save_output else folder_paths.get_temp_directory()
|
||
# 展开 %date:yyyy-MM-dd% 之类的标记，否则会生成非法目录名（Windows: WinError 267）
|
||
filename_prefix = resolve_filename_prefix(filename_prefix)
|
||
full_output_folder, filename, _, subfolder, _ = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
||
output_files = []
|
||
|
||
metadata = _PngInfo()
|
||
video_metadata = {}
|
||
if prompt is not None:
|
||
metadata.add_text("prompt", _json.dumps(prompt))
|
||
video_metadata["prompt"] = prompt
|
||
if extra_pnginfo is not None:
|
||
for x in extra_pnginfo:
|
||
metadata.add_text(x, _json.dumps(extra_pnginfo[x]))
|
||
video_metadata[x] = extra_pnginfo[x]
|
||
extra_options = extra_pnginfo.get('workflow', {}).get('extra', {})
|
||
else:
|
||
extra_options = {}
|
||
metadata.add_text("CreationTime", _dt.datetime.now().isoformat(" ")[:19])
|
||
|
||
max_counter = 0
|
||
matcher = re.compile(f"{re.escape(filename)}_(\\d+)\\D*\\..+", re.IGNORECASE)
|
||
for existing_file in os.listdir(full_output_folder):
|
||
match = matcher.fullmatch(existing_file)
|
||
if match:
|
||
max_counter = max(max_counter, int(match.group(1)))
|
||
counter = max_counter + 1
|
||
|
||
first_image_file = f"{filename}_{counter:05}.png"
|
||
file_path = os.path.join(full_output_folder, first_image_file)
|
||
if extra_options.get('VHS_MetadataImage', True) != False:
|
||
_PILImage.fromarray(_tensor_to_bytes(first_image)).save(file_path, pnginfo=metadata, compress_level=4)
|
||
output_files.append(file_path)
|
||
|
||
format_type, format_ext = format.split("/")
|
||
if format_type == "image":
|
||
image_kwargs = {}
|
||
if format_ext == "gif":
|
||
image_kwargs['disposal'] = 2
|
||
if format_ext == "webp":
|
||
exif = _PILImage.Exif()
|
||
exif[_ExifTags.IFD.Exif] = {36867: _dt.datetime.now().isoformat(" ")[:19]}
|
||
image_kwargs['exif'] = exif
|
||
image_kwargs['lossless'] = kwargs.get("lossless", True)
|
||
file = f"{filename}_{counter:05}.{format_ext}"
|
||
file_path = os.path.join(full_output_folder, file)
|
||
if pingpong:
|
||
images = _to_pingpong(images)
|
||
def frames_gen(images):
|
||
for i in images:
|
||
pbar.update(1)
|
||
yield _PILImage.fromarray(_tensor_to_bytes(i))
|
||
frames = frames_gen(images)
|
||
next(frames).save(file_path, format=format_ext.upper(), save_all=True,
|
||
append_images=frames, duration=round(1000/frame_rate),
|
||
loop=loop_count, compress_level=4, **image_kwargs)
|
||
output_files.append(file_path)
|
||
else:
|
||
if _ffmpeg is None:
|
||
raise RuntimeError("ffmpeg 未找到")
|
||
|
||
if manual_format_widgets is not None:
|
||
kwargs.update(manual_format_widgets)
|
||
|
||
has_alpha = first_image.shape[-1] == 4
|
||
kwargs["has_alpha"] = has_alpha
|
||
video_format = _apply_format_widgets(format_ext, kwargs)
|
||
dim_alignment = video_format.get("dim_alignment", 2)
|
||
if (first_image.shape[1] % dim_alignment) or (first_image.shape[0] % dim_alignment):
|
||
to_pad = (-first_image.shape[1] % dim_alignment, -first_image.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):
|
||
image = image.permute((2,0,1))
|
||
padded = padfunc(image.to(dtype=torch.float32))
|
||
return padded.permute((1,2,0))
|
||
images = map(pad, images)
|
||
dimensions = (-first_image.shape[1] % dim_alignment + first_image.shape[1],
|
||
-first_image.shape[0] % dim_alignment + first_image.shape[0])
|
||
else:
|
||
dimensions = (first_image.shape[1], first_image.shape[0])
|
||
if pingpong:
|
||
images = _to_pingpong(images)
|
||
if num_frames > 2:
|
||
num_frames += num_frames -2
|
||
pbar.total = num_frames
|
||
if loop_count > 0:
|
||
loop_args = ["-vf", "loop=loop=" + str(loop_count)+":size=" + str(num_frames)]
|
||
else:
|
||
loop_args = []
|
||
if video_format.get('input_color_depth', '8bit') == '16bit':
|
||
images = map(_tensor_to_shorts, images)
|
||
i_pix_fmt = 'rgba64' if has_alpha else 'rgb48'
|
||
else:
|
||
images = map(_tensor_to_bytes, images)
|
||
i_pix_fmt = 'rgba' if has_alpha else 'rgb24'
|
||
file = f"{filename}_{counter:05}.{video_format['extension']}"
|
||
file_path = os.path.join(full_output_folder, file)
|
||
bitrate_arg = []
|
||
bitrate = video_format.get('bitrate')
|
||
if bitrate is not None:
|
||
bitrate_arg = ["-b:v", str(bitrate) + "M" if video_format.get('megabit') == 'True' else str(bitrate) + "K"]
|
||
args = [_ffmpeg, "-v", "error", "-f", "rawvideo", "-pix_fmt", i_pix_fmt,
|
||
"-color_range", "pc", "-colorspace", "rgb", "-color_primaries", "bt709",
|
||
"-color_trc", video_format.get("fake_trc", "iec61966-2-1"),
|
||
"-s", f"{dimensions[0]}x{dimensions[1]}", "-r", str(frame_rate), "-i", "-"] + loop_args
|
||
|
||
images = map(lambda x: x.tobytes(), images)
|
||
env = os.environ.copy()
|
||
if "environment" in video_format:
|
||
env.update(video_format["environment"])
|
||
|
||
if "pre_pass" in video_format:
|
||
images = [b''.join(images)]
|
||
os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)
|
||
in_args_len = args.index("-i") + 2
|
||
pre_pass_args = args[:in_args_len] + video_format['pre_pass']
|
||
_merge_filter_args(pre_pass_args)
|
||
try:
|
||
subprocess.run(pre_pass_args, input=images[0], env=env, capture_output=True, check=True)
|
||
except subprocess.CalledProcessError as e:
|
||
raise Exception("An error occurred in the ffmpeg prepass:\n" + e.stderr.decode(*ENCODE_ARGS))
|
||
if "inputs_main_pass" in video_format:
|
||
in_args_len = args.index("-i") + 2
|
||
args = args[:in_args_len] + video_format['inputs_main_pass'] + args[in_args_len:]
|
||
|
||
if 'gifski_pass' in video_format:
|
||
format = 'image/gif'
|
||
output_process = _gifski_process(args, dimensions, frame_rate, video_format, file_path, env)
|
||
audio = None
|
||
else:
|
||
args += video_format['main_pass'] + bitrate_arg
|
||
_merge_filter_args(args)
|
||
output_process = _ffmpeg_process(args, video_format, video_metadata, file_path, env)
|
||
output_process.send(None)
|
||
|
||
for image in images:
|
||
pbar.update(1)
|
||
output_process.send(image)
|
||
try:
|
||
total_frames_output = output_process.send(None)
|
||
output_process.send(None)
|
||
except StopIteration:
|
||
pass
|
||
output_files.append(file_path)
|
||
|
||
a_waveform = None
|
||
if audio is not None:
|
||
try:
|
||
a_waveform = audio['waveform']
|
||
except:
|
||
pass
|
||
if a_waveform is not None:
|
||
output_file_with_audio = f"{filename}_{counter:05}-audio.{video_format['extension']}"
|
||
output_file_with_audio_path = os.path.join(full_output_folder, output_file_with_audio)
|
||
if "audio_pass" not in video_format:
|
||
video_format["audio_pass"] = ["-c:a", "libopus"]
|
||
channels = audio['waveform'].size(1)
|
||
min_audio_dur = total_frames_output / frame_rate + 1
|
||
if video_format.get('trim_to_audio', 'False') != 'False':
|
||
apad = []
|
||
else:
|
||
apad = ["-af", "apad=whole_dur="+str(min_audio_dur)]
|
||
mux_args = [_ffmpeg, "-v", "error", "-n", "-i", file_path,
|
||
"-ar", str(audio['sample_rate']), "-ac", str(channels),
|
||
"-f", "f32le", "-i", "-", "-c:v", "copy"] \
|
||
+ video_format["audio_pass"] + apad + ["-shortest", output_file_with_audio_path]
|
||
audio_data = audio['waveform'].squeeze(0).transpose(0,1).numpy().tobytes()
|
||
_merge_filter_args(mux_args, '-af')
|
||
try:
|
||
res = subprocess.run(mux_args, input=audio_data, env=env, capture_output=True, check=True)
|
||
except subprocess.CalledProcessError as e:
|
||
raise Exception("An error occured in the ffmpeg subprocess:\n" + e.stderr.decode(*ENCODE_ARGS))
|
||
output_files.append(output_file_with_audio_path)
|
||
file = output_file_with_audio
|
||
if extra_options.get('VHS_KeepIntermediate', True) == False:
|
||
for intermediate in output_files[1:-1]:
|
||
if os.path.exists(intermediate):
|
||
os.remove(intermediate)
|
||
preview = {
|
||
"filename": file,
|
||
"subfolder": subfolder,
|
||
"type": "output" if save_output else "temp",
|
||
"format": format,
|
||
"frame_rate": frame_rate,
|
||
"workflow": first_image_file,
|
||
"fullpath": output_files[-1],
|
||
}
|
||
if num_frames == 1 and 'png' in format and '%03d' in file:
|
||
preview['format'] = 'image/png'
|
||
preview['filename'] = file.replace('%03d', '001')
|
||
return {"ui": {"gifs": [preview]}, "result": ((save_output, output_files),)} |