From aa6d088e91cd0290761edb073343440e7bab3dbd Mon Sep 17 00:00:00 2001 From: root <2267330597@qq.com> Date: Tue, 30 Apr 2024 15:23:43 +0000 Subject: [PATCH] train vis --- .gitignore | 4 +- opensora/dataset/t2v_datasets.py | 2 + opensora/sample/pipeline_videogen.py | 6 +- opensora/sample/sample_t2v.py | 1 + opensora/train/train_t2v.py | 165 ++++++------------ scripts/causalvae/reconstruction.sh | 4 +- scripts/text_condition/sample_video.sh | 6 +- .../train_videoae_65x512x512.sh | 9 +- scripts/train_data/video_data.txt | 4 +- 9 files changed, 75 insertions(+), 126 deletions(-) diff --git a/.gitignore b/.gitignore index 95d0f4204..6bb760f1f 100644 --- a/.gitignore +++ b/.gitignore @@ -28,4 +28,6 @@ sample_videos/ 1024* debug* private* -caption* \ No newline at end of file +caption* +*deepspeed* +revised* \ No newline at end of file diff --git a/opensora/dataset/t2v_datasets.py b/opensora/dataset/t2v_datasets.py index 7fa829990..c2f99b7d6 100644 --- a/opensora/dataset/t2v_datasets.py +++ b/opensora/dataset/t2v_datasets.py @@ -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) diff --git a/opensora/sample/pipeline_videogen.py b/opensora/sample/pipeline_videogen.py index 303414b8c..c48943a25 100644 --- a/opensora/sample/pipeline_videogen.py +++ b/opensora/sample/pipeline_videogen.py @@ -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 diff --git a/opensora/sample/sample_t2v.py b/opensora/sample/sample_t2v.py index bfb743b65..d9a3069cd 100644 --- a/opensora/sample/sample_t2v.py +++ b/opensora/sample/sample_t2v.py @@ -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 diff --git a/opensora/train/train_t2v.py b/opensora/train/train_t2v.py index 58bc2ba5e..c2a1f5ecf 100644 --- a/opensora/train/train_t2v.py +++ b/opensora/train/train_t2v.py @@ -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, diff --git a/scripts/causalvae/reconstruction.sh b/scripts/causalvae/reconstruction.sh index 3e89d2c28..0c71d27ea 100644 --- a/scripts/causalvae/reconstruction.sh +++ b/scripts/causalvae/reconstruction.sh @@ -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 \ diff --git a/scripts/text_condition/sample_video.sh b/scripts/text_condition/sample_video.sh index fb6b7bd10..6c17894b4 100644 --- a/scripts/text_condition/sample_video.sh +++ b/scripts/text_condition/sample_video.sh @@ -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 \ diff --git a/scripts/text_condition/train_videoae_65x512x512.sh b/scripts/text_condition/train_videoae_65x512x512.sh index db823c3cb..411affde2 100644 --- a/scripts/text_condition/train_videoae_65x512x512.sh +++ b/scripts/text_condition/train_videoae_65x512x512.sh @@ -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 diff --git a/scripts/train_data/video_data.txt b/scripts/train_data/video_data.txt index 89fff89d5..f95abc801 100644 --- a/scripts/train_data/video_data.txt +++ b/scripts/train_data/video_data.txt @@ -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 \ No newline at end of file +/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 \ No newline at end of file