nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor dynamic training

Browse Source
This commit is contained in:
LinB203
2024-04-14 11:37:16 +00:00
parent 8fa3e614a1
commit 8a6db0b399
6 changed files with 103 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4 -1
View File
@@ -18,4 +18,7 @@ wandb/
*.png
*.gif
*.pth
*.pt
*.pt
cache_dir/
wandb/
test/
opensora/dataset/t2v_datasets.py
+49 -34
View File
@@ -14,6 +14,10 @@ from opensora.utils.dataset_utils import DecordInit
from opensora.utils.utils import text_preprocessing
def random_video_noise(t, c, h, w):
vid = torch.rand(t, c, h, w) * 255.0
vid = vid.to(torch.uint8)
return vid
class T2V_dataset(Dataset):
def __init__(self, args, transform, temporal_sample, tokenizer):
@@ -40,47 +44,58 @@ class T2V_dataset(Dataset):
def __getitem__(self, idx):
try:
# video = torch.randn(3, 16, 128, 128)
# input_ids = torch.ones(1, 120).to(torch.long).squeeze(0)
# cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0)
# return video, input_ids, cond_mask
video_path = self.samples[idx]['path']
video = self.decord_read(video_path)
video = self.transform(video) # T C H W -> T C H W
video = video.transpose(0, 1) # T C H W -> C T H W
text = self.samples[idx]['cap'][0]
text = text_preprocessing(text)
text_tokens_and_mask = self.tokenizer(
text,
max_length=self.model_max_length,
padding='max_length',
truncation=True,
return_attention_mask=True,
add_special_tokens=True,
return_tensors='pt'
)
input_ids = text_tokens_and_mask['input_ids'].squeeze(0)
cond_mask = text_tokens_and_mask['attention_mask'].squeeze(0)
video_data = self.get_video(idx)
image_data = {}
if self.use_image_num != 0 and self.use_img_from_vid:
select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int)
assert self.num_frames >= self.use_image_num
images = video[:, select_image_idx] # c, num_img, h, w
video = torch.cat([video, images], dim=1) # c, num_frame+num_img, h, w
input_ids = torch.stack([input_ids] * (1+self.use_image_num)) # 1+self.use_image_num, l
cond_mask = torch.stack([cond_mask] * (1+self.use_image_num)) # 1+self.use_image_num, l
image_data = self.get_image_from_video(video_data)
elif self.use_image_num != 0 and not self.use_img_from_vid:
images, captions = self.img_cap_list[idx]
raise NotImplementedError
image_data = self.get_image(idx)
else:
pass
return video, input_ids, cond_mask
raise NotImplementedError
return dict(video_data=video_data, image_data=image_data)
except Exception as e:
print(f'Error with {e}, {self.samples[idx]}')
return self.__getitem__(random.randint(0, self.__len__() - 1))
def get_video(self, idx):
# video = random.choice([random_video_noise(65, 3, 720, 360) * 255, random_video_noise(65, 3, 1024, 1024), random_video_noise(65, 3, 360, 720)])
# print('random shape', video.shape)
# input_ids = torch.ones(1, 120).to(torch.long).squeeze(0)
# cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0)
video_path = self.samples[idx]['path'].replace('/remote-home1/dataset/data_split_tt', '/yzwldata01/dataset/split_1024')
video = self.decord_read(video_path)
video = self.transform(video) # T C H W -> T C H W
video = video.transpose(0, 1) # T C H W -> C T H W
text = self.samples[idx]['cap'][0]
text = text_preprocessing(text)
text_tokens_and_mask = self.tokenizer(
text,
max_length=self.model_max_length,
padding='max_length',
truncation=True,
return_attention_mask=True,
add_special_tokens=True,
return_tensors='pt'
)
input_ids = text_tokens_and_mask['input_ids']
cond_mask = text_tokens_and_mask['attention_mask']
return dict(video=video, input_ids=input_ids, cond_mask=cond_mask)
def get_image_from_video(self, video_data):
select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int)
assert self.num_frames >= self.use_image_num
image = [video_data['video'][:, i:i+1] for i in select_image_idx] # num_img [c, 1, h, w]
input_ids = video_data['input_ids'].repeat(self.use_image_num, 1) # self.use_image_num, l
cond_mask = video_data['cond_mask'].repeat(self.use_image_num, 1) # self.use_image_num, l
return dict(image=image, input_ids=input_ids, cond_mask=cond_mask)
def get_image(self, idx):
raise NotImplementedError
image_data = self.img_cap_list[idx]
return dict(image=image, input_ids=input_ids, cond_mask=cond_mask)
def tv_read(self, path):
vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
total_frames = len(vframes)
opensora/models/diffusion/latte/modules.py
+11 -7
View File
@@ -263,10 +263,6 @@ class Attention(nn.Module):
self.dropout = dropout
self.use_rope = use_rope
if self.use_rope:
self.rope2d = RoPE2D(freq=100)
self.rope1d = RoPE1D(freq=100)
# we make use of this private variable to know whether this class is loaded
# with an deprecated state dict so that we can convert it on the fly
self._from_deprecated_attn_block = _from_deprecated_attn_block
@@ -350,7 +346,7 @@ class Attention(nn.Module):
# but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
if processor is None:
processor = (
AttnProcessor2_0(attention_mode) if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
AttnProcessor2_0(attention_mode, use_rope) if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
)
self.set_processor(processor)
@@ -843,8 +839,13 @@ class AttnProcessor2_0:
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
"""
def __init__(self, attention_mode='xformers'):
def __init__(self, attention_mode='xformers', use_rope=False):
self.attention_mode = attention_mode
self.use_rope = use_rope
if self.use_rope:
self.rope2d = RoPE2D(freq=100)
self.rope1d = RoPE1D(freq=100)
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
@@ -904,7 +905,10 @@ class AttnProcessor2_0:
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
if self.rope is not None:
if self.use_rope:
print('use rope')
print('use rope')
print('use rope')
# require the shape of (batch_size x nheads x ntokens x dim)
if position_q.ndim == 3:
query = self.rope2d(query, position_q)
opensora/utils/dataset_utils.py
+32 -22
View File
@@ -1,5 +1,5 @@
import math
from einops import rearrange
import decord
from torch.nn import functional as F
import torch
@@ -58,45 +58,55 @@ class Collate:
self.num_frames = args.num_frames
self.use_image_num = args.use_image_num
self.max_thw = (self.num_frames, self.max_image_size, self.max_image_size)
self.max_ihw = (self.use_image_num, self.max_image_size, self.max_image_size)
self.max_1hw = (1, self.max_image_size, self.max_image_size)
def package(self, batch):
batch_tubes_vid = [i['video_data']['video'] for i in batch] # b [c t h w]
input_ids_vid = torch.stack([i['video_data']['input_ids'] for i in batch]) # b 1 l
cond_mask_vid = torch.stack([i['video_data']['cond_mask'] for i in batch]) # b 1 l
batch_tubes_img, input_ids_img, cond_mask_img = None, None, None
if self.use_image_num != 0:
import ipdb;ipdb.set_trace()
batch_tubes_img = [j for i in batch for j in i['image_data']['image']] # b*num_img [c 1 h w]
input_ids_img = torch.stack([i['image_data']['input_ids'] for i in batch]) # b image_num l
cond_mask_img = torch.stack([i['image_data']['cond_mask'] for i in batch]) # b image_num l
return batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img
def __call__(self, batch):
unzip = tuple(zip(*batch))
batch_tubes, input_ids, cond_mask = unzip
input_ids = torch.stack(input_ids).squeeze(1)
cond_mask = torch.stack(cond_mask).squeeze(1)
batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img = self.package(batch)
ds_stride = self.ae_stride * self.patch_size
t_ds_stride = self.ae_stride_t * self.patch_size_t
if self.use_image_num == 0:
pad_batch_tubes, attention_mask = process(batch_tubes, t_ds_stride, ds_stride, self.max_thw, self.ae_stride_thw, self.patch_size_thw)
pad_batch_tubes, attention_mask = process(batch_tubes_vid, t_ds_stride, ds_stride,
self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
# attention_mask: b t hw, use first frame
attention_mask = attention_mask[:, 0] # b hw
input_ids, cond_mask = input_ids_vid.squeeze(1), cond_mask_vid.squeeze(1) # b 1 l -> b l
else:
batch_tubes_vid = [i[:, :-self.use_image_num] for i in batch_tubes]
pad_batch_tubes_vid, attention_mask_vid = process(batch_tubes_vid, t_ds_stride, ds_stride, self.max_thw, self.ae_stride_thw, self.patch_size_thw)
print('pad_batch_tubes_vid.shape, attention_mask_vid.shape', pad_batch_tubes_vid.shape, attention_mask_vid.shape)
pad_batch_tubes_vid, attention_mask_vid = self.process(batch_tubes_vid, t_ds_stride, ds_stride,
self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
# attention_mask_vid: b t hw, use first frame
attention_mask_vid = attention_mask_vid[:, :1] # b 1 hw
batch_tubes_img = [i[:, -self.use_image_num:] for i in batch_tubes]
pad_batch_tubes_img, attention_mask_img = process(batch_tubes_img, 1, ds_stride, self.max_ihw, self.ae_stride_1hw, self.patch_size_1hw)
print('pad_batch_tubes_img.shape, attention_mask_img.shape', pad_batch_tubes_img.shape, attention_mask_img.shape)
pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=3)
pad_batch_tubes_img, attention_mask_img = self.process(batch_tubes_img, 1, ds_stride,
self.max_1hw, self.ae_stride_1hw, self.patch_size_1hw, extra_1=False)
pad_batch_tubes_img = rearrange(pad_batch_tubes_img, '(b i) c 1 h w -> b c i h w', i=self.use_image_num)
attention_mask_img = rearrange(attention_mask_img, '(b i) 1 hw -> b i hw', i=self.use_image_num)
pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=2) # concat at temporal, video first
# attention_mask_img: b num_img hw
attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1) # b 1+num_img hw
input_ids = torch.cat([input_ids_vid, input_ids_img], dim=1) # b 1+num_img hw
cond_mask = torch.cat([cond_mask_vid, cond_mask_img], dim=1) # b 1+num_img hw
return pad_batch_tubes, attention_mask, input_ids, cond_mask
def process(batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw):
def process(self, batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw, extra_1):
# 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)]
max_t, max_h, max_w = max_thw
# max_t, max_h, max_w = max([i[1] for i in batch_input_size]), \
# max([i[2] for i in batch_input_size]), \
# max([i[3] for i in batch_input_size])
pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t, t_ds_stride), \
pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1 if extra_1 else max_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 if extra_1 else pad_max_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]
@@ -108,13 +118,13 @@ class Collate:
# make attention_mask
max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
max_latent_size = [max_tube_size[0] // ae_stride_thw[0],
max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
max_tube_size[1] // ae_stride_thw[1],
max_tube_size[2] // ae_stride_thw[2]]
max_patchify_latent_size = [max_latent_size[0] // patch_size_thw[0],
max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]),
max_latent_size[1] // patch_size_thw[1],
max_latent_size[2] // patch_size_thw[2]]
valid_patchify_latent_size = [[int(math.ceil(i[1] / t_ds_stride)),
valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)),
int(math.ceil(i[2] / ds_stride)),
int(math.ceil(i[3] / ds_stride))] for i in batch_input_size]
attention_mask = [F.pad(torch.ones(i),
scripts/accelerate_configs/ddp_config.yaml
+3 -3
View File
@@ -1,11 +1,11 @@
compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
distributed_type: NO
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: 29501
main_training_function: main
num_machines: 1
num_processes: 8
gpu_ids: 0,1,2,3,4,5,6,7
num_processes: 1
gpu_ids: 0,
use_cpu: false
scripts/text_condition/train_videoae_65x512x512.sh
+4 -4
View File
@@ -2,7 +2,7 @@ export WANDB_KEY=""
export ENTITY=""
export PROJECT="test"
accelerate launch \
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
--config_file scripts/accelerate_configs/ddp_config.yaml \
opensora/train/train_t2v.py \
--model LatteT2V-XL/122 \
--text_encoder_name DeepFloyd/t5-v1_1-xxl \
@@ -17,15 +17,15 @@ accelerate launch \
--max_image_size 512 \
--gradient_checkpointing \
--attention_mode xformers \
--train_batch_size=2 \
--dataloader_num_workers 10 \
--train_batch_size=4 \
--dataloader_num_workers 0 \
--gradient_accumulation_steps=1 \
--max_train_steps=1000000 \
--learning_rate=2e-05 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--mixed_precision="bf16" \
--report_to="wandb" \
--report_to="tensorboard" \
--checkpointing_steps=500 \
--output_dir="test" \
--allow_tf32 \