449 lines
16 KiB
Python
449 lines
16 KiB
Python
"""
|
|
MetaHub Save Video Node for ComfyUI
|
|
Adds metadata to video files generated by VHS (Video Helper Suite)
|
|
"""
|
|
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
try:
|
|
from . import metadata_utils as utils
|
|
from . import video_metadata_utils as video_utils
|
|
from .workflow_extractor import WorkflowExtractor
|
|
except ImportError:
|
|
import metadata_utils as utils
|
|
import video_metadata_utils as video_utils
|
|
from workflow_extractor import WorkflowExtractor
|
|
|
|
|
|
class MetaHubSaveVideoNode:
|
|
"""
|
|
ComfyUI custom node for adding MetaHub metadata to video files.
|
|
|
|
Accepts VHS_FILENAMES from VHS Video Combine node and injects metadata
|
|
into the video container (MP4, WebM, MKV) without re-encoding.
|
|
|
|
Features:
|
|
- A1111/Civitai compatible metadata (description field)
|
|
- Video MetaHub metadata (comment field as JSON)
|
|
- Auto-detection of workflow parameters
|
|
- Frame rate and frame count tracking
|
|
- Pass-through design for workflow chaining
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"filenames": ("VHS_FILENAMES",),
|
|
},
|
|
"optional": {
|
|
# Generation parameter overrides
|
|
"seed": ("INT", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override seed"
|
|
}),
|
|
"steps": ("INT", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override steps"
|
|
}),
|
|
"cfg": ("FLOAT", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override CFG scale"
|
|
}),
|
|
"sampler_name": ("STRING", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override sampler name"
|
|
}),
|
|
"scheduler": ("STRING", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override scheduler"
|
|
}),
|
|
"model_name": ("STRING", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override model filename"
|
|
}),
|
|
"positive": ("STRING", {
|
|
"multiline": True,
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override positive prompt"
|
|
}),
|
|
"negative": ("STRING", {
|
|
"multiline": True,
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override negative prompt"
|
|
}),
|
|
"denoise": ("FLOAT", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override denoise strength"
|
|
}),
|
|
"vae_name": ("STRING", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Override VAE filename"
|
|
}),
|
|
# Video-specific fields
|
|
"frame_rate": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 120.0,
|
|
"tooltip": "Video frame rate (FPS). 0 = auto-detect from file"
|
|
}),
|
|
"frame_count": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"tooltip": "Number of frames. 0 = auto-detect from file"
|
|
}),
|
|
"motion_model_name": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Motion model name (AnimateDiff, etc.)"
|
|
}),
|
|
# IMH Pro fields
|
|
"user_tags": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "User tags (comma-separated)"
|
|
}),
|
|
"notes": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Notes"
|
|
}),
|
|
"project_name": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Project name"
|
|
}),
|
|
# Timer integration
|
|
"generation_time_override": ("FLOAT", {
|
|
"default": None,
|
|
"forceInput": True,
|
|
"tooltip": "Timestamp from MetaHub Timer Node"
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"unique_id": "UNIQUE_ID"
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("VHS_FILENAMES",)
|
|
RETURN_NAMES = ("filenames",)
|
|
FUNCTION = "add_metadata"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "video/save"
|
|
DESCRIPTION = "Add A1111/Civitai and Video MetaHub metadata to video files from VHS"
|
|
|
|
def add_metadata(
|
|
self,
|
|
filenames,
|
|
seed=None,
|
|
steps=None,
|
|
cfg=None,
|
|
sampler_name=None,
|
|
scheduler=None,
|
|
model_name=None,
|
|
positive=None,
|
|
negative=None,
|
|
denoise=None,
|
|
vae_name=None,
|
|
frame_rate=0.0,
|
|
frame_count=0,
|
|
motion_model_name="",
|
|
user_tags="",
|
|
notes="",
|
|
project_name="",
|
|
generation_time_override=None,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
unique_id=None,
|
|
):
|
|
"""
|
|
Adds metadata to video files from VHS_FILENAMES.
|
|
|
|
1. Extracts video paths from VHS_FILENAMES tuple
|
|
2. Uses WorkflowExtractor to get generation parameters
|
|
3. Builds A1111 and MetaHub metadata structures
|
|
4. Injects metadata into video using ffmpeg
|
|
5. Returns original filenames for chaining
|
|
"""
|
|
# Check ffmpeg availability first
|
|
ffmpeg = video_utils.find_ffmpeg_binary()
|
|
if not ffmpeg:
|
|
print("[MetaHub Video] Warning: FFmpeg not found. Metadata injection skipped.")
|
|
print("[MetaHub Video] Install FFmpeg and ensure it's in PATH, or set FFMPEG_PATH.")
|
|
return (filenames,)
|
|
|
|
# Extract video files from VHS_FILENAMES
|
|
video_paths = video_utils.extract_video_files(filenames)
|
|
|
|
if not video_paths:
|
|
print("[MetaHub Video] No video files found in VHS_FILENAMES")
|
|
return (filenames,)
|
|
|
|
# Extract workflow data
|
|
try:
|
|
workflow_json = utils.get_workflow_json(extra_pnginfo)
|
|
prompt_data = prompt if isinstance(prompt, dict) else workflow_json.get("prompt", {})
|
|
if not isinstance(prompt_data, dict):
|
|
prompt_data = {}
|
|
|
|
workflow_json = utils.ensure_prompt_in_workflow(workflow_json, prompt_data)
|
|
save_node_id = str(unique_id) if unique_id is not None else None
|
|
imh_attribution = utils.extract_workflow_attribution(workflow_json, save_node_id)
|
|
|
|
# Use WorkflowExtractor to get generation params
|
|
extractor = WorkflowExtractor(prompt_data)
|
|
extracted, missing_fields = extractor.extract(save_node_id=save_node_id)
|
|
lora_list = extracted.get("lora_list") or utils.extract_loras_from_workflow(workflow_json)
|
|
|
|
except Exception as e:
|
|
print(f"[MetaHub Video] Warning: Could not extract workflow data: {e}")
|
|
extracted = {}
|
|
lora_list = []
|
|
workflow_json = {}
|
|
imh_attribution = None
|
|
|
|
# Resolve values (manual override > extracted > default)
|
|
def resolve_value(manual_value, extracted_value, default_value):
|
|
if manual_value is not None:
|
|
return manual_value
|
|
if extracted_value is not None:
|
|
return extracted_value
|
|
return default_value
|
|
|
|
def normalize_int(value, default_value):
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return default_value
|
|
|
|
def normalize_float(value, default_value):
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return default_value
|
|
|
|
seed_value = normalize_int(resolve_value(seed, extracted.get("seed"), 0), 0)
|
|
steps_value = normalize_int(resolve_value(steps, extracted.get("steps"), 20), 20)
|
|
cfg_value = normalize_float(resolve_value(cfg, extracted.get("cfg"), 7.0), 7.0)
|
|
sampler_value = resolve_value(sampler_name, extracted.get("sampler_name"), "euler")
|
|
scheduler_value = resolve_value(scheduler, extracted.get("scheduler"), "normal")
|
|
model_name_value = resolve_value(model_name, extracted.get("model_name"), "")
|
|
positive_value = resolve_value(positive, extracted.get("positive"), "")
|
|
negative_value = resolve_value(negative, extracted.get("negative"), "")
|
|
denoise_value = normalize_float(resolve_value(denoise, extracted.get("denoise"), 1.0), 1.0)
|
|
vae_name_value = resolve_value(vae_name, extracted.get("vae_name"), "")
|
|
metadata_fields = [
|
|
"seed",
|
|
"steps",
|
|
"cfg",
|
|
"sampler_name",
|
|
"scheduler",
|
|
"model_name",
|
|
"positive",
|
|
"negative",
|
|
"denoise",
|
|
"vae_name",
|
|
]
|
|
manual_inputs = {
|
|
"seed": seed,
|
|
"steps": steps,
|
|
"cfg": cfg,
|
|
"sampler_name": sampler_name,
|
|
"scheduler": scheduler,
|
|
"model_name": model_name,
|
|
"positive": positive,
|
|
"negative": negative,
|
|
"denoise": denoise,
|
|
"vae_name": vae_name,
|
|
}
|
|
metadata_sources = utils.build_metadata_sources(manual_inputs, extracted, metadata_fields)
|
|
metadata_status = utils.build_metadata_status(metadata_sources)
|
|
|
|
# Calculate model hash
|
|
if model_name_value:
|
|
model_hash = utils.calculate_model_hash(model_name_value, model_type="checkpoint")
|
|
else:
|
|
model_hash = "0000000000"
|
|
lora_hashes = utils.calculate_lora_hashes(lora_list) if lora_list else {}
|
|
|
|
# Calculate generation time
|
|
if generation_time_override is not None and generation_time_override > 0:
|
|
final_time = time.time() - generation_time_override
|
|
else:
|
|
final_time = 0.0
|
|
|
|
generation_time_ms = int(final_time * 1000) if final_time > 0 else None
|
|
steps_per_second = None
|
|
if generation_time_ms and generation_time_ms > 0 and steps_value > 0:
|
|
steps_per_second = round((steps_value / (generation_time_ms / 1000)), 2)
|
|
|
|
# Collect GPU metrics
|
|
gpu_metrics = utils.collect_gpu_metrics()
|
|
version_info = utils.collect_version_info()
|
|
|
|
# Process each video file
|
|
for video_path in video_paths:
|
|
try:
|
|
self._process_video(
|
|
video_path=video_path,
|
|
seed=seed_value,
|
|
steps=steps_value,
|
|
cfg=cfg_value,
|
|
sampler=sampler_value,
|
|
scheduler=scheduler_value,
|
|
model_name=model_name_value,
|
|
model_hash=model_hash,
|
|
positive=positive_value,
|
|
negative=negative_value,
|
|
denoise=denoise_value,
|
|
vae_name=vae_name_value,
|
|
frame_rate=frame_rate,
|
|
frame_count=frame_count,
|
|
motion_model_name=motion_model_name,
|
|
lora_list=lora_list,
|
|
lora_hashes=lora_hashes,
|
|
user_tags=user_tags,
|
|
notes=notes,
|
|
project_name=project_name,
|
|
generation_time_ms=generation_time_ms,
|
|
steps_per_second=steps_per_second,
|
|
gpu_metrics=gpu_metrics,
|
|
version_info=version_info,
|
|
workflow_json=workflow_json,
|
|
imh_attribution=imh_attribution,
|
|
metadata_sources=metadata_sources,
|
|
metadata_status=metadata_status,
|
|
)
|
|
except Exception as e:
|
|
print(f"[MetaHub Video] Warning: Failed to add metadata to {video_path.name}: {e}")
|
|
|
|
return (filenames,)
|
|
|
|
def _process_video(
|
|
self,
|
|
video_path: Path,
|
|
seed: int,
|
|
steps: int,
|
|
cfg: float,
|
|
sampler: str,
|
|
scheduler: str,
|
|
model_name: str,
|
|
model_hash: str,
|
|
positive: str,
|
|
negative: str,
|
|
denoise: float,
|
|
vae_name: str,
|
|
frame_rate: float,
|
|
frame_count: int,
|
|
motion_model_name: str,
|
|
lora_list: List[Dict],
|
|
lora_hashes: Dict[str, str],
|
|
user_tags: str,
|
|
notes: str,
|
|
project_name: str,
|
|
generation_time_ms: Optional[int],
|
|
steps_per_second: Optional[float],
|
|
gpu_metrics: Dict[str, Any],
|
|
version_info: Dict[str, str],
|
|
workflow_json: Dict,
|
|
imh_attribution: Optional[Dict[str, Any]],
|
|
metadata_sources: Dict[str, str],
|
|
metadata_status: str,
|
|
) -> None:
|
|
"""
|
|
Processes a single video file and injects metadata.
|
|
"""
|
|
# Get video info from file if not provided
|
|
video_info = video_utils.get_video_info(video_path)
|
|
|
|
if video_info:
|
|
if frame_rate <= 0:
|
|
frame_rate = video_info.get('frame_rate') or 0
|
|
if frame_count <= 0:
|
|
frame_count = video_info.get('frame_count') or 0
|
|
width = video_info.get('width', 512)
|
|
height = video_info.get('height', 512)
|
|
video_codec = video_info.get('codec', 'h264')
|
|
else:
|
|
width = 512
|
|
height = 512
|
|
video_codec = 'h264'
|
|
|
|
video_format = video_utils.detect_video_format(video_path) or 'mp4'
|
|
|
|
# Build params dict
|
|
params = {
|
|
"positive": positive,
|
|
"negative": negative,
|
|
"steps": steps,
|
|
"sampler": sampler,
|
|
"scheduler": scheduler,
|
|
"cfg": cfg,
|
|
"seed": seed,
|
|
"width": width,
|
|
"height": height,
|
|
"model_name": model_name,
|
|
"model_hash": model_hash,
|
|
"vae_name": vae_name,
|
|
"denoise": denoise,
|
|
# Video-specific
|
|
"frame_rate": frame_rate,
|
|
"frame_count": frame_count,
|
|
"video_format": video_format,
|
|
"video_codec": video_codec,
|
|
"motion_model_name": motion_model_name,
|
|
# LoRAs
|
|
"lora_list": lora_list,
|
|
"lora_hashes": lora_hashes,
|
|
# IMH Pro
|
|
"user_tags": user_tags,
|
|
"notes": notes,
|
|
"project_name": project_name,
|
|
# Performance metrics
|
|
"vram_peak_mb": gpu_metrics.get("vram_peak_mb"),
|
|
"gpu_device": gpu_metrics.get("gpu_device"),
|
|
"generation_time_ms": generation_time_ms,
|
|
"steps_per_second": steps_per_second,
|
|
"comfyui_version": version_info.get("comfyui_version"),
|
|
"torch_version": version_info.get("torch_version"),
|
|
"python_version": version_info.get("python_version"),
|
|
"imh_attribution": imh_attribution,
|
|
"metadata_sources": metadata_sources,
|
|
"metadata_status": metadata_status,
|
|
}
|
|
|
|
# Build metadata
|
|
a1111_metadata = video_utils.build_video_a1111_metadata(params)
|
|
metahub_metadata = video_utils.build_video_metahub_metadata(params, workflow_json)
|
|
|
|
# Inject metadata into video
|
|
video_utils.inject_video_metadata(
|
|
video_path=video_path,
|
|
a1111_metadata=a1111_metadata,
|
|
metahub_metadata=metahub_metadata,
|
|
)
|
|
|
|
print(f"[MetaHub Video] Added metadata to: {video_path.name}")
|
|
|
|
# Verify metadata was injected
|
|
verification = video_utils.verify_video_metadata(video_path)
|
|
if verification and verification.get('has_metahub_data'):
|
|
print(f"[MetaHub Video] Metadata verified successfully")
|
|
else:
|
|
print(f"[MetaHub Video] Warning: Metadata verification failed")
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"MetaHubSaveVideoNode": MetaHubSaveVideoNode}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"MetaHubSaveVideoNode": "MetaHub Save Video"}
|