56 lines
2.9 KiB
Python
56 lines
2.9 KiB
Python
"""PromptForge Load Video: a video from the input folder as 24 fps frames, its sound and a description; nothing at all when no file is chosen or the node is switched off."""
|
|
import os
|
|
|
|
import torch
|
|
|
|
import folder_paths
|
|
from comfy_api.latest import InputImpl
|
|
|
|
FPS = 24
|
|
|
|
|
|
class PROMPTFORGE_LoadVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
files = folder_paths.filter_files_content_types(os.listdir(folder_paths.get_input_directory()), ["video"])
|
|
return {
|
|
"required": {
|
|
"video": ([""] + sorted(files), {"tooltip": "A video from the input folder (put the files there). Empty = no video."}),
|
|
"description": ("STRING", {"default": "", "multiline": True, "placeholder": "What this video shows. It goes to the prompt writer.", "tooltip": "Output as DESCRIPTION."}),
|
|
"enabled": ("BOOLEAN", {"default": True, "tooltip": "Off = the node outputs nothing, as if it were bypassed, so it can stay in a workflow and be switched on when needed."}),
|
|
"max_seconds": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 3600.0, "step": 0.5, "tooltip": "Only the first seconds are loaded. 0 = the whole video. Reference videos work best at 2 to 15 seconds."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "AUDIO", "STRING")
|
|
RETURN_NAMES = ("IMAGES", "AUDIO", "DESCRIPTION")
|
|
FUNCTION = "load"
|
|
CATEGORY = "\U0001f52e PromptForge/Load"
|
|
DESCRIPTION = "Load a video from the input folder as 24 fps frames with its sound and a description. With no file chosen, or switched off, the outputs are empty."
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, video, description, enabled, max_seconds):
|
|
return f"{os.path.getmtime(folder_paths.get_annotated_filepath(video)) if video and enabled else 0}|{description}|{enabled}|{max_seconds}"
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, video, enabled=True):
|
|
return not video or not enabled or folder_paths.exists_annotated_filepath(video) or f"Invalid video file: {video}"
|
|
|
|
def load(self, video, description, enabled, max_seconds):
|
|
if not video or not enabled:
|
|
return (None, None, "")
|
|
components = InputImpl.VideoFromFile(folder_paths.get_annotated_filepath(video)).get_components()
|
|
frames, rate = components.images, float(components.frame_rate)
|
|
count = round(frames.shape[0] * FPS / rate)
|
|
if max_seconds:
|
|
count = min(count, round(max_seconds * FPS))
|
|
frames = frames[torch.arange(count).mul(rate / FPS).floor().long().clamp(max=frames.shape[0] - 1)]
|
|
audio = components.audio
|
|
if audio is not None and max_seconds:
|
|
audio = {"waveform": audio["waveform"][..., :round(max_seconds * audio["sample_rate"])], "sample_rate": audio["sample_rate"]}
|
|
return (frames, audio, description.strip())
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"PROMPTFORGE_LoadVideo": PROMPTFORGE_LoadVideo}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_LoadVideo": "PromptForge Load Video"}
|