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

[feat]: Add batchy data preprocess (#53)

Browse Source 
Co-authored-by:Runlong <rlsu9@ucsd.edu>
This commit is contained in:
rlsu9
2024-11-25 21:52:36 -08:00
committed by GitHub
co-authored by Runlong
parent 9c45a15f53
commit faef037f2e
4 changed files with 86 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/dataset/latent_datasets.py
+1 -1
View File
@@ -27,7 +27,7 @@ class LatentDataset(Dataset):
self.num_latent_t = num_latent_t
self.uncond_prompt_embed = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "embed.pt"), map_location="cpu", weights_only=True)
self.uncond_prompt_mask = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "mask.pt"), map_location="cpu", weights_only=True)
self.lengths = [torch.load(os.path.join(self.latent_dir, data_item["latent_path"]), map_location="cpu").shape[1] for data_item in self.data_anno]
self.lengths = [data_item['length'] if "length" in data_item else 1 for data_item in self.data_anno]
def __getitem__(self, idx):
latent_file = self.data_anno[idx]["latent_path"]
prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
fastvideo/dataset/t2v_datasets.py
-2
View File
@@ -141,7 +141,6 @@ class T2V_dataset(Dataset):
video = torchvision_video[frame_indices]
video = self.transform(video)
video = rearrange(video, 't c h w -> c t h w')
video = video.unsqueeze(0)
video = video.to(torch.uint8)
assert video.dtype == torch.uint8
@@ -181,7 +180,6 @@ class T2V_dataset(Dataset):
image = self.transform_topcrop(image) if 'human_images' in image_data['path'] else self.transform(image) # [1 C H W] -> num_img [1 C H W]
image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W]
image = image.unsqueeze(0)
image = image.float() / 127.5 - 1.0
fastvideo/utils/data_preprocess/finetune_data_T5.py
+71 -39
View File
@@ -7,7 +7,35 @@ import json
import os
import torch.distributed as dist
logger = get_logger(__name__)
from torch.utils.data import Dataset
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
class T5dataset(Dataset):
def __init__(
self,
json_path,
vae_debug,
):
self.json_path = json_path
self.vae_debug = vae_debug
with open(self.json_path, "r") as f:
train_dataset = json.load(f)
self.train_dataset = sorted(train_dataset, key=lambda x: x['latent_path'])
def __getitem__(self, idx):
caption = self.train_dataset[idx]['caption']
filename = self.train_dataset[idx]['latent_path'].split('.')[0]
length = self.train_dataset[idx]['length']
if self.vae_debug:
latents = torch.load(os.path.join(args.output_dir, 'latent', self.train_dataset[idx]['latent_path']), map_location="cpu")
else:
latents = []
return dict(caption=caption, latents=latents, filename=filename, length=length)
def __len__(self):
return len(self.train_dataset)
def main(args):
local_rank = int(os.getenv('RANK', 0))
world_size = int(os.getenv('WORLD_SIZE', 1))
@@ -27,45 +55,47 @@ def main(args):
os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
with open(latents_json_path, "r") as f:
train_dataset = json.load(f)
train_dataset = sorted(train_dataset, key=lambda x: x['latent_path'])
train_dataset = T5dataset(latents_json_path, args.vae_debug)
sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
train_dataloader = DataLoader(
train_dataset,
sampler=sampler,
batch_size=args.train_batch_size,
num_workers=args.dataloader_num_workers,
)
json_data = []
for _, data in enumerate(train_dataset):
video_name =data['latent_path'].split(".")[0]
if int(video_name) % world_size != local_rank:
continue
try:
with torch.inference_mode():
with torch.autocast("cuda", dtype=torch.bfloat16):
latent = torch.load(os.path.join(args.output_dir, 'latent', data['latent_path']))
prompt_embeds, prompt_attention_mask, _, _ = pipe.encode_prompt(
prompt=data['caption'],
)
prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask", video_name + ".pt")
# save latent
torch.save(prompt_embeds[0], prompt_embed_path)
torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
print(f"sample {video_name} saved")
video = pipe.vae.decode(latent.unsqueeze(0).to(device), return_dict=False)[0]
video = pipe.video_processor.postprocess_video(video)
export_to_video(video[0], video_path, fps=30)
item = {}
item["latent_path"] = video_name + ".pt"
item["prompt_embed_path"] = video_name + ".pt"
item["prompt_attention_mask"] = video_name + ".pt"
item["caption"] = data['caption']
json_data.append(item)
except:
print("video out of memory")
continue
dist.barrier()
local_data = json_data
gathered_data = [None] * world_size
dist.all_gather_object(gathered_data, local_data)
json_data = []
for _, data in enumerate(train_dataloader):
with torch.inference_mode():
with torch.autocast("cuda", dtype=torch.bfloat16):
prompt_embeds, prompt_attention_mask, _, _ = pipe.encode_prompt(
prompt=data['caption'],
)
if args.vae_debug:
latents = data['latents']
video = pipe.vae.decode(latents.to(device), return_dict=False)[0]
video = pipe.video_processor.postprocess_video(video)
for idx, video_name in enumerate(data['filename']):
prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask", video_name + ".pt")
# save latent
torch.save(prompt_embeds[idx], prompt_embed_path)
torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
print(f"sample {video_name} saved")
if args.vae_debug:
export_to_video(video[idx], video_path, fps=30)
item = {}
item['length'] = int(data['length'][idx])
item["latent_path"] = video_name + ".pt"
item["prompt_embed_path"] = video_name + ".pt"
item["prompt_attention_mask"] = video_name + ".pt"
item["caption"] = data['caption'][idx]
json_data.append(item)
dist.barrier()
local_data = json_data
gathered_data = [None] * world_size
dist.all_gather_object(gathered_data, local_data)
if local_rank == 0:
# os.remove(latents_json_path)
all_json_data = [item for sublist in gathered_data for item in sublist]
@@ -77,9 +107,11 @@ if __name__ == "__main__":
# dataset & dataloader
parser.add_argument("--model_path", type=str, default="data/mochi")
# text encoder & vae & diffusion model
parser.add_argument("--dataloader_num_workers", type=int, default=1, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.")
parser.add_argument("--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader.")
parser.add_argument("--text_encoder_name", type=str, default='google/t5-v1_1-xxl')
parser.add_argument("--cache_dir", type=str, default='./cache_dir')
parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.")
parser.add_argument("--vae_debug",action="store_true")
args = parser.parse_args()
main(args)
fastvideo/utils/data_preprocess/finetune_data_VAE.py
+14 -12
View File
@@ -10,6 +10,7 @@ import json
import os
from diffusers import AutoencoderKLMochi
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
logger = get_logger(__name__)
@@ -27,9 +28,10 @@ def main(args):
project_config=accelerator_project_config,
)
train_dataset = getdataset(args)
sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
train_dataloader = DataLoader(
train_dataset,
shuffle=True,
sampler=sampler,
batch_size=args.train_batch_size,
num_workers=args.dataloader_num_workers,
)
@@ -46,19 +48,19 @@ def main(args):
json_data = []
for _, data in enumerate(train_dataloader):
video_name = os.path.basename(data['path'][0]).split(".")[0]
if int(video_name) % world_size != local_rank:
continue
with torch.inference_mode():
with torch.autocast("cuda", dtype=torch.bfloat16):
latents = vae.encode(data['pixel_values'][0].to(encoder_device))['latent_dist'].sample()
latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
torch.save(latents[0].to(torch.bfloat16), latent_path)
item = {}
item["latent_path"] = video_name + ".pt"
item["caption"] = data['text']
json_data.append(item)
print(f"{video_name} processed")
latents = vae.encode(data['pixel_values'].to(encoder_device))['latent_dist'].sample()
for idx, video_path in enumerate(data['path']):
video_name = os.path.basename(video_path).split(".")[0]
latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
torch.save(latents[idx].to(torch.bfloat16), latent_path)
item = {}
item["length"] = latents[idx].shape[1]
item["latent_path"] = video_name + ".pt"
item["caption"] = data['text'][idx]
json_data.append(item)
print(f"{video_name} processed")
dist.barrier()
local_data = json_data
gathered_data = [None] * world_size