483 lines
16 KiB
Python
483 lines
16 KiB
Python
import os
|
|
import subprocess
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
import cv2
|
|
|
|
# Import ComfyUI-specific modules
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
# For standalone testing
|
|
class FolderPaths:
|
|
def get_input_directory(self):
|
|
return "input"
|
|
def get_output_directory(self):
|
|
return "output"
|
|
def get_temp_directory(self):
|
|
return "temp"
|
|
folder_paths = FolderPaths()
|
|
|
|
# Define supported video extensions
|
|
video_extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov']
|
|
|
|
# Check for ffmpeg availability
|
|
def get_ffmpeg_path():
|
|
"""Get the path to ffmpeg executable"""
|
|
# Try to find ffmpeg in the system path
|
|
try:
|
|
if os.name == 'nt': # Windows
|
|
ffmpeg_command = "where ffmpeg"
|
|
else: # Unix/Linux/Mac
|
|
ffmpeg_command = "which ffmpeg"
|
|
|
|
result = subprocess.run(ffmpeg_command, shell=True, capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
return result.stdout.strip()
|
|
except:
|
|
pass
|
|
|
|
# Check if ffmpeg is in the current directory
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if os.name == 'nt': # Windows
|
|
ffmpeg_path = os.path.join(current_dir, "ffmpeg.exe")
|
|
else: # Unix/Linux/Mac
|
|
ffmpeg_path = os.path.join(current_dir, "ffmpeg")
|
|
|
|
if os.path.exists(ffmpeg_path):
|
|
return ffmpeg_path
|
|
|
|
# Try to use imageio's ffmpeg
|
|
try:
|
|
import imageio
|
|
return imageio.plugins.ffmpeg.get_exe()
|
|
except:
|
|
pass
|
|
|
|
return None
|
|
|
|
ffmpeg_path = get_ffmpeg_path()
|
|
|
|
def target_size(width, height, custom_width, custom_height):
|
|
"""Calculate target size while maintaining aspect ratio"""
|
|
if custom_width == 0 and custom_height == 0:
|
|
return width, height
|
|
|
|
if custom_width == 0:
|
|
# Calculate width based on height
|
|
new_width = int(width * (custom_height / height))
|
|
return new_width, custom_height
|
|
|
|
if custom_height == 0:
|
|
# Calculate height based on width
|
|
new_height = int(height * (custom_width / width))
|
|
return custom_width, new_height
|
|
|
|
# Both dimensions specified
|
|
return custom_width, custom_height
|
|
|
|
def extract_frames_from_video(video_path, output_tensors=True, max_width=512, max_height=512):
|
|
"""Extract frames from a video file and return as tensors or PIL images"""
|
|
if not os.path.exists(video_path):
|
|
raise FileNotFoundError(f"Video file not found: {video_path}")
|
|
|
|
# Open the video file
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
raise ValueError(f"Could not open video file: {video_path}")
|
|
|
|
# Get video properties
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
|
# Calculate target size
|
|
target_w, target_h = target_size(width, height, max_width, max_height)
|
|
|
|
frames = []
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Convert from BGR to RGB
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# Resize frame if needed
|
|
if width != target_w or height != target_h:
|
|
frame = cv2.resize(frame, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
if output_tensors:
|
|
# Convert to tensor
|
|
frame_tensor = torch.from_numpy(frame).float() / 255.0
|
|
frames.append(frame_tensor)
|
|
else:
|
|
# Convert to PIL Image
|
|
frame_pil = Image.fromarray(frame)
|
|
frames.append(frame_pil)
|
|
|
|
cap.release()
|
|
|
|
if output_tensors:
|
|
# Stack tensors into a batch
|
|
if frames:
|
|
return torch.stack(frames)
|
|
return torch.zeros((0, target_h, target_w, 3))
|
|
|
|
return frames
|
|
|
|
def resize_image(image, max_width=512, max_height=512):
|
|
"""Resize an image while maintaining aspect ratio"""
|
|
width, height = image.size
|
|
target_w, target_h = target_size(width, height, max_width, max_height)
|
|
|
|
if width != target_w or height != target_h:
|
|
return image.resize((target_w, target_h), Image.LANCZOS)
|
|
|
|
return image
|
|
|
|
def combine_frames_to_video(frames, output_path, fps=30, audio_path=None):
|
|
"""Combine frames into a video file"""
|
|
if ffmpeg_path is None:
|
|
raise RuntimeError("ffmpeg not found. Please install ffmpeg to use video functionality.")
|
|
|
|
if not frames:
|
|
raise ValueError("No frames provided")
|
|
|
|
# Create temporary directory for frames if it doesn't exist
|
|
temp_dir = os.path.join(folder_paths.get_temp_directory(), "frames")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
# Save frames as images
|
|
for i, frame in enumerate(frames):
|
|
if isinstance(frame, torch.Tensor):
|
|
# Convert tensor to PIL Image
|
|
if frame.dim() == 4: # Batch of images
|
|
frame = frame[0]
|
|
|
|
frame = (frame * 255).byte().cpu().numpy()
|
|
frame = Image.fromarray(frame)
|
|
|
|
frame.save(os.path.join(temp_dir, f"frame_{i:05d}.png"))
|
|
|
|
# Construct ffmpeg command
|
|
ffmpeg_cmd = [
|
|
ffmpeg_path,
|
|
"-y", # Overwrite output file if it exists
|
|
"-framerate", str(fps),
|
|
"-i", os.path.join(temp_dir, "frame_%05d.png"),
|
|
"-c:v", "libx264",
|
|
"-pix_fmt", "yuv420p",
|
|
"-crf", "23", # Quality setting (lower is better)
|
|
]
|
|
|
|
# Add audio if provided
|
|
if audio_path and os.path.exists(audio_path):
|
|
ffmpeg_cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])
|
|
|
|
ffmpeg_cmd.append(output_path)
|
|
|
|
# Run ffmpeg
|
|
try:
|
|
subprocess.run(ffmpeg_cmd, check=True, capture_output=True)
|
|
except subprocess.CalledProcessError as e:
|
|
raise RuntimeError(f"Error creating video: {e.stderr.decode()}")
|
|
|
|
# Clean up temporary files
|
|
for i in range(len(frames)):
|
|
os.remove(os.path.join(temp_dir, f"frame_{i:05d}.png"))
|
|
|
|
return output_path
|
|
|
|
class SBSVideoUploader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
try:
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = []
|
|
|
|
if os.path.exists(input_dir):
|
|
for f in os.listdir(input_dir):
|
|
if os.path.isfile(os.path.join(input_dir, f)):
|
|
file_parts = f.split('.')
|
|
if len(file_parts) > 1 and (file_parts[-1].lower() in video_extensions):
|
|
files.append(f)
|
|
|
|
if len(files) == 0:
|
|
files = ["none"]
|
|
except Exception:
|
|
files = ["none"]
|
|
|
|
return {
|
|
"required": {
|
|
"video": (sorted(files),),
|
|
"max_width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
|
|
"max_height": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8, "display": "Max height (0 = auto)"}),
|
|
"frame_load_cap": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1, "display": "Max frames (0 = all)"}),
|
|
"skip_first_frames": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
|
"select_every_nth": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT")
|
|
RETURN_NAMES = ("images", "frame_count")
|
|
FUNCTION = "load_video"
|
|
CATEGORY = "👀 SamSeen"
|
|
|
|
def load_video(self, video, max_width, max_height, frame_load_cap=0, skip_first_frames=0, select_every_nth=1):
|
|
if video == "none":
|
|
# Return an empty tensor if no video is selected
|
|
return (torch.zeros((0, 512, 512, 3)), 0)
|
|
|
|
try:
|
|
video_path = os.path.join(folder_paths.get_input_directory(), video)
|
|
|
|
# Open the video file
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
raise ValueError(f"Could not open video file: {video_path}")
|
|
except Exception as e:
|
|
print(f"Error loading video: {e}")
|
|
return (torch.zeros((0, 512, 512, 3)), 0)
|
|
|
|
# Get video properties
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
|
# Calculate target size
|
|
target_w, target_h = target_size(width, height, max_width, max_height)
|
|
|
|
# Skip frames if needed
|
|
if skip_first_frames > 0:
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, skip_first_frames)
|
|
|
|
frames = []
|
|
frame_count = 0
|
|
frame_index = 0
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Process every nth frame
|
|
if frame_index % select_every_nth == 0:
|
|
# Convert from BGR to RGB
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# Resize frame if needed
|
|
if width != target_w or height != target_h:
|
|
frame = cv2.resize(frame, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
# Convert to tensor
|
|
frame_tensor = torch.from_numpy(frame).float() / 255.0
|
|
frames.append(frame_tensor)
|
|
|
|
frame_count += 1
|
|
|
|
# Stop if we've reached the frame cap
|
|
if frame_load_cap > 0 and frame_count >= frame_load_cap:
|
|
break
|
|
|
|
frame_index += 1
|
|
|
|
cap.release()
|
|
|
|
# Stack tensors into a batch
|
|
if frames:
|
|
return (torch.stack(frames), frame_count)
|
|
|
|
return (torch.zeros((0, target_h, target_w, 3)), 0)
|
|
|
|
class SBSImageUploader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
try:
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = []
|
|
|
|
if os.path.exists(input_dir):
|
|
for f in os.listdir(input_dir):
|
|
if os.path.isfile(os.path.join(input_dir, f)):
|
|
file_parts = f.split('.')
|
|
if len(file_parts) > 1 and (file_parts[-1].lower() in ['png', 'jpg', 'jpeg', 'webp', 'bmp']):
|
|
files.append(f)
|
|
|
|
if len(files) == 0:
|
|
files = ["none"]
|
|
except Exception:
|
|
files = ["none"]
|
|
|
|
return {
|
|
"required": {
|
|
"image": (sorted(files),),
|
|
"max_width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
|
|
"max_height": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8, "display": "Max height (0 = auto)"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "load_image"
|
|
CATEGORY = "👀 SamSeen"
|
|
|
|
def load_image(self, image, max_width, max_height):
|
|
if image == "none":
|
|
# Return an empty tensor if no image is selected
|
|
return (torch.zeros((1, 512, 512, 3)),)
|
|
|
|
try:
|
|
image_path = os.path.join(folder_paths.get_input_directory(), image)
|
|
|
|
if not os.path.exists(image_path):
|
|
raise FileNotFoundError(f"Image file not found: {image_path}")
|
|
except Exception as e:
|
|
print(f"Error loading image: {e}")
|
|
return (torch.zeros((1, 512, 512, 3)),)
|
|
|
|
# Load the image
|
|
img = Image.open(image_path).convert('RGB')
|
|
|
|
# Resize if needed
|
|
width, height = img.size
|
|
target_w, target_h = target_size(width, height, max_width, max_height)
|
|
|
|
if width != target_w or height != target_h:
|
|
img = img.resize((target_w, target_h), Image.LANCZOS)
|
|
|
|
# Convert to tensor
|
|
img_tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
return (img_tensor,)
|
|
|
|
class SBSVideoCombiner:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"frame_rate": ("INT", {"default": 30, "min": 1, "max": 120, "step": 1}),
|
|
"filename_prefix": ("STRING", {"default": "SBS_Video"}),
|
|
"format": (["mp4", "webm", "gif"], {"default": "mp4"}),
|
|
"save_output": ("BOOLEAN", {"default": True})
|
|
},
|
|
"optional": {
|
|
"audio": ("AUDIO",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("video_path",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "combine_video"
|
|
CATEGORY = "👀 SamSeen"
|
|
|
|
def combine_video(self, images, frame_rate, filename_prefix, format, save_output, audio=None):
|
|
if ffmpeg_path is None:
|
|
raise RuntimeError("ffmpeg not found. Please install ffmpeg to use video functionality.")
|
|
|
|
if images.size(0) == 0:
|
|
raise ValueError("No images provided")
|
|
|
|
# Get output directory
|
|
output_dir = folder_paths.get_output_directory() if save_output else folder_paths.get_temp_directory()
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Generate a unique filename
|
|
counter = 1
|
|
while True:
|
|
filename = f"{filename_prefix}_{counter:05d}.{format}"
|
|
file_path = os.path.join(output_dir, filename)
|
|
if not os.path.exists(file_path):
|
|
break
|
|
counter += 1
|
|
|
|
# Create temporary directory for frames
|
|
temp_dir = os.path.join(folder_paths.get_temp_directory(), "frames")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
# Save frames as images
|
|
for i in range(images.size(0)):
|
|
frame = images[i].cpu().numpy() * 255
|
|
frame = Image.fromarray(frame.astype(np.uint8))
|
|
frame.save(os.path.join(temp_dir, f"frame_{i:05d}.png"))
|
|
|
|
# Construct ffmpeg command
|
|
ffmpeg_cmd = [
|
|
ffmpeg_path,
|
|
"-y", # Overwrite output file if it exists
|
|
"-framerate", str(frame_rate),
|
|
"-i", os.path.join(temp_dir, "frame_%05d.png")
|
|
]
|
|
|
|
# Add format-specific options
|
|
if format == "mp4":
|
|
ffmpeg_cmd.extend([
|
|
"-c:v", "libx264",
|
|
"-pix_fmt", "yuv420p",
|
|
"-crf", "23" # Quality setting (lower is better)
|
|
])
|
|
elif format == "webm":
|
|
ffmpeg_cmd.extend([
|
|
"-c:v", "libvpx-vp9",
|
|
"-pix_fmt", "yuv420p",
|
|
"-crf", "30" # Quality setting for VP9
|
|
])
|
|
elif format == "gif":
|
|
ffmpeg_cmd.extend([
|
|
"-filter_complex", "[0:v] split [a][b];[a] palettegen [p];[b][p] paletteuse"
|
|
])
|
|
|
|
# Add audio if provided
|
|
if audio is not None:
|
|
# Create a temporary WAV file for the audio
|
|
audio_path = os.path.join(folder_paths.get_temp_directory(), "temp_audio.wav")
|
|
|
|
# Get audio data
|
|
waveform = audio['waveform']
|
|
sample_rate = audio['sample_rate']
|
|
|
|
# Save audio to WAV file using scipy
|
|
try:
|
|
import scipy.io.wavfile
|
|
# Convert from torch tensor to numpy array
|
|
audio_data = waveform.squeeze(0).transpose(0, 1).cpu().numpy()
|
|
scipy.io.wavfile.write(audio_path, sample_rate, audio_data)
|
|
|
|
# Add audio to ffmpeg command
|
|
ffmpeg_cmd.extend([
|
|
"-i", audio_path,
|
|
"-c:a", "aac" if format != "webm" else "libopus",
|
|
"-shortest" # End when the shortest input stream ends
|
|
])
|
|
except Exception as e:
|
|
print(f"Warning: Could not add audio to video: {e}")
|
|
|
|
# Add output file
|
|
ffmpeg_cmd.append(file_path)
|
|
|
|
# Run ffmpeg
|
|
try:
|
|
subprocess.run(ffmpeg_cmd, check=True, capture_output=True)
|
|
except subprocess.CalledProcessError as e:
|
|
raise RuntimeError(f"Error creating video: {e.stderr.decode()}")
|
|
|
|
# Clean up temporary files
|
|
for i in range(images.size(0)):
|
|
try:
|
|
os.remove(os.path.join(temp_dir, f"frame_{i:05d}.png"))
|
|
except:
|
|
pass
|
|
|
|
if audio is not None and os.path.exists(audio_path):
|
|
try:
|
|
os.remove(audio_path)
|
|
except:
|
|
pass
|
|
|
|
return (file_path,)
|