latenightlabs-ComfyUI-LNL/modules/video_utils.py

import subprocess
import shutil
import os
import re
import json
from collections.abc import Mapping

import cv2
import numpy as np
import torch

from PIL import Image

from folder_paths import base_path

"""
Attribution: ComfyUI-VideoHelperSuite

Portions of this code are adapted from GitHub repository `https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite`,
which is licensed under the GNU General Public License version 3 (GPL-3.0):

"""

_AUDIO_STREAM_PROBE_CACHE = {}

def __lnl_ffmpeg_suitability(path):
    try:
        version = subprocess.run([path, "-version"], check=True,
                                 capture_output=True).stdout.decode("utf-8")
    except:
        return 0
    score = 0
    #rough layout of the importance of various features
    simple_criterion = [("libvpx", 20),("264",10), ("265",3),
                        ("svtav1",5),("libopus", 1)]
    for criterion in simple_criterion:
        if version.find(criterion[0]) >= 0:
            score += criterion[1]
    #obtain rough compile year from copyright information
    copyright_index = version.find('2000-2')
    if copyright_index >= 0:
        copyright_year = version[copyright_index+6:copyright_index+9]
        if copyright_year.isnumeric():
            score += int(copyright_year)
    return score

def lnl_get_audio(file, start_time=0, duration=0):
    if ffmpeg_path is None:
        return b""
    args = [ffmpeg_path, "-v", "error", "-nostdin", "-i", file, "-map", "0:a:0", "-vn"]
    if start_time > 0:
        args += ["-ss", str(start_time)]
    if duration > 0:
        args += ["-t", str(duration)]
    try:
        return subprocess.run(args + ["-f", "wav", "-"],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True).stdout
    except subprocess.CalledProcessError as e:
        stderr = e.stderr.decode("utf-8", errors="ignore") if e.stderr else ""
        stdout = e.stdout.decode("utf-8", errors="ignore") if e.stdout else ""
        combined = (stderr + "\n" + stdout).strip()
        if _lnl_is_no_audio_error(combined):
            return b""
        raise
    except OSError:
        return b""

def lnl_lazy_eval(func):
    class Cache:
        def __init__(self, func):
            self.res = None
            self.func = func
        def get(self):
            if self.res is None:
                self.res = self.func()
            return self.res
    cache = Cache(func)
    return lambda : cache.get()

def _lnl_probe_audio_stream_params(file):
    cached = _AUDIO_STREAM_PROBE_CACHE.get(file)
    if cached is not None:
        return cached

    ffprobe_cmd = shutil.which("ffprobe")
    if ffprobe_cmd is None and ffmpeg_path is not None:
        ffprobe_name = "ffprobe.exe" if os.name == "nt" else "ffprobe"
        ffprobe_candidate = os.path.join(os.path.dirname(ffmpeg_path), ffprobe_name)
        if os.path.exists(ffprobe_candidate):
            ffprobe_cmd = ffprobe_candidate

    if ffprobe_cmd is None:
        _AUDIO_STREAM_PROBE_CACHE[file] = (None, None)
        return (None, None)

    cmd = [
        ffprobe_cmd,
        "-v", "error",
        "-select_streams", "a:0",
        "-show_entries", "stream=sample_rate,channels",
        "-of", "json",
        file,
    ]
    try:
        process = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
        data = json.loads(process.stdout) if process.stdout else {}
        stream = (data.get("streams") or [None])[0]
        if not isinstance(stream, dict):
            _AUDIO_STREAM_PROBE_CACHE[file] = (None, None)
            return (None, None)
        sample_rate_raw = stream.get("sample_rate")
        channels_raw = stream.get("channels")
        sample_rate = int(sample_rate_raw) if str(sample_rate_raw).isdigit() else None
        channels = int(channels_raw) if isinstance(channels_raw, int) else None
        _AUDIO_STREAM_PROBE_CACHE[file] = (sample_rate, channels)
        return (sample_rate, channels)
    except Exception:
        _AUDIO_STREAM_PROBE_CACHE[file] = (None, None)
        return (None, None)

def _lnl_get_audio(file, start_time=0, duration=0):
    if ffmpeg_path is None:
        return lnl_empty_audio_dict()
    args = [ffmpeg_path, "-v", "error", "-nostdin", "-i", file, "-map", "0:a:0", "-vn"]
    if start_time > 0:
        args += ["-ss", str(start_time)]
    if duration > 0:
        args += ["-t", str(duration)]
    try:
        #TODO: scan for sample rate and maintain
        res =  subprocess.run(args + ["-f", "f32le", "-"],
                              capture_output=True, check=True)
        stderr_text = res.stderr.decode("utf-8", errors="ignore")
        raw_audio = res.stdout or b""
        if len(raw_audio) == 0:
            return lnl_empty_audio_dict()
        # f32le must be 4-byte aligned; truncate trailing partial bytes defensively.
        remainder = len(raw_audio) % 4
        if remainder:
            raw_audio = raw_audio[: len(raw_audio) - remainder]
        if len(raw_audio) == 0:
            return lnl_empty_audio_dict()
        audio = torch.frombuffer(bytearray(raw_audio), dtype=torch.float32)
        if audio.numel() == 0:
            return lnl_empty_audio_dict()
        match = re.search(', (\\d+) Hz, (\\w+), ', stderr_text)
    except subprocess.CalledProcessError as e:
        stderr = e.stderr.decode("utf-8", errors="ignore") if e.stderr else ""
        stdout = e.stdout.decode("utf-8", errors="ignore") if e.stdout else ""
        combined = (stderr + "\n" + stdout).strip()
        if _lnl_is_no_audio_error(combined):
            return lnl_empty_audio_dict()
        raise Exception(f"VHS failed to extract audio from {file}:\n" \
                + (combined or stderr))
    except OSError:
        return lnl_empty_audio_dict()
    if match:
        ar = int(match.group(1))
        # NOTE: Just throwing an error for other channel types right now
        # Will deal with issues if they come
        ac = {"mono": 1, "stereo": 2}.get(match.group(2), 2)
    else:
        probed_ar, probed_ac = _lnl_probe_audio_stream_params(file)
        ar = probed_ar if probed_ar and probed_ar > 0 else 44100
        ac = probed_ac if probed_ac and probed_ac > 0 else 2
    if ac <= 0:
        ac = 2
    usable_values = (audio.numel() // ac) * ac
    if usable_values <= 0:
        return lnl_empty_audio_dict(ar)
    if usable_values != audio.numel():
        audio = audio[:usable_values]
    audio = audio.reshape((-1,ac)).transpose(0,1).unsqueeze(0)
    return {'waveform': audio, 'sample_rate': ar}

def lnl_empty_audio_dict(sample_rate=44100):
    return {'waveform': torch.zeros((1, 1, 0), dtype=torch.float32), 'sample_rate': sample_rate}

def _lnl_is_no_audio_error(stderr):
    if not stderr:
        return False
    text = stderr.lower()
    return ("audio" not in text and "video" in text) \
        or "matches no streams" in text \
        or "no audio" in text \
        or "does not contain any stream" in text \
        or "output file does not contain any stream" in text \
        or ("error opening output file" in text and "pipe:" in text)

class LNLLazyAudioMap(Mapping):
    def __init__(self, file, start_time, duration):
        self.file = file
        self.start_time=start_time
        self.duration=duration
        self._dict=None
    def __getitem__(self, key):
        if self._dict is None:
            try:
                self._dict = _lnl_get_audio(self.file, self.start_time, self.duration)
            except Exception:
                self._dict = lnl_empty_audio_dict()
        return self._dict[key]
    def __iter__(self):
        if self._dict is None:
            try:
                self._dict = _lnl_get_audio(self.file, self.start_time, self.duration)
            except Exception:
                self._dict = lnl_empty_audio_dict()
        return iter(self._dict)
    def __len__(self):
        if self._dict is None:
            try:
                self._dict = _lnl_get_audio(self.file, self.start_time, self.duration)
            except Exception:
                self._dict = lnl_empty_audio_dict()
        return len(self._dict)

def lnl_lazy_get_audio(file, start_time=0, duration=0):
    return LNLLazyAudioMap(file, start_time, duration)

def lnl_cv_frame_generator(video, number_of_frames_to_process, skip_first_frames, select_every_nth):
    try:
        video_cap = cv2.VideoCapture(video)
        if not video_cap.isOpened():
            raise ValueError(f"{video} could not be loaded with cv.")
        # set video_cap to look at start_index frame
        total_frame_count = 0
        total_frames_evaluated = -1
        frames_added = 0
        base_frame_time = 1/video_cap.get(cv2.CAP_PROP_FPS)
        width = video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        height = video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        prev_frame = None

        target_frame_time = base_frame_time
        yield (width, height, target_frame_time)
        
        time_offset=target_frame_time - base_frame_time
        while video_cap.isOpened():
            if time_offset < target_frame_time:
                is_returned = video_cap.grab()
                # if didn't return frame, video has ended
                if not is_returned:
                    break
                time_offset += base_frame_time
            if time_offset < target_frame_time:
                continue
            time_offset -= target_frame_time
            # if not at start_index, skip doing anything with frame
            total_frame_count += 1
            if total_frame_count < skip_first_frames:
                continue
            else:
                total_frames_evaluated += 1

            # if should not be selected, skip doing anything with frame
            if total_frames_evaluated%select_every_nth != 0:
                frames_added += 1
                if total_frame_count >= number_of_frames_to_process + skip_first_frames:
                    break
                continue

            # opencv loads images in BGR format (yuck), so need to convert to RGB for ComfyUI use
            # follow up: can videos ever have an alpha channel?
            # To my testing: No. opencv has no support for alpha
            unused, frame = video_cap.retrieve()
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # convert frame to comfyui's expected format
            # TODO: frame contains no exif information. Check if opencv2 has already applied
            frame = np.array(frame, dtype=np.float32) / 255.0
            if prev_frame is not None:
                inp  = yield prev_frame
                if inp is not None:
                    #ensure the finally block is called
                    return
            prev_frame = frame
            frames_added += 1
            # if cap exists and we've reached it, stop processing frames
            if number_of_frames_to_process > 0 and frames_added >= number_of_frames_to_process:
                break
        if prev_frame is not None:
            yield prev_frame
    finally:
        video_cap.release()

def lnl_bislerp(samples, width, height):
    def slerp(b1, b2, r):
        '''slerps batches b1, b2 according to ratio r, batches should be flat e.g. NxC'''
        
        c = b1.shape[-1]

        #norms
        b1_norms = torch.norm(b1, dim=-1, keepdim=True)
        b2_norms = torch.norm(b2, dim=-1, keepdim=True)

        #normalize
        b1_normalized = b1 / b1_norms
        b2_normalized = b2 / b2_norms

        #zero when norms are zero
        b1_normalized[b1_norms.expand(-1,c) == 0.0] = 0.0
        b2_normalized[b2_norms.expand(-1,c) == 0.0] = 0.0

        #slerp
        dot = (b1_normalized*b2_normalized).sum(1)
        omega = torch.acos(dot)
        so = torch.sin(omega)

        #technically not mathematically correct, but more pleasing?
        res = (torch.sin((1.0-r.squeeze(1))*omega)/so).unsqueeze(1)*b1_normalized + (torch.sin(r.squeeze(1)*omega)/so).unsqueeze(1) * b2_normalized
        res *= (b1_norms * (1.0-r) + b2_norms * r).expand(-1,c)

        #edge cases for same or polar opposites
        res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5] 
        res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1]
        return res
    
    def generate_bilinear_data(length_old, length_new, device):
        coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1))
        coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear")
        ratios = coords_1 - coords_1.floor()
        coords_1 = coords_1.to(torch.int64)
        
        coords_2 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1)) + 1
        coords_2[:,:,:,-1] -= 1
        coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear")
        coords_2 = coords_2.to(torch.int64)
        return ratios, coords_1, coords_2

    orig_dtype = samples.dtype
    samples = samples.float()
    n,c,h,w = samples.shape
    h_new, w_new = (height, width)
    
    #linear w
    ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device)
    coords_1 = coords_1.expand((n, c, h, -1))
    coords_2 = coords_2.expand((n, c, h, -1))
    ratios = ratios.expand((n, 1, h, -1))

    pass_1 = samples.gather(-1,coords_1).movedim(1, -1).reshape((-1,c))
    pass_2 = samples.gather(-1,coords_2).movedim(1, -1).reshape((-1,c))
    ratios = ratios.movedim(1, -1).reshape((-1,1))

    result = slerp(pass_1, pass_2, ratios)
    result = result.reshape(n, h, w_new, c).movedim(-1, 1)

    #linear h
    ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device)
    coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
    coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
    ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new))

    pass_1 = result.gather(-2,coords_1).movedim(1, -1).reshape((-1,c))
    pass_2 = result.gather(-2,coords_2).movedim(1, -1).reshape((-1,c))
    ratios = ratios.movedim(1, -1).reshape((-1,1))

    result = slerp(pass_1, pass_2, ratios)
    result = result.reshape(n, h_new, w_new, c).movedim(-1, 1)
    return result.to(orig_dtype)

def lnl_lanczos(samples, width, height):
    images = [Image.fromarray(np.clip(255. * image.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8)) for image in samples]
    images = [image.resize((width, height), resample=Image.Resampling.LANCZOS) for image in images]
    images = [torch.from_numpy(np.array(image).astype(np.float32) / 255.0).movedim(-1, 0) for image in images]
    result = torch.stack(images)
    return result.to(samples.device, samples.dtype)

def lnl_common_upscale(samples, width, height, upscale_method, crop):
        if crop == "center":
            old_width = samples.shape[3]
            old_height = samples.shape[2]
            old_aspect = old_width / old_height
            new_aspect = width / height
            x = 0
            y = 0
            if old_aspect > new_aspect:
                x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)
            elif old_aspect < new_aspect:
                y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)
            s = samples[:,:,y:old_height-y,x:old_width-x]
        else:
            s = samples

        if upscale_method == "bislerp":
            return lnl_bislerp(s, width, height)
        elif upscale_method == "lanczos":
            return lnl_lanczos(s, width, height)
        else:
            return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)

def lnl_target_size(width, height, force_size, custom_width, custom_height) -> tuple[int, int]:
    if force_size == "Custom":
        return (custom_width, custom_height)
    elif force_size == "Custom Height":
        force_size = "?x"+str(custom_height)
    elif force_size == "Custom Width":
        force_size = str(custom_width)+"x?"

    if force_size != "Disabled":
        force_size = force_size.split("x")
        if force_size[0] == "?":
            width = (width*int(force_size[1]))//height
            #Limit to a multple of 8 for latent conversion
            width = int(width)+4 & ~7
            height = int(force_size[1])
        elif force_size[1] == "?":
            height = (height*int(force_size[0]))//width
            height = int(height)+4 & ~7
            width = int(force_size[0])
        else:
            width = int(force_size[0])
            height = int(force_size[1])
    return (width, height)

def get_video_info(video_path):
    full_video_path = video_path
    if not os.path.isabs(full_video_path) and not os.path.exists(full_video_path):
        full_video_path = os.path.join(base_path, video_path)
    if not os.path.exists(full_video_path):
        raise Exception(f"Video path does not exist: {full_video_path}")

    frame_rate = None
    total_frames = None
    duration = None

    ffprobe_cmd = shutil.which("ffprobe")
    if ffprobe_cmd is None and ffmpeg_path is not None:
        ffprobe_name = "ffprobe.exe" if os.name == "nt" else "ffprobe"
        ffprobe_candidate = os.path.join(os.path.dirname(ffmpeg_path), ffprobe_name)
        if os.path.exists(ffprobe_candidate):
            ffprobe_cmd = ffprobe_candidate

    if ffprobe_cmd is not None:
        cmd = [
            ffprobe_cmd, '-v', 'error', '-select_streams', 'v:0',
            '-count_frames',
            '-show_entries', 'stream=avg_frame_rate,r_frame_rate,nb_frames,nb_read_frames,duration',
            '-show_entries', 'format=duration',
            '-of', 'json',
            full_video_path,
        ]
        process = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        try:
            data = json.loads(process.stdout) if process.stdout else {}
        except json.JSONDecodeError:
            data = {}

        stream = None
        streams = data.get("streams") or []
        if streams:
            stream = streams[0]

        def _parse_rate(rate_value):
            if rate_value is None:
                return None
            if isinstance(rate_value, (int, float)):
                return float(rate_value)
            if isinstance(rate_value, str):
                if "/" in rate_value:
                    try:
                        num, den = map(float, rate_value.split("/", 1))
                        if den != 0:
                            return num / den
                    except ValueError:
                        return None
                try:
                    return float(rate_value)
                except ValueError:
                    return None
            return None

        if stream:
            frame_rate = _parse_rate(stream.get("avg_frame_rate")) or _parse_rate(stream.get("r_frame_rate"))
            nb_frames = stream.get("nb_frames")
            nb_read_frames = stream.get("nb_read_frames")
            stream_duration = stream.get("duration")

            def _parse_int(value):
                if isinstance(value, (int, float)):
                    return int(value)
                if isinstance(value, str) and value.isdigit():
                    return int(value)
                return None

            parsed_read = _parse_int(nb_read_frames)
            parsed_meta = _parse_int(nb_frames)
            if parsed_read is not None:
                total_frames = parsed_read
            elif parsed_meta is not None:
                total_frames = parsed_meta

            try:
                if duration is None and stream_duration is not None:
                    duration = float(stream_duration)
            except ValueError:
                duration = None

        if duration is None:
            try:
                duration = float(data.get("format", {}).get("duration"))
            except (TypeError, ValueError):
                duration = None

        if total_frames is None and frame_rate and duration:
            total_frames = int(round(duration * frame_rate))

    if frame_rate is None or total_frames is None or duration is None:
        cap = cv2.VideoCapture(full_video_path)
        if cap.isOpened():
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
            if frame_rate is None and fps > 0:
                frame_rate = fps
            if total_frames is None and frame_count > 0:
                total_frames = int(frame_count)
            if duration is None and fps > 0 and frame_count > 0:
                duration = frame_count / fps

            if ffprobe_cmd is None:
                try:
                    cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                    precise_count = 0
                    while True:
                        ok = cap.grab()
                        if not ok:
                            break
                        precise_count += 1
                    if precise_count > 0:
                        total_frames = int(precise_count)
                        if frame_rate and frame_rate > 0:
                            duration = total_frames / frame_rate
                except Exception:
                    pass
        cap.release()

    if frame_rate is None or frame_rate <= 0:
        frame_rate = 1.0
    if total_frames is None or total_frames <= 0:
        total_frames = 1
    if duration is None or duration <= 0:
        duration = total_frames / frame_rate

    return frame_rate, total_frames, duration

ffmpeg_paths = []
try:
    from imageio_ffmpeg import get_ffmpeg_exe
    imageio_ffmpeg_path = get_ffmpeg_exe()
    ffmpeg_paths.append(imageio_ffmpeg_path)
except:
    print("Failed to import imageio_ffmpeg")
system_ffmpeg = shutil.which("ffmpeg")
if system_ffmpeg is not None:
    ffmpeg_paths.append(system_ffmpeg)

if len(ffmpeg_paths) == 0:
    print("No valid ffmpeg found.")
    ffmpeg_path = None
elif len(ffmpeg_paths) == 1:
    ffmpeg_path = ffmpeg_paths[0]
else:
    ffmpeg_path = max(ffmpeg_paths, key=__lnl_ffmpeg_suitability)