Add nodes for flexible typing and conversion to Any
This commit adds several new nodes to enable flexible typing and conversion between types: - ValueToAny - Converts primitive values like ints, floats, strings, bools to Any type - IterableToAny - Converts iterable values like lists, dicts to Any - CallableToAny - Converts callables like functions to Any - ObjectToAny - Converts tensor objects like images, latents etc to Any The motivation is to allow passing different types flexibly into functions using Any typing. This enables writing functions in a more generic way by accepting Any instead of specific types. The nodes convert the input values to Any and return debug info on what was converted. This allows tracing what types entered a function even if the function uses generic Any types. No breaking changes, new nodes added: - ValueToAny - IterableToAny - CallableToAny - ObjectToAny These provide building blocks for more flexible typing, opening up possibilities like: - Writing utility and processing functions accepting Any instead of specific types - Calling functions dynamically with values of different types - Tracing and logging what types passed through a function Next steps could include AnyToValue nodes, more specialized Any converters, and exploring uses for generic functions.
This commit is contained in:
Tasha Upchurch
@@ -121,4 +121,56 @@ class JoinAudioFiles:
|
||||
with open(output_path, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
return (output_path, file_bytes, metadata,)
|
||||
return (output_path, file_bytes, metadata,)
|
||||
|
||||
@ETK_audio_base
|
||||
class ExtractAudioFromVideo:
|
||||
""" Extracts audio from a video file and saves it as an audio file with a specified format and bitrate """
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"video_path": ("STRING", {"multiline": False, "default": ""}),
|
||||
"format":(["wav", "mp3"], {"default": "wav"}),
|
||||
"output_path": ("STRING", {"multiline": False, "default": ""}),
|
||||
},
|
||||
"optional": {
|
||||
"bitrate (MP3)": ("INT", {"default": 192}) # Default bitrate value (change as needed)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "BYTES")
|
||||
RETURN_NAMES = ("output_wave_file_path", "audio_bytes")
|
||||
FUNCTION = 'extract_audio_from_video'
|
||||
CATEGORY = 'ETK/audio'
|
||||
|
||||
def extract_audio_from_video(self, **kwargs):
|
||||
import moviepy.editor as mp
|
||||
import io
|
||||
|
||||
kwargs = copy.deepcopy(kwargs)
|
||||
video_path = kwargs.get('video_path')
|
||||
output_path = kwargs.get('output_path',"tmp.wav")
|
||||
audio_format = kwargs.get('format', 'wav')
|
||||
bitrate = kwargs.get('bitrate', None) # Default bitrate is None
|
||||
|
||||
# Load the video clip
|
||||
video_clip = mp.VideoFileClip(video_path)
|
||||
|
||||
# Determine the appropriate codec based on the audio format
|
||||
if audio_format == 'wav':
|
||||
codec = "pcm_s16le" # For WAV format
|
||||
elif audio_format == 'mp3':
|
||||
codec = "mp3" # For MP3 format
|
||||
|
||||
# Extract audio and save it with the specified format and bitrate
|
||||
audio_clip = video_clip.audio
|
||||
audio_clip.write_audiofile(output_path, codec=codec, bitrate=bitrate)
|
||||
|
||||
# Read the generated audio file as bytes
|
||||
audio_bytes = None
|
||||
with open(output_path, 'rb') as audio_file:
|
||||
audio_bytes = audio_file.read()
|
||||
|
||||
return (output_path, audio_bytes,)
|
||||
|
||||
@@ -97,6 +97,7 @@ class StrBreakout:
|
||||
pass
|
||||
return ("\n".join(notes), out_bool, out_int, out_float, out_list, out_lines, out_comma,)
|
||||
|
||||
|
||||
class IntBreakout:
|
||||
""" node that breaks out an int into a:
|
||||
- Boolean
|
||||
@@ -172,8 +173,136 @@ class IntBreakout:
|
||||
return ("\n".join(notes), out_bool, out_str, out_float, out_list, out_range, out_none,)
|
||||
|
||||
|
||||
class ValueToAny:
|
||||
"""takes INT,STRING,FLOAT,BOOL, and returns the value coded as *"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"optional": {
|
||||
"INT": ("INT", {"default": None, }),
|
||||
"FLOAT": ("FLOAT", {"default": None, }),
|
||||
"STRING": ("STRING", {"default": None, }),
|
||||
"BOOL": ("BOOL", {"default": None, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*", "STRING",)
|
||||
RETURN_NAMES = ("*", "debug info",)
|
||||
FUNCTION = "get_value"
|
||||
CATEGORY = "conversion"
|
||||
|
||||
def get_value(self, INT=None, FLOAT=None, STRING=None, BOOL=None):
|
||||
if INT:
|
||||
dbg = f"found INT: {INT}"
|
||||
return (INT, dbg,)
|
||||
if FLOAT:
|
||||
dbg = f"found FLOAT: {FLOAT}"
|
||||
return (FLOAT, dbg,)
|
||||
if STRING:
|
||||
dbg = f"found STRING: {STRING}"
|
||||
return (STRING, dbg,)
|
||||
if BOOL:
|
||||
dbg = f"found BOOL: {BOOL}"
|
||||
return (BOOL, dbg,)
|
||||
dbg = f"found nothing, will return the first value that is not None; from INT,FLOAT,STRING,BOOL"
|
||||
return (None, dbg)
|
||||
|
||||
|
||||
class IterableToAny:
|
||||
"""takes LIST,DICT, and returns the value coded as *"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"optional": {
|
||||
"LIST": ("LIST", {"default": None, }),
|
||||
"DICT": ("DICT", {"default": None, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*", "STRING",)
|
||||
RETURN_NAMES = ("*", "debug info",)
|
||||
FUNCTION = "get_value"
|
||||
CATEGORY = "conversion"
|
||||
|
||||
def get_value(self, LIST=None, DICT=None):
|
||||
if LIST:
|
||||
dbg = f"found LIST: {LIST}"
|
||||
return (LIST, dbg,)
|
||||
if DICT:
|
||||
dbg = f"found DICT: {DICT}"
|
||||
return (DICT, dbg,)
|
||||
dbg = f"found nothing, will return the first value that is not None; from LIST,DICT"
|
||||
return (None, dbg)
|
||||
|
||||
|
||||
class CallableToAny:
|
||||
"""takes FUNC, (others soon) and returns the value coded as *"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"optional": {
|
||||
"FUNC": ("FUNC", {"default": None, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*", "STRING",)
|
||||
RETURN_NAMES = ("*", "debug info",)
|
||||
FUNCTION = "get_value"
|
||||
CATEGORY = "conversion"
|
||||
|
||||
def get_value(self, FUNC=None):
|
||||
if FUNC:
|
||||
dbg = f"found FUNC: {FUNC}"
|
||||
return (FUNC, dbg,)
|
||||
dbg = f"found nothing, will return the first value that is not None; from FUNC"
|
||||
return (None, dbg)
|
||||
|
||||
|
||||
class ObjectToAny:
|
||||
""" takes non-primitive objects like TENSORS and returns the value coded as *"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"optional": {
|
||||
"LATENT": ("LATENT", {"default": None, }),
|
||||
"CONDITIONING": ("CONDITIONING", {"default": None, }),
|
||||
"IMAGE": ("IMAGE", {"default": None, }),
|
||||
"MASK": ("MASK", {"default": None, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*", "STRING",)
|
||||
RETURN_NAMES = ("*", "debug info",)
|
||||
FUNCTION = "get_value"
|
||||
CATEGORY = "conversion"
|
||||
|
||||
def get_value(self, LATENT=None, CONDITIONING=None, IMAGE=None, MASK=None):
|
||||
if LATENT:
|
||||
dbg = f"found LATENT: {LATENT}"
|
||||
return (LATENT, dbg,)
|
||||
if CONDITIONING:
|
||||
dbg = f"found CONDITIONING: {CONDITIONING}"
|
||||
return (CONDITIONING, dbg,)
|
||||
if IMAGE:
|
||||
dbg = f"found IMAGE: {IMAGE}"
|
||||
return (IMAGE, dbg,)
|
||||
if MASK:
|
||||
dbg = f"found MASK: {MASK}"
|
||||
return (MASK, dbg,)
|
||||
dbg = f"found nothing, will return the first value that is not None; from LATENT,CONDITIONING,IMAGE,MASK"
|
||||
return (None, dbg)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_CLASS_MAPPINGS["MultilineStringNode"] = MultilineStringNode
|
||||
NODE_CLASS_MAPPINGS["StrBreakout"] = StrBreakout
|
||||
NODE_CLASS_MAPPINGS["IntBreakout"] = IntBreakout
|
||||
|
||||
NODE_CLASS_MAPPINGS["ValueToAny"] = ValueToAny
|
||||
NODE_CLASS_MAPPINGS["IterableToAny"] = IterableToAny
|
||||
NODE_CLASS_MAPPINGS["CallableToAny"] = CallableToAny
|
||||
NODE_CLASS_MAPPINGS["ObjectToAny"] = ObjectToAny
|
||||
|
||||
+46
-1
@@ -5,9 +5,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
_list = {"list": ("LIST",)}
|
||||
_func = {"func": ("FUNC", {"default": None})}
|
||||
_func_extra = {"func_extra": ("FUNC",)}
|
||||
_text = {"code": ("STRING", {"multiline": True, "default": "y=x()"})}
|
||||
_text = {"code": ("STRING", {"multiline": True, "default": "y mapped to output"})}
|
||||
_image = {"image": ("IMAGE",)}
|
||||
_string = {"string": ("STRING", {"default": None})}
|
||||
_any = {"any": ("*", {"default": None})}
|
||||
|
||||
optional = lambda x: {"optional": {k: v for k, v in x.items()}}
|
||||
required = lambda x: {"required": {k: v for k, v in x.items()}}
|
||||
@@ -440,6 +441,50 @@ class GetFirstCodeBlock():
|
||||
return (text,)
|
||||
|
||||
|
||||
@ETK_functional_base
|
||||
class FuncAnysToImage:
|
||||
"""runs eval on the given text, with 5 anys as input and outputs an image"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
req = optional(
|
||||
both(
|
||||
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
|
||||
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
|
||||
),
|
||||
)
|
||||
req["required"] = {}
|
||||
req["required"].update(_text)
|
||||
|
||||
return req
|
||||
|
||||
CATEGORY = "ETK/func"
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "func"
|
||||
|
||||
def func(self, **kwargs):
|
||||
code = kwargs.get("code", None)
|
||||
any1 = kwargs.get("any1", None)
|
||||
any2 = kwargs.get("any2", None)
|
||||
any3 = kwargs.get("any3", None)
|
||||
any4 = kwargs.get("any4", None)
|
||||
|
||||
my_globals = globals()
|
||||
my_locals = locals()
|
||||
|
||||
my_locals["any1"] = any1
|
||||
my_locals["any2"] = any2
|
||||
my_locals["any3"] = any3
|
||||
my_locals["any4"] = any4
|
||||
|
||||
exec(code, my_globals, my_locals)
|
||||
|
||||
return (my_locals["y"],)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test = {}
|
||||
|
||||
|
||||
@@ -337,7 +337,7 @@ class TinyTxtToImg:
|
||||
if self.latent_image == None:
|
||||
self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0]
|
||||
|
||||
if "latent" in render_what or "all" in render_what:
|
||||
if "latent" in render_what or "all" in render_what or "image" in render_what:
|
||||
self.KSampler = KSampler()
|
||||
self.samples = self.KSampler.sample(
|
||||
model=self.mdl,
|
||||
@@ -384,6 +384,9 @@ class TinyTxtToImg:
|
||||
|
||||
self.FUNC = self.tinytxt2img
|
||||
self.ARGS = kwargs
|
||||
|
||||
self.samples["samples"] = self.samples["samples"].to("cpu", non_blocking=True)
|
||||
|
||||
ret = (image,
|
||||
self,
|
||||
self.samples,
|
||||
@@ -1653,33 +1656,159 @@ class TextRender:
|
||||
}),
|
||||
"color": ("STRING", {
|
||||
"default": "#000000"
|
||||
})
|
||||
}),
|
||||
"func_only": ([True, False], {"default": False}),
|
||||
"stroke fill": ("STRING", {"default": "#000000"}),
|
||||
"stroke width": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
CATEGORY = "ETK/text"
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
||||
RETURN_NAMES = ("image", "image rgba",)
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", "FUNC",)
|
||||
RETURN_NAMES = ("image", "image rgba", "FUNC(**kwargs)",)
|
||||
|
||||
FUNCTION = "render_text"
|
||||
|
||||
def render_text(self, text, x, y, width, height, font='Arial', size=16, color='#000000'):
|
||||
@profile
|
||||
def render_text(self, text, x, y, width, height, font='Arial', size=16, color='#888888', func_only=False, **kwargs):
|
||||
"""
|
||||
This function renders the provided text at specified location, with the given width and height.
|
||||
The text is rendered in the provided font, size, and color.
|
||||
|
||||
>>> tr = TextRender()
|
||||
>>> result = tr.render_text('Hello, world!', 128, 128, 512, 512,"Arial",16,"#FF11000")
|
||||
>>> torch_image_show(result[0][0])
|
||||
"""
|
||||
from PIL import ImageFont, ImageDraw
|
||||
import gc
|
||||
import psutil
|
||||
process = psutil.Process();
|
||||
memory_info = process.memory_info();
|
||||
start_memory = memory_info.rss / (1024 * 1024)
|
||||
|
||||
font_name = font
|
||||
sw = kwargs.get("stroke width", None)
|
||||
sf = kwargs.get("stroke fill", None)
|
||||
|
||||
@profile
|
||||
def _wrap_text(font, line, max_width, d):
|
||||
words = line.split(' ')
|
||||
new_line = ''
|
||||
lines = []
|
||||
for word in words:
|
||||
temp_line = new_line + word + ' '
|
||||
|
||||
|
||||
#w, _ = d.textsize(temp_line, font=font)
|
||||
|
||||
# w, h = d.textsize(line, font=font)
|
||||
(left, top, right, bottom) = d.textbbox((0, 0), text=temp_line, font=font)
|
||||
w = right - left
|
||||
h = bottom - top
|
||||
previous_height = h # Update the previous height
|
||||
|
||||
|
||||
|
||||
|
||||
if w > max_width:
|
||||
lines.append(new_line.strip())
|
||||
new_line = word + ' '
|
||||
else:
|
||||
new_line = temp_line
|
||||
|
||||
lines.append(new_line.strip())
|
||||
wrapped_line = '\n'.join(lines)
|
||||
return wrapped_line
|
||||
|
||||
@profile
|
||||
def _render_text(font_name, image, size, text, x, y, allow_wrap=True, allow_shrink=True, sw=None, sf=None):
|
||||
width = image.size[0]
|
||||
d = ImageDraw.Draw(image)
|
||||
font = ImageFont.truetype(font_name, size)
|
||||
lines = text.split('\n')
|
||||
|
||||
shrink_threshold = 24 # The minimum font size before we start wrapping text
|
||||
line_spacing = 5 # The constant space between lines
|
||||
y_offset = y
|
||||
previous_height = 0 # Store the height of the previous line
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
if line.strip() == '' or line.strip() == '\n':
|
||||
# For blank lines, add a blank line of the previous line's height
|
||||
y_offset += previous_height + line_spacing
|
||||
else:
|
||||
# w, h = d.textsize(line, font=font)
|
||||
(left, top, right, bottom) = d.textbbox((0, 0), text=line, font=font)
|
||||
w = right - left
|
||||
h = bottom - top
|
||||
previous_height = h # Update the previous height
|
||||
|
||||
if w > width:
|
||||
if allow_shrink:
|
||||
new_font_size = int(size * width / w)
|
||||
if new_font_size >= shrink_threshold:
|
||||
font = ImageFont.truetype(font_name, new_font_size)
|
||||
else:
|
||||
# The font size is too small, start wrapping text
|
||||
wrapped_line = _wrap_text(font, line, width, d)
|
||||
y_offset = _render_text(font_name,
|
||||
image,
|
||||
size,
|
||||
wrapped_line, x, y_offset + line_spacing,
|
||||
allow_wrap=False,
|
||||
allow_shrink=False)
|
||||
continue # Skip to the next line
|
||||
|
||||
if allow_wrap:
|
||||
# Wrap the text
|
||||
wrapped_line = _wrap_text(font, line, width, d)
|
||||
y_offset = _render_text(font_name,
|
||||
image,
|
||||
new_font_size,
|
||||
wrapped_line,
|
||||
x,
|
||||
y_offset + line_spacing,
|
||||
allow_wrap=False,
|
||||
allow_shrink=False)
|
||||
continue # Skip to the next line
|
||||
|
||||
# Draw the line
|
||||
x_offset = x + (width - w) / 2
|
||||
if sw or sw:
|
||||
d.text((x_offset, y_offset), line, font=font, fill="#FFFFFF", stroke_width=sw, stroke_fill=sf)
|
||||
else:
|
||||
d.text((x_offset, y_offset), line, font=font, fill="#FFFFFF")
|
||||
|
||||
y_offset += h + line_spacing # Move y_offset to the bottom of the last line drawn plus the line spacing
|
||||
# debug print the amount of free memory in torch
|
||||
|
||||
return y_offset
|
||||
|
||||
if func_only:
|
||||
def lll(new_kwargs):
|
||||
if "font" not in new_kwargs:
|
||||
new_kwargs["font"] = kwargs.get("font_name", font)
|
||||
if "color" not in new_kwargs:
|
||||
new_kwargs["color"] = color
|
||||
if "size" not in new_kwargs:
|
||||
new_kwargs["size"] = size
|
||||
if "text" not in new_kwargs:
|
||||
new_kwargs["text"] = text
|
||||
if "x" not in new_kwargs:
|
||||
new_kwargs["x"] = x
|
||||
if "y" not in new_kwargs:
|
||||
new_kwargs["y"] = y
|
||||
if "width" not in new_kwargs:
|
||||
new_kwargs["width"] = width
|
||||
if "height" not in new_kwargs:
|
||||
new_kwargs["height"] = height
|
||||
|
||||
return TextRender.render_text(None, **new_kwargs)
|
||||
|
||||
return (None, None, lll,)
|
||||
|
||||
font_name = kwargs.get("font_name", font)
|
||||
|
||||
# Create an empty image with RGBA channels
|
||||
image = Image.new('RGBA', (width, height))
|
||||
|
||||
self._render_text(font_name, image, size, text, x, y, allow_shrink=True)
|
||||
_render_text(font_name, image, size, text, x, y, allow_shrink=True, sf=sf, sw=sw)
|
||||
|
||||
# Convert the image to numpy array
|
||||
image_array = np.array(image)
|
||||
@@ -1700,87 +1829,19 @@ class TextRender:
|
||||
|
||||
# make a copy for the RGB image
|
||||
image_rgb = image_tensor[..., :3]
|
||||
print(torch.cuda.memory_stats())
|
||||
return (image_rgb, image_tensor,)
|
||||
# print(torch.cuda.memory_stats())
|
||||
del image
|
||||
del image_array
|
||||
del color_arr
|
||||
|
||||
from PIL import ImageFont, ImageDraw
|
||||
from PIL import ImageFont, ImageDraw
|
||||
process = psutil.Process();
|
||||
memory_info = process.memory_info();
|
||||
end_memory = memory_info.rss / (1024 * 1024)
|
||||
print("Memory used: {} MB".format(end_memory - start_memory))
|
||||
|
||||
def _wrap_text(self, font, line, max_width, d):
|
||||
words = line.split(' ')
|
||||
new_line = ''
|
||||
lines = []
|
||||
for word in words:
|
||||
temp_line = new_line + word + ' '
|
||||
w, _ = d.textsize(temp_line, font=font)
|
||||
if w > max_width:
|
||||
lines.append(new_line.strip())
|
||||
new_line = word + ' '
|
||||
else:
|
||||
new_line = temp_line
|
||||
|
||||
lines.append(new_line.strip())
|
||||
wrapped_line = '\n'.join(lines)
|
||||
return wrapped_line
|
||||
|
||||
def _render_text(self, font_name, image, size, text, x, y, allow_wrap=True, allow_shrink=True):
|
||||
width = image.size[0]
|
||||
d = ImageDraw.Draw(image)
|
||||
font = ImageFont.truetype(font_name, size)
|
||||
lines = text.split('\n')
|
||||
|
||||
shrink_threshold = 24 # The minimum font size before we start wrapping text
|
||||
line_spacing = 5 # The constant space between lines
|
||||
y_offset = y
|
||||
previous_height = 0 # Store the height of the previous line
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
if line.strip() == '' or line.strip() == '\n':
|
||||
# For blank lines, add a blank line of the previous line's height
|
||||
y_offset += previous_height + line_spacing
|
||||
else:
|
||||
#w, h = d.textsize(line, font=font)
|
||||
(left, top, right, bottom) = d.textbbox((0, 0), text=line, font=font)
|
||||
w = right - left
|
||||
h = bottom - top
|
||||
previous_height = h # Update the previous height
|
||||
|
||||
if w > width:
|
||||
if allow_shrink:
|
||||
new_font_size = int(size * width / w)
|
||||
if new_font_size >= shrink_threshold:
|
||||
font = ImageFont.truetype(font_name, new_font_size)
|
||||
else:
|
||||
# The font size is too small, start wrapping text
|
||||
wrapped_line = self._wrap_text(font, line, width, d)
|
||||
y_offset = self._render_text(font_name,
|
||||
image,
|
||||
size,
|
||||
wrapped_line, x, y_offset + line_spacing,
|
||||
allow_wrap=False,
|
||||
allow_shrink=False)
|
||||
continue # Skip to the next line
|
||||
|
||||
if allow_wrap:
|
||||
# Wrap the text
|
||||
wrapped_line = self._wrap_text(font, line, width, d)
|
||||
y_offset = self._render_text(font_name,
|
||||
image,
|
||||
new_font_size,
|
||||
wrapped_line,
|
||||
x,
|
||||
y_offset + line_spacing,
|
||||
allow_wrap=False,
|
||||
allow_shrink=False)
|
||||
continue # Skip to the next line
|
||||
|
||||
# Draw the line
|
||||
x_offset = x + (width - w) / 2
|
||||
d.text((x_offset, y_offset), line, font=font, fill="#FFFFFF")
|
||||
y_offset += h + line_spacing # Move y_offset to the bottom of the last line drawn plus the line spacing
|
||||
# debug print the amount of free memory in torch
|
||||
|
||||
return y_offset
|
||||
image_rgb = image_rgb.to("cpu")
|
||||
image_tensor = image_tensor.to("cpu")
|
||||
return (image_rgb, image_tensor, None,)
|
||||
|
||||
|
||||
@ETK_image_base
|
||||
|
||||
+1
-1
@@ -72,7 +72,7 @@ app.registerExtension({
|
||||
// Set widget value to filename
|
||||
// filter to the widget that has the name "vide_in"
|
||||
const widget = node.widgets.find(w => w.name === "video_in");
|
||||
widget.value = filename.filename;
|
||||
widget.value = file.name;
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -11,7 +11,9 @@ from abc import ABC, ABCMeta, abstractmethod
|
||||
from typing import List, Tuple, Union
|
||||
import cv2
|
||||
import torch
|
||||
import torchaudio
|
||||
from cv2 import VideoWriter, VideoWriter_fourcc, VideoCapture
|
||||
from line_profiler import profile
|
||||
from torch import dtype
|
||||
from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
|
||||
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights, raft_small, Raft_Small_Weights
|
||||
@@ -163,12 +165,9 @@ class ABCABCVideoFileToImage(ABCVideoWidget, metaclass=ABCWidgetMetaclass):
|
||||
"idx_slice_stop": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 1}),
|
||||
"slice_idx_step": ("INT", {"min": 0, "max": 9999, "step": 1, "default": 1}),
|
||||
"idx": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0}),
|
||||
"func_only": ([True, False], {"default": False}),
|
||||
},
|
||||
|
||||
"optional":
|
||||
{"text":
|
||||
("STRING", {"multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
@@ -189,6 +188,8 @@ class VideoToFramesFolder(ABCVideoFileToABCVideoFolder, metaclass=ABCWidgetMetac
|
||||
# self.INPUT_TYPES = lambda s: self.old().update({})
|
||||
|
||||
def handler(self, video_in, folder_out, text):
|
||||
import time
|
||||
start = time.time()
|
||||
video_in_path = os.path.join(self.input_dir, video_in)
|
||||
video_out_path = os.path.join(self.output_dir, folder_out)
|
||||
|
||||
@@ -200,6 +201,64 @@ class VideoToFramesFolder(ABCVideoFileToABCVideoFolder, metaclass=ABCWidgetMetac
|
||||
success, image = vidcap.read()
|
||||
count += 1
|
||||
|
||||
vidcap.release()
|
||||
end = time.time()
|
||||
print("time elapsed: ", end - start)
|
||||
return (folder_out,)
|
||||
|
||||
|
||||
class VideoToFramesFolderFFMPEG(ABCVideoFileToABCVideoFolder, metaclass=ABCWidgetMetaclass):
|
||||
"""use ffmpeg to open the video and use it as a source for image stacks, use PNG or jpeg with quality 100"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required":
|
||||
{
|
||||
"video_in": ("STRING", {"default": "full_path_to_folder"}),
|
||||
"folder_out": ("STRING", {"default": "full_path_to_folder"}),
|
||||
"format": (["PNG", "JPEG"], {"default": "PNG"}),
|
||||
},
|
||||
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("folder_out",)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.old = self.INPUT_TYPES
|
||||
# self.INPUT_TYPES = lambda s: self.old().update({})
|
||||
|
||||
def handler(self, video_in, folder_out, format):
|
||||
import time
|
||||
start = time.time()
|
||||
video_in_path = os.path.join(self.input_dir, video_in)
|
||||
video_out_path = os.path.join(self.output_dir, folder_out)
|
||||
|
||||
if not os.path.exists(video_out_path):
|
||||
os.makedirs(video_out_path)
|
||||
|
||||
# specifically this calls ffmpeg with the following arguments:
|
||||
# -i means input file
|
||||
# -qscale:v means quality scale video, default is 2, which is a good compromise between size and quality
|
||||
# 1 is the quality scale value, which is a value between 1 and 31, where 1 is the highest quality
|
||||
|
||||
start = time.time()
|
||||
|
||||
if format == "JPEG":
|
||||
subprocess.call(
|
||||
["ffmpeg", "-i", video_in_path, "-qscale:v", "1", os.path.join(video_out_path, "frame%05d.jpg")])
|
||||
end = time.time()
|
||||
print("ffmpeg jpg process took {} seconds".format(end - start))
|
||||
|
||||
elif format == "PNG":
|
||||
subprocess.call(["ffmpeg", "-i", video_in_path, os.path.join(video_out_path, "frame%05d.png")])
|
||||
end = time.time()
|
||||
print("ffmpeg png process took {} seconds".format(end - start))
|
||||
|
||||
else:
|
||||
raise ValueError("format must be either PNG or JPEG")
|
||||
|
||||
return (folder_out,)
|
||||
|
||||
|
||||
@@ -532,17 +591,30 @@ class GetImageStackStatistics(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetac
|
||||
return (out_str, out_hist_image,)
|
||||
|
||||
|
||||
@profile
|
||||
class VideoFileToImageStack(ABCABCVideoFileToImage, metaclass=ABCWidgetMetaclass):
|
||||
"""Use cv2 to open a video and save the videos individual frames"""
|
||||
RETURN_TYPES = ("IMAGE", "FUNC",)
|
||||
RETURN_NAMES = ("image_stack", "FUNC(**kwargs)",)
|
||||
|
||||
def handler(self, video_in, text, idx_slice_start, idx_slice_stop, slice_idx_step, idx):
|
||||
def handler(self, video_in, idx_slice_start, idx_slice_stop, slice_idx_step, idx, func_only=False):
|
||||
"""
|
||||
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
||||
- torch tensors (T,H,W,C)
|
||||
"""
|
||||
import sys
|
||||
import gc
|
||||
# print memory usage
|
||||
print(f"Memory usage: {sys.getsizeof(self)} bytes")
|
||||
|
||||
if func_only:
|
||||
return (None, VideoFileToImageStack.handler,)
|
||||
|
||||
image_list = []
|
||||
video_in_fp = os.path.join(self.input_dir, video_in)
|
||||
video_in_fp = video_in
|
||||
vidcap = cv2.VideoCapture(video_in_fp)
|
||||
max_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
if not vidcap.isOpened():
|
||||
print("Error: Unable to open video file")
|
||||
return
|
||||
@@ -551,6 +623,9 @@ class VideoFileToImageStack(ABCABCVideoFileToImage, metaclass=ABCWidgetMetaclass
|
||||
stop = (idx * slice_idx_step) + idx_slice_stop
|
||||
slice_length = stop - start
|
||||
|
||||
if (start + slice_length) > max_frames:
|
||||
raise ValueError(f"idx_slice_stop ({idx_slice_stop}) is greater than the max frames ({max_frames})")
|
||||
|
||||
for frame_count in range(start, start + slice_length):
|
||||
# Seek to frame_count frame
|
||||
vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
|
||||
@@ -563,12 +638,24 @@ class VideoFileToImageStack(ABCABCVideoFileToImage, metaclass=ABCWidgetMetaclass
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
image_list.append(image)
|
||||
print(f"Memory usage: {sys.getsizeof(self)} bytes")
|
||||
|
||||
vidcap.release()
|
||||
# Convert list of images to tensor (T,H,W,C)
|
||||
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
||||
tensor = tensor / 255.0
|
||||
try:
|
||||
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
||||
except RuntimeError as e:
|
||||
print("Error: Unable to convert image list to tensor")
|
||||
raise e
|
||||
|
||||
return (tensor,)
|
||||
tensor = tensor / 255.0
|
||||
del image_list
|
||||
del vidcap
|
||||
del image
|
||||
gc.collect()
|
||||
# print memory usage
|
||||
print(f"Memory usage: {sys.getsizeof(self)} bytes")
|
||||
return (tensor, VideoFileToImageStack.handler,)
|
||||
|
||||
|
||||
class ImageStackToVideoFramesFolder(metaclass=ABCWidgetMetaclass):
|
||||
@@ -1177,6 +1264,9 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
|
||||
import imageio
|
||||
import audioread
|
||||
import cv2
|
||||
import torchvision
|
||||
|
||||
use_torchvision = True
|
||||
|
||||
# video_dir = kwargs["video_dir"]
|
||||
|
||||
@@ -1222,52 +1312,54 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
|
||||
else:
|
||||
video_fps = frames_length / fps_length_value
|
||||
|
||||
folder_out = "temp_frames_folder"
|
||||
if use_torchvision ==False:
|
||||
writer = imageio.get_writer(video_out_path, fps=float(video_fps))
|
||||
|
||||
# Now, read the frames from disk as numpy data
|
||||
import tempfile
|
||||
# use tempfile to create a temporary folder
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
frames_out_path = temp_dir
|
||||
|
||||
# Convert the image stack to frames on disk
|
||||
frame_folder_handler = ImageStackToVideoFramesFolder()
|
||||
frame_folder_handler.handler(image_stack=image_stack, folder_out=frames_out_path)
|
||||
|
||||
frames = []
|
||||
for fl in sorted(os.listdir(frames_out_path)):
|
||||
# if the fl is an actual image file
|
||||
if not fl.lower().endswith(('.jpg', '.jpeg', '.png')):
|
||||
continue
|
||||
# create the numpy array
|
||||
frame = cv2.imread(os.path.join(frames_out_path, fl))
|
||||
# convert to RGB
|
||||
for frame_data in image_stack:
|
||||
frame = np.array(frame_data) # Convert torch tensor to numpy array if needed
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
# append to the list of frames
|
||||
frames.append(frame)
|
||||
writer.append_data(frame)
|
||||
|
||||
writer = imageio.get_writer(video_out_path, fps=video_fps)
|
||||
for frame in frames:
|
||||
writer.append_data(frame)
|
||||
writer.close()
|
||||
if audio_file is not None:
|
||||
difference = video_length - audio_length
|
||||
|
||||
if audio_file is not None:
|
||||
if audio_length > video_length:
|
||||
while video_length < audio_length:
|
||||
image_stack = np.concatenate((image_stack, image_stack[-1:]), axis=0)
|
||||
video_length += 1 / video_fps
|
||||
frames_length += 1
|
||||
needed_extra_frames = int(difference * video_fps)
|
||||
|
||||
args = [
|
||||
"-i", video_out_path,
|
||||
"-i", audio_file,
|
||||
"-c:v", codec,
|
||||
"-b:v", str(bitrate),
|
||||
"-c:a", "aac",
|
||||
"-strict", "experimental",
|
||||
video_out_path
|
||||
]
|
||||
run_ffmpeg_command(args)
|
||||
if needed_extra_frames > 0:
|
||||
for i in range(needed_extra_frames):
|
||||
writer.append_data(frame)
|
||||
|
||||
writer.close()
|
||||
|
||||
args = [
|
||||
"-i", video_out_path,
|
||||
"-i", audio_file,
|
||||
"-c:v", codec,
|
||||
"-b:v", str(bitrate),
|
||||
"-c:a", "aac",
|
||||
"-strict", "experimental",
|
||||
"-y",
|
||||
video_out_path.replace(".mp4", "_with_audio.mp4")
|
||||
]
|
||||
run_ffmpeg_command(args)
|
||||
else:
|
||||
writer.close()
|
||||
|
||||
if use_torchvision:
|
||||
audio_tensor = None
|
||||
if audio_file is not None:
|
||||
aud_sr = torchaudio.io.StreamReader(audio_file)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
torchvision.io.write_video(filename=video_out_path,
|
||||
video_array=image_stack,
|
||||
fps=video_fps,
|
||||
audio_array=audio_tensor,
|
||||
)
|
||||
|
||||
return (self, video_out_path,)
|
||||
|
||||
@@ -1459,43 +1551,293 @@ def combine_videos(video_paths, base_input_directory, output_path):
|
||||
os.remove(temp_file.name)
|
||||
|
||||
|
||||
def transcribe_to_segs(input_filename):
|
||||
import whisper as wh
|
||||
import folder_paths as fp
|
||||
|
||||
if not os.sep in input_filename:
|
||||
input_dir = fp.input_directory
|
||||
full_path = os.path.join(input_dir, input_filename)
|
||||
else:
|
||||
full_path = input_filename
|
||||
# check that the file exists
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
raise ValueError(f"File {full_path} does not exist.")
|
||||
# check that it is a file
|
||||
if not os.path.isfile(full_path):
|
||||
raise ValueError(f"{full_path} is not a file.")
|
||||
|
||||
model = wh.load_model("medium.en")
|
||||
|
||||
res = wh.transcribe(model=model, audio=full_path, word_timestamps=True)
|
||||
segs = res["segments"]
|
||||
|
||||
return segs
|
||||
|
||||
|
||||
class WhisperTranscription(metaclass=ABCWidgetMetaclass):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"file_path": ("STRING", {"default": "fullpath"})}}
|
||||
return {"required": {"input_filename": ("STRING", {"default": "fullpath"})}}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_TYPES = ("STRING", "STRING",)
|
||||
RETURN_NAMES = ("lines_out", "words_out",)
|
||||
FUNCTION = "transcribe_movie"
|
||||
CATEGORY = "transcription"
|
||||
|
||||
def transcribe_movie(self, file_path):
|
||||
import os
|
||||
import time
|
||||
import numpy as np
|
||||
import whisper as wh
|
||||
import moviepy as mp
|
||||
def transcribe_movie(self, input_filename):
|
||||
|
||||
import folder_paths as fp
|
||||
input_dir = fp.input_directory
|
||||
full_path = os.path.join(input_dir, input_filename)
|
||||
|
||||
try:
|
||||
model = wh.load_model("medium.en")
|
||||
filename = os.path.basename(file_path)
|
||||
audio_file_path = f"audio_{filename}.wav"
|
||||
segs = transcribe_to_segs(full_path)
|
||||
|
||||
video = mp.VideoFileClip(file_path)
|
||||
video.audio.write_audiofile(audio_file_path)
|
||||
|
||||
res = wh.transcribe(model=model, audio=audio_file_path)
|
||||
segs = res["segments"]
|
||||
out = []
|
||||
lines_out = []
|
||||
for seg in segs:
|
||||
out.append(f"{seg['start']:3.3f} - {seg['end']:3.3f}: {seg['text']}")
|
||||
ret = '\n'.join(out)
|
||||
lines_out.append(f"{seg['start']:3.3f} - {seg['end']:3.3f}: {seg['text']}")
|
||||
|
||||
os.remove(audio_file_path)
|
||||
lines_out = '\n'.join(lines_out)
|
||||
|
||||
return (ret,)
|
||||
words_out = []
|
||||
for seg in segs:
|
||||
words_out.append([])
|
||||
for word in seg['words']:
|
||||
line = f"{word['start']:3.3f} - {word['end']:3.3f}: {word['word']}"
|
||||
words_out[-1].append(line)
|
||||
|
||||
words_out = '\n'.join(['\n'.join(x) for x in words_out])
|
||||
|
||||
return (lines_out, words_out,)
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return ("Error",)
|
||||
return ("Error",)
|
||||
|
||||
|
||||
class WhisperTranscribeToFrameIDXandWords(metaclass=ABCWidgetMetaclass):
|
||||
"""Just like WhisperTranscription, but returns a list of frame numbers and words instead of a string."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# needs full path to video file
|
||||
# needs fps for conversion
|
||||
ret = {"required": {"input_filename": ("STRING", {"default": "fullpath"}),
|
||||
"fps": ("INT", {"default": 30})}}
|
||||
return ret
|
||||
|
||||
RETURN_TYPES = ("LIST", "LIST", "LIST",)
|
||||
RETURN_NAMES = ("frame_idx", "words_out", "combined",)
|
||||
FUNCTION = "transcribe_movie"
|
||||
CATEGORY = "transcription"
|
||||
|
||||
def transcribe_movie(self, input_filename, fps):
|
||||
import folder_paths as fp
|
||||
input_dir = fp.input_directory
|
||||
full_path = os.path.join(input_dir, input_filename)
|
||||
try:
|
||||
segs = transcribe_to_segs(full_path)
|
||||
frame_idx = []
|
||||
words_out = []
|
||||
combined = []
|
||||
for seg in segs:
|
||||
for word in seg['words']:
|
||||
start_frame = int(word['start'] * fps)
|
||||
end_frame = int(word['end'] * fps)
|
||||
word_text = word['word']
|
||||
|
||||
frame_idx.append((start_frame, end_frame))
|
||||
words_out.append(word_text)
|
||||
combined.append((start_frame, end_frame, word_text))
|
||||
|
||||
# now we fill in the gaps so the entire video is covered by some word
|
||||
# we do this by finding if there is a gap then changing the frame stamps of the last word and the next word
|
||||
# to cover the gap
|
||||
for i in range(len(frame_idx) - 1):
|
||||
if frame_idx[i][1] < frame_idx[i + 1][0]:
|
||||
# there is a gap
|
||||
gap_len = frame_idx[i + 1][0] - frame_idx[i][1]
|
||||
to_previous_word = gap_len // 2
|
||||
to_next_word = gap_len - to_previous_word
|
||||
|
||||
frame_idx[i] = (frame_idx[i][0], frame_idx[i][1] + to_previous_word)
|
||||
frame_idx[i + 1] = (frame_idx[i + 1][0] - to_next_word, frame_idx[i + 1][1])
|
||||
|
||||
combined[i] = (combined[i][0], combined[i][1] + to_previous_word, combined[i][2])
|
||||
combined[i + 1] = (combined[i + 1][0] - to_next_word, combined[i + 1][1], combined[i + 1][2])
|
||||
|
||||
# leave words_out alone
|
||||
|
||||
return (frame_idx, words_out, combined,)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return ("Error",)
|
||||
|
||||
|
||||
@profile
|
||||
def get_text_frames(word_idx, words_data, x: callable, y: callable, size: (), text_kwargs: dict, style="default"):
|
||||
from .image import TextRender as text_render_class
|
||||
text_render = text_render_class().render_text
|
||||
|
||||
start = words_data[word_idx][0]
|
||||
end = words_data[word_idx][1]
|
||||
word_text = words_data[word_idx][2]
|
||||
frame_length = end - start
|
||||
|
||||
# if x and y are not callables
|
||||
if not callable(x):
|
||||
x = lambda i: x
|
||||
if not callable(y):
|
||||
y = lambda i: y
|
||||
|
||||
out_frames = []
|
||||
for i in range(start, end):
|
||||
it_dict = dict()
|
||||
it_dict["text"] = word_text
|
||||
it_dict["x"] = x(i / float(frame_length))
|
||||
it_dict["y"] = y(i / float(frame_length))
|
||||
it_dict["width"] = size[0]
|
||||
it_dict["height"] = size[1]
|
||||
|
||||
it_dict.update(text_kwargs)
|
||||
# it_dict["size"] = 128
|
||||
# it_dict["color"] = "#ffffff"
|
||||
# it_dict["func_only"] = False
|
||||
# it_dict["stroke fill"] = "#050505"
|
||||
# it_dict["stroke width"] = 5
|
||||
|
||||
text_img = text_render(**it_dict)
|
||||
out_frames.append(text_img[0])
|
||||
return out_frames
|
||||
|
||||
|
||||
class AddTranscriptionToVideo(metaclass=ABCWidgetMetaclass):
|
||||
"""
|
||||
Add a transcription to a video. Requires a video and transcription data.
|
||||
transcription data must be [(start_frame, end_frame, word), ...]
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# needs full path to video file
|
||||
# needs fps for conversion
|
||||
ret = {"required": {"video_file": ("STRING", {"default": "fullpath"}),
|
||||
"words": ("LIST", {"default": None}),
|
||||
}
|
||||
}
|
||||
return ret
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
||||
RETURN_NAMES = ("1 fps preview", "Full Video")
|
||||
FUNCTION = "add_transcription_to_video"
|
||||
CATEGORY = "transcription"
|
||||
|
||||
@profile
|
||||
def add_transcription_to_video(self, **kwargs):
|
||||
import torch
|
||||
import gc
|
||||
|
||||
video_file = kwargs["video_file"]
|
||||
words = kwargs["words"]
|
||||
|
||||
def create_video_frames_folder(video_file):
|
||||
fnc = VideoToFramesFolderFFMPEG().handler
|
||||
call_dict = dict()
|
||||
call_dict["video_in"] = video_file
|
||||
call_dict["folder_out"] = os.path.join(video_file, os.sep, "frames_tmp")
|
||||
call_dict["format"] = "PNG"
|
||||
|
||||
tmp_video_folder = fnc(**call_dict)[0]
|
||||
del fnc
|
||||
del call_dict
|
||||
gc.collect()
|
||||
return tmp_video_folder
|
||||
|
||||
def get_video_frames(start, end):
|
||||
fnc = VideoFramesFolderToImageStack().handler
|
||||
call_dict = dict()
|
||||
call_dict["use_subfolders"] = False
|
||||
call_dict["by_date_or_name"] = "NAME"
|
||||
call_dict["folder_in"] = os.path.join(video_file, os.sep, "frames_tmp")
|
||||
call_dict["idx_slice_start"] = 0
|
||||
call_dict["idx_slice_stop"] = end - start
|
||||
call_dict["slice_idx_step"] = 1
|
||||
call_dict["idx"] = start
|
||||
call_dict["use_float16"] = False
|
||||
|
||||
ret = fnc(**call_dict)[0]
|
||||
|
||||
return ret
|
||||
|
||||
def overlay_text_on_video(video_frames, text_frames):
|
||||
import torch
|
||||
# Create a mask for the text frames (where text is black)
|
||||
text_mask = (text_frames == 0) * 1 # Assuming black is represented as 0 in text_frames
|
||||
|
||||
# Use the ~ operator to invert the boolean mask
|
||||
# inverted_mask = text_mask
|
||||
|
||||
# Overlay text_frames on video_frames using the inverted_mask
|
||||
overlaid_frames = video_frames * text_mask + text_frames
|
||||
|
||||
return overlaid_frames
|
||||
|
||||
frames_out = []
|
||||
num_frames = 0
|
||||
frames_folder = create_video_frames_folder(video_file)
|
||||
for w_idx in range(len(words)):
|
||||
word_start = words[w_idx][0]
|
||||
word_end = words[w_idx][1]
|
||||
word = words[w_idx][2].strip()
|
||||
print(f"word_start: {word_start}, word_end: {word_end}, word: {word}")
|
||||
|
||||
words_video_frames = get_video_frames(word_start, word_end)
|
||||
i_width = words_video_frames.shape[2]
|
||||
i_height = words_video_frames.shape[1]
|
||||
text_frames = get_text_frames(word_idx=w_idx,
|
||||
words_data=words,
|
||||
x=lambda i: 0,
|
||||
y=lambda i: 0,
|
||||
size=(i_width, i_height),
|
||||
text_kwargs={"size": 128, "color": "#ffffff", "func_only": False,
|
||||
"stroke fill": "#050505", "stroke width": 5,
|
||||
"font_name": "c:/Windows/Fonts/arial.ttf"})
|
||||
|
||||
text_frames = torch.cat(text_frames)
|
||||
|
||||
# Overlay text on video frames
|
||||
video_frames_with_text = overlay_text_on_video(words_video_frames, text_frames)
|
||||
frames_out.append(video_frames_with_text)
|
||||
num_frames += video_frames_with_text.shape[0]
|
||||
gc.collect()
|
||||
|
||||
# free memory
|
||||
del words_video_frames
|
||||
del text_frames
|
||||
del video_frames_with_text
|
||||
|
||||
num_stks = len(frames_out)
|
||||
# write to disk
|
||||
for i, stk in enumerate(frames_out):
|
||||
torch.save(stk, f"tmp_str_{i}.pt")
|
||||
del frames_out
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# pre allocate tensor
|
||||
y = torch.zeros((num_frames, i_height, i_width, 3), dtype=torch.float32, device="cpu")
|
||||
# load each save and put into tensor
|
||||
new_start = 0
|
||||
for i in range(num_stks):
|
||||
stk = torch.load(f"tmp_str_{i}.pt")
|
||||
y[new_start:new_start + stk.shape[0]] = stk
|
||||
new_start = new_start + stk.shape[0]
|
||||
del stk
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
preview = y[::30].clone()
|
||||
|
||||
return (preview, y,)
|
||||
|
||||
Reference in New Issue
Block a user