vendor/VideoX-Fun/videox_fun/data/dataset_video.py

import csv
import json
import math
import os
import random

import cv2
import librosa
import numpy as np
import torch
import torchvision.transforms as transforms
from einops import rearrange
from func_timeout import FunctionTimedOut, func_timeout
from PIL import Image
from torch.utils.data.dataset import Dataset

try:
    from decord import VideoReader
except ImportError:
    from .utils import AVVideoReader as VideoReader

from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager,
                    get_random_mask, get_video_reader_batch, resize_frame)


class WebVid10M(Dataset):
    def __init__(
        self,
        csv_path, 
        video_folder,
        sample_size=256, 
        sample_stride=4, 
        sample_n_frames=16,
        enable_bucket=False, 
        enable_inpaint=False, 
        is_image=False,
    ):
        print(f"loading annotations from {csv_path} ...")
        with open(csv_path, 'r') as csvfile:
            self.dataset = list(csv.DictReader(csvfile))
        self.length = len(self.dataset)
        print(f"data scale: {self.length}")

        self.video_folder    = video_folder
        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.enable_bucket   = enable_bucket
        self.enable_inpaint  = enable_inpaint
        self.is_image        = is_image
        
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(sample_size[0]),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])
    
    def get_batch(self, idx):
        video_dict = self.dataset[idx]
        videoid, name, page_dir = video_dict['videoid'], video_dict['name'], video_dict['page_dir']
        
        video_dir    = os.path.join(self.video_folder, f"{videoid}.mp4")
        video_reader = VideoReader(video_dir)
        video_length = len(video_reader)
        
        if not self.is_image:
            clip_length = min(video_length, (self.sample_n_frames - 1) * self.sample_stride + 1)
            start_idx   = random.randint(0, video_length - clip_length)
            batch_index = np.linspace(start_idx, start_idx + clip_length - 1, self.sample_n_frames, dtype=int)
        else:
            batch_index = [random.randint(0, video_length - 1)]

        if not self.enable_bucket:
            pixel_values = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous()
            pixel_values = pixel_values / 255.
            del video_reader
        else:
            pixel_values = video_reader.get_batch(batch_index).asnumpy()

        if self.is_image:
            pixel_values = pixel_values[0]
        return pixel_values, name

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        while True:
            try:
                pixel_values, name = self.get_batch(idx)
                break

            except Exception as e:
                print("Error info:", e)
                idx = random.randint(0, self.length-1)

        if not self.enable_bucket:
            pixel_values = self.pixel_transforms(pixel_values)
        if self.enable_inpaint:
            mask = get_random_mask(pixel_values.size())
            mask_pixel_values = pixel_values * (1 - mask) + torch.ones_like(pixel_values) * -1 * mask
            sample = dict(pixel_values=pixel_values, mask_pixel_values=mask_pixel_values, mask=mask, text=name)
        else:
            sample = dict(pixel_values=pixel_values, text=name)
        return sample


class VideoDataset(Dataset):
    """Dataset for video training with inpainting support."""
    def __init__(
        self,
        ann_path, 
        data_root=None,
        sample_size=256, 
        sample_stride=4, 
        sample_n_frames=16,
        enable_bucket=False, 
        enable_inpaint=False,
        inpaint_mask_fill_value=0,
        video_length_drop_start=0.0,
        video_length_drop_end=1.0,
        text_drop_ratio=0.1,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
        self.dataset = json.load(open(ann_path, 'r'))
        self.length = len(self.dataset)
        print(f"data scale: {self.length}")

        self.data_root = data_root
        self.sample_stride = sample_stride
        self.sample_n_frames = sample_n_frames
        self.enable_bucket = enable_bucket
        self.enable_inpaint = enable_inpaint
        self.inpaint_mask_fill_value = inpaint_mask_fill_value
        self.video_length_drop_start = video_length_drop_start
        self.video_length_drop_end = video_length_drop_end
        self.text_drop_ratio = text_drop_ratio
        
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size)
        self.pixel_transforms = transforms.Compose(
            [
                transforms.Resize(sample_size[0]),
                transforms.CenterCrop(sample_size),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
            ]
        )
    
    def get_batch(self, idx):
        """Load and preprocess a single video sample."""
        video_dict = self.dataset[idx]
        video_id, text = video_dict['file_path'], video_dict['text']

        # Resolve video path
        if self.data_root is None:
            video_dir = video_id
        else:
            video_dir = os.path.join(self.data_root, video_id)

        with VideoReader_contextmanager(video_dir, num_threads=2) as video_reader:
            # Calculate frame sampling range with length dropout
            min_sample_n_frames = min(
                self.sample_n_frames, 
                int(len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) // self.sample_stride)
            )
            if min_sample_n_frames == 0:
                raise ValueError(f"No Frames in video.")

            # Select contiguous clip with random start position
            video_length = int(self.video_length_drop_end * len(video_reader))
            clip_length = min(video_length, (min_sample_n_frames - 1) * self.sample_stride + 1)
            start_idx   = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
            batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

            try:
                sample_args = (video_reader, batch_index)
                pixel_values = func_timeout(
                    VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                )
            except FunctionTimedOut:
                raise ValueError(f"Read {idx} timeout.")
            except Exception as e:
                raise ValueError(f"Failed to extract frames from video. Error is {e}.")

            # Convert to tensor, normalize to [-1, 1], apply transforms
            if not self.enable_bucket:
                pixel_values = torch.from_numpy(pixel_values).permute(0, 3, 1, 2).contiguous()
                pixel_values = pixel_values / 255.
                del video_reader
                pixel_values = self.pixel_transforms(pixel_values)
            
            # Random text dropout for classifier-free guidance
            if random.random() < self.text_drop_ratio:
                text = ''
            return pixel_values, text

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        """Get a sample with retry on failure."""
        while True:
            sample = {}
            try:
                pixel_values, name = self.get_batch(idx)
                sample["pixel_values"] = pixel_values
                sample["text"] = name
                sample["idx"] = idx
                if len(sample) > 0:
                    break

            except Exception as e:
                print(e, self.dataset[idx % len(self.dataset)])
                idx = random.randint(0, self.length-1)

        if self.enable_inpaint and not self.enable_bucket:
            mask = get_random_mask(pixel_values.size())
            # Fill masked regions with configurable value (default -1.0, some models use 0.0)
            mask_pixel_values = torch.where(mask.bool(), torch.tensor(self.inpaint_mask_fill_value), pixel_values)
            sample["mask_pixel_values"] = mask_pixel_values
            sample["mask"] = mask

            # Prepare CLIP pixel values for first frame
            sample["clip_pixel_values"] = (sample["pixel_values"][0].permute(1, 2, 0).contiguous() * 0.5 + 0.5) * 255

        return sample


class VideoSpeechDataset(Dataset):
    """Dataset for video-speech paired training with motion and inpainting support."""
    def __init__(
        self,
        ann_path, 
        data_root=None,
        video_sample_size=512,
        video_sample_stride=4,
        video_sample_n_frames=16,
        enable_bucket=False, 
        enable_inpaint=False,
        inpaint_mask_fill_value=0,
        audio_sr=16000,
        text_drop_ratio=0.1,
        enable_motion_info=False,
        motion_frames=73,
        return_file_name=False,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
        self.dataset = json.load(open(ann_path, 'r'))
        self.length = len(self.dataset)
        print(f"data scale: {self.length}")

        self.data_root = data_root
        self.enable_bucket = enable_bucket
        self.enable_inpaint = enable_inpaint
        self.inpaint_mask_fill_value = inpaint_mask_fill_value
        self.audio_sr = audio_sr
        self.text_drop_ratio = text_drop_ratio
        self.enable_motion_info = enable_motion_info
        self.motion_frames = motion_frames
        self.return_file_name = return_file_name
        
        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride = video_sample_stride
        self.video_sample_n_frames = video_sample_n_frames
        self.video_sample_size = tuple(video_sample_size) if not isinstance(video_sample_size, int) else (video_sample_size, video_sample_size)
        self.pixel_transforms = transforms.Compose(
            [
                transforms.Resize(self.video_sample_size[0]),
                transforms.CenterCrop(self.video_sample_size),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
            ]
        )
    
    def get_batch(self, idx):
        """Load and preprocess a single video sample with corresponding audio."""
        video_dict = self.dataset[idx]
        video_id, text = video_dict['file_path'], video_dict['text']
        audio_id = video_dict['audio_path']

        # Resolve video and audio paths
        if self.data_root is None:
            video_path = video_id
            audio_path = audio_id
        else:
            video_path = os.path.join(self.data_root, video_id)
            audio_path = os.path.join(self.data_root, audio_id)

        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"Audio file not found for {video_path}")

        with VideoReader_contextmanager(video_path, num_threads=2) as video_reader:
            total_frames = len(video_reader)
            fps = video_reader.get_avg_fps()

            # Adjust stride to avoid fps > 30
            local_video_sample_stride = self.video_sample_stride
            new_fps = int(fps // local_video_sample_stride)
            while new_fps > 30:
                local_video_sample_stride = local_video_sample_stride + 1
                new_fps = int(fps // local_video_sample_stride)

            # Calculate the actual number of sampled frames (considering boundaries)
            max_possible_frames = (total_frames - 1) // local_video_sample_stride + 1
            actual_n_frames = min(self.video_sample_n_frames, max_possible_frames)
            if actual_n_frames <= 0:
                raise ValueError(f"Video too short: {video_path}")

            # Randomly select the starting frame
            max_start = total_frames - (actual_n_frames - 1) * local_video_sample_stride - 1
            start_frame = random.randint(0, max_start) if max_start > 0 else 0
            frame_indices = [start_frame + i * local_video_sample_stride for i in range(actual_n_frames)]

            # Read video frames
            try:
                sample_args = (video_reader, frame_indices)
                raw_frames = func_timeout(
                    VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                )
                # Resize each frame and free the original array early to reduce peak memory
                resized_frames = []
                for i in range(len(raw_frames)):
                    resized_frames.append(resize_frame(raw_frames[i], max(self.video_sample_size)))
                del raw_frames
                pixel_values = np.array(resized_frames)
                del resized_frames
            except FunctionTimedOut:
                raise ValueError(f"Read {idx} timeout.")
            except Exception as e:
                raise ValueError(f"Failed to extract frames from video. Error is {e}.")

            # Motion video processing
            _, height, width, channel = np.shape(pixel_values)
            if self.enable_motion_info:
                motion_pixel_values = np.ones([self.motion_frames, height, width, channel]) * 127.5
                if start_frame > 0:
                    # Collect motion frames before start_frame (from start_frame-stride towards 0)
                    motion_frame_indices = []
                    current_idx = start_frame - local_video_sample_stride
                    while current_idx >= 0 and len(motion_frame_indices) < self.motion_frames:
                        motion_frame_indices.append(current_idx)
                        current_idx -= local_video_sample_stride
                    motion_frame_indices = motion_frame_indices[::-1]  # Reverse to ascending order

                    _motion_sample_args = (video_reader, motion_frame_indices)
                    motion_raw_frames = func_timeout(
                        VIDEO_READER_TIMEOUT, get_video_reader_batch, args=_motion_sample_args
                    )
                    # Resize each frame and free the original array early
                    motion_resized_frames = []
                    for i in range(len(motion_raw_frames)):
                        motion_resized_frames.append(resize_frame(motion_raw_frames[i], max(self.video_sample_size)))
                    del motion_raw_frames
                    if len(motion_resized_frames) > 0:
                        motion_pixel_values[-len(motion_resized_frames):] = motion_resized_frames
                    del motion_resized_frames

                if not self.enable_bucket:
                    motion_pixel_values = torch.from_numpy(motion_pixel_values).permute(0, 3, 1, 2).contiguous()
                    motion_pixel_values = motion_pixel_values / 255.
                    motion_pixel_values = self.pixel_transforms(motion_pixel_values)
            else:
                motion_pixel_values = None

            # Video post-processing: convert to tensor, normalize to [-1, 1], apply transforms
            if not self.enable_bucket:
                pixel_values = torch.from_numpy(pixel_values).permute(0, 3, 1, 2).contiguous()
                pixel_values = pixel_values / 255.
                pixel_values = self.pixel_transforms(pixel_values)

        # Load and extract the corresponding audio segment
        # Calculate start and end times (in seconds) of the video clip
        start_time = start_frame / fps
        end_time = (start_frame + (actual_n_frames - 1) * local_video_sample_stride) / fps
        duration = end_time - start_time

        # Load entire audio and resample to target sample rate
        audio_input, sample_rate = librosa.load(audio_path, sr=self.audio_sr)

        # Convert time to sample indices
        start_sample = round(start_time * self.audio_sr)
        target_len = round(duration * self.audio_sr)
        end_sample = start_sample + target_len

        # Extract audio segment with validation
        if start_sample >= len(audio_input):
            raise ValueError(f"Audio file too short: {audio_path}")
        else:
            audio_segment = audio_input[start_sample:end_sample]
            if len(audio_segment) < target_len:
                raise ValueError(f"Audio file too short: {audio_path}")

        # Random text dropout for classifier-free guidance
        if random.random() < self.text_drop_ratio:
            text = ''

        return pixel_values, motion_pixel_values, text, audio_segment, sample_rate, new_fps

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        """Get a sample with retry on failure."""
        data_info = self.dataset[idx % len(self.dataset)]
        while True:
            sample = {}
            try:
                pixel_values, motion_pixel_values, text, audio, sample_rate, fps = self.get_batch(idx)
                sample["pixel_values"] = pixel_values
                sample["motion_pixel_values"] = motion_pixel_values
                sample["text"] = text
                sample["audio"] = torch.from_numpy(audio).float()
                sample["sample_rate"] = sample_rate
                sample["fps"] = fps
                sample["idx"] = idx
                
                if self.return_file_name:
                    sample["file_name"] = os.path.basename(data_info['file_path'])

                if len(sample) > 0:
                    break
            except Exception as e:
                print(e, self.dataset[idx % len(self.dataset)])
                idx = random.randint(0, self.length - 1)

        if self.enable_inpaint and not self.enable_bucket:
            mask = get_random_mask(pixel_values.size(), image_start_only=True)
            # Fill masked regions with configurable value (default -1.0, some models use 0.0)
            mask_pixel_values = torch.where(mask.bool(), torch.tensor(self.inpaint_mask_fill_value), pixel_values)
            sample["mask_pixel_values"] = mask_pixel_values
            sample["mask"] = mask

            clip_pixel_values = sample["pixel_values"][0].permute(1, 2, 0).contiguous()
            clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255
            sample["clip_pixel_values"] = clip_pixel_values

        return sample


class VideoSpeechControlDataset(Dataset):
    """Dataset for video-speech-control paired training with motion and inpainting support."""
    def __init__(
        self,
        ann_path, 
        data_root=None,
        video_sample_size=512, 
        video_sample_stride=4, 
        video_sample_n_frames=16,
        enable_bucket=False, 
        enable_inpaint=False,
        inpaint_mask_fill_value=0,
        audio_sr=16000,
        text_drop_ratio=0.1,
        enable_motion_info=False,
        motion_frames=73,
        return_file_name=False,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
        self.dataset = json.load(open(ann_path, 'r'))
        self.length = len(self.dataset)
        print(f"data scale: {self.length}")

        self.data_root = data_root
        self.enable_bucket = enable_bucket
        self.enable_inpaint = enable_inpaint
        self.inpaint_mask_fill_value = inpaint_mask_fill_value
        self.audio_sr = audio_sr
        self.text_drop_ratio = text_drop_ratio
        self.enable_motion_info = enable_motion_info
        self.motion_frames = motion_frames
        self.return_file_name = return_file_name
        
        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride = video_sample_stride
        self.video_sample_n_frames = video_sample_n_frames
        self.video_sample_size = tuple(video_sample_size) if not isinstance(video_sample_size, int) else (video_sample_size, video_sample_size)
        self.pixel_transforms = transforms.Compose(
            [
                transforms.Resize(self.video_sample_size[0]),
                transforms.CenterCrop(self.video_sample_size),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
            ]
        )
    
    def get_batch(self, idx):
        """Load and preprocess a single video sample with control and audio."""
        video_dict = self.dataset[idx]
        video_id, text = video_dict['file_path'], video_dict['text']
        audio_id = video_dict['audio_path']
        control_video_id = video_dict['control_file_path']

        # Resolve video, audio, and control paths
        if self.data_root is None:
            video_path = video_id
            audio_path = audio_id
            control_path = control_video_id
        else:
            video_path = os.path.join(self.data_root, video_id)
            audio_path = os.path.join(self.data_root, audio_id)
            control_path = os.path.join(self.data_root, control_video_id)

        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"Audio file not found for {video_path}")

        # Video information
        with VideoReader_contextmanager(video_path, num_threads=2) as video_reader:
            total_frames = len(video_reader)
            fps = video_reader.get_avg_fps()  # Get the original video frame rate
            if fps <= 0:
                raise ValueError(f"Video has negative fps: {video_path}")
            
            # Avoid fps > 30
            local_video_sample_stride = self.video_sample_stride
            new_fps = int(fps // local_video_sample_stride)
            while new_fps > 30:
                local_video_sample_stride = local_video_sample_stride + 1
                new_fps = int(fps // local_video_sample_stride)

            # Calculate the actual number of sampled video frames (considering boundaries)
            max_possible_frames = (total_frames - 1) // local_video_sample_stride + 1
            actual_n_frames = min(self.video_sample_n_frames, max_possible_frames)
            if actual_n_frames <= 0:
                raise ValueError(f"Video too short: {video_path}")

            # Randomly select the starting frame
            max_start = total_frames - (actual_n_frames - 1) * local_video_sample_stride - 1
            start_frame = random.randint(0, max_start) if max_start > 0 else 0
            frame_indices = [start_frame + i * local_video_sample_stride for i in range(actual_n_frames)]

            # Read video frames
            try:
                sample_args = (video_reader, frame_indices)
                raw_frames = func_timeout(
                    VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                )
                # Resize each frame and free the original array early to reduce peak memory
                resized_frames = []
                for i in range(len(raw_frames)):
                    resized_frames.append(resize_frame(raw_frames[i], max(self.video_sample_size)))
                del raw_frames
                pixel_values = np.array(resized_frames)
                del resized_frames
            except FunctionTimedOut:
                raise ValueError(f"Read {idx} timeout.")
            except Exception as e:
                raise ValueError(f"Failed to extract frames from video. Error is {e}.")

            # Motion video processing
            _, height, width, channel = np.shape(pixel_values)
            if self.enable_motion_info:
                motion_pixel_values = np.ones([self.motion_frames, height, width, channel]) * 127.5
                if start_frame > 0:
                    # Collect motion frames before start_frame (from start_frame-stride towards 0)
                    motion_frame_indices = []
                    current_idx = start_frame - local_video_sample_stride
                    while current_idx >= 0 and len(motion_frame_indices) < self.motion_frames:
                        motion_frame_indices.append(current_idx)
                        current_idx -= local_video_sample_stride
                    motion_frame_indices = motion_frame_indices[::-1]  # Reverse to ascending order

                    _motion_sample_args = (video_reader, motion_frame_indices)
                    motion_raw_frames = func_timeout(
                        VIDEO_READER_TIMEOUT, get_video_reader_batch, args=_motion_sample_args
                    )
                    # Resize each frame and free the original array early
                    motion_resized_frames = []
                    for i in range(len(motion_raw_frames)):
                        motion_resized_frames.append(resize_frame(motion_raw_frames[i], max(self.video_sample_size)))
                    del motion_raw_frames
                    if len(motion_resized_frames) > 0:
                        motion_pixel_values[-len(motion_resized_frames):] = motion_resized_frames
                    del motion_resized_frames

                if not self.enable_bucket:
                    motion_pixel_values = torch.from_numpy(motion_pixel_values).permute(0, 3, 1, 2).contiguous()
                    motion_pixel_values = motion_pixel_values / 255.
                    motion_pixel_values = self.pixel_transforms(motion_pixel_values)
            else:
                motion_pixel_values = None

            # Video post-processing: convert to tensor, normalize to [-1, 1], apply transforms
            if not self.enable_bucket:
                pixel_values = torch.from_numpy(pixel_values).permute(0, 3, 1, 2).contiguous()
                pixel_values = pixel_values / 255.
                pixel_values = self.pixel_transforms(pixel_values)

        # Control information
        with VideoReader_contextmanager(control_path, num_threads=2) as control_video_reader:
            try:
                sample_args = (control_video_reader, frame_indices)
                control_raw_frames = func_timeout(
                    VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                )
                # Resize each frame and free the original array early
                resized_frames = []
                for i in range(len(control_raw_frames)):
                    resized_frames.append(resize_frame(control_raw_frames[i], max(self.video_sample_size)))
                del control_raw_frames
                control_pixel_values = np.stack(resized_frames)
                del resized_frames
            except FunctionTimedOut:
                raise ValueError(f"Read {idx} timeout.")
            except Exception as e:
                raise ValueError(f"Failed to extract frames from video. Error is {e}.")

            if not self.enable_bucket:
                control_pixel_values = torch.from_numpy(control_pixel_values).permute(0, 3, 1, 2).contiguous()
                control_pixel_values = control_pixel_values / 255.
                control_pixel_values = self.pixel_transforms(control_pixel_values)

        # Load and extract the corresponding audio segment
        # Calculate start and end times (in seconds) of the video clip
        start_time = start_frame / fps
        end_time = (start_frame + (actual_n_frames - 1) * local_video_sample_stride) / fps
        duration = end_time - start_time

        # Load entire audio and resample to target sample rate
        audio_input, sample_rate = librosa.load(audio_path, sr=self.audio_sr)

        # Convert time to sample indices
        start_sample = round(start_time * self.audio_sr)
        target_len = round(duration * self.audio_sr)
        end_sample = start_sample + target_len

        # Extract audio segment with validation
        if start_sample >= len(audio_input):
            raise ValueError(f"Audio file too short: {audio_path}")
        else:
            audio_segment = audio_input[start_sample:end_sample]
            if len(audio_segment) < target_len:
                raise ValueError(f"Audio file too short: {audio_path}")

        # Random text dropout for classifier-free guidance
        if random.random() < self.text_drop_ratio:
            text = ''

        return pixel_values, motion_pixel_values, control_pixel_values, text, audio_segment, sample_rate, new_fps

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        """Get a sample with retry on failure."""
        data_info = self.dataset[idx % len(self.dataset)]
        while True:
            sample = {}
            try:
                pixel_values, motion_pixel_values, control_pixel_values, text, audio, sample_rate, fps = self.get_batch(idx)
                sample["pixel_values"] = pixel_values
                sample["motion_pixel_values"] = motion_pixel_values
                sample["control_pixel_values"] = control_pixel_values
                sample["text"] = text
                sample["audio"] = torch.from_numpy(audio).float()
                sample["sample_rate"] = sample_rate
                sample["fps"] = fps
                sample["idx"] = idx
                
                if self.return_file_name:
                    sample["file_name"] = os.path.basename(data_info['file_path'])

                if len(sample) > 0:
                    break
            except Exception as e:
                print(e, self.dataset[idx % len(self.dataset)])
                idx = random.randint(0, self.length-1)

        if self.enable_inpaint and not self.enable_bucket:
            mask = get_random_mask(pixel_values.size(), image_start_only=True)
            # Fill masked regions with configurable value (default -1.0, some models use 0.0)
            mask_pixel_values = torch.where(mask.bool(), torch.tensor(self.inpaint_mask_fill_value), pixel_values)
            sample["mask_pixel_values"] = mask_pixel_values
            sample["mask"] = mask

            clip_pixel_values = sample["pixel_values"][0].permute(1, 2, 0).contiguous()
            clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255
            sample["clip_pixel_values"] = clip_pixel_values

        return sample


class VideoAnimateDataset(Dataset):
    """Dataset for video animation training with control, face, background, and mask support."""
    def __init__(
        self,
        ann_path, 
        data_root=None,
        video_sample_size=512, 
        video_sample_stride=4, 
        video_sample_n_frames=16,
        video_repeat=0,
        text_drop_ratio=0.1,
        enable_bucket=False,
        video_length_drop_start=0.1, 
        video_length_drop_end=0.9,
        return_file_name=False,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
        if ann_path.endswith('.csv'):
            with open(ann_path, 'r') as csvfile:
                dataset = list(csv.DictReader(csvfile))
        elif ann_path.endswith('.json'):
            dataset = json.load(open(ann_path))
    
        self.data_root = data_root

        # Balance image/video ratio by duplicating video entries
        if video_repeat > 0:
            self.dataset = []
            for data in dataset:
                if data.get('type', 'image') != 'video':
                    self.dataset.append(data)
                    
            for _ in range(video_repeat):
                for data in dataset:
                    if data.get('type', 'image') == 'video':
                        self.dataset.append(data)
        else:
            self.dataset = dataset
        del dataset

        self.length = len(self.dataset)
        print(f"data scale: {self.length}")
        
        self.enable_bucket = enable_bucket
        self.text_drop_ratio = text_drop_ratio
        self.return_file_name = return_file_name

        self.video_length_drop_start = video_length_drop_start
        self.video_length_drop_end = video_length_drop_end

        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride = video_sample_stride
        self.video_sample_n_frames = video_sample_n_frames
        self.video_sample_size = tuple(video_sample_size) if not isinstance(video_sample_size, int) else (video_sample_size, video_sample_size)
        self.video_transforms = transforms.Compose(
            [
                transforms.Resize(min(self.video_sample_size)),
                transforms.CenterCrop(self.video_sample_size),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
            ]
        )

        self.larger_side_of_image_and_video = min(self.video_sample_size)
    
    def get_batch(self, idx):
        """Load and preprocess a single video sample with control, face, background, and mask."""
        data_info = self.dataset[idx % len(self.dataset)]
        video_id, text = data_info['file_path'], data_info['text']

        # Resolve video path
        if self.data_root is None:
            video_dir = video_id
        else:
            video_dir = os.path.join(self.data_root, video_id)

        with VideoReader_contextmanager(video_dir, num_threads=2) as video_reader:
            # Calculate frame sampling range with length dropout
            min_sample_n_frames = min(
                self.video_sample_n_frames, 
                int(len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) // self.video_sample_stride)
            )
            if min_sample_n_frames == 0:
                raise ValueError(f"No Frames in video.")

            # Select contiguous clip with random start position
            video_length = int(self.video_length_drop_end * len(video_reader))
            clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
            start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
            batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

            try:
                sample_args = (video_reader, batch_index)
                raw_frames = func_timeout(
                    VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                )
                # Resize each frame and free the original array early
                resized_frames = []
                for i in range(len(raw_frames)):
                    resized_frames.append(resize_frame(raw_frames[i], self.larger_side_of_image_and_video))
                del raw_frames
                pixel_values = np.stack(resized_frames)
                del resized_frames
            except FunctionTimedOut:
                raise ValueError(f"Read {idx} timeout.")
            except Exception as e:
                raise ValueError(f"Failed to extract frames from video. Error is {e}.")

            # Release video reader early
            del video_reader

            # Convert to tensor and apply transforms
            if not self.enable_bucket:
                pixel_values = torch.from_numpy(pixel_values).permute(0, 3, 1, 2).contiguous()
                pixel_values = pixel_values / 255.
                pixel_values = self.video_transforms(pixel_values)
            
            # Random text dropout for classifier-free guidance
            if random.random() < self.text_drop_ratio:
                text = ''

        # Load control video
        control_video_id = data_info['control_file_path']
        if control_video_id is not None:
            control_video_path = control_video_id if self.data_root is None else os.path.join(self.data_root, control_video_id)
        else:
            control_video_path = None
        
        if control_video_path is not None:
            with VideoReader_contextmanager(control_video_path, num_threads=2) as control_video_reader:
                try:
                    sample_args = (control_video_reader, batch_index)
                    control_raw_frames = func_timeout(
                        VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                    )
                    # Resize each frame and free the original array early
                    resized_frames = []
                    for i in range(len(control_raw_frames)):
                        resized_frames.append(resize_frame(control_raw_frames[i], self.larger_side_of_image_and_video))
                    del control_raw_frames
                    control_pixel_values = np.stack(resized_frames)
                    del resized_frames
                except FunctionTimedOut:
                    raise ValueError(f"Read {idx} timeout.")
                except Exception as e:
                    raise ValueError(f"Failed to extract frames from video. Error is {e}.")

                # Release control video reader early
                del control_video_reader

                # Convert to tensor and apply transforms
                if not self.enable_bucket:
                    control_pixel_values = torch.from_numpy(control_pixel_values).permute(0, 3, 1, 2).contiguous()
                    control_pixel_values = control_pixel_values / 255.
                    control_pixel_values = self.video_transforms(control_pixel_values)
        else:
            control_pixel_values = torch.zeros_like(pixel_values) if not self.enable_bucket else np.zeros_like(pixel_values)

        # Load face video
        face_video_id = data_info['face_file_path']
        if face_video_id is not None:
            face_video_path = face_video_id if self.data_root is None else os.path.join(self.data_root, face_video_id)
        else:
            face_video_path = None
        
        if face_video_path is not None:
            with VideoReader_contextmanager(face_video_path, num_threads=2) as face_video_reader:
                try:
                    sample_args = (face_video_reader, batch_index)
                    face_raw_frames = func_timeout(
                        VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                    )
                    # Resize each frame and free the original array early
                    resized_frames = []
                    for i in range(len(face_raw_frames)):
                        resized_frames.append(resize_frame(face_raw_frames[i], self.larger_side_of_image_and_video))
                    del face_raw_frames
                    face_pixel_values = np.stack(resized_frames)
                    del resized_frames
                except FunctionTimedOut:
                    raise ValueError(f"Read {idx} timeout.")
                except Exception as e:
                    raise ValueError(f"Failed to extract frames from video. Error is {e}.")

                # Release face video reader early
                del face_video_reader

                # Convert to tensor and apply transforms
                if not self.enable_bucket:
                    face_pixel_values = torch.from_numpy(face_pixel_values).permute(0, 3, 1, 2).contiguous()
                    face_pixel_values = face_pixel_values / 255.
                    face_pixel_values = self.video_transforms(face_pixel_values)
        else:
            face_pixel_values = torch.zeros_like(pixel_values) if not self.enable_bucket else np.zeros_like(pixel_values)

        # Load background video
        background_video_id = data_info.get('background_file_path', None)
        if background_video_id is not None:
            background_video_path = background_video_id if self.data_root is None else os.path.join(self.data_root, background_video_id)
        else:
            background_video_path = None
        
        if background_video_path is not None:
            with VideoReader_contextmanager(background_video_path, num_threads=2) as background_video_reader:
                try:
                    sample_args = (background_video_reader, batch_index)
                    background_raw_frames = func_timeout(
                        VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                    )
                    # Resize each frame and free the original array early
                    resized_frames = []
                    for i in range(len(background_raw_frames)):
                        resized_frames.append(resize_frame(background_raw_frames[i], self.larger_side_of_image_and_video))
                    del background_raw_frames
                    background_pixel_values = np.stack(resized_frames)
                    del resized_frames
                except FunctionTimedOut:
                    raise ValueError(f"Read {idx} timeout.")
                except Exception as e:
                    raise ValueError(f"Failed to extract frames from video. Error is {e}.")

                # Release background video reader early
                del background_video_reader

                # Convert to tensor and apply transforms
                if not self.enable_bucket:
                    background_pixel_values = torch.from_numpy(background_pixel_values).permute(0, 3, 1, 2).contiguous()
                    background_pixel_values = background_pixel_values / 255.
                    background_pixel_values = self.video_transforms(background_pixel_values)
        else:
            background_pixel_values = torch.ones_like(pixel_values) * 127.5 if not self.enable_bucket else np.ones_like(pixel_values) * 127.5

        # Load mask video
        mask_video_id = data_info.get('mask_file_path', None)
        if mask_video_id is not None:
            mask_video_path = mask_video_id if self.data_root is None else os.path.join(self.data_root, mask_video_id)
        else:
            mask_video_path = None
        
        if mask_video_path is not None:
            with VideoReader_contextmanager(mask_video_path, num_threads=2) as mask_video_reader:
                try:
                    sample_args = (mask_video_reader, batch_index)
                    mask_raw_frames = func_timeout(
                        VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
                    )
                    # Resize each frame and free the original array early
                    resized_frames = []
                    for i in range(len(mask_raw_frames)):
                        resized_frames.append(resize_frame(mask_raw_frames[i], self.larger_side_of_image_and_video))
                    del mask_raw_frames
                    mask = np.stack(resized_frames)
                    del resized_frames
                except FunctionTimedOut:
                    raise ValueError(f"Read {idx} timeout.")
                except Exception as e:
                    raise ValueError(f"Failed to extract frames from video. Error is {e}.")

                # Release mask video reader early
                del mask_video_reader

                # Convert to tensor (no transforms for mask)
                if not self.enable_bucket:
                    mask = torch.from_numpy(mask).permute(0, 3, 1, 2).contiguous()
                    mask = mask / 255.
        else:
            mask = torch.ones_like(pixel_values) if not self.enable_bucket else np.ones_like(pixel_values) * 255
        
        # Extract only the first channel
        mask = mask[:, :, :, :1]
        
        # Load reference image
        ref_pixel_values_path = data_info.get('ref_file_path', [])
        if self.data_root is not None:
            ref_pixel_values_path = os.path.join(self.data_root, ref_pixel_values_path)
        ref_pixel_values = Image.open(ref_pixel_values_path).convert('RGB')

        if not self.enable_bucket:
            raise ValueError("Not enable_bucket is not supported now. ")
        else:
            ref_pixel_values = np.array(ref_pixel_values)
    
        return pixel_values, control_pixel_values, face_pixel_values, background_pixel_values, mask, ref_pixel_values, text, "video"

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        """Get a sample with retry on failure."""
        data_info = self.dataset[idx % len(self.dataset)]
        data_type = data_info.get('type', 'image')
        while True:
            sample = {}
            try:
                data_info_local = self.dataset[idx % len(self.dataset)]
                data_type_local = data_info_local.get('type', 'image')
                if data_type_local != data_type:
                    raise ValueError("data_type_local != data_type")

                pixel_values, control_pixel_values, face_pixel_values, background_pixel_values, mask, ref_pixel_values, name, data_type = \
                    self.get_batch(idx)

                sample["pixel_values"] = pixel_values
                sample["control_pixel_values"] = control_pixel_values
                sample["face_pixel_values"] = face_pixel_values
                sample["background_pixel_values"] = background_pixel_values
                sample["mask"] = mask
                sample["ref_pixel_values"] = ref_pixel_values
                sample["clip_pixel_values"] = ref_pixel_values
                sample["text"] = name
                sample["data_type"] = data_type
                sample["idx"] = idx

                if self.return_file_name:
                    sample["file_name"] = os.path.basename(data_info['file_path'])

                if len(sample) > 0:
                    break
            except Exception as e:
                print(e, self.dataset[idx % len(self.dataset)])
                idx = random.randint(0, self.length-1)

        return sample


if __name__ == "__main__":
    if 1:
        dataset = VideoDataset(
            json_path="./webvidval/results_2M_val.json",
            sample_size=256,
            sample_stride=4, sample_n_frames=16,
        )

    if 0:
        dataset = WebVid10M(
            csv_path="./webvid/results_2M_val.csv",
            video_folder="./webvid/2M_val",
            sample_size=256,
            sample_stride=4, sample_n_frames=16,
            is_image=False,
        )

    dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=0,)
    for idx, batch in enumerate(dataloader):
        print(batch["pixel_values"].shape, len(batch["text"]))