138 lines
5.0 KiB
Python
138 lines
5.0 KiB
Python
import subprocess
|
|
import shutil
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
"""
|
|
Attribution: ComfyUI-VideoHelperSuite
|
|
|
|
Portions of this code are adapted from GitHub repository `https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite`,
|
|
which is licensed under the GNU General Public License version 3 (GPL-3.0):
|
|
|
|
"""
|
|
|
|
def __lnl_ffmpeg_suitability(path):
|
|
try:
|
|
version = subprocess.run([path, "-version"], check=True,
|
|
capture_output=True).stdout.decode("utf-8")
|
|
except:
|
|
return 0
|
|
score = 0
|
|
#rough layout of the importance of various features
|
|
simple_criterion = [("libvpx", 20),("264",10), ("265",3),
|
|
("svtav1",5),("libopus", 1)]
|
|
for criterion in simple_criterion:
|
|
if version.find(criterion[0]) >= 0:
|
|
score += criterion[1]
|
|
#obtain rough compile year from copyright information
|
|
copyright_index = version.find('2000-2')
|
|
if copyright_index >= 0:
|
|
copyright_year = version[copyright_index+6:copyright_index+9]
|
|
if copyright_year.isnumeric():
|
|
score += int(copyright_year)
|
|
return score
|
|
|
|
def lnl_get_audio(file, start_time=0, duration=0):
|
|
args = [ffmpeg_path, "-v", "error", "-i", file]
|
|
if start_time > 0:
|
|
args += ["-ss", str(start_time)]
|
|
if duration > 0:
|
|
args += ["-t", str(duration)]
|
|
return subprocess.run(args + ["-f", "wav", "-"],
|
|
stdout=subprocess.PIPE, check=True).stdout
|
|
|
|
def lnl_lazy_eval(func):
|
|
class Cache:
|
|
def __init__(self, func):
|
|
self.res = None
|
|
self.func = func
|
|
def get(self):
|
|
if self.res is None:
|
|
self.res = self.func()
|
|
return self.res
|
|
cache = Cache(func)
|
|
return lambda : cache.get()
|
|
|
|
def lnl_cv_frame_generator(video, frame_load_cap, skip_first_frames, select_every_nth):
|
|
try:
|
|
video_cap = cv2.VideoCapture(video)
|
|
if not video_cap.isOpened():
|
|
raise ValueError(f"{video} could not be loaded with cv.")
|
|
# set video_cap to look at start_index frame
|
|
total_frame_count = 0
|
|
total_frames_evaluated = -1
|
|
frames_added = 0
|
|
base_frame_time = 1/video_cap.get(cv2.CAP_PROP_FPS)
|
|
width = video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)
|
|
height = video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
|
|
prev_frame = None
|
|
target_frame_time = base_frame_time
|
|
yield (width, height, target_frame_time)
|
|
time_offset=target_frame_time - base_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:
|
|
frames_added += 1
|
|
if total_frame_count >= frame_load_cap + skip_first_frames:
|
|
break
|
|
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) / 255.0
|
|
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 cap exists and we've reached it, stop processing frames
|
|
if frame_load_cap > 0 and frames_added >= frame_load_cap:
|
|
break
|
|
if prev_frame is not None:
|
|
yield prev_frame
|
|
finally:
|
|
video_cap.release()
|
|
|
|
ffmpeg_paths = []
|
|
try:
|
|
from imageio_ffmpeg import get_ffmpeg_exe
|
|
imageio_ffmpeg_path = get_ffmpeg_exe()
|
|
ffmpeg_paths.append(imageio_ffmpeg_path)
|
|
except:
|
|
print("Failed to import imageio_ffmpeg")
|
|
system_ffmpeg = shutil.which("ffmpeg")
|
|
if system_ffmpeg is not None:
|
|
ffmpeg_paths.append(system_ffmpeg)
|
|
|
|
if len(ffmpeg_paths) == 0:
|
|
print("No valid ffmpeg found.")
|
|
ffmpeg_path = None
|
|
elif len(ffmpeg_paths) == 1:
|
|
ffmpeg_path = ffmpeg_paths[0]
|
|
else:
|
|
ffmpeg_path = max(ffmpeg_paths, key=__lnl_ffmpeg_suitability)
|