metahub_save_video_node.py

"""
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"}