Alvidi-timeline_viewer_AV_comfyUI/nodes.py

import os
import shutil
import subprocess
import tempfile
import numpy as np
import folder_paths


def get_ffmpeg():
    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:
        pass
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg:
        return ffmpeg
    raise RuntimeError("[TimelineViewer] ffmpeg no encontrado.")


class TimelineViewer:

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images":          ("IMAGE",),
                "frame_rate":      ("FLOAT",   {"default": 24.0, "min": 1.0, "max": 120.0, "step": 0.1}),
                "loop":            ("BOOLEAN", {"default": False}),
                "filename_prefix": ("STRING",  {"default": "TimelineViewer"}),
                "save_output":     ("BOOLEAN", {"default": True}),
                "frame_in":        ("INT",     {"default": 0, "min": 0, "max": 99999}),
                "frame_out":       ("INT",     {"default": 0, "min": 0, "max": 99999,
                                                "tooltip": "0 = hasta el final"}),
                "max_width":       ("INT",     {"default": 1920, "min": 256, "max": 4096, "step": 2,
                                                "tooltip": "Ancho maximo del preview"}),
            },
            "optional": {
                "audio": ("AUDIO",),
            }
        }

    RETURN_TYPES   = ()
    OUTPUT_NODE    = True
    FUNCTION       = "combine_video"
    CATEGORY       = "Video"

    def combine_video(self, images, frame_rate, loop, filename_prefix,
                      save_output, frame_in, frame_out, max_width, audio=None):

        ffmpeg = get_ffmpeg()

        frames = (images.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
        total_frames = len(frames)

        f_in  = max(0, frame_in)
        f_out = frame_out if frame_out > 0 else total_frames
        f_out = min(f_out, total_frames)
        if f_in >= f_out:
            f_in, f_out = 0, total_frames

        frames = frames[f_in:f_out]
        frame_count = len(frames)
        H, W = frames.shape[1], frames.shape[2]

        if W > max_width:
            from PIL import Image as PILImage
            scale = max_width / W
            new_W = max_width
            new_H = int(H * scale)
            new_H = new_H if new_H % 2 == 0 else new_H + 1
            new_W = new_W if new_W % 2 == 0 else new_W + 1
            print(f"[TimelineViewer] Scaling {W}x{H} -> {new_W}x{new_H}")
            resized = []
            for f in frames:
                img = PILImage.fromarray(f).resize((new_W, new_H), PILImage.LANCZOS)
                resized.append(np.array(img))
            frames = np.stack(resized)
            W, H = new_W, new_H

        print(f"[TimelineViewer] {frame_count} frames ({f_in}-{f_out}) @ {frame_rate}fps  {W}x{H}")

        if save_output:
            out_dir  = folder_paths.get_output_directory()
            out_type = "output"
        else:
            out_dir  = folder_paths.get_temp_directory()
            out_type = "temp"

        os.makedirs(out_dir, exist_ok=True)

        counter = 0
        while True:
            filename   = f"{filename_prefix}_{counter:05d}.mp4"
            final_path = os.path.join(out_dir, filename)
            if not os.path.exists(final_path):
                break
            counter += 1

        tmp_dir  = tempfile.mkdtemp(dir="/tmp")
        tmp_path = os.path.join(tmp_dir, filename)

        audio_path    = None
        waveform_data = []

        if audio is not None:
            try:
                import torchaudio
                wf          = audio["waveform"].squeeze(0)
                sample_rate = audio["sample_rate"]
                wf_cut      = wf
                if f_in > 0 or f_out < total_frames:
                    s_in  = int((f_in  / frame_rate) * sample_rate)
                    s_out = int((f_out / frame_rate) * sample_rate)
                    wf_cut = wf[:, s_in:s_out]

                audio_path = os.path.join(tmp_dir, "audio.wav")
                torchaudio.save(audio_path, wf_cut, sample_rate)
                print(f"[TimelineViewer] Audio OK")

                # Calcular waveform para el frontend
                mono       = wf_cut.mean(dim=0).abs().cpu().numpy()
                samples    = 300
                block_size = max(1, len(mono) // samples)
                raw_wave   = []
                for i in range(samples):
                    chunk = mono[i * block_size:(i + 1) * block_size]
                    raw_wave.append(float(chunk.mean()) if len(chunk) > 0 else 0.0)
                max_val       = max(raw_wave) if max(raw_wave) > 0 else 1.0
                waveform_data = [round(v / max_val, 4) for v in raw_wave]
                print(f"[TimelineViewer] Waveform OK ({len(waveform_data)} samples)")

            except Exception as e:
                print(f"[TimelineViewer] Audio error: {e}")
                audio_path = None

        cmd = [
            ffmpeg, "-y",
            "-f", "rawvideo",
            "-vcodec", "rawvideo",
            "-s", f"{W}x{H}",
            "-pix_fmt", "rgb24",
            "-r", str(frame_rate),
            "-i", "pipe:0",
        ]
        if audio_path:
            cmd += ["-i", audio_path, "-shortest"]
        cmd += ["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", tmp_path]

        try:
            raw_data = frames.tobytes()
            result   = subprocess.run(
                cmd,
                input=raw_data,
                capture_output=True,
                timeout=300
            )
            if result.returncode != 0:
                raise RuntimeError(f"[TimelineViewer] ffmpeg error:\n{result.stderr.decode()}")
        except Exception as e:
            shutil.rmtree(tmp_dir, ignore_errors=True)
            raise e

        shutil.move(tmp_path, final_path)
        shutil.rmtree(tmp_dir, ignore_errors=True)

        print(f"[TimelineViewer] Saved: {final_path}")

        return {
            "ui": {
                "videos": [{
                    "filename":    filename,
                    "type":        out_type,
                    "format":      "video/mp4",
                    "frame_rate":  frame_rate,
                    "frame_count": frame_count,
                    "frame_in":    f_in,
                    "loop":        loop,
                    "waveform":    waveform_data,
                }]
            }
        }


NODE_CLASS_MAPPINGS = {
    "TimelineViewer": TimelineViewer
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TimelineViewer": "Timeline Viewer 🎬"
}