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

update

Browse Source
This commit is contained in:
Will Lin
2025-05-28 11:58:03 -07:00
parent 0ec4e8ecf0
commit dbfccfef40
3 changed files with 0 additions and 465 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/data_preprocess/preprocess_text_embeddings_v1.py
-199
View File
@@ -1,199 +0,0 @@
import argparse
import json
import os
import torch
import torch.distributed as dist
from accelerate.logging import get_logger
from diffusers.utils import export_to_video
from diffusers.video_processor import VideoProcessor
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm
# from fastvideo.utils.load import load_text_encoder, load_vae
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
from fastvideo.v1.utils import maybe_download_model
from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.configs.models.vaes import WanVAEConfig
from fastvideo.v1.configs.models.encoders.t5 import T5Config
logger = get_logger(__name__)
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))
rank = int(os.getenv("RANK", 0))
init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
print("world_size", world_size, "local rank", local_rank)
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
world_group = get_world_group()
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# torch.cuda.set_device(local_rank)
# if not dist.is_initialized():
# dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
videoprocessor = VideoProcessor(vae_scale_factor=8)
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True)
os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
vae_precision = "fp16"
text_encoder_precision = "fp32"
fastvideo_args = FastVideoArgs(model_path=args.model_path,
use_cpu_offload=False,
vae_precision=vae_precision,
text_encoder_precisions=(text_encoder_precision,))
fastvideo_args.device = device
fastvideo_args.device_str = f"cuda:{local_rank}"
# fastvideo_args.dit_config = HunyuanVideoConfig()
fastvideo_args.vae_config = WanVAEConfig()
fastvideo_args.text_encoder_configs = (T5Config(),)
# vae_loader = VAELoader()
# vae = vae_loader.load_vae()
text_encoder_loader = TextEncoderLoader()
tokenizer_loader = TokenizerLoader()
model_path = args.model_path
path = maybe_download_model(model_path)
# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
ENCODER_PATH = os.path.join(path, "text_encoder")
TOKENIZER_PATH = os.path.join(path, "tokenizer")
print(ENCODER_PATH)
text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
train_dataset = T5dataset(latents_json_path, args.vae_debug)
# text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
# vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
# vae.enable_tiling()
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 tqdm(enumerate(train_dataloader), disable=local_rank != 0):
with torch.inference_mode():
# with torch.autocast("cuda", dtype=torch.float32):
print(data["caption"])
text_inputs = tokenizer(data["caption"], **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
fastvideo_args.device)
input_ids = text_inputs["input_ids"]
attention_mask = text_inputs["attention_mask"]
outputs = text_encoder(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
)
from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
post_process_func = t5_postprocess_text
prompt_embeds = post_process_func(outputs)
prompt_attention_mask = attention_mask
if args.vae_debug:
latents = data["latents"]
video = vae.decode(latents.to(device), return_dict=False)[0]
video = videoprocessor.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=16)
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]
with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f:
json.dump(all_json_data, f, indent=4)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# dataset & dataloader
parser.add_argument("--model_path", type=str, default="data/mochi")
# parser.add_argument("--model_type", type=str, default="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/data_preprocess/preprocess_vae_latents_v1.py
-151
View File
@@ -1,151 +0,0 @@
import argparse
import json
import os
import torch
# import torch.distributed as dist
# from accelerate.logging import get_logger
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm
from fastvideo.dataset import getdataset
# from fastvideo.utils.load import load_vae
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.loader.component_loader import VAELoader
from fastvideo.v1.utils import maybe_download_model
from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.configs.models.vaes import WanVAEConfig
logger = init_logger(__name__)
model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
path = maybe_download_model(model_path)
# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
VAE_PATH = os.path.join(path, "vae")
print(VAE_PATH)
def main(args):
local_rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))
rank = int(os.getenv("RANK", 0))
init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
print("world_size", world_size, "local rank", local_rank)
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
world_group = get_world_group()
vae_precision = "fp16"
fastvideo_args = FastVideoArgs(model_path=VAE_PATH,
use_cpu_offload=False,
vae_precision=vae_precision)
fastvideo_args.device = device
# fastvideo_args.dit_config = HunyuanVideoConfig()
fastvideo_args.vae_config = WanVAEConfig()
train_dataset = getdataset(args)
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,
)
# encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# torch.cuda.set_device(local_rank)
# if not dist.is_initialized():
# dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
vae_loader = VAELoader()
vae = vae_loader.load(VAE_PATH, "vae", fastvideo_args)
# vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
# vae.enable_tiling()
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
json_data = []
for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
with torch.inference_mode():
with torch.autocast("cuda", dtype=torch.float16):
latents = vae.encode(data["pixel_values"].to(device)).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")
world_group.barrier()
local_data = json_data
gathered_data = [None] * world_size
for i in range(world_size):
if local_rank == i:
world_group.broadcast_object(local_data, src=i)
else:
gathered_data[i] = world_group.broadcast_object(None, src=i)
gathered_data[local_rank] = json_data
print(gathered_data)
if local_rank == 0:
all_json_data = [item for sublist in gathered_data for item in sublist]
with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f:
json.dump(all_json_data, f, indent=4)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# dataset & dataloader
parser.add_argument("--model_path", type=str, default="data/mochi")
# parser.add_argument("--model_type", type=str, default="mochi")
parser.add_argument("--data_merge_path", type=str, required=True)
parser.add_argument("--num_frames", type=int, default=163)
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=16,
help="Batch size (per device) for the training dataloader.",
)
parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
parser.add_argument("--max_height", type=int, default=480)
parser.add_argument("--max_width", type=int, default=848)
parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
parser.add_argument("--group_frame", action="store_true") # TODO
parser.add_argument("--group_resolution", action="store_true") # TODO
parser.add_argument("--dataset", default="t2v")
parser.add_argument("--train_fps", type=int, default=30)
parser.add_argument("--use_image_num", type=int, default=0)
parser.add_argument("--text_max_length", type=int, default=256)
parser.add_argument("--speed_factor", type=float, default=1.0)
parser.add_argument("--drop_short_ratio", type=float, default=1.0)
# text encoder & vae & diffusion model
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("--cfg", type=float, default=0.0)
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(
"--logging_dir",
type=str,
default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
)
args = parser.parse_args()
main(args)
fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py
-115
View File
@@ -1,115 +0,0 @@
import argparse
import os
import torch
# import torch.distributed as dist
from accelerate.logging import get_logger
# from fastvideo.utils.load import load_text_encoder
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
from fastvideo.v1.utils import maybe_download_model
from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.configs.models.vaes import WanVAEConfig
from fastvideo.v1.configs.models.encoders.t5 import T5Config
logger = get_logger(__name__)
def main(args):
local_rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))
rank = int(os.getenv("RANK", 0))
init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
print("world_size", world_size, "local rank", local_rank)
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
world_group = get_world_group()
vae_precision = "fp16"
text_encoder_precision = "fp32"
fastvideo_args = FastVideoArgs(model_path=args.model_path,
use_cpu_offload=False,
vae_precision=vae_precision,
text_encoder_precisions=(text_encoder_precision,))
fastvideo_args.device = device
fastvideo_args.device_str = f"cuda:{local_rank}"
# fastvideo_args.dit_config = HunyuanVideoConfig()
fastvideo_args.vae_config = WanVAEConfig()
fastvideo_args.text_encoder_configs = (T5Config(),)
# vae_loader = VAELoader()
# vae = vae_loader.load_vae()
text_encoder_loader = TextEncoderLoader()
tokenizer_loader = TokenizerLoader()
model_path = args.model_path
path = maybe_download_model(model_path)
# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
ENCODER_PATH = os.path.join(path, "text_encoder")
TOKENIZER_PATH = os.path.join(path, "tokenizer")
print(ENCODER_PATH)
text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
# text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
# autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16
# output_dir/validation/prompt_attention_mask
# output_dir/validation/prompt_embed
os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True)
os.makedirs(
os.path.join(args.output_dir, "validation", "prompt_attention_mask"),
exist_ok=True,
)
os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True)
with open(args.validation_prompt_txt, "r", encoding="utf-8") as file:
lines = file.readlines()
prompts = [line.strip() for line in lines]
for prompt in prompts:
with torch.inference_mode():
# with torch.autocast("cuda", dtype=autocast_type):
text_inputs = tokenizer(prompt, **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
fastvideo_args.device)
input_ids = text_inputs["input_ids"]
attention_mask = text_inputs["attention_mask"]
outputs = text_encoder(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
)
from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
post_process_func = t5_postprocess_text
prompt_embeds = post_process_func(outputs)
prompt_attention_mask = attention_mask
file_name = prompt.split(".")[0]
prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt")
prompt_attention_mask_path = os.path.join(
args.output_dir,
"validation",
"prompt_attention_mask",
f"{file_name}.pt",
)
torch.save(prompt_embeds[0], prompt_embed_path)
torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
print(f"sample {file_name} saved")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# dataset & dataloader
parser.add_argument("--model_path", type=str, default="data/mochi")
parser.add_argument("--validation_prompt_txt", type=str)
parser.add_argument(
"--output_dir",
type=str,
default=None,
help="The output directory where the model predictions and checkpoints will be written.",
)
args = parser.parse_args()
main(args)