Update wan22nodes.py

wan22nodes.py

@@ -1,21 +1,28 @@
-import os, json, hashlib, glob
+from __future__ import annotations
+import os, re, glob, json, hashlib, uuid, math
 from typing import Any, Dict, Tuple, Optional, List, Union
+
 import torch
-from safetensors import safe_open
-from comfy.cli_args import args
 import numpy as np
 from PIL import Image
+from safetensors import safe_open
+
 import folder_paths
 import comfy.utils
-from nodes import KSamplerAdvanced
-import nodes
 import comfy.model_management
+from comfy.cli_args import args
+from nodes import KSamplerAdvanced
 import node_helpers
-from comfy_api.latest import ComfyExtension, io, ui
+import nodes
+
+# Comfy API
+from comfy_api.latest import io, ui
+from comfy_api.input import VideoInput
+from comfy_api.input_impl import VideoFromFile, VideoFromComponents
+from comfy_api.util import VideoComponents, VideoContainer, VideoCodec
+
 import imageio.v3 as iio
 
-
-
 # ---------- SaveLatent (Comfy-only; saves into input/latents) ----------
 class SaveLatentMXD:
     DESCRIPTION = """
@@ -401,12 +408,16 @@ class LoadLatent_WithParams:
         return 5.0
 
     def load(self, latent):
-        latent_path = folder_paths.get_annotated_filepath(f"latents/{latent}")
+        # ✅ Ensure we prepend "latents/" if missing, but don't duplicate it
+        if not latent.startswith("latents/"):
+            latent_path = folder_paths.get_annotated_filepath(f"latents/{latent}")
+        else:
+            latent_path = folder_paths.get_annotated_filepath(latent)
+
         sample_dict, meta, _ = _load_latent_file(latent_path)
         t = sample_dict["samples"]
 
         if isinstance(t, torch.Tensor) and t.dim() >= 4 and t.size(0) > 1:
-            # for safety, only take first slice (multi-batch handling is folder loader’s job)
             samples = {"samples": t[0:1].contiguous()}
         elif isinstance(t, torch.Tensor) and t.dim() >= 4 and t.size(0) == 1:
             samples = {"samples": t}
@@ -459,7 +470,10 @@ class LoadLatent_WithParams:
 
     @classmethod
     def VALIDATE_INPUTS(s, latent):
-        if not folder_paths.exists_annotated_filepath(latent):
+        check_path = latent if latent.startswith("latents/") else f"latents/{latent}"
+        try:
+            folder_paths.get_annotated_filepath(check_path)
+        except Exception:
             return f"Invalid latent file: {latent}"
         return True
 
@@ -1382,6 +1396,457 @@ class WAN22_I2V_Image_Scaler_MXD:
             out, _, _ = _resize_fit_inside(image, bw, bh)
 
         return (out,)
+    
+class Frames_Select_End_MXD:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "frames": ("IMAGE",),
+                "count": ("INT", {"default": 10, "min": 1, "max": 10000, "tooltip": "Number of frames to select from the end"}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION     = "main"
+    CATEGORY     = "MXD/images"
+
+    def main(self, frames=None, count=10):
+        total = frames.shape[0]
+        start = max(0, total - count)
+        frames_end = frames[start:].clone()
+        return (frames_end,)
+    
+class Frames_Remove_From_Start_MXD:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "frames": ("IMAGE",),
+                "count": ("INT", {
+                    "default": 10,
+                    "min": 1,
+                    "max": 10000,
+                    "tooltip": "Number of frames to remove from the start"
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION     = "main"
+    CATEGORY     = "MXD/images"
+
+    def main(self, frames=None, count=10):
+        # ✅ Skip the first `count` frames instead of keeping them
+        frames_after = frames[count:].clone()
+        return (frames_after,)
+
+
+class CombineVideos_MXD:
+    """
+    Combine two VIDEO inputs end-to-end (sequentially).
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video_a": ("VIDEO", {"tooltip": "The first video (plays first)"}),
+                "video_b": ("VIDEO", {"tooltip": "The second video (plays after video_a)"}),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "combine"
+    CATEGORY = "MXD/video"
+
+    def combine(self, video_a, video_b):
+        comp_a = video_a.get_components()
+        comp_b = video_b.get_components()
+
+        # Check frame rate consistency
+        if comp_a.frame_rate != comp_b.frame_rate:
+            raise ValueError(f"FPS mismatch: {comp_a.frame_rate} vs {comp_b.frame_rate}")
+
+        # ✅ Correct way: concatenate frame tensors along batch/time dimension (dim=0)
+        frames_a = torch.stack(comp_a.images) if isinstance(comp_a.images, list) else comp_a.images
+        frames_b = torch.stack(comp_b.images) if isinstance(comp_b.images, list) else comp_b.images
+        combined_images = torch.cat([frames_a, frames_b], dim=0)
+
+        # ✅ Combine audio sequentially
+        combined_audio = None
+        if comp_a.audio is not None or comp_b.audio is not None:
+            audio_a = comp_a.audio if comp_a.audio is not None else torch.zeros((1, 0))
+            audio_b = comp_b.audio if comp_b.audio is not None else torch.zeros((1, 0))
+            combined_audio = torch.cat([audio_a, audio_b], dim=1)
+
+
+
+        combined_video = VideoFromComponents(
+            VideoComponents(
+                images=combined_images,
+                audio=combined_audio,
+                frame_rate=comp_a.frame_rate,
+            )
+        )
+
+        return (combined_video,)
+    
+
+class LoadVideoMXD(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+        files = folder_paths.filter_files_content_types(files, ["video"])
+        return io.Schema(
+            node_id="LoadVideoMXD",
+            display_name="Load Video MXD (Auto-Reload)",
+            category="image/video",
+            description="Always reloads the newest video with the same base name each time you run.",
+            inputs=[
+                io.Combo.Input(
+                    "file",
+                    options=sorted(files),
+                    upload=io.UploadType.video,
+                    tooltip="Pick or upload the base video. The newest version will be auto-loaded on next run."
+                ),
+            ],
+            outputs=[
+                io.Video.Output("video"),
+                io.String.Output("video_path"),
+            ],
+        )
+
+    @classmethod
+    def _find_latest(cls, base_path: str) -> str:
+        base_dir, base_filename = os.path.split(base_path)
+        base_name, _ = os.path.splitext(base_filename)
+
+        related = [
+            os.path.join(base_dir, f)
+            for f in os.listdir(base_dir)
+            if f.startswith(base_name) and os.path.isfile(os.path.join(base_dir, f))
+        ]
+        if not related:
+            return base_path
+        newest = max(related, key=os.path.getmtime)
+        return newest
+
+    @classmethod
+    def execute(cls, file):
+        video_path = folder_paths.get_annotated_filepath(file)
+        latest = cls._find_latest(video_path)
+        if latest != video_path:
+            print(f"[LoadVideoMXD] Reloading latest: {os.path.basename(latest)}")
+        return io.NodeOutput(VideoFromFile(latest), latest)
+
+    # 🔥 This is the missing piece — forces reload whenever a newer file exists
+    @classmethod
+    def fingerprint_inputs(cls, file):
+        video_path = folder_paths.get_annotated_filepath(file)
+        latest = cls._find_latest(video_path)
+        return os.path.getmtime(latest)
+
+
+class LoadVideoMXD(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+        files = folder_paths.filter_files_content_types(files, ["video"])
+        return io.Schema(
+            node_id="LoadVideoMXD",
+            display_name="Load Video MXD",
+            category="image/video",
+            description="Always reloads the newest video with the same base name each time you run.",
+            inputs=[
+                io.Combo.Input(
+                    "file",
+                    options=sorted(files),
+                    upload=io.UploadType.video,
+                    tooltip="Pick or upload the base video. The newest version will be auto-loaded on next run."
+                ),
+            ],
+            outputs=[
+                io.Video.Output("video"),
+                io.String.Output("video_path"),
+            ],
+        )
+
+    @classmethod
+    def _find_latest(cls, base_path: str) -> str:
+        base_dir, base_filename = os.path.split(base_path)
+        base_name, _ = os.path.splitext(base_filename)
+        related = [
+            os.path.join(base_dir, f)
+            for f in os.listdir(base_dir)
+            if f.startswith(base_name) and os.path.isfile(os.path.join(base_dir, f))
+        ]
+        if not related:
+            return base_path
+        return max(related, key=os.path.getmtime)
+
+    @classmethod
+    def execute(cls, file):
+        video_path = folder_paths.get_annotated_filepath(file)
+        latest = cls._find_latest(video_path)
+        if latest != video_path:
+            print(f"[LoadVideoMXD] Reloading latest: {os.path.basename(latest)}")
+        return io.NodeOutput(VideoFromFile(latest), latest)
+
+    @classmethod
+    def fingerprint_inputs(cls, file):
+        video_path = folder_paths.get_annotated_filepath(file)
+        latest = cls._find_latest(video_path)
+        return os.path.getmtime(latest)
+
+
+class SaveVideoMXD(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SaveVideoMXD",
+            display_name="Save Video MXD",
+            category="image/video",
+            description="Saves a new version of the video next to the original, auto-incrementing filenames cleanly.",
+            inputs=[
+                io.Video.Input("video"),
+                io.String.Input("video_path"),
+                io.Combo.Input("save_to_outputs", options=[False, True], default=False),
+                io.Combo.Input("format", options=VideoContainer.as_input(), default="auto"),
+                io.Combo.Input("codec", options=VideoCodec.as_input(), default="auto"),
+            ],
+            outputs=[],
+            hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, video: VideoInput, video_path: str, save_to_outputs: bool, format: str, codec: str):
+        base_dir, base_filename = os.path.split(video_path)
+        base_name, ext = os.path.splitext(base_filename)
+
+        # 🧹 Clean trailing counters like "__001__002" → remove them all
+        base_clean = re.sub(r'(__\d+)+$', '', base_name)
+
+        # 🧮 Now find the next available counter
+        pattern = re.compile(rf"^{re.escape(base_clean)}__(\d+){re.escape(ext)}$")
+        existing = [
+            int(m.group(1))
+            for f in os.listdir(base_dir)
+            if (m := pattern.match(f))
+        ]
+        next_counter = max(existing, default=0) + 1
+
+        new_filename = f"{base_clean}__{next_counter:03d}{ext}"
+        save_path = os.path.join(base_dir, new_filename)
+
+        # 💾 Metadata
+        saved_metadata = None
+        if not args.disable_metadata:
+            metadata = {}
+            if cls.hidden.extra_pnginfo is not None:
+                metadata.update(cls.hidden.extra_pnginfo)
+            if cls.hidden.prompt is not None:
+                metadata["prompt"] = cls.hidden.prompt
+            if metadata:
+                saved_metadata = metadata
+
+        # 🚀 Save main copy
+        video.save_to(save_path, format=format, codec=codec, metadata=saved_metadata)
+
+        # 🪣 Optional copy to outputs folder
+        if save_to_outputs:
+            out_dir = folder_paths.get_output_directory()
+            os.makedirs(out_dir, exist_ok=True)
+            alt_path = os.path.join(out_dir, new_filename)
+            video.save_to(alt_path, format=format, codec=codec, metadata=saved_metadata)
+            print(f"[SaveVideoMXD] Also saved copy to outputs: {alt_path}")
+
+        print(f"[SaveVideoMXD] Saved clean new version: {new_filename}")
+        rel_folder = os.path.relpath(base_dir, folder_paths.get_output_directory())
+        return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(new_filename, rel_folder, io.FolderType.output)]))
+    
+class GroupVideoFramesMXD:
+    CATEGORY = "MXD/Video"
+    TITLE = "Group Video Frames (MXD)"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("IMAGE_GROUPS",)
+    OUTPUT_IS_LIST = (True,)
+    FUNCTION = "group_frames"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "frames": ("IMAGE",),
+                "group_size": ("INT", {"default": 81, "min": 1, "max": 5000, "step": 1}),
+            }
+        }
+
+    def group_frames(self, frames, group_size):
+        import math, torch
+
+        all_frames = list(frames)
+        total = len(all_frames)
+        num_groups = math.ceil(total / group_size)
+        grouped_tensors = []
+
+        for i in range(num_groups):
+            start = i * group_size
+            end = min(start + group_size, total)
+            group = all_frames[start:end]
+
+            clean = []
+            for f in group:
+                # ✅ drop redundant singleton batch dim if present
+                if f.ndim == 4 and f.shape[0] == 1:
+                    f = f.squeeze(0)  # (H,W,C)
+                # ✅ ensure shape (H,W,C)
+                if f.ndim != 3:
+                    print(f"[GroupVideoFramesMXD] weird frame shape {f.shape}")
+                    continue
+                clean.append(f)
+
+            # ✅ stack back to (N,H,W,C)
+            if len(clean) == 0:
+                continue
+            stacked = torch.stack(clean, dim=0)
+            grouped_tensors.append(stacked)
+
+        print(f"[GroupVideoFramesMXD] Split {total} frames into {len(grouped_tensors)} groups of up to {group_size}.")
+        return (grouped_tensors,)
+
+import os
+import torch
+from comfy_api.latest import io, ui
+from comfy_api.util import VideoComponents, VideoContainer, VideoCodec
+from comfy_api.input_impl import VideoFromFile, VideoFromComponents
+import folder_paths
+from comfy.cli_args import args
+
+
+class SaveAndMergeWhenComplete_MXD(io.ComfyNode):
+    """
+    Saves batched video parts and merges them automatically once all expected parts are present.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SaveAndMergeWhenComplete_MXD",
+            display_name="Save & Merge When Complete (MXD)",
+            category="MXD/video",
+            description="Saves each incoming video part, and merges all parts once the expected count is reached.",
+            inputs=[
+                io.Video.Input("video", tooltip="Video to save (one per batch item)."),
+                io.String.Input(
+                    "folder_name",
+                    default="temp_batch_merge",
+                    tooltip="Subfolder under output/video/ to save temporary parts."
+                ),
+                io.Int.Input(
+                    "expected_parts",
+                    default=2,
+                    min=1,
+                    max=9999,
+                    tooltip="Total number of parts to wait for before merging."
+                ),
+                io.Combo.Input(
+                    "format",
+                    options=VideoContainer.as_input(),
+                    default="auto",
+                    tooltip="Container format for the saved videos."
+                ),
+                io.Combo.Input(
+                    "codec",
+                    options=VideoCodec.as_input(),
+                    default="auto",
+                    tooltip="Codec to use for the video."
+                ),
+            ],
+            outputs=[
+                io.String.Output("final_video_path", tooltip="Full path of the merged video (only once complete)."),
+            ],
+            hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
+            is_output_node=False,
+        )
+
+    @classmethod
+    def execute(cls, video, folder_name, expected_parts, format, codec) -> io.NodeOutput:
+        output_dir = os.path.join(folder_paths.get_output_directory(), "video", folder_name)
+        os.makedirs(output_dir, exist_ok=True)
+
+        # --- Save incoming video part ---
+        part_index = len([f for f in os.listdir(output_dir) if f.endswith(".mp4")])
+        part_path = os.path.join(output_dir, f"part_{part_index+1:03d}.mp4")
+
+        # --- Metadata (same as SaveVideo) ---
+        saved_metadata = None
+        if not args.disable_metadata:
+            metadata = {}
+            if cls.hidden.extra_pnginfo is not None:
+                metadata.update(cls.hidden.extra_pnginfo)
+            if cls.hidden.prompt is not None:
+                metadata["prompt"] = cls.hidden.prompt
+            if metadata:
+                saved_metadata = metadata
+
+        video.save_to(part_path, format=format, codec=codec, metadata=saved_metadata)
+        print(f"[SaveAndMergeWhenComplete_MXD] 💾 Saved part {part_index+1}/{expected_parts} → {part_path}")
+
+        # --- Check how many parts exist ---
+        part_files = sorted([
+            os.path.join(output_dir, f)
+            for f in os.listdir(output_dir)
+            if f.lower().endswith(".mp4")
+        ])
+
+        if len(part_files) < expected_parts:
+            print(f"[SaveAndMergeWhenComplete_MXD] Waiting for all parts ({len(part_files)}/{expected_parts})...")
+            return io.NodeOutput(final_video_path="")  # Not ready yet
+
+        # ✅ All parts present → merge them
+        print(f"[SaveAndMergeWhenComplete_MXD] All {expected_parts} parts found. Starting merge...")
+
+        comp_ref = VideoFromFile(part_files[0]).get_components()
+        all_frames = []
+        all_audio = []
+        frame_rate = comp_ref.frame_rate
+
+        for path in part_files:
+            vid = VideoFromFile(path)
+            comp = vid.get_components()
+            frames = torch.stack(comp.images) if isinstance(comp.images, list) else comp.images
+            all_frames.append(frames)
+            if comp.audio is not None:
+                all_audio.append(comp.audio)
+
+        merged_frames = torch.cat(all_frames, dim=0)
+        merged_audio = torch.cat(all_audio, dim=1) if all_audio else None
+
+        combined_video = VideoFromComponents(
+            VideoComponents(images=merged_frames, audio=merged_audio, frame_rate=frame_rate)
+        )
+
+        final_path = os.path.join(output_dir, "merged_final.mp4")
+        combined_video.save_to(final_path, format=format, codec=codec, metadata=saved_metadata)
+
+        print(f"[SaveAndMergeWhenComplete_MXD] ✅ Merged {expected_parts} parts → {final_path}")
+
+        return io.NodeOutput(final_path)
+
+
 
 # ---------- Node registration ----------
 NODE_CLASS_MAPPINGS = {
@@ -1395,6 +1860,12 @@ NODE_CLASS_MAPPINGS = {
     "LoadLatents_FromFolder_I2V_MXD": LoadLatents_FromFolder_I2V_MXD,
     "WanImageToVideoMXD": WanImageToVideoMXD,
     "WAN22_I2V_Image_Scaler_MXD": WAN22_I2V_Image_Scaler_MXD,
+    "Frames_Select_End_MXD": Frames_Select_End_MXD,
+    "Frames_Remove_From_Start_MXD": Frames_Remove_From_Start_MXD,
+    "CombineVideos_MXD": CombineVideos_MXD,
+    "LoadVideoMXD": LoadVideoMXD,
+    "SaveVideoMXD": SaveVideoMXD,
+    "GroupVideoFramesMXD": GroupVideoFramesMXD,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -1408,4 +1879,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "LoadLatents_FromFolder_I2V_MXD": "Load Latent Batch I2V MXD",
     "WanImageToVideoMXD": "WAN Image to Video MXD",
     "WAN22_I2V_Image_Scaler_MXD": "WAN 2.2 I2V Image Scaler MXD",
+    "Frames_Select_End_MXD": "Frames Select End MXD",
+    "Frames_Remove_From_Start_MXD": "Frames Remove From Start MXD",
+    "CombineVideos_MXD": "Combine Videos MXD",
+    "LoadVideoMXD": "Load Video MXD",
+    "SaveVideoMXD": "Save Video MXD",
+    "GroupVideoFramesMXD": "Group Video Frames MXD",
 }