import os import decord import numpy as np import random import json import torchvision import torchvision.transforms as T import torch from glob import glob from PIL import Image from itertools import islice from pathlib import Path from .bucketing import sensible_buckets decord.bridge.set_bridge('torch') from torch.utils.data import Dataset from einops import rearrange, repeat def get_prompt_ids(prompt, tokenizer): prompt_ids = tokenizer( prompt, truncation=True, padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt", ).input_ids return prompt_ids def read_caption_file(caption_file): with open(caption_file, 'r', encoding="utf8") as t: return t.read() def get_text_prompt( text_prompt: str = '', fallback_prompt: str= '', file_path:str = '', ext_types=['.mp4'], use_caption=False ): try: if use_caption: if len(text_prompt) > 1: return text_prompt caption_file = '' # Use caption on per-video basis (One caption PER video) for ext in ext_types: maybe_file = file_path.replace(ext, '.txt') if maybe_file.endswith(ext_types): continue if os.path.exists(maybe_file): caption_file = maybe_file break if os.path.exists(caption_file): return read_caption_file(caption_file) # Return fallback prompt if no conditions are met. return fallback_prompt return text_prompt except: print(f"Couldn't read prompt caption for {file_path}. Using fallback.") return fallback_prompt def get_video_frames(vr, start_idx, sample_rate=1, max_frames=24): max_range = len(vr) frame_number = sorted((0, start_idx, max_range))[1] frame_range = range(frame_number, max_range, sample_rate) frame_range_indices = list(frame_range)[:max_frames] return frame_range_indices def process_video(vid_path, use_bucketing, w, h, get_frame_buckets, get_frame_batch): if use_bucketing: vr = decord.VideoReader(vid_path) resize = get_frame_buckets(vr) video = get_frame_batch(vr, resize=resize) else: vr = decord.VideoReader(vid_path, width=w, height=h) video = get_frame_batch(vr) return video, vr # https://github.com/ExponentialML/Video-BLIP2-Preprocessor class VideoJsonDataset(Dataset): def __init__( self, tokenizer = None, width: int = 256, height: int = 256, n_sample_frames: int = 4, sample_start_idx: int = 1, frame_step: int = 1, json_path: str ="", json_data = None, vid_data_key: str = "video_path", preprocessed: bool = False, use_bucketing: bool = False, **kwargs ): self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg") self.use_bucketing = use_bucketing self.tokenizer = tokenizer self.preprocessed = preprocessed self.vid_data_key = vid_data_key self.train_data = self.load_from_json(json_path, json_data) self.width = width self.height = height self.n_sample_frames = n_sample_frames self.sample_start_idx = sample_start_idx self.frame_step = frame_step def build_json(self, json_data): extended_data = [] for data in json_data['data']: for nested_data in data['data']: self.build_json_dict( data, nested_data, extended_data ) json_data = extended_data return json_data def build_json_dict(self, data, nested_data, extended_data): clip_path = nested_data['clip_path'] if 'clip_path' in nested_data else None extended_data.append({ self.vid_data_key: data[self.vid_data_key], 'frame_index': nested_data['frame_index'], 'prompt': nested_data['prompt'], 'clip_path': clip_path }) def load_from_json(self, path, json_data): try: with open(path) as jpath: print(f"Loading JSON from {path}") json_data = json.load(jpath) return self.build_json(json_data) except: self.train_data = [] print("Non-existant JSON path. Skipping.") def validate_json(self, base_path, path): return os.path.exists(f"{base_path}/{path}") def get_frame_range(self, vr): return get_video_frames( vr, self.sample_start_idx, self.frame_step, self.n_sample_frames ) def get_vid_idx(self, vr, vid_data=None): frames = self.n_sample_frames if vid_data is not None: idx = vid_data['frame_index'] else: idx = self.sample_start_idx return idx def get_frame_buckets(self, vr): _, h, w = vr[0].shape width, height = sensible_buckets(self.width, self.height, h, w) # width, height = self.width, self.height resize = T.transforms.Resize((height, width), antialias=True) return resize def get_frame_batch(self, vr, resize=None): frame_range = self.get_frame_range(vr) frames = vr.get_batch(frame_range) video = rearrange(frames, "f h w c -> f c h w") if resize is not None: video = resize(video) return video def process_video_wrapper(self, vid_path): video, vr = process_video( vid_path, self.use_bucketing, self.width, self.height, self.get_frame_buckets, self.get_frame_batch ) return video, vr def train_data_batch(self, index): # If we are training on individual clips. if 'clip_path' in self.train_data[index] and \ self.train_data[index]['clip_path'] is not None: vid_data = self.train_data[index] clip_path = vid_data['clip_path'] # Get video prompt prompt = vid_data['prompt'] video, _ = self.process_video_wrapper(clip_path) prompt_ids = get_prompt_ids(prompt, self.tokenizer) return video, prompt, prompt_ids # Assign train data train_data = self.train_data[index] # Get the frame of the current index. self.sample_start_idx = train_data['frame_index'] # Initialize resize resize = None video, vr = self.process_video_wrapper(train_data[self.vid_data_key]) # Get video prompt prompt = train_data['prompt'] vr.seek(0) prompt_ids = get_prompt_ids(prompt, self.tokenizer) return video, prompt, prompt_ids @staticmethod def __getname__(): return 'json' def __len__(self): if self.train_data is not None: return len(self.train_data) else: return 0 def __getitem__(self, index): # Initialize variables video = None prompt = None prompt_ids = None # Use default JSON training if self.train_data is not None: video, prompt, prompt_ids = self.train_data_batch(index) example = { "pixel_values": (video / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt, 'dataset': self.__getname__() } return example class SingleVideoDataset(Dataset): def __init__( self, tokenizer = None, width: int = 256, height: int = 256, n_sample_frames: int = 4, frame_step: int = 1, single_video_path: str = "", single_video_prompt: str = "", use_caption: bool = False, use_bucketing: bool = False, **kwargs ): self.tokenizer = tokenizer self.use_bucketing = use_bucketing self.frames = [] self.index = 1 self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg") self.n_sample_frames = n_sample_frames self.frame_step = frame_step self.single_video_path = single_video_path self.single_video_prompt = single_video_prompt self.width = width self.height = height def create_video_chunks(self): vr = decord.VideoReader(self.single_video_path) vr_range = range(0, len(vr), self.frame_step) self.frames = list(self.chunk(vr_range, self.n_sample_frames)) return self.frames def chunk(self, it, size): it = iter(it) return iter(lambda: tuple(islice(it, size)), ()) def get_frame_batch(self, vr, resize=None): index = self.index frames = vr.get_batch(self.frames[self.index]) video = rearrange(frames, "f h w c -> f c h w") if resize is not None: video = resize(video) return video def get_frame_buckets(self, vr): _, h, w = vr[0].shape # width, height = sensible_buckets(self.width, self.height, h, w) width, height = self.width, self.height resize = T.transforms.Resize((height, width), antialias=True) return resize def process_video_wrapper(self, vid_path): video, vr = process_video( vid_path, self.use_bucketing, self.width, self.height, self.get_frame_buckets, self.get_frame_batch ) return video, vr def single_video_batch(self, index): train_data = self.single_video_path self.index = index if train_data.endswith(self.vid_types): video, _ = self.process_video_wrapper(train_data) prompt = self.single_video_prompt prompt_ids = get_prompt_ids(prompt, self.tokenizer) return video, prompt, prompt_ids else: raise ValueError(f"Single video is not a video type. Types: {self.vid_types}") @staticmethod def __getname__(): return 'single_video' def __len__(self): return len(self.create_video_chunks()) def __getitem__(self, index): video, prompt, prompt_ids = self.single_video_batch(index) example = { "pixel_values": (video / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt, 'dataset': self.__getname__() } return example class ImageDataset(Dataset): def __init__( self, tokenizer = None, width: int = 256, height: int = 256, base_width: int = 256, base_height: int = 256, use_caption: bool = False, image_dir: str = '', single_img_prompt: str = '', use_bucketing: bool = False, fallback_prompt: str = '', **kwargs ): self.tokenizer = tokenizer self.img_types = (".png", ".jpg", ".jpeg", '.bmp') self.use_bucketing = use_bucketing self.image_dir = self.get_images_list(image_dir) self.fallback_prompt = fallback_prompt self.use_caption = use_caption self.single_img_prompt = single_img_prompt self.width = width self.height = height def get_images_list(self, image_dir): if os.path.exists(image_dir): imgs = [x for x in os.listdir(image_dir) if x.endswith(self.img_types)] full_img_dir = [] for img in imgs: full_img_dir.append(f"{image_dir}/{img}") return sorted(full_img_dir) return [''] def image_batch(self, index): train_data = self.image_dir[index] img = train_data try: img = torchvision.io.read_image(img, mode=torchvision.io.ImageReadMode.RGB) except: img = T.transforms.PILToTensor()(Image.open(img).convert("RGB")) width = self.width height = self.height if self.use_bucketing: _, h, w = img.shape width, height = sensible_buckets(width, height, w, h) resize = T.transforms.Resize((height, width), antialias=True) img = resize(img) img = repeat(img, 'c h w -> f c h w', f=1) prompt = get_text_prompt( file_path=train_data, text_prompt=self.single_img_prompt, fallback_prompt=self.fallback_prompt, ext_types=self.img_types, use_caption=True ) prompt_ids = get_prompt_ids(prompt, self.tokenizer) return img, prompt, prompt_ids @staticmethod def __getname__(): return 'image' def __len__(self): # Image directory if os.path.exists(self.image_dir[0]): return len(self.image_dir) else: return 0 def __getitem__(self, index): img, prompt, prompt_ids = self.image_batch(index) example = { "pixel_values": (img / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt, 'dataset': self.__getname__() } return example class VideoFolderDataset(Dataset): def __init__( self, tokenizer=None, width: int = 256, height: int = 256, n_sample_frames: int = 16, fps: int = 8, path: str = "./data", fallback_prompt: str = "", use_bucketing: bool = False, **kwargs ): self.tokenizer = tokenizer self.use_bucketing = use_bucketing self.fallback_prompt = fallback_prompt self.video_files = glob(f"{path}/*.mp4") self.width = width self.height = height self.n_sample_frames = n_sample_frames self.fps = fps def get_frame_buckets(self, vr): _, h, w = vr[0].shape width, height = sensible_buckets(self.width, self.height, h, w) # width, height = self.width, self.height resize = T.transforms.Resize((height, width), antialias=True) return resize def get_frame_batch(self, vr, resize=None): n_sample_frames = self.n_sample_frames native_fps = vr.get_avg_fps() every_nth_frame = max(1, round(native_fps / self.fps)) every_nth_frame = min(len(vr), every_nth_frame) effective_length = len(vr) // every_nth_frame if effective_length < n_sample_frames: n_sample_frames = effective_length effective_idx = random.randint(0, (effective_length - n_sample_frames)) idxs = every_nth_frame * np.arange(effective_idx, effective_idx + n_sample_frames) video = vr.get_batch(idxs) video = rearrange(video, "f h w c -> f c h w") if resize is not None: video = resize(video) return video, vr def process_video_wrapper(self, vid_path): video, vr = process_video( vid_path, self.use_bucketing, self.width, self.height, self.get_frame_buckets, self.get_frame_batch ) return video, vr def get_prompt_ids(self, prompt): return self.tokenizer( prompt, truncation=True, padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt", ).input_ids @staticmethod def __getname__(): return 'folder' def __len__(self): return len(self.video_files) def __getitem__(self, index): video, _ = self.process_video_wrapper(self.video_files[index]) prompt = self.fallback_prompt prompt_ids = self.get_prompt_ids(prompt) return {"pixel_values": (video[0] / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt, 'dataset': self.__getname__()} class CachedDataset(Dataset): def __init__(self,cache_dir: str = ''): self.cache_dir = cache_dir self.cached_data_list = self.get_files_list() def get_files_list(self): tensors_list = [f"{self.cache_dir}/{x}" for x in os.listdir(self.cache_dir) if x.endswith('.pt')] return sorted(tensors_list) def __len__(self): return len(self.cached_data_list) def __getitem__(self, index): cached_latent = torch.load(self.cached_data_list[index], map_location='cuda:0') return cached_latent