Peiyuan Zhang
8ec82cbc37
Delete Open Sora Plan Modeling amend log validation Change name to fast video Remove OSP modeling clean up name changing remove files Deleted unnecessary files commit first commit first training ok Update generate_synthetic.sh and deepspeed_zero2_config.yaml 14 debug overfitting .... debugging .. Debug successful! Add zero3 OK update optimizer load small bug random seed args sp enable & still has bug rename update inference sp code / can run / still has bug / don't output normal mp4 switch to deepspeed dummyoptim fix some bugs; still output green SP inference done! fix typos in readme and latent dataset debug file
330 lines
14 KiB
Python
330 lines
14 KiB
Python
import math
|
|
from einops import rearrange
|
|
import decord
|
|
from torch.nn import functional as F
|
|
import torch
|
|
from typing import Optional
|
|
import torch.utils
|
|
import torch.utils.data
|
|
import torch
|
|
from torch.utils.data import Sampler
|
|
from typing import List
|
|
from collections import Counter
|
|
import random
|
|
|
|
|
|
IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
|
|
|
|
def is_image_file(filename):
|
|
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
|
|
|
class DecordInit(object):
|
|
"""Using Decord(https://github.com/dmlc/decord) to initialize the video_reader."""
|
|
|
|
def __init__(self, num_threads=1):
|
|
self.num_threads = num_threads
|
|
self.ctx = decord.cpu(0)
|
|
|
|
def __call__(self, filename):
|
|
"""Perform the Decord initialization.
|
|
Args:
|
|
results (dict): The resulting dict to be modified and passed
|
|
to the next transform in pipeline.
|
|
"""
|
|
reader = decord.VideoReader(filename,
|
|
ctx=self.ctx,
|
|
num_threads=self.num_threads)
|
|
return reader
|
|
|
|
def __repr__(self):
|
|
repr_str = (f'{self.__class__.__name__}('
|
|
f'sr={self.sr},'
|
|
f'num_threads={self.num_threads})')
|
|
return repr_str
|
|
|
|
def pad_to_multiple(number, ds_stride):
|
|
remainder = number % ds_stride
|
|
if remainder == 0:
|
|
return number
|
|
else:
|
|
padding = ds_stride - remainder
|
|
return number + padding
|
|
|
|
class Collate:
|
|
def __init__(self, args):
|
|
self.batch_size = args.train_batch_size
|
|
self.group_frame = args.group_frame
|
|
self.group_resolution = args.group_resolution
|
|
|
|
self.max_height = args.max_height
|
|
self.max_width = args.max_width
|
|
self.ae_stride = args.ae_stride
|
|
|
|
self.ae_stride_t = args.ae_stride_t
|
|
self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride)
|
|
|
|
self.patch_size = args.patch_size
|
|
self.patch_size_t = args.patch_size_t
|
|
|
|
self.num_frames = args.num_frames
|
|
self.use_image_num = args.use_image_num
|
|
self.max_thw = (self.num_frames, self.max_height, self.max_width)
|
|
|
|
def package(self, batch):
|
|
batch_tubes = [i['pixel_values'] for i in batch] # b [c t h w]
|
|
input_ids = [i['input_ids'] for i in batch] # b [1 l]
|
|
cond_mask = [i['cond_mask'] for i in batch] # b [1 l]
|
|
return batch_tubes, input_ids, cond_mask
|
|
|
|
def __call__(self, batch):
|
|
batch_tubes, input_ids, cond_mask = self.package(batch)
|
|
|
|
ds_stride = self.ae_stride * self.patch_size
|
|
t_ds_stride = self.ae_stride_t * self.patch_size_t
|
|
|
|
pad_batch_tubes, attention_mask, input_ids, cond_mask = self.process(batch_tubes, input_ids, cond_mask, t_ds_stride, ds_stride, self.max_thw, self.ae_stride_thw)
|
|
assert not torch.any(torch.isnan(pad_batch_tubes)), 'after pad_batch_tubes'
|
|
return pad_batch_tubes, attention_mask, input_ids, cond_mask
|
|
|
|
|
|
def process(self, batch_tubes, input_ids, cond_mask, t_ds_stride, ds_stride, max_thw, ae_stride_thw):
|
|
# pad to max multiple of ds_stride
|
|
batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)]
|
|
assert len(batch_input_size) == self.batch_size
|
|
if self.group_frame or self.group_resolution or self.batch_size == 1: #
|
|
len_each_batch = batch_input_size
|
|
idx_length_dict = dict([*zip(list(range(self.batch_size)), len_each_batch)])
|
|
count_dict = Counter(len_each_batch)
|
|
if len(count_dict) != 1:
|
|
sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
|
|
# import ipdb;ipdb.set_trace()
|
|
# print(batch, idx_length_dict, count_dict, sorted_by_value)
|
|
pick_length = sorted_by_value[-1][0] # the highest frequency
|
|
candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length]
|
|
random_select_batch = [random.choice(candidate_batch) for _ in range(len(len_each_batch) - len(candidate_batch))]
|
|
print(batch_input_size, idx_length_dict, count_dict, sorted_by_value, pick_length, candidate_batch, random_select_batch)
|
|
pick_idx = candidate_batch + random_select_batch
|
|
|
|
batch_tubes = [batch_tubes[i] for i in pick_idx]
|
|
batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)]
|
|
input_ids = [input_ids[i] for i in pick_idx] # b [1, l]
|
|
cond_mask = [cond_mask[i] for i in pick_idx] # b [1, l]
|
|
|
|
for i in range(1, self.batch_size):
|
|
assert batch_input_size[0] == batch_input_size[i]
|
|
max_t = max([i[1] for i in batch_input_size])
|
|
max_h = max([i[2] for i in batch_input_size])
|
|
max_w = max([i[3] for i in batch_input_size])
|
|
else:
|
|
max_t, max_h, max_w = max_thw
|
|
pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1+self.ae_stride_t, t_ds_stride), \
|
|
pad_to_multiple(max_h, ds_stride), \
|
|
pad_to_multiple(max_w, ds_stride)
|
|
pad_max_t = pad_max_t + 1 - self.ae_stride_t
|
|
each_pad_t_h_w = [
|
|
[
|
|
pad_max_t - i.shape[1],
|
|
pad_max_h - i.shape[2],
|
|
pad_max_w - i.shape[3]
|
|
] for i in batch_tubes
|
|
]
|
|
pad_batch_tubes = [
|
|
F.pad(im, (0, pad_w, 0, pad_h, 0, pad_t), value=0)
|
|
for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)
|
|
]
|
|
pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0)
|
|
|
|
|
|
max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
|
|
max_latent_size = [
|
|
((max_tube_size[0]-1) // ae_stride_thw[0] + 1),
|
|
max_tube_size[1] // ae_stride_thw[1],
|
|
max_tube_size[2] // ae_stride_thw[2]
|
|
]
|
|
valid_latent_size = [
|
|
[
|
|
int(math.ceil((i[1]-1) / ae_stride_thw[0])) + 1,
|
|
int(math.ceil(i[2] / ae_stride_thw[1])),
|
|
int(math.ceil(i[3] / ae_stride_thw[2]))
|
|
] for i in batch_input_size]
|
|
attention_mask = [
|
|
F.pad(torch.ones(i, dtype=pad_batch_tubes.dtype), (0, max_latent_size[2] - i[2],
|
|
0, max_latent_size[1] - i[1],
|
|
0, max_latent_size[0] - i[0]), value=0) for i in valid_latent_size]
|
|
attention_mask = torch.stack(attention_mask) # b t h w
|
|
if self.batch_size == 1 or self.group_frame or self.group_resolution:
|
|
assert torch.all(attention_mask.bool())
|
|
|
|
input_ids = torch.stack(input_ids) # b 1 l
|
|
cond_mask = torch.stack(cond_mask) # b 1 l
|
|
|
|
return pad_batch_tubes, attention_mask, input_ids, cond_mask
|
|
|
|
|
|
def split_to_even_chunks(indices, lengths, num_chunks, batch_size):
|
|
"""
|
|
Split a list of indices into `chunks` chunks of roughly equal lengths.
|
|
"""
|
|
|
|
if len(indices) % num_chunks != 0:
|
|
chunks = [indices[i::num_chunks] for i in range(num_chunks)]
|
|
else:
|
|
num_indices_per_chunk = len(indices) // num_chunks
|
|
|
|
chunks = [[] for _ in range(num_chunks)]
|
|
chunks_lengths = [0 for _ in range(num_chunks)]
|
|
for index in indices:
|
|
shortest_chunk = chunks_lengths.index(min(chunks_lengths))
|
|
chunks[shortest_chunk].append(index)
|
|
chunks_lengths[shortest_chunk] += lengths[index]
|
|
if len(chunks[shortest_chunk]) == num_indices_per_chunk:
|
|
chunks_lengths[shortest_chunk] = float("inf")
|
|
# return chunks
|
|
|
|
pad_chunks = []
|
|
for idx, chunk in enumerate(chunks):
|
|
if batch_size != len(chunk):
|
|
assert batch_size > len(chunk)
|
|
if len(chunk) != 0:
|
|
chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))]
|
|
else:
|
|
chunk = random.choice(pad_chunks)
|
|
print(chunks[idx], '->', chunk)
|
|
pad_chunks.append(chunk)
|
|
return pad_chunks
|
|
|
|
def group_frame_fun(indices, lengths):
|
|
# sort by num_frames
|
|
indices.sort(key=lambda i: lengths[i], reverse=True)
|
|
return indices
|
|
|
|
def group_resolution_fun(indices):
|
|
raise NotImplementedError
|
|
return indices
|
|
|
|
def group_frame_and_resolution_fun(indices):
|
|
raise NotImplementedError
|
|
return indices
|
|
|
|
def last_group_frame_fun(shuffled_megabatches, lengths):
|
|
re_shuffled_megabatches = []
|
|
# print('shuffled_megabatches', len(shuffled_megabatches))
|
|
for i_megabatch, megabatch in enumerate(shuffled_megabatches):
|
|
re_megabatch = []
|
|
for i_batch, batch in enumerate(megabatch):
|
|
assert len(batch) != 0
|
|
len_each_batch = [lengths[i] for i in batch]
|
|
idx_length_dict = dict([*zip(batch, len_each_batch)])
|
|
count_dict = Counter(len_each_batch)
|
|
if len(count_dict) != 1:
|
|
sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
|
|
# print(batch, idx_length_dict, count_dict, sorted_by_value)
|
|
pick_length = sorted_by_value[-1][0] # the highest frequency
|
|
candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length]
|
|
random_select_batch = [random.choice(candidate_batch) for i in range(len(len_each_batch) - len(candidate_batch))]
|
|
# print(batch, idx_length_dict, count_dict, sorted_by_value, pick_length, candidate_batch, random_select_batch)
|
|
batch = candidate_batch + random_select_batch
|
|
# print(batch)
|
|
re_megabatch.append(batch)
|
|
re_shuffled_megabatches.append(re_megabatch)
|
|
|
|
|
|
# for megabatch, re_megabatch in zip(shuffled_megabatches, re_shuffled_megabatches):
|
|
# for batch, re_batch in zip(megabatch, re_megabatch):
|
|
# for i, re_i in zip(batch, re_batch):
|
|
# if i != re_i:
|
|
# print(i, re_i)
|
|
return re_shuffled_megabatches
|
|
|
|
|
|
def last_group_resolution_fun(indices):
|
|
raise NotImplementedError
|
|
return indices
|
|
|
|
def last_group_frame_and_resolution_fun(indices):
|
|
raise NotImplementedError
|
|
return indices
|
|
|
|
def get_length_grouped_indices(lengths, batch_size, world_size, generator=None, group_frame=False, group_resolution=False, seed=42):
|
|
# We need to use torch for the random part as a distributed sampler will set the random seed for torch.
|
|
if generator is None:
|
|
generator = torch.Generator().manual_seed(seed) # every rank will generate a fixed order but random index
|
|
# print('lengths', lengths)
|
|
|
|
indices = torch.randperm(len(lengths), generator=generator).tolist()
|
|
# print('indices', len(indices))
|
|
|
|
if group_frame and not group_resolution:
|
|
indices = group_frame_fun(indices, lengths)
|
|
elif not group_frame and group_resolution:
|
|
indices = group_resolution_fun(indices)
|
|
elif group_frame and group_resolution:
|
|
indices = group_frame_and_resolution_fun(indices)
|
|
# print('sort indices', len(indices))
|
|
# print('sort indices', indices)
|
|
# print('sort lengths', [lengths[i] for i in indices])
|
|
|
|
megabatch_size = world_size * batch_size
|
|
megabatches = [indices[i: i + megabatch_size] for i in range(0, len(lengths), megabatch_size)]
|
|
# print('megabatches', len(megabatches))
|
|
# print('\nmegabatches', megabatches)
|
|
megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]
|
|
# print('sort megabatches', len(megabatches))
|
|
# megabatches_len = [[lengths[i] for i in megabatch] for megabatch in megabatches]
|
|
# print('\nsorted megabatches', megabatches)
|
|
# print('\nsorted megabatches_len', megabatches_len)
|
|
megabatches = [split_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches]
|
|
# print('nsplit_to_even_chunks megabatches', len(megabatches))
|
|
# print('\nsplit_to_even_chunks megabatches', megabatches)
|
|
# print('\nsplit_to_even_chunks len', [lengths[i] for megabatch in megabatches for batch in megabatch for i in batch])
|
|
# return [i for megabatch in megabatches for batch in megabatch for i in batch]
|
|
|
|
indices = torch.randperm(len(megabatches), generator=generator).tolist()
|
|
shuffled_megabatches = [megabatches[i] for i in indices]
|
|
# print('shuffled_megabatches', len(shuffled_megabatches))
|
|
if group_frame and not group_resolution:
|
|
shuffled_megabatches = last_group_frame_fun(shuffled_megabatches, lengths)
|
|
elif not group_frame and group_resolution:
|
|
shuffled_megabatches = last_group_resolution_fun(shuffled_megabatches, indices)
|
|
elif group_frame and group_resolution:
|
|
shuffled_megabatches = last_group_frame_and_resolution_fun(shuffled_megabatches, indices)
|
|
# print('\nshuffled_megabatches', shuffled_megabatches)
|
|
# import ipdb;ipdb.set_trace()
|
|
# print('\nshuffled_megabatches len', [lengths[i] for megabatch in shuffled_megabatches for batch in megabatch for i in batch])
|
|
|
|
return [i for megabatch in shuffled_megabatches for batch in megabatch for i in batch]
|
|
|
|
|
|
class LengthGroupedSampler(Sampler):
|
|
r"""
|
|
Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
|
|
keeping a bit of randomness.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
batch_size: int,
|
|
world_size: int,
|
|
lengths: Optional[List[int]] = None,
|
|
group_frame=False,
|
|
group_resolution=False,
|
|
generator=None,
|
|
):
|
|
if lengths is None:
|
|
raise ValueError("Lengths must be provided.")
|
|
|
|
self.batch_size = batch_size
|
|
self.world_size = world_size
|
|
self.lengths = lengths
|
|
self.group_frame = group_frame
|
|
self.group_resolution = group_resolution
|
|
self.generator = generator
|
|
|
|
def __len__(self):
|
|
return len(self.lengths)
|
|
|
|
def __iter__(self):
|
|
indices = get_length_grouped_indices(self.lengths, self.batch_size, self.world_size, group_frame=self.group_frame,
|
|
group_resolution=self.group_resolution, generator=self.generator)
|
|
return iter(indices)
|