modelscope-scepter/scepter/modules/data/dataset/video_gen_dataset.py

import io
import random
import sys
import os
import warnings

import torch
import numpy as np
from tqdm import tqdm

from scepter.modules.utils.distribute import we
from scepter.modules.data.dataset import DATASETS, BaseDataset
from scepter.modules.utils.file_system import FS

try:
    import decord
    decord.bridge.set_bridge("torch")
except ImportError:
    warnings.warn(
        "The `decord` package is required for loading the video dataset. Install with `pip install decord`"
    )


@DATASETS.register_class()
class VideoGenDataset(BaseDataset):
    def __init__(self, cfg, logger = None):
        super().__init__(cfg, logger=logger)
        self.prompt_prefix = cfg.get('PROMPT_PREFIX', '')
        self.path_prefix = cfg.get('PATH_PREFIX', '')
        self.p_zero = cfg.get('P_ZERO', 0.0)
        self.max_num_frames = cfg.get("NUM_FRAMES", 49)
        self.fps = cfg.get("FPS", 8)
        self.height = cfg.get("HEIGHT", 480)
        self.width = cfg.get("WIDTH", 720)
        self.skip_frames_start = cfg.get("SKIP_FRAMES_START", 0)
        self.skip_frames_end = cfg.get("SKIP_FRAMES_END", 0)
        self.data_type = cfg.get('DATA_TYPE', 't2v')

    def worker_init_fn(self, worker_id, num_workers=1):
        super().worker_init_fn(worker_id, num_workers=num_workers)
        randseed = np.random.randint(0, 2 ** 32 - num_workers - 1)
        workerseed = randseed + worker_id
        random.seed(workerseed)
        np.random.seed(workerseed)

    def _preprocess_video_data(self, video_path):

        with FS.get_object(video_path) as video_data:
            video_reader = decord.VideoReader(io.BytesIO(video_data), width=self.width, height=self.height)
        video_num_frames = len(video_reader)

        start_frame = min(self.skip_frames_start, video_num_frames)
        end_frame = max(0, video_num_frames - self.skip_frames_end)
        if end_frame <= start_frame:
            frames = video_reader.get_batch([start_frame])
        elif end_frame - start_frame <= self.max_num_frames:
            frames = video_reader.get_batch(list(range(start_frame, end_frame)))
        else:
            indices = list(range(start_frame, end_frame, (end_frame - start_frame) // self.max_num_frames))
            frames = video_reader.get_batch(indices)

        # Ensure that we don't go over the limit
        frames = frames[: self.max_num_frames]
        selected_num_frames = frames.shape[0]

        # Choose first (4k + 1) frames as this is how many is required by the VAE
        remainder = (3 + (selected_num_frames % 4)) % 4
        if remainder != 0:
            frames = frames[:-remainder]
        selected_num_frames = frames.shape[0]

        assert (selected_num_frames - 1) % 4 == 0

        # Training transforms
        frames = frames.float().div_(127.5).sub_(1.)
        frames = frames.permute(3, 0, 1, 2).contiguous() # [C, F, H, W]
        return frames

    def _parse_index(self, index):
        meta = dict()
        for key, value in zip(index[-1], index[:-1]):
            if key in ['oss_key', 'path', 'video_path']:
                meta['video_path'] = value
            elif key in ['prompt', 'caption', 'text']:
                meta['prompt'] = value
            elif key in ['width', 'height']:
                meta[key] = int(value)
            else:
                meta[key] = value
        return meta

    def _get(self, index):
        meta = self._parse_index(index)

        video_path = os.path.join(self.path_prefix, meta.get('video_path', ''))
        video = self._preprocess_video_data(video_path)

        prompt = self.prompt_prefix + meta.get('prompt', '')
        if self.mode == 'train' and np.random.uniform() < self.p_zero:
            prompt = ''

        item = {
            'video': video,
            'prompt': prompt,
            'meta': meta,
        }
        if self.data_type == 'i2v':
            item['image'] = item['video'][:, :1, :, :]
        return item

    def __len__(self):
        return sys.maxsize

    @staticmethod
    def collate_fn(batch):
        collect = {}
        for sample in batch:
            for k, v in sample.items():
                if k not in collect:
                    collect[k] = []
                collect[k].append(v)
        return collect



@DATASETS.register_class()
class VideoGenDatasetOTF(VideoGenDataset):
    def __init__(self, cfg, logger=None):
        super().__init__(cfg, logger)
        self.data_file = cfg.DATA_FILE
        self.delimiter = cfg.get('DELIMITER', '#;#')
        self.fields = cfg.get('FIELDS', ['video_path', 'prompt'])
        self.use_num = cfg.get('USE_NUM', -1)

        from scepter.modules.model.registry import MODELS
        model_cfg = cfg.get('MODEL', None)
        if model_cfg is not None:
            self.model = MODELS.build(cfg.MODEL, logger=logger).eval().requires_grad_(False).to(we.device_id)
        self.items = self.parse_data(self.data_file, self.delimiter, self.fields)
        if self.use_num and self.use_num > 0:
            self.items = self.items[:self.use_num]
        self.data = self.encode(self.items)
        self.real_number = len(self.data)
        if model_cfg is not None:
            self.model.to('cpu')
            del self.model
            torch.cuda.empty_cache()

    def parse_data(self, data_file, delimiter, fields):
        items = list()
        with FS.get_object(data_file) as local_data:
            rows = [
                i.split(delimiter,
                        len(fields) - 1)
                for i in local_data.decode('utf-8').strip().split('\n')
            ]
            for i, row in enumerate(rows):
                item = {}
                for key, value in zip(self.fields, row):
                    if key in ['oss_key', 'path', 'video_path']:
                        item['video_path'] = value
                    elif key in ['prompt', 'caption', 'text']:
                        item['prompt'] = value
                    elif key in ['width', 'height']:
                        item[key] = int(value)
                    else:
                        item[key] = value
                items.append(item)
        return items

    def encode(self, items):
        self.logger.info("Start to encode video data [{}]!".format(len(items)))
        for item in tqdm(items):
            video_path = os.path.join(self.path_prefix, item.get('video_path', ''))
            video = self._preprocess_video_data(video_path)
            latent = self.model.encode_first_stage(video.unsqueeze(0).to(we.device_id)).squeeze(0)
            item['video_latent'] = latent.detach().cpu()
            item['video'] = video
            if self.data_type == 'i2v':
                item['image'] = item['video'][:, :1, :, :]
        return items

    def _get(self, index):
        return self.data[index % self.real_number]