mirror of
https://github.com/justUmen/Bjornulf_custom_nodes.git
synced 2026-03-24 05:52:11 -03:00
120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
import subprocess
|
|
from PIL import Image
|
|
import soundfile as sf
|
|
import glob
|
|
|
|
class imagesToVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"fps": ("INT", {"default": 24, "min": 1, "max": 60}),
|
|
"name_prefix": ("STRING", {"default": "output/imgs2video/me"}),
|
|
"format": (["mp4", "webm"],),
|
|
"crf": ("INT", {"default": 19, "min": 0, "max": 63}),
|
|
},
|
|
"optional": {
|
|
"audio": ("AUDIO",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("comment",)
|
|
FUNCTION = "image_to_video"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Bjornulf"
|
|
|
|
def image_to_video(self, images, fps, name_prefix, format, crf, audio=None):
|
|
# Remove any existing extension
|
|
name_prefix = os.path.splitext(name_prefix)[0]
|
|
|
|
# Find the next available number
|
|
existing_files = glob.glob(f"{name_prefix}_*.{format}")
|
|
if existing_files:
|
|
max_num = max([int(f.split('_')[-1].split('.')[0]) for f in existing_files])
|
|
next_num = max_num + 1
|
|
else:
|
|
next_num = 1
|
|
|
|
# Create the new filename with the incremented number
|
|
output_file = f"{name_prefix}_{next_num:04d}.{format}"
|
|
|
|
temp_dir = "temp_images_imgs2video"
|
|
# Clean up temp dir
|
|
if os.path.exists(temp_dir) and os.path.isdir(temp_dir):
|
|
for file in os.listdir(temp_dir):
|
|
os.remove(os.path.join(temp_dir, file))
|
|
os.rmdir(temp_dir)
|
|
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
# Ensure the output directory exists
|
|
os.makedirs(os.path.dirname(output_file) if os.path.dirname(output_file) else ".", exist_ok=True)
|
|
|
|
# Save the tensor images as PNG files
|
|
for i, img_tensor in enumerate(images):
|
|
img = Image.fromarray((img_tensor.cpu().numpy() * 255).astype(np.uint8))
|
|
if format == "webm":
|
|
img = img.convert("RGBA") # Ensure alpha channel for WebM
|
|
img.save(os.path.join(temp_dir, f"frame_{i:04d}.png"))
|
|
|
|
# Handle audio
|
|
temp_audio_file = None
|
|
if audio is not None:
|
|
temp_audio_file = os.path.join(temp_dir, "temp_audio.wav")
|
|
waveform = audio['waveform'].squeeze().numpy()
|
|
sample_rate = audio['sample_rate']
|
|
sf.write(temp_audio_file, waveform, sample_rate)
|
|
|
|
# Construct the FFmpeg command based on the selected format
|
|
ffmpeg_cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-framerate", str(fps),
|
|
"-i", os.path.join(temp_dir, "frame_%04d.png"),
|
|
]
|
|
|
|
if temp_audio_file:
|
|
ffmpeg_cmd.extend(["-i", temp_audio_file])
|
|
|
|
if format == "mp4":
|
|
ffmpeg_cmd.extend([
|
|
"-filter_complex", "[0:v]scale=iw:ih,format=rgba,split[s0][s1];[s0]lutrgb=r=0:g=0:b=0:a=0[transparent];[transparent][s1]overlay",
|
|
"-crf", str(crf),
|
|
"-c:v", "libx264",
|
|
"-pix_fmt", "yuv420p",
|
|
])
|
|
comment = "MP4 format: Widely compatible, efficient compression, no transparency support."
|
|
elif format == "webm":
|
|
ffmpeg_cmd.extend([
|
|
"-filter_complex", "[0:v]scale=iw:ih,format=rgba,split[s0][s1];[s0]lutrgb=r=0:g=0:b=0:a=0[transparent];[transparent][s1]overlay",
|
|
"-c:v", "libvpx-vp9",
|
|
"-pix_fmt", "yuva420p",
|
|
"-b:v", "0",
|
|
"-crf", str(crf),
|
|
"-auto-alt-ref", "0",
|
|
])
|
|
comment = "WebM format: Supports transparency, open format, smaller file size, but less compatible than MP4."
|
|
|
|
if temp_audio_file:
|
|
ffmpeg_cmd.extend(["-c:a", "libvorbis", "-shortest"])
|
|
|
|
ffmpeg_cmd.append(output_file)
|
|
|
|
# Run FFmpeg
|
|
try:
|
|
subprocess.run(ffmpeg_cmd, check=True)
|
|
print(f"Video created successfully: {output_file}")
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error creating video: {e}")
|
|
finally:
|
|
# Clean up temporary files
|
|
# for file in os.listdir(temp_dir):
|
|
# os.remove(os.path.join(temp_dir, file))
|
|
# os.rmdir(temp_dir)
|
|
print("Temporary files not removed for debugging purposes.")
|
|
|
|
return (comment,) |