From 6c723e75f34fd108cd71a3ec09f2bac818a65e97 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Mon, 28 Oct 2024 17:50:34 +0000 Subject: [PATCH] amend log validation --- opensora/train/train_t2v_diffusers.py | 91 ++++++++++----------------- 1 file changed, 34 insertions(+), 57 deletions(-) diff --git a/opensora/train/train_t2v_diffusers.py b/opensora/train/train_t2v_diffusers.py index 0050da71f..2405fac12 100644 --- a/opensora/train/train_t2v_diffusers.py +++ b/opensora/train/train_t2v_diffusers.py @@ -14,12 +14,10 @@ import os import shutil from pathlib import Path import gc -import numpy as np +import wandb from einops import rearrange from tqdm import tqdm -from opensora.adaptor.modules import replace_with_fp32_forwards - from opensora.utils.parallel_states import initialize_sequence_parallel_state, \ destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state from opensora.utils.communications import prepare_parallel_data, broadcast @@ -37,7 +35,12 @@ from packaging import version from tqdm.auto import tqdm import diffusers -from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler +from diffusers import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + MochiTransformer3DModel, + MochiPipeline, +) from diffusers.optimization import get_scheduler from diffusers.training_utils import compute_snr from diffusers.utils import check_min_version, is_wandb_available @@ -59,42 +62,24 @@ logger = get_logger(__name__) @torch.inference_mode() def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, ema=False): - positive_prompt = "(masterpiece), (best quality), (ultra-detailed), {}. emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, sharp focus, high budget, cinemascope, moody, epic, gorgeous" - negative_prompt = """nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, - """ - validation_prompt = [ - "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.", - "A serene underwater scene featuring a sea turtle swimming through a coral reef. The turtle, with its greenish-brown shell, is the main focus of the video, swimming gracefully towards the right side of the frame. The coral reef, teeming with life, is visible in the background, providing a vibrant and colorful backdrop to the turtle's journey. Several small fish, darting around the turtle, add a sense of movement and dynamism to the scene." - ] - if 'mt5' in args.text_encoder_name: - validation_prompt_cn = [ - "一只戴着墨镜在泳池当救生员的猫咪。", - "这是一个宁静的水下场景，一只海龟游过珊瑚礁。海龟带着绿褐色的龟壳，优雅地游向画面右侧，成为视频的焦点。背景中的珊瑚礁生机盎然，为海龟的旅程提供了生动多彩的背景。几条小鱼在海龟周围穿梭，为画面增添了动感和活力。" - ] - validation_prompt += validation_prompt_cn logger.info(f"Running validation....\n") model = accelerator.unwrap_model(model) - scheduler = DPMSolverMultistepScheduler() - opensora_pipeline = OpenSoraPipeline(vae=vae, + scheduler = FlowMatchEulerDiscreteScheduler() + mochi_pipeline = MochiPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler, transformer=model).to(device=accelerator.device) videos = [] - for prompt in validation_prompt: + for prompt in args.validaiton_prompt: logger.info('Processing the ({}) prompt'.format(prompt)) - video = opensora_pipeline( - positive_prompt.format(prompt), - negative_prompt=negative_prompt, + video = mochi_pipeline( + args.prompt, num_frames=args.num_frames, height=args.max_height, width=args.max_width, - num_inference_steps=args.num_sampling_steps, + num_inference_steps=args.validation_sampling_steps, guidance_scale=args.guidance_scale, - enable_temporal_attentions=True, - num_images_per_prompt=1, - mask_feature=True, - max_sequence_length=args.model_max_length, ).images videos.append(video[0]) # import ipdb;ipdb.set_trace() @@ -103,35 +88,24 @@ def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weigh 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": - if videos.shape[1] == 1: - assert args.num_frames == 1 - images = rearrange(videos, 'b 1 c h w -> (b 1) h w c') - np_images = np.stack([np.asarray(img) for img in images]) - tracker.writer.add_images(f"{'ema_' if ema else ''}validation", np_images, global_step, dataformats="NHWC") - else: - np_videos = np.stack([np.asarray(vid) for vid in videos]) - tracker.writer.add_video(f"{'ema_' if ema else ''}validation", np_videos, global_step, fps=24) - if tracker.name == "wandb": - import wandb - if videos.shape[1] == 1: - images = rearrange(videos, 'b 1 c h w -> (b 1) h w c') - logs = { - f"{'ema_' if ema else ''}validation": [ - wandb.Image(image, caption=f"{i}: {prompt}") - for i, (image, prompt) in enumerate(zip(images, validation_prompt)) - ] - } - else: - logs = { - f"{'ema_' if ema else ''}validation": [ - wandb.Video(video, caption=f"{i}: {prompt}", fps=24) - for i, (video, prompt) in enumerate(zip(videos, validation_prompt)) - ] - } - tracker.log(logs, step=global_step) + if videos.shape[1] == 1: + images = rearrange(videos, 'b 1 c h w -> (b 1) h w c') + logs = { + f"{'ema_' if ema else ''}validation": [ + wandb.Image(image, caption=f"{i}: {prompt}") + for i, (image, prompt) in enumerate(zip(images, validation_prompt)) + ] + } + else: + logs = { + f"{'ema_' if ema else ''}validation": [ + wandb.Video(video, caption=f"{i}: {prompt}", fps=24) + for i, (video, prompt) in enumerate(zip(videos, validation_prompt)) + ] + } + tracker.log(logs, step=global_step) - del opensora_pipeline + del mochi_pipeline gc.collect() torch.cuda.empty_cache() @@ -150,6 +124,7 @@ def main(args): logging_dir = Path(args.output_dir, args.logging_dir) # use LayerNorm, GeLu, SiLu always as fp32 mode + # TODO: if args.enable_stable_fp32: replace_with_fp32_forwards() @@ -160,6 +135,7 @@ def main(args): mixed_precision=args.mixed_precision, log_with=args.report_to, project_config=accelerator_project_config, + log_with="wandb" ) if args.num_frames != 1 and args.use_image_num == 0: @@ -813,7 +789,8 @@ if __name__ == "__main__": parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.") # validation & logs - parser.add_argument("--num_sampling_steps", type=int, default=20) + parser.add_argument("--validaiton_prompt", nargs='+', help="List of prompts to use for validation.") + parser.add_argument("--validation_sampling_steps", type=int, default=64) parser.add_argument('--guidance_scale', type=float, default=4.5) parser.add_argument("--enable_tracker", action="store_true") parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")