styletransfer-ComfyUI-TRELLIS2_Motion/utils/video_utils.py

"""Video frame utilities for extraction and interpolation."""

import logging
from typing import List, Tuple

import torch

logger = logging.getLogger(__name__)


def extract_frames(
    video_tensor: torch.Tensor,
    skip: int = 1,
    max_frames: int = 0
) -> Tuple[List[torch.Tensor], List[int]]:
    """
    Extract frames from a video tensor with optional frame skipping.

    Args:
        video_tensor: Video frames as [N, H, W, C] tensor (ComfyUI format)
        skip: Process every Nth frame (1 = all frames, 2 = every other, etc.)
        max_frames: Maximum number of frames to extract (0 = no limit)

    Returns:
        Tuple of:
            - frames: List of [H, W, C] tensors for each extracted frame
            - indices: List of original frame indices that were extracted
    """
    total_frames = video_tensor.shape[0]

    # Generate indices with skipping
    indices = list(range(0, total_frames, skip))

    # Apply max_frames limit
    if max_frames > 0:
        indices = indices[:max_frames]

    # Extract frames
    frames = [video_tensor[i] for i in indices]

    logger.info(
        f"Extracted {len(frames)} frames from {total_frames} total "
        f"(skip={skip}, max={max_frames})"
    )

    return frames, indices


def interpolate_missing_frames(
    rendered_frames: List[torch.Tensor],
    original_indices: List[int],
    total_frames: int
) -> torch.Tensor:
    """
    Interpolate frames that were skipped during processing.

    Uses linear interpolation between the nearest processed frames to fill in gaps.

    Args:
        rendered_frames: List of rendered [H, W, C] tensors for processed frames
        original_indices: The original frame indices that were processed
        total_frames: Total number of frames in the original video

    Returns:
        Full video tensor [N, H, W, C] with all frames filled in
    """
    if len(original_indices) == total_frames:
        # No interpolation needed
        return torch.stack(rendered_frames)

    if len(rendered_frames) == 0:
        raise ValueError("No rendered frames provided")

    # Get shape from first frame
    H, W, C = rendered_frames[0].shape
    device = rendered_frames[0].device
    dtype = rendered_frames[0].dtype

    # Create output tensor
    output = torch.zeros(total_frames, H, W, C, dtype=dtype, device=device)

    # Create a set for faster lookup
    processed_set = set(original_indices)

    # Fill in processed frames
    for i, idx in enumerate(original_indices):
        output[idx] = rendered_frames[i]

    # Interpolate gaps
    for frame_idx in range(total_frames):
        if frame_idx in processed_set:
            continue

        # Find nearest processed frames before and after
        prev_idx = None
        next_idx = None

        for idx in original_indices:
            if idx < frame_idx:
                prev_idx = idx
            elif idx > frame_idx:
                next_idx = idx
                break

        # Handle edge cases
        if prev_idx is None:
            # Before first processed frame - use first frame
            output[frame_idx] = output[original_indices[0]]
        elif next_idx is None:
            # After last processed frame - use last frame
            output[frame_idx] = output[original_indices[-1]]
        else:
            # Interpolate between prev and next
            # Find which indices in original_indices correspond to prev/next
            prev_frame = output[prev_idx]
            next_frame = output[next_idx]

            # Linear interpolation factor
            t = (frame_idx - prev_idx) / (next_idx - prev_idx)

            # Interpolate pixel values
            output[frame_idx] = prev_frame * (1 - t) + next_frame * t

    logger.info(
        f"Interpolated {total_frames - len(original_indices)} missing frames "
        f"(total: {total_frames})"
    )

    return output


def batch_to_list(batch_tensor: torch.Tensor) -> List[torch.Tensor]:
    """
    Convert a batch tensor to a list of individual tensors.

    Args:
        batch_tensor: [N, H, W, C] tensor

    Returns:
        List of N tensors, each [H, W, C]
    """
    return [batch_tensor[i] for i in range(batch_tensor.shape[0])]


def list_to_batch(frame_list: List[torch.Tensor]) -> torch.Tensor:
    """
    Convert a list of frame tensors to a batch tensor.

    Args:
        frame_list: List of [H, W, C] tensors

    Returns:
        Stacked tensor [N, H, W, C]
    """
    return torch.stack(frame_list)


def resize_frames(
    frames: torch.Tensor,
    width: int,
    height: int,
    mode: str = "bilinear"
) -> torch.Tensor:
    """
    Resize a batch of frames to a target resolution.

    Args:
        frames: [N, H, W, C] tensor in ComfyUI format
        width: Target width
        height: Target height
        mode: Interpolation mode ("nearest", "bilinear", "bicubic")

    Returns:
        Resized [N, height, width, C] tensor
    """
    # ComfyUI format: [N, H, W, C]
    # PyTorch interpolate expects: [N, C, H, W]

    n_frames = frames.shape[0]

    # Convert to [N, C, H, W]
    frames_nchw = frames.permute(0, 3, 1, 2)

    # Resize
    resized = torch.nn.functional.interpolate(
        frames_nchw,
        size=(height, width),
        mode=mode,
        align_corners=False if mode != "nearest" else None
    )

    # Convert back to [N, H, W, C]
    resized_nhwc = resized.permute(0, 2, 3, 1)

    return resized_nhwc


def get_video_info(video_tensor: torch.Tensor) -> dict:
    """
    Get information about a video tensor.

    Args:
        video_tensor: [N, H, W, C] tensor

    Returns:
        Dictionary with video properties
    """
    n, h, w, c = video_tensor.shape

    return {
        "num_frames": n,
        "height": h,
        "width": w,
        "channels": c,
        "dtype": str(video_tensor.dtype),
        "device": str(video_tensor.device),
        "min_value": video_tensor.min().item(),
        "max_value": video_tensor.max().item(),
    }