mirror of
https://github.com/justUmen/Bjornulf_custom_nodes.git
synced 2026-03-21 12:42:11 -03:00
92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
import subprocess
|
|
from PIL import Image
|
|
|
|
class imgs2vid:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"fps": ("INT", {"default": 30, "min": 1, "max": 60}),
|
|
"video_name_NO_format": ("STRING", {"default": "output"}),
|
|
"format": (["mp4", "webm"],),
|
|
"audio_path": ("STRING", {"default": "/home/umen/6sec.wav"}), # New audio input
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("comment",)
|
|
FUNCTION = "create_video"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Bjornulf"
|
|
|
|
def create_video(self, images, fps, video_name_NO_format, format, audio_path):
|
|
# Remove any existing extension
|
|
video_name_NO_format = os.path.splitext(video_name_NO_format)[0]
|
|
# Add the correct extension
|
|
output_file = f"{video_name_NO_format}.{format}"
|
|
temp_dir = "temp_images"
|
|
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"))
|
|
|
|
# Construct the FFmpeg command based on the selected format
|
|
if format == "mp4":
|
|
ffmpeg_cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-framerate", str(fps),
|
|
"-i", os.path.join(temp_dir, "frame_%04d.png"),
|
|
"-i", str(audio_path),
|
|
"-crf", "19",
|
|
"-c:v", "libx264",
|
|
"-pix_fmt", "yuv420p",
|
|
output_file
|
|
]
|
|
comment = "MP4 format: Widely compatible, efficient compression, no transparency support."
|
|
elif format == "webm":
|
|
ffmpeg_cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-framerate", str(fps),
|
|
"-i", os.path.join(temp_dir, "frame_%04d.png"),
|
|
"-i", str(audio_path),
|
|
"-crf", "19",
|
|
"-c:v", "libvpx",
|
|
"-b:v", "1M", # Set video bitrate
|
|
"-auto-alt-ref", "0", # Disable auto alt ref
|
|
"-c:a", "libvorbis",
|
|
"-pix_fmt", "yuva420p",
|
|
"-shortest",
|
|
output_file
|
|
]
|
|
comment = "WebM format: Supports transparency, open format, smaller file size, but less compatible than MP4."
|
|
|
|
# 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)
|
|
|
|
return (comment,)
|
|
|
|
# Example usage
|
|
# images = [torch.rand(256, 256, 3) for _ in range(10)] # Replace with actual image tensors
|
|
# imgs2vid().create_video(images, 30, "output", "webm", "/home/
|