Skip to content

Commit

Permalink
Merge pull request #7 from Sxela/v0.3.0
Browse files Browse the repository at this point in the history 
V0.3.0
  • Loading branch information
@Sxela
Sxela authored Oct 6, 2024
2 parents aeae7e0 + 60fd528 commit f8c7371
Show file tree
Hide file tree
Showing 7 changed files with 8,836 additions and 17 deletions.
3 changes: 3 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,6 @@ Save frame to a folder, using current frame number.

## MixConsistencyMaps
Mix consistency maps, blur, dilate.

## RenderVideo
Trigger output video render at a given frame
45 changes: 41 additions & 4 deletions frame_nodes.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import torch
import numpy as np
import folder_paths
from .frame_utils import FrameDataset, StylizedFrameDataset, get_size
from .frame_utils import FrameDataset, StylizedFrameDataset, get_size, save_video

class LoadFrameSequence:
@classmethod
Expand Down Expand Up @@ -264,14 +264,49 @@ def INPUT_TYPES(self):

def save_img(self, image, output_dir, batch_name, frame_number):
os.makedirs(output_dir, exist_ok=True)
fname = f'{batch_name}_{frame_number}.png'
fname = f'{batch_name}_{frame_number:06}.png'
out_fname = os.path.join(output_dir, fname)
print('image.shape', image.shape, image.max(), image.min())
image = (image[0].clip(0,1)*255.).cpu().numpy().astype('uint8')
image = Image.fromarray(image)
image.save(out_fname)
print('fname', out_fname)
return ()

class RenderVideo:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"output_dir": ("STRING", {"multiline": True,
"default":''}),
"frames_input_dir": ("STRING", {"multiline": True,
"default":''}),
"batch_name": ("STRING", {"default":'ComfyWarp'}),
"first_frame":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"last_frame":("INT",{"default": -1, "min": -1, "max": 9999999999}),
"render_at_frame":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"current_frame":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"fps":("FLOAT",{"default": 24, "min": 0, "max": 9999999999}),
"output_format":(["h264_mp4", "qtrle_mov", "prores_mov"],),
"use_deflicker": ("BOOLEAN", {"default": False})
}
}

CATEGORY = "WarpFusion"

RETURN_TYPES = ()
FUNCTION = "export_video"
OUTPUT_NODE = True

def export_video(self, output_dir, frames_input_dir, batch_name, first_frame=1, last_frame=-1,
render_at_frame=999999, current_frame=0, fps=30, output_format='h264_mp4', use_deflicker=False):
if current_frame>=render_at_frame:
print('Exporting video.')
save_video(indir=frames_input_dir, video_out=output_dir, batch_name=batch_name, start_frame=first_frame,
last_frame=last_frame, fps=fps, output_format=output_format, use_deflicker=use_deflicker)
return ()


NODE_CLASS_MAPPINGS = {
"LoadFrameSequence": LoadFrameSequence,
Expand All @@ -281,7 +316,8 @@ def save_img(self, image, output_dir, batch_name, frame_number):
"LoadFramePairFromDataset":LoadFramePairFromDataset,
"LoadFrameFromFolder":LoadFrameFromFolder,
"ResizeToFit":ResizeToFit,
"SaveFrame":SaveFrame
"SaveFrame":SaveFrame,
"RenderVideo": RenderVideo
}

NODE_DISPLAY_NAME_MAPPINGS = {
Expand All @@ -292,5 +328,6 @@ def save_img(self, image, output_dir, batch_name, frame_number):
"LoadFramePairFromDataset":"Load Frame Pair From Dataset",
"LoadFrameFromFolder": "Maybe Load Frame From Folder",
"ResizeToFit":"Resize To Fit",
"SaveFrame":"SaveFrame"
"SaveFrame":"SaveFrame",
"RenderVideo": "RenderVideo"
}
119 changes: 106 additions & 13 deletions frame_utils.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,13 @@
#(c) Alex Spirin 2023

import hashlib, os, sys, glob, subprocess, pathlib

import hashlib, os, sys, glob, subprocess, pathlib, platform
import zipfile
from shutil import copy
import requests

def generate_file_hash(input_file):
"""Generates has for video vile based on name, size, creation time
"""
# Get file name and metadata
file_name = os.path.basename(input_file)
file_size = os.path.getsize(input_file)
Expand Down Expand Up @@ -50,7 +54,7 @@ def __init__(self, source_path, outdir_prefix='', videoframes_root='', update_on
if len(glob.glob(source_path))>0:
self.frame_paths = sorted(glob.glob(source_path))
else:
raise Exception(f'Frame source for {outdir_prefix} not found at {source_path}\nPlease specify an existing source path.')
raise FileNotFoundError(f'Frame source for {outdir_prefix} not found at {source_path}\nPlease specify an existing source path.')
if os.path.exists(source_path):
if os.path.isfile(source_path):
if os.path.splitext(source_path)[1][1:].lower() in image_extenstions:
Expand All @@ -63,11 +67,11 @@ def __init__(self, source_path, outdir_prefix='', videoframes_root='', update_on
self.frame_paths = glob.glob(os.path.join(out_path, '*.*'))
self.source_path = out_path
if len(self.frame_paths)<1:
raise Exception(f'Couldn`t extract frames from {source_path}\nPlease specify an existing source path.')
raise FileNotFoundError(f'Couldn`t extract frames from {source_path}\nPlease specify an existing source path.')
elif os.path.isdir(source_path):
self.frame_paths = glob.glob(os.path.join(source_path, '*.*'))
if len(self.frame_paths)<1:
raise Exception(f'Found 0 frames in {source_path}\nPlease specify an existing source path.')
raise FileNotFoundError(f'Found 0 frames in {source_path}\nPlease specify an existing source path.')
extensions = []
if self.frame_paths is not None:
for f in self.frame_paths:
Expand All @@ -81,7 +85,7 @@ def __init__(self, source_path, outdir_prefix='', videoframes_root='', update_on

self.frame_paths = sorted(self.frame_paths)

else: raise Exception(f'Frame source for {outdir_prefix} not found at {source_path}\nPlease specify an existing source path.')
else: raise FileNotFoundError(f'Frame source for {outdir_prefix} not found at {source_path}\nPlease specify an existing source path.')
print(f'Found {len(self.frame_paths)} frames at {source_path}')

def __getitem__(self, idx):
Expand All @@ -100,10 +104,10 @@ def __init__(self, source_path):
self.frame_paths = None
self.source_path = source_path
if not os.path.exists(source_path):
raise Exception(f'Frame source not found at {source_path}\nPlease specify an existing source path.')
raise FileNotFoundError(f'Frame source not found at {source_path}\nPlease specify an existing source path.')
if os.path.exists(source_path):
if os.path.isfile(source_path):
raise Exception(f'{source_path} is a file. Please specify path to a folder.')
raise NotADirectoryError(f'{source_path} is a file. Please specify path to a folder.')
elif os.path.isdir(source_path):
self.frame_paths = glob.glob(os.path.join(source_path, '*.*'))

Expand All @@ -123,8 +127,97 @@ def get_size(size, max_size, divisible_by=8):
divisible_by = int(divisible_by)
x,y = size
max_dim = max(size)
if max_dim>max_size:
ratio = max_size/max_dim
new_size = ((int(x*ratio)//divisible_by*divisible_by),(int(y*ratio)//divisible_by*divisible_by))
return new_size
return size
# if max_dim>max_size:
ratio = max_size/max_dim
new_size = ((int(x*ratio)//divisible_by*divisible_by),(int(y*ratio)//divisible_by*divisible_by))
return new_size
# return size

def find_ffmpeg(start_dir):
files = glob.glob(f"{start_dir}/**/*.*", recursive=True)
for f in files:
if platform.system() == 'Linux':
if f.endswith('ffmpeg'): return f
elif f.endswith('ffmpeg.exe'): return f
return None

def download_ffmpeg(start_dir):
ffmpeg_url = 'https://github.com/GyanD/codexffmpeg/releases/download/6.0/ffmpeg-6.0-full_build.zip'
print('ffmpeg.exe not found, downloading...')
r = requests.get(ffmpeg_url, allow_redirects=True, timeout=60)
print('downloaded, extracting')
open('ffmpeg-6.0-full_build.zip', 'wb').write(r.content)

with zipfile.ZipFile('ffmpeg-6.0-full_build.zip', 'r') as zip_ref:
zip_ref.extractall(f'{start_dir}/')

copy(f'{start_dir}/ffmpeg-6.0-full_build/bin/ffmpeg.exe', f'{start_dir}/')
return f'{start_dir}/ffmpeg.exe'


def get_ffmpeg():
start_dir = os.getcwd()
ffmpeg_path = find_ffmpeg(start_dir)
if ffmpeg_path is None or not os.path.exists(ffmpeg_path):
ffmpeg_path = download_ffmpeg(start_dir)
return ffmpeg_path

def save_video(indir, video_out, batch_name='', start_frame=1, last_frame=-1, fps=30, output_format='h264_mp4', use_deflicker=False):
ffmpeg_path = get_ffmpeg()
print('Found ffmpeg at: ', ffmpeg_path)
os.makedirs(video_out, exist_ok=True)
indir = indir.replace('\\','/')
image_path = f"{indir}/{batch_name}_%06d.png"

postfix = ''
if use_deflicker:
postfix+='_dfl'

indir_stem = indir.replace('\\','/').split('/')[-1]
out_filepath = f"{video_out}/{indir_stem}_{postfix}.{output_format.split('_')[-1]}"
if last_frame == -1:
last_frame = len(glob.glob(f"{indir}/*.png"))

cmd = [ffmpeg_path,
'-y',
'-vcodec',
'png',
'-r',
str(fps),
'-start_number',
str(start_frame),
'-i',
image_path,
'-frames:v',
str(last_frame+1),
'-c:v']

if output_format == 'h264_mp4':
cmd+=['libx264',
'-pix_fmt',
'yuv420p']
elif output_format == 'qtrle_mov':
cmd+=['qtrle',
'-vf',
f'fps={fps}']
elif output_format == 'prores_mov':
cmd+=['prores_aw',
'-profile:v',
'2',
'-pix_fmt',
'yuv422p10',
'-vf',
f'fps={fps}']

if use_deflicker:
cmd+=['-vf','deflicker=mode=pm:size=10']
cmd+=[out_filepath]

process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate()
if process.returncode != 0:
print(stderr)
raise RuntimeError(stderr)
else:
print(f"The video is ready and saved to {out_filepath}")
return 0
Loading

0 comments on commit f8c7371

Please sign in to comment.