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

train vis

Browse Source
This commit is contained in:
root
2024-04-30 15:23:43 +00:00
parent 594d98265b
commit aa6d088e91
9 changed files with 75 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3 -1
View File
@@ -28,4 +28,6 @@ sample_videos/
1024*
debug*
private*
caption*
caption*
*deepspeed*
revised*
opensora/dataset/t2v_datasets.py
+2
View File
@@ -69,6 +69,7 @@ class T2V_dataset(Dataset):
video = self.decord_read(video_path, frame_idx)
video = self.transform(video) # T C H W -> T C H W
# video = torch.rand(65, 3, 512, 512)
video = video.transpose(0, 1) # T C H W -> C T H W
text = self.vid_cap_list[idx]['cap'][0]
@@ -147,6 +148,7 @@ class T2V_dataset(Dataset):
start_frame_ind, end_frame_ind = int(start_frame_ind), int(end_frame_ind)
# assert end_frame_ind - start_frame_ind >= self.num_frames
frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
# frame_indice = np.linspace(0, end_frame_ind - 2, self.num_frames, dtype=int)
video_data = decord_vr.get_batch(frame_indice).asnumpy()
video_data = torch.from_numpy(video_data)
opensora/sample/pipeline_videogen.py
+2 -4
View File
@@ -751,9 +751,7 @@ class VideoGenPipeline(DiffusionPipeline):
return VideoPipelineOutput(video=video)
def decode_latents(self, latents):
video = self.vae.decode(latents)
# video = self.vae.decode(latents / 0.18215)
# video = rearrange(video, 'b c t h w -> b t c h w').contiguous()
video = self.vae.decode(latents) # b t c h w
# b t c h w -> b t h w c
video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous()
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
return video
opensora/sample/sample_t2v.py
+1
View File
@@ -110,6 +110,7 @@ def main(args):
num_images_per_prompt=1,
mask_feature=True,
).video
print(videos.shape)
try:
if args.force_images:
videos = videos[:, 0].permute(0, 3, 1, 2) # b t h w c -> b c h w
opensora/train/train_t2v.py
+55 -110
View File
@@ -14,14 +14,13 @@ import os
import shutil
from pathlib import Path
from typing import Optional
import gc
import numpy as np
from einops import rearrange
from tqdm import tqdm
from dataclasses import field, dataclass
from torch.utils.data import DataLoader
from copy import deepcopy
import accelerate
import torch
from torch.nn import functional as F
@@ -49,53 +48,64 @@ from opensora.models.text_encoder import get_text_enc, get_text_warpper
from opensora.utils.dataset_utils import Collate
from opensora.models.ae import ae_stride_config, ae_channel_config
from opensora.models.diffusion import Diffusion_models
from opensora.sample.pipeline_videogen import VideoGenPipeline
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
check_min_version("0.24.0")
logger = get_logger(__name__)
def generate_timestep_weights(args, num_timesteps):
weights = torch.ones(num_timesteps)
# Determine the indices to bias
num_to_bias = int(args.timestep_bias_portion * num_timesteps)
if args.timestep_bias_strategy == "later":
bias_indices = slice(-num_to_bias, None)
elif args.timestep_bias_strategy == "earlier":
bias_indices = slice(0, num_to_bias)
elif args.timestep_bias_strategy == "range":
# Out of the possible 1000 timesteps, we might want to focus on eg. 200-500.
range_begin = args.timestep_bias_begin
range_end = args.timestep_bias_end
if range_begin < 0:
raise ValueError(
"When using the range strategy for timestep bias, you must provide a beginning timestep greater or equal to zero."
@torch.inference_mode()
def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step):
validation_prompt = [
"A quiet beach at dawn, the waves gently lapping at the shore and the sky painted in pastel hues.",
"The majestic beauty of a waterfall cascading down a cliff into a serene lake."
]
logger.info(f"Running validation....\n")
model = accelerator.unwrap_model(model)
scheduler = PNDMScheduler()
videogen_pipeline = VideoGenPipeline(vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
scheduler=scheduler,
transformer=model).to(device=accelerator.device)
videos = []
for prompt in validation_prompt:
logger.info('Processing the ({}) prompt'.format(prompt))
video = videogen_pipeline(prompt,
video_length=args.video_length,
height=args.max_image_size,
width=args.max_image_size,
num_inference_steps=args.num_sampling_steps,
guidance_scale=args.guidance_scale,
enable_temporal_attentions=True,
num_images_per_prompt=1,
mask_feature=True,
).video
videos.append(video[0])
# import ipdb;ipdb.set_trace()
gc.collect()
torch.cuda.empty_cache()
videos = torch.stack(videos).numpy()
videos = rearrange(videos, 'b t h w c -> b t c h w')
for tracker in accelerator.trackers:
if tracker.name == "tensorboard":
np_videos = np.stack([np.asarray(vid) for vid in videos])
tracker.writer.add_video("validation", np_videos, global_step, fps=10)
if tracker.name == "wandb":
import wandb
tracker.log(
{
"validation": [
wandb.Video(video, caption=f"{i}: {prompt}", fps=10)
for i, (video, prompt) in enumerate(zip(videos, validation_prompt))
]
}
)
if range_end > num_timesteps:
raise ValueError(
"When using the range strategy for timestep bias, you must provide an ending timestep smaller than the number of timesteps."
)
bias_indices = slice(range_begin, range_end)
else: # 'none' or any other string
return weights
if args.timestep_bias_multiplier <= 0:
return ValueError(
"The parameter --timestep_bias_multiplier is not intended to be used to disable the training of specific timesteps."
" If it was intended to disable timestep bias, use `--timestep_bias_strategy none` instead."
" A timestep bias multiplier less than or equal to 0 is not allowed."
)
# Apply the bias
weights[bias_indices] *= args.timestep_bias_multiplier
# Normalize
weights /= weights.sum()
return weights
del videogen_pipeline
gc.collect()
torch.cuda.empty_cache()
#################################################################################
# Training Loop #
#################################################################################
@@ -115,7 +125,6 @@ def main(args):
if args.report_to == "wandb":
if not is_wandb_available():
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
import wandb
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
@@ -179,6 +188,7 @@ def main(args):
args.stride_t = ae_stride_t * patch_size_t
args.stride = ae_stride_h * patch_size_h
latent_size = (args.max_image_size // ae_stride_h, args.max_image_size // ae_stride_w)
ae.latent_size = latent_size
if getae_wrapper(args.ae) == CausalVQVAEModelWrapper or getae_wrapper(args.ae) == CausalVAEModelWrapper:
args.video_length = video_length = args.num_frames // ae_stride_t + 1
@@ -534,68 +544,7 @@ def main(args):
ema_model.copy_to(model.parameters())
if args.enable_tracker:
validation_prompt = "The majestic beauty of a waterfall cascading down a cliff into a serene lake."
logger.info(f"Running validation...Using DDPM naive sampling...\n"
f"Generating {args.num_validation_videos} videos with prompt: {validation_prompt}")
with torch.no_grad():
# create pipeline
ae_ = getae_wrapper(args.ae)(args.ae_path).to(accelerator.device).eval()
if args.enable_tiling:
ae_.vae.enable_tiling()
ae_.vae.tile_overlap_factor = args.tile_overlap_factor
# text_enc_ = get_text_enc(args).to(accelerator.device).eval()
model_ = LatteT2V.from_pretrained(save_path, subfolder="model").to(accelerator.device).eval()
diffusion_ = create_diffusion(str(250))
tokenizer_ = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir='./cache_dir')
videos = []
for idx in range(args.num_validation_videos):
with torch.autocast(device_type='cuda', dtype=weight_dtype):
z = torch.randn(1, model_.in_channels, video_length,
latent_size[0], latent_size[1], device=accelerator.device)
text_tokens_and_mask = tokenizer_(
validation_prompt,
max_length=args.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'].to(accelerator.device)
cond_mask = text_tokens_and_mask['attention_mask'].to(accelerator.device)
# cond = text_enc_(input_ids, cond_mask) # B L D
cond = text_enc(input_ids, cond_mask) # B L D
model_kwargs = dict(encoder_hidden_states=cond, attention_mask=None, encoder_attention_mask=cond_mask)
sample_fn = model_.forward
# Sample images:
samples = diffusion_.p_sample_loop(
sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True,
device=accelerator.device
)
samples = ae_.decode(samples)
# Save and display images:
video = (ae_denorm[args.ae](samples[0]) * 255).add_(0.5).clamp_(0, 255).to(
dtype=torch.uint8).cpu().contiguous() # t c h w
videos.append(video)
videos = torch.stack(videos).numpy()
for tracker in accelerator.trackers:
if tracker.name == "tensorboard":
np_videos = np.stack([np.asarray(vid) for vid in videos])
tracker.writer.add_video("validation", np_videos, global_step, fps=10)
if tracker.name == "wandb":
tracker.log(
{
"validation": [
wandb.Video(video, caption=f"{i}: {validation_prompt}", fps=10)
for i, video in enumerate(videos)
]
}
)
# del ae_, text_enc_, model_, diffusion_, tokenizer_
del ae_, model_, diffusion_, tokenizer_
torch.cuda.empty_cache()
log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, weight_dtype, global_step)
accelerator.wait_for_everyone()
accelerator.end_training()
@@ -628,15 +577,11 @@ if __name__ == "__main__":
parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
parser.add_argument("--cache_dir", type=str, default='./cache_dir')
parser.add_argument("--num_sampling_steps", type=int, default=50)
parser.add_argument('--guidance_scale', type=float, default=10.0)
parser.add_argument("--enable_tracker", action="store_true")
parser.add_argument("--use_deepspeed", action="store_true")
parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
parser.add_argument(
"--num_validation_videos",
type=int,
default=2,
help="Number of images that should be generated during validation with `validation_prompt`.",
)
parser.add_argument(
"--output_dir",
type=str,
scripts/causalvae/reconstruction.sh
+2 -2
View File
@@ -1,6 +1,6 @@
CUDA_VISIBLE_DEVICES=0 python examples/rec_imvi_vae.py \
--ae_path "../../Open-Sora-Plan/CausalVAEModel_4x8x8/" \
--video_path origin_tiger.mp4 \
--ae_path "/path/to/CausalVAEModel_4x8x8/" \
--video_path input.mp4 \
--rec_path rec.mp4 \
--device cuda \
--sample_rate 1 \
scripts/text_condition/sample_video.sh
+3 -3
View File
@@ -1,12 +1,12 @@
CUDA_VISIBLE_DEVICES=6 python opensora/sample/sample_t2v.py \
--model_path /remote-home1/yeyang/dev/Open-Sora-Plan/debug_mask_loss/checkpoint-500/model \
--model_path 512/checkpoint-20/model \
--version 65x512x512 \
--image_size 512 \
--cache_dir "/remote-home1/yeyang/sora4.8/Open-Sora-Plan/cache_dir" \
--cache_dir "./cache_dir" \
--text_encoder_name DeepFloyd/t5-v1_1-xxl \
--text_prompt examples/prompt_list_0.txt \
--ae CausalVAEModel_4x8x8 \
--ae_path "/remote-home1/yeyang/sora4.8/Open-Sora-Plan/CausalVAEModel_4x8x8/" \
--ae_path "/dxyl_data02/CausalVAEModel_4x8x8/" \
--save_img_path "./sample_videos" \
--fps 24 \
--guidance_scale 10.0 \
scripts/text_condition/train_videoae_65x512x512.sh
+5 -4
View File
@@ -9,14 +9,14 @@ accelerate launch \
--cache_dir "./cache_dir" \
--dataset t2v \
--ae CausalVAEModel_4x8x8 \
--ae_path "/dxyl_code02/dev/CausalVAEModel_4x8x8/" \
--ae_path "/dxyl_data02/CausalVAEModel_4x8x8/" \
--video_data "/dxyl_code02/dev/Open-Sora-Plan/scripts/train_data/video_data.txt" \
--sample_rate 1 \
--num_frames 65 \
--max_image_size 512 \
--gradient_checkpointing \
--attention_mode xformers \
--train_batch_size=4 \
--train_batch_size=2 \
--dataloader_num_workers 10 \
--gradient_accumulation_steps=1 \
--max_train_steps=1000000 \
@@ -24,7 +24,7 @@ accelerate launch \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--mixed_precision="bf16" \
--report_to="tensorboard" \
--report_to="wandb" \
--checkpointing_steps=100 \
--output_dir="512" \
--allow_tf32 \
@@ -33,4 +33,5 @@ accelerate launch \
--use_image_num 4 \
--enable_tiling \
--use_img_from_vid \
--pretrained t2v.pt
--pretrained t2v.pt \
--enable_tracker
scripts/train_data/video_data.txt
+2 -2
View File
@@ -1,2 +1,2 @@
/dxyl_data02/mixkit,/dxyl_data02/anno_jsons/video_mixkit_65f_54735.json
/dxyl_data02/mixkit,/dxyl_data02/anno_jsons/video_mixkit_65f_54735.json
/dxyl_data02/mixkit,/dxyl_data02/anno_jsons/tmp_video_mixkit_65f_54735.json
/dxyl_data02/mixkit,/dxyl_data02/anno_jsons/tmp_video_mixkit_65f_54735.json