| 3D animation of a small, round, fluffy creature with big, expressive eyes explores ... | +A single drop of liquid metal falls from a floating orb, landing on a mirror-like ... | +The video presents an abstract composition centered around a hexagonal shape adorned ... | +
| + | + | + |
| A drone camera circles around a beautiful historic church built on a rocky outcropping ... | +Aerial view of Santorini during the blue hour, showcasing the stunning architecture ... | +An aerial shot of a lighthouse standing tall on a rocky cliff, its beacon cutting ... | +
| + | + | + |
| A snowy forest landscape with a dirt road running through it. The road is flanked by ... | +Drone shot along the Hawaii jungle coastline, sunny day. Kayaks in the water. | +The camera rotates around a large stack of vintage televisions all showing different ... | +
| + | + | + |
| In an ornate, historical hall, a massive tidal wave peaks and begins to crash. Two ... | +A Shiba Inu dog wearing a beret and black turtleneck. | +A painting of a boat on water comes to life, with waves crashing and the boat becoming ... | +
| + | + | + |
| A person clad in a space suit with a helmet and equipped with a chest light and arm ... | +3D animation of a small, round, fluffy creature with big, expressive eyes explores a ... | +In a studio, there is a painting depicting a ship sailing through the rough sea. | +
| + | + | + |
| A robot dog trots down a deserted alley at night, its metallic paws clinking softly ... | +A lone surfer rides a massive wave, skillfully maneuvering through the surf. The water ... | +A solitary cheetah sprints across the savannah, its powerful muscles propelling it ... | +
| + | + | + |
| Generated | ++ | + | + |
| Edited | ++ | + | + |
+Highly recommend trying out our web demo by the following command. We also provide [online demo](https://huggingface.co/spaces/LanguageBind/Open-Sora-Plan-v1.1.0) [](https://huggingface.co/spaces/LanguageBind/Open-Sora-Plan-v1.1.0).
+
+
+
+ # Open-Sora Plan 93x720p
+ #### [Open-Sora Plan 93x720p](https://github.com/PKU-YuanGroup/Open-Sora-Plan) is a transformer-based text-to-video diffusion system trained on text embeddings from mT5. This demo uses the []() checkpoint.
+ #### Multilingual prompts SUPPORT; 支持多语言
+ #### Welcome to Star🌟 our [GitHub](https://github.com/PKU-YuanGroup/Open-Sora-Plan)
+ ### You may change frame to 29 if you're not satisfied with the speed, as generating short video is significantly faster than long one.
+ """
+NEG_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.
+ """
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
- if randomize_seed:
- seed = random.randint(0, 203279)
- return seed
-title_markdown = ("""
-
-
-# 
-# Running on CPU 🥶 This demo does not work on CPU.
" -@torch.inference_mode() -def generate_img(prompt, sample_steps, scale, seed=0, randomize_seed=False, force_images=False): - seed = int(randomize_seed_fn(seed, randomize_seed)) - set_env(seed) - video_length = transformer_model.config.video_length if not force_images else 1 - height, width = int(args.version.split('x')[1]), int(args.version.split('x')[2]) - num_frames = 1 if video_length == 1 else int(args.version.split('x')[0]) - videos = videogen_pipeline(prompt, - video_length=video_length, - height=height, - width=width, - num_inference_steps=sample_steps, - guidance_scale=scale, - enable_temporal_attentions=not force_images, - num_images_per_prompt=1, - mask_feature=True, - ).video - - torch.cuda.empty_cache() - videos = videos[0] - tmp_save_path = 'tmp.mp4' - imageio.mimwrite(tmp_save_path, videos, fps=24, quality=9) # highest quality is 10, lowest is 0 - display_model_info = f"Video size: {num_frames}×{height}×{width}, \nSampling Step: {sample_steps}, \nGuidance Scale: {scale}" - return tmp_save_path, prompt, display_model_info, seed - -if __name__ == '__main__': - args = type('args', (), { - 'ae': 'CausalVAEModel_4x8x8', - 'force_images': False, - 'model_path': 'LanguageBind/Open-Sora-Plan-v1.0.0', - 'text_encoder_name': 'DeepFloyd/t5-v1_1-xxl', - 'version': '65x512x512' - }) - device = torch.device('cuda:0') - - # Load model: - transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, torch_dtype=torch.float16, cache_dir='cache_dir').to(device) - - vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir').to(device) - vae = vae.half() +CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1" +CACHE_EXAMPLES = False +MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "6000")) +MAX_VIDEO_FRAME = int(os.getenv("MAX_IMAGE_SIZE", "93")) +SPEED_UP_T5 = os.getenv("USE_TORCH_COMPILE", "0") == "1" +USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1" +ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1" +PORT = int(os.getenv("DEMO_PORT", "15432")) + + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list} +STYLE_NAMES = list(styles.keys()) +DEFAULT_STYLE_NAME = "(Default)" +SCHEDULE_NAME = [ + "PNDM-Solver", "EulerA-Solver", "DPM-Solver", "SA-Solver", + "DDIM-Solver", "Euler-Solver", "DDPM-Solver", "DEISM-Solver"] +DEFAULT_SCHEDULE_NAME = "PNDM-Solver" +NUM_IMAGES_PER_PROMPT = 1 + +def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]: + p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) + if not negative: + negative = "" + return p.replace("{prompt}", positive), n + negative + +if torch.cuda.is_available(): + weight_dtype = torch.bfloat16 + T5_token_max_length = 512 + + vae = getae_wrapper('CausalVAEModel_4x8x8')("/storage/dataset/test140k") + vae.vae = vae.vae.to(device=device, dtype=weight_dtype) vae.vae.enable_tiling() - image_size = int(args.version.split('x')[1]) - latent_size = (image_size // ae_stride_config[args.ae][1], image_size // ae_stride_config[args.ae][2]) - vae.latent_size = latent_size - transformer_model.force_images = args.force_images - tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir="cache_dir") - text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir="cache_dir", - torch_dtype=torch.float16).to(device) - - # set eval mode - transformer_model.eval() - vae.eval() - text_encoder.eval() + vae.vae.tile_overlap_factor = 0.125 + vae.vae.tile_sample_min_size = 256 + vae.vae.tile_latent_min_size = 32 + vae.vae.tile_sample_min_size_t = 29 + vae.vae.tile_latent_min_size_t = 8 + vae.vae_scale_factor = ae_stride_config['CausalVAEModel_4x8x8'] + + text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", + low_cpu_mem_usage=True, torch_dtype=weight_dtype) + tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl") + transformer = OpenSoraT2V.from_pretrained("/storage/dataset/hw29/model_ema", low_cpu_mem_usage=False, + device_map=None, torch_dtype=weight_dtype) scheduler = PNDMScheduler() - videogen_pipeline = VideoGenPipeline(vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - transformer=transformer_model).to(device) - - - demo = gr.Interface( - fn=generate_img, - inputs=[Textbox(label="", - placeholder="Please enter your prompt. \n"), - gr.Slider( - label='Sample Steps', + pipe = OpenSoraPipeline(vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer) + pipe.to(device) + print("Loaded on Device!") + + # speed-up T5 + if SPEED_UP_T5: + pipe.text_encoder.to_bettertransformer() + + if USE_TORCH_COMPILE: + pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True) + print("Model Compiled!") + +# @spaces.GPU(duration=120) +@torch.no_grad() +@torch.inference_mode() +def generate( + prompt: str, + negative_prompt: str = "", + style: str = DEFAULT_STYLE_NAME, + use_negative_prompt: bool = False, + seed: int = 0, + frame: int = 29, + schedule: str = 'DPM-Solver', + guidance_scale: float = 4.5, + num_inference_steps: int = 25, + randomize_seed: bool = False, + progress=gr.Progress(track_tqdm=True), +): + seed = int(randomize_seed_fn(seed, randomize_seed)) + generator = torch.Generator().manual_seed(seed) + + if schedule == 'DPM-Solver': + if not isinstance(pipe.scheduler, DPMSolverMultistepScheduler): + pipe.scheduler = DPMSolverMultistepScheduler() + elif schedule == "PNDM-Solver": + if not isinstance(pipe.scheduler, PNDMScheduler): + pipe.scheduler = PNDMScheduler() + elif schedule == "DDIM-Solver": + if not isinstance(pipe.scheduler, DDIMScheduler): + pipe.scheduler = DDIMScheduler() + elif schedule == "Euler-Solver": + if not isinstance(pipe.scheduler, EulerDiscreteScheduler): + pipe.scheduler = EulerDiscreteScheduler() + elif schedule == "DDPM-Solver": + if not isinstance(pipe.scheduler, DDPMScheduler): + pipe.scheduler = DDPMScheduler() + elif schedule == "EulerA-Solver": + if not isinstance(pipe.scheduler, EulerAncestralDiscreteScheduler): + pipe.scheduler = EulerAncestralDiscreteScheduler() + elif schedule == "DEISM-Solver": + if not isinstance(pipe.scheduler, DEISMultistepScheduler): + pipe.scheduler = DEISMultistepScheduler() + elif schedule == "SA-Solver": + if not isinstance(pipe.scheduler, SASolverScheduler): + pipe.scheduler = SASolverScheduler.from_config(pipe.scheduler.config, algorithm_type='data_prediction', tau_func=lambda t: 1 if 200 <= t <= 800 else 0, predictor_order=2, corrector_order=2) + else: + raise ValueError(f"Unknown schedule: {schedule}") + + if not use_negative_prompt: + negative_prompt = None # type: ignore + prompt, negative_prompt = apply_style(style, prompt, negative_prompt) + print(prompt, negative_prompt) + videos = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + num_frames=frame, + # width=1280, + # height=720, + width=640, + height=480, + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + generator=generator, + num_images_per_prompt=1, # num_imgs + max_sequence_length=T5_token_max_length, + ).images + + video_paths = [save_video(vid) for vid in videos] + print(video_paths) + return video_paths[0], seed + + +examples = [ + "A small cactus with a happy face in the Sahara desert.", + "Eiffel Tower was Made up of more than 2 million translucent straws to look like a cloud, with the bell tower at the top of the building, Michel installed huge foam-making machines in the forest to blow huge amounts of unpredictable wet clouds in the building's classic architecture.", + "3D animation of a small, round, fluffy creature with big, expressive eyes explores a vibrant, enchanted forest. The creature, a whimsical blend of a rabbit and a squirrel, has soft blue fur and a bushy, striped tail. It hops along a sparkling stream, its eyes wide with wonder. The forest is alive with magical elements: flowers that glow and change colors, trees with leaves in shades of purple and silver, and small floating lights that resemble fireflies. The creature stops to interact playfully with a group of tiny, fairy-like beings dancing around a mushroom ring. The creature looks up in awe at a large, glowing tree that seems to be the heart of the forest.", + "Color photo of a corgi made of transparent glass, standing on the riverside in Yosemite National Park.", + "A close-up photo of a person. The subject is a woman. She wore a blue coat with a gray dress underneath. She has blue eyes and blond hair, and wears a pair of earrings. Behind are blurred city buildings and streets.", + "A litter of golden retriever puppies playing in the snow. Their heads pop out of the snow, covered in.", + "a handsome young boy in the middle with sky color background wearing eye glasses, it's super detailed with anime style, it's a portrait with delicated eyes and nice looking face", + "an astronaut sitting in a diner, eating fries, cinematic, analog film", + "Pirate ship trapped in a cosmic maelstrom nebula, rendered in cosmic beach whirlpool engine, volumetric lighting, spectacular, ambient lights, light pollution, cinematic atmosphere, art nouveau style, illustration art artwork by SenseiJaye, intricate detail.", + "professional portrait photo of an anthropomorphic cat wearing fancy gentleman hat and jacket walking in autumn forest.", + "The parametric hotel lobby is a sleek and modern space with plenty of natural light. The lobby is spacious and open with a variety of seating options. The front desk is a sleek white counter with a parametric design. The walls are a light blue color with parametric patterns. The floor is a light wood color with a parametric design. There are plenty of plants and flowers throughout the space. The overall effect is a calm and relaxing space. occlusion, moody, sunset, concept art, octane rendering, 8k, highly detailed, concept art, highly detailed, beautiful scenery, cinematic, beautiful light, hyperreal, octane render, hdr, long exposure, 8K, realistic, fog, moody, fire and explosions, smoke, 50mm f2.8", +] + +with gr.Blocks(css="style.css") as demo: + gr.Markdown(DESCRIPTION) + gr.DuplicateButton( + value="Duplicate Space for private use", + elem_id="duplicate-button", + visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1", + ) + with gr.Row(equal_height=False): + with gr.Group(): + with gr.Row(): + use_negative_prompt = gr.Checkbox(label="Use additional negative prompt", value=False, visible=True) + negative_prompt = gr.Text( + label="Negative prompt", + max_lines=1, + placeholder="Enter a additional negative prompt", + visible=True, + ) + with gr.Row(visible=True): + schedule = gr.Radio( + show_label=True, + container=True, + interactive=True, + choices=SCHEDULE_NAME, + value=DEFAULT_SCHEDULE_NAME, + label="Sampler Schedule", + visible=True, + ) + style_selection = gr.Radio( + show_label=True, + container=True, + interactive=True, + choices=STYLE_NAMES, + value=DEFAULT_STYLE_NAME, + label="Video Style", + ) + seed = gr.Slider( + label="Seed", + minimum=0, + maximum=MAX_SEED, + step=1, + value=0, + ) + randomize_seed = gr.Checkbox(label="Randomize seed", value=True) + with gr.Row(visible=True): + frame = gr.Slider( + label="Frame", + minimum=29, + maximum=MAX_VIDEO_FRAME, + step=16, + value=29, + ) + with gr.Row(): + guidance_scale = gr.Slider( + label="Guidance scale", minimum=1, - maximum=500, - value=50, - step=10 - ), - gr.Slider( - label='Guidance Scale', - minimum=0.1, - maximum=30.0, - value=10.0, - step=0.1 - ), - gr.Slider( - label="Seed", - minimum=0, - maximum=203279, + maximum=10, + step=0.1, + value=5.0, + ) + inference_steps = gr.Slider( + label="inference steps", + minimum=10, + maximum=200, step=1, - value=0, - ), - gr.Checkbox(label="Randomize seed", value=True), - gr.Checkbox(label="Generate image (1 frame video)", value=False), - ], - outputs=[Video(label="Vid", width=512, height=512), - Textbox(label="input prompt"), - Textbox(label="model info"), - gr.Slider(label='seed')], - title=title_markdown, description=DESCRIPTION, theme=gr.themes.Default(), css=block_css, + value=50, + ) + with gr.Group(): + with gr.Row(): + prompt = gr.Text( + label="Prompt", + show_label=False, + max_lines=1, + placeholder="Enter your prompt", + container=False, + ) + run_button = gr.Button("Run", scale=0) + result = gr.Video(label="Result") + + gr.Examples( examples=examples, + inputs=prompt, + outputs=[result, seed], + fn=generate, + cache_examples=CACHE_EXAMPLES, ) - demo.launch() \ No newline at end of file + + use_negative_prompt.change( + fn=lambda x: gr.update(visible=x), + inputs=use_negative_prompt, + outputs=negative_prompt, + api_name=False, + ) + + gr.on( + triggers=[ + prompt.submit, + negative_prompt.submit, + run_button.click, + ], + fn=generate, + inputs=[ + prompt, + negative_prompt, + style_selection, + use_negative_prompt, + seed, + frame, + schedule, + guidance_scale, + inference_steps, + randomize_seed, + ], + outputs=[result, seed], + api_name="run", + ) + +if __name__ == "__main__": + # demo.queue(max_size=20).launch(server_name='0.0.0.0', share=True) + demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=11900, debug=True) \ No newline at end of file diff --git a/opensora/serve/style.css b/opensora/serve/style.css new file mode 100644 index 000000000..73054b14b --- /dev/null +++ b/opensora/serve/style.css @@ -0,0 +1 @@ +.gradio-container{width:1280px!important} \ No newline at end of file diff --git a/opensora/train/train_inpaint.py b/opensora/train/train_inpaint.py new file mode 100644 index 000000000..596d73cf9 --- /dev/null +++ b/opensora/train/train_inpaint.py @@ -0,0 +1,1063 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +A minimal training script for DiT using PyTorch DDP. +""" +import argparse +import logging +import math +import os +import shutil +from pathlib import Path +from typing import Optional +import gc +import numpy as np +from einops import rearrange +import torch.utils +import torch.utils.data +from tqdm import tqdm +import torch.multiprocessing as mp +import dill + + +from opensora.adaptor.modules import replace_with_fp32_forwards + +try: + import torch_npu + from opensora.npu_config import npu_config + from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, \ + destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state + from opensora.acceleration.communications import prepare_parallel_data, broadcast +except: + torch_npu = None + npu_config = None + 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 + pass +import time +# import cv2 +from ultralytics import YOLO + +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 +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedType, ProjectConfiguration, set_seed +from packaging import version +from tqdm.auto import tqdm + +import diffusers +from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import EMAModel, compute_snr +from diffusers.utils import check_min_version, is_wandb_available +# from torchvision import transforms + +from opensora.models.causalvideovae import ae_stride_config, ae_channel_config +from opensora.models.causalvideovae import ae_norm, ae_denorm +from opensora.models.text_encoder import get_text_enc, get_text_warpper +from opensora.dataset import getdataset +from opensora.models.diffusion import Diffusion_models, Diffusion_models_class +from opensora.utils.dataset_utils import Collate, LengthGroupedSampler +from opensora.sample.pipeline_opensora import OpenSoraPipeline +from opensora.models.causalvideovae import ae_stride_config, ae_wrapper +from opensora.utils.ema import EMAModel +# from opensora.sample.pipeline_inpaint import OpenSoraInpaintPipeline +# from opensora.dataset.transform import CenterCropResizeVideo + +# 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 get_video(video_path): +# cap = cv2.VideoCapture(video_path) +# frames = [] +# while True: +# ret, frame = cap.read() +# if not ret: +# break # 视频读取结束 + +# # frame 是一个 (H, W, C) 的 numpy 数组，代表一帧 +# frames.append(frame) + +# # 释放视频文件 +# cap.release() + +# # 将帧列表转换为 numpy 数组，形状为 (T, H, W, C) +# video_numpy = np.array(frames) + +# return video_numpy + +class ProgressInfo: + def __init__(self, global_step, train_loss=0.0): + self.global_step = global_step + self.train_loss = train_loss + + +# def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, motion_score, 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, +# """ + + +# input_prompts = ["A man is riding a horse"] +# videos_path = ["/home/image_data/hxy/data/video/output_video.mp4"] +# masks_path = ["/home/image_data/hxy/data/video/output_mask.mp4"] +# model = accelerator.unwrap_model(model) +# scheduler = DPMSolverMultistepScheduler() + +# inpaint_pipeline = OpenSoraInpaintPipeline( +# vae=vae, +# text_encoder=text_encoder, +# tokenizer=tokenizer, +# scheduler=scheduler, +# transformer=model +# ).to(device=accelerator.device) + +# resize = [CenterCropResizeVideo((args.height, args.width)), ] +# resize_transform = transforms.Compose([*resize]) + +# inference_videos = [] + + +# for index, (prompt, video_path, mask_path) in enumerate(zip(input_prompts, videos_path, masks_path)): +# video = get_video((video_path)) +# mask = get_video(mask_path) + +# video_tensor = resize_transform(torch.from_numpy(video.transpose(0,3,1,2))) +# mask_tensor = resize_transform(torch.from_numpy(mask.transpose(0,3,1,2))) + +# video_resize = video_tensor.numpy() +# mask_resize = mask_tensor.numpy() + + +# inference_video = inpaint_pipeline( +# video = video_resize, +# masks = mask_resize, +# prompt=prompt, +# negative_prompt=negative_prompt, +# num_frames=args.num_frames, +# height=args.height, +# width=args.width, +# motion_score=motion_score, +# num_inference_steps=args.num_sampling_steps, +# guidance_scale=args.guidance_scale, +# num_images_per_prompt=1, +# mask_feature=True, +# device=accelerator.device, +# max_sequence_length=args.max_sequence_length, +# ).images +# inference_videos.append(inference_video[0]) +# videos = torch.stack(inference_videos).numpy() +# videos = rearrange(videos, 'b t h w c -> b t c h w') + +# for tracker in accelerator.trackers: +# 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, input_prompts)) +# ] +# } +# 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, input_prompts)) +# ] +# } +# tracker.log(logs, step=global_step) + + +################################################################################# +# Training Loop # +################################################################################# + +def main(args): + logging_dir = Path(args.output_dir, args.logging_dir) + + # use LayerNorm, GeLu, SiLu always as fp32 mode + if args.enable_stable_fp32: + replace_with_fp32_forwards() + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(args) + npu_config.seed_everything(args.seed) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + if args.num_frames != 1 and args.use_image_num == 0: + initialize_sequence_parallel_state(args.sp_size) + + 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.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed, device_specific=True) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Create model: + kwargs = {} + ae = ae_wrapper[args.ae](args.ae_path, cache_dir=args.cache_dir, **kwargs).eval() + if args.enable_tiling: + ae.vae.enable_tiling() + ae.vae.tile_overlap_factor = args.tile_overlap_factor + + kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True} + text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval() + + ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] + ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w) + assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" + args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w + args.ae_stride = args.ae_stride_h + patch_size = args.model[-3:] + patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + args.patch_size = patch_size_h + args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w + assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" + # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." + assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})." + assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})." + + args.stride_t = ae_stride_t * patch_size_t + args.stride = ae_stride_h * patch_size_h + latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w) + ae.latent_size = latent_size + + if args.num_frames % 2 == 1: + args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1 + else: + latent_size_t = args.num_frames // ae_stride_t + + model_kwargs = {'vae_scale_factor_t': ae_stride_t} + + model = Diffusion_models[args.model]( + in_channels=ae_channel_config[args.ae], + out_channels=ae_channel_config[args.ae], + # caption_channels=4096, + # cross_attention_dim=1152, + attention_bias=True, + sample_size=latent_size, + sample_size_t=latent_size_t, + num_vector_embeds=None, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + use_linear_projection=False, + only_cross_attention=False, + double_self_attention=False, + upcast_attention=False, + # norm_type="ada_norm_single", + norm_elementwise_affine=False, + norm_eps=1e-6, + attention_type='default', + attention_mode=args.attention_mode, + interpolation_scale_h=args.interpolation_scale_h, + interpolation_scale_w=args.interpolation_scale_w, + interpolation_scale_t=args.interpolation_scale_t, + downsampler=args.downsampler, + # compress_kv_factor=args.compress_kv_factor, + use_rope=args.use_rope, + # model_max_length=args.model_max_length, + use_stable_fp32=args.enable_stable_fp32, + sparse1d=args.sparse1d, + sparse2d=args.sparse2d, + sparse_n=args.sparse_n, + use_motion=args.use_motion, + **model_kwargs, + ) + model.gradient_checkpointing = args.gradient_checkpointing + + pretrained_transformer_model_path = args.pretrained_transformer_model_path + pretrained_model_path = dict(transformer_model=pretrained_transformer_model_path) + if pretrained_transformer_model_path is not None: + model.custom_load_state_dict(pretrained_model_path) + + # Freeze main models + ae.vae.requires_grad_(False) + text_enc.requires_grad_(False) + # Set model as trainable. + model.train() + + ae.vae.tile_sample_min_size = args.tile_sample_min_size + ae.vae.tile_sample_min_size_t = args.tile_sample_min_size_t + + noise_scheduler = DDPMScheduler(rescale_betas_zero_snr=args.rescale_betas_zero_snr) + # Move unet, vae and text_encoder to device and cast to weight_dtype + # The VAE is in float32 to avoid NaN losses. + ae.vae.to(accelerator.device, dtype=torch.float32) + # ae.vae.to(accelerator.device, dtype=weight_dtype) + text_enc.to(accelerator.device, dtype=weight_dtype) + + # Create EMA for the unet. + if args.use_ema: + ema_model = deepcopy(model) + ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step, + model_cls=Diffusion_models_class[args.model], model_config=ema_model.config) + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + if args.use_ema: + ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) + + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "model")) + if weights: # Don't pop if empty + # make sure to pop weight so that corresponding model is not saved again + weights.pop() + + + def load_model_hook(models, input_dir): + # loading ema with customed 'from_pretrained' function + if args.use_ema: + load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model]) + ema_model.load_state_dict(load_model.state_dict()) + ema_model.to(accelerator.device) + del load_model + + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model") + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + params_to_optimize = list(filter(lambda p: p.requires_grad, model.parameters())) + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + logger.info(f"optimizer: {optimizer}") + + if args.trained_data_global_step is not None: + initial_global_step_for_sampler = args.trained_data_global_step + else: + initial_global_step_for_sampler = 0 + + train_dataset = getdataset(args) + sampler = LengthGroupedSampler( + args.train_batch_size, + world_size=accelerator.num_processes, + gradient_accumulation_size=args.gradient_accumulation_steps, + initial_global_step=initial_global_step_for_sampler, + lengths=train_dataset.lengths, + group_data=args.group_data, + ) + + Yolomodel = YOLO(args.yolomodel_pathorname).to(accelerator.device) + + train_dataloader = DataLoader( + train_dataset, + shuffle=False, + # pin_memory=True, + collate_fn=Collate(args,Yolomodel), + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + sampler=sampler, + drop_last=True, + # prefetch_factor=4 + ) + logger.info(f'after train_dataloader') + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, + num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, + ) + + # Prepare everything with our `accelerator`. + # model.requires_grad_(False) + # model.pos_embed.requires_grad_(True) + if args.adapt_vae: + model.requires_grad_(False) + for name, param in model.named_parameters(): + if 'pos_embed' in name or 'proj_out' in name: + param.requires_grad = True + logger.info(f'before accelerator.prepare') + model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, lr_scheduler + ) + logger.info(f'after accelerator.prepare') + if args.use_ema: + ema_model.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + # NOTE wandb + if accelerator.is_main_process: + logger.info("init trackers...") + project_name = os.getenv('PROJECT', os.path.basename(args.output_dir)) + entity = os.getenv('ENTITY', None) + run_name = os.getenv('WANDB_NAME', None) + init_kwargs = { + "entity": entity, + "run_name": run_name, + } + accelerator.init_trackers(project_name=project_name, config=vars(args), init_kwargs=init_kwargs) + + # Setup data: + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size + + # Train! + logger.info("***** Running training *****") + logger.info(f" Model = {model}") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + logger.info(f" Total optimization steps (num_update_steps_per_epoch) = {num_update_steps_per_epoch}") + logger.info(f" Total trainable parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B") + + global_step = 0 + first_epoch = 0 + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + progress_info = ProgressInfo(global_step, train_loss=0.0) + + def sync_gradients_info(loss): + # Checks if the accelerator has performed an optimization step behind the scenes + if args.use_ema: + ema_model.step(model.parameters()) + progress_bar.update(1) + progress_info.global_step += 1 + end_time = time.time() + one_step_duration = end_time - start_time + accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step) + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(f"Step: [{progress_info.global_step}], local_loss={loss.detach().item()}, " + f"train_loss={progress_info.train_loss}, time_cost={one_step_duration}", + rank=0) + progress_info.train_loss = 0.0 + + # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues. + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: + if progress_info.global_step % args.checkpointing_steps == 0 or progress_info.global_step == args.after_one_epoch_global_step: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if accelerator.is_main_process and args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + def run(model_input, model_kwargs, prof): + global start_time + start_time = time.time() + + try: + in_channels = ae_channel_config[args.ae] + model_input, masked_input, video_mask = model_input[:, 0:in_channels], model_input[:, in_channels:2 * in_channels], model_input[:, 2 * in_channels:] + except: + raise ValueError("masked_x and video_mask is None!") + + noise = torch.randn_like(model_input) + if args.noise_offset: + # https://www.crosslabs.org//blog/diffusion-with-offset-noise + noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1), + device=model_input.device) + + bsz = model_input.shape[0] + current_step_frame = model_input.shape[2] + # Sample a random timestep for each image without bias. + timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device) + # print('accelerator.process_index, timesteps', accelerator.process_index, timesteps) + if current_step_frame != 1 and get_sequence_parallel_state(): # image do not need sp + broadcast(timesteps) + motion_score = model_kwargs.pop('motion_score', None) + if motion_score is not None: + raise NotImplementedError + + # Add noise to the model input according to the noise magnitude at each timestep + # (this is the forward diffusion process) + + noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) + + model_pred = model( + torch.cat([noisy_model_input, masked_input, video_mask], dim=1), + timesteps, + **model_kwargs, + )[0] + + # Get the target for loss depending on the prediction type + if args.prediction_type is not None: + # set prediction_type of scheduler if defined + noise_scheduler.register_to_config(prediction_type=args.prediction_type) + + if noise_scheduler.config.prediction_type == "epsilon": + target = noise + elif noise_scheduler.config.prediction_type == "v_prediction": + target = noise_scheduler.get_velocity(model_input, noise, timesteps) + elif noise_scheduler.config.prediction_type == "sample": + # We set the target to latents here, but the model_pred will return the noise sample prediction. + target = model_input + # We will have to subtract the noise residual from the prediction to get the target sample. + model_pred = model_pred - noise + else: + raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") + + mask = model_kwargs.get('attention_mask', None) + if torch.all(mask.bool()): + mask = None + if get_sequence_parallel_state(): + assert mask is None + b, c, _, _, _ = model_pred.shape + if mask is not None: + mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float() # b t h w -> b c t h w + mask = mask.reshape(b, -1) + if args.snr_gamma is None: + # model_pred: b c t h w, attention_mask: b t h w + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.reshape(b, -1) + if mask is not None: + loss = (loss * mask).sum() / mask.sum() # mean loss on unpad patches + else: + loss = loss.mean() + else: + # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. + # Since we predict the noise instead of x_0, the original formulation is slightly changed. + # This is discussed in Section 4.2 of the same paper. + snr = compute_snr(noise_scheduler, timesteps) + mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min( + dim=1 + )[0] + if noise_scheduler.config.prediction_type == "epsilon": + mse_loss_weights = mse_loss_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + mse_loss_weights = mse_loss_weights / (snr + 1) + else: + raise NameError(f'{noise_scheduler.config.prediction_type}') + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.reshape(b, -1) + mse_loss_weights = mse_loss_weights.reshape(b, 1) + if mask is not None: + loss = (loss * mask * mse_loss_weights).sum() / mask.sum() # mean loss on unpad patches + else: + loss = (loss * mse_loss_weights).mean() + + # Gather the losses across all processes for logging (if we use distributed training). + avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() + progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps + + # Backpropagate + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = params_to_optimize + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if accelerator.sync_gradients: + sync_gradients_info(loss) + + # if accelerator.is_main_process: + + # if progress_info.global_step % 1 == 0: + + # if args.enable_tracker: + # log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + # weight_dtype, progress_info.global_step) + + # if args.use_ema and npu_config is None: + # # Store the UNet parameters temporarily and load the EMA parameters to perform inference. + # ema_model.store(model.parameters()) + # ema_model.copy_to(model.parameters()) + # log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + # weight_dtype, progress_info.global_step, motion_score,ema=True) + # # Switch back to the original UNet parameters. + # ema_model.restore(model.parameters()) + + if prof is not None: + prof.step() + + + return loss + + def train_one_step(step_, data_item_, prof_=None): + train_loss = 0.0 + x, attn_mask, input_ids, cond_mask, motion_score = data_item_ + assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))' + x = x.to(accelerator.device, dtype=ae.vae.dtype) # B C T+num_images H W, 16 + 4 + + attn_mask = attn_mask.to(accelerator.device) # B T+num_images H W + input_ids = input_ids.to(accelerator.device) # B 1+num_images L + cond_mask = cond_mask.to(accelerator.device) # B 1+num_images L + motion_score = motion_score.to(accelerator.device) if motion_score is not None else motion_score # B 1 + # if accelerator.process_index == 0: + # logger.info(f'rank: {accelerator.process_index}, x: {x.shape}, attn_mask: {attn_mask.shape}') + + with torch.no_grad(): + # import ipdb;ipdb.set_trace() + # use for loop to avoid OOM, because T5 is too huge... + B, N, L = input_ids.shape # B 1 L + # cond_ = torch.stack([text_enc(input_ids[i], cond_mask[i]) for i in range(B)]) # B 1 L D + + # use batch inference + input_ids_ = input_ids.reshape(-1, L) + cond_mask_ = cond_mask.reshape(-1, L) + cond = text_enc(input_ids_, cond_mask_) # B 1 L D + cond = cond.reshape(B, N, L, -1) + + + def preprocess_x_for_inpaint(x): + + x, masked_x, mask = x[:, :3], x[:, 3:6], x[:, 6:7] + x, masked_x = ae.encode(x), ae.encode(masked_x) + batch_size, channels, frame = mask.shape[:3] + new_frame, new_height, new_width = x.shape[2:] + mask = rearrange(mask, 'b c t h w -> (b c t) 1 h w') + mask = F.interpolate(mask, size=(new_height, new_width), mode='bilinear') + mask = rearrange(mask, '(b c t) 1 h w -> b c t h w', t=frame, b=batch_size) + mask_first_frame = mask[:, :, 0:1].repeat(1, 1, ae_stride_t, 1, 1).contiguous() + mask = torch.cat([mask_first_frame, mask[:, :, 1:]], dim=2) + mask = mask.view(batch_size, new_frame, ae_stride_t, new_height, new_width) + mask = mask.transpose(1, 2).contiguous() + + return x, masked_x, mask + + + # Map input images to latent space + normalize latents + x, masked_x, mask = preprocess_x_for_inpaint(x) + x = torch.cat([x, masked_x, mask], dim=1) + + current_step_frame = x.shape[2] + current_step_sp_state = get_sequence_parallel_state() + if args.sp_size != 1: # enable sp + if current_step_frame == 1: # but image do not need sp + set_sequence_parallel_state(False) + else: + set_sequence_parallel_state(True) + if get_sequence_parallel_state(): + x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask, + args.use_image_num) + for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size): + with accelerator.accumulate(model): + st_idx = iter * args.train_sp_batch_size + ed_idx = (iter + 1) * args.train_sp_batch_size + model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx], + attention_mask=attn_mask[st_idx: ed_idx], + encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num) + run(x[st_idx: ed_idx], model_kwargs, prof_) + + else: + with accelerator.accumulate(model): + assert not torch.any(torch.isnan(x)), 'after vae' + x = x.to(weight_dtype) + model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, motion_score=motion_score, + encoder_attention_mask=cond_mask, use_image_num=args.use_image_num,) + run(x, model_kwargs, prof_) + + set_sequence_parallel_state(current_step_sp_state) # in case the next step use sp, which need broadcast(timesteps) + + if progress_info.global_step >= args.max_train_steps: + return True + + return False + + def train_one_epoch(prof_=None): + # for epoch in range(first_epoch, args.num_train_epochs): + progress_info.train_loss = 0.0 + if progress_info.global_step >= args.max_train_steps: + return True + + args.after_one_epoch_global_step = progress_info.global_step + len(train_dataloader) // args.gradient_accumulation_steps - 1 + + for step, data_item in enumerate(train_dataloader): + if train_one_step(step, data_item, prof_): + break + + if step >= 2 and torch_npu is not None and npu_config is not None: + npu_config.free_mm() + + if npu_config is not None and npu_config.on_npu and npu_config.profiling: + experimental_config = torch_npu.profiler._ExperimentalConfig( + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization + ) + profile_output_path = f"/home/image_data/npu_profiling_t2v/{os.getenv('PROJECT_NAME', 'local')}" + os.makedirs(profile_output_path, exist_ok=True) + + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU], + with_stack=True, + record_shapes=True, + profile_memory=True, + experimental_config=experimental_config, + schedule=torch_npu.profiler.schedule(wait=npu_config.profiling_step, warmup=0, active=1, repeat=1, + skip_first=0), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/") + ) as prof: + train_one_epoch(prof) + else: + train_one_epoch() + accelerator.wait_for_everyone() + accelerator.end_training() + if npu_config is not None and get_sequence_parallel_state(): + destroy_sequence_parallel_group() + + +if __name__ == "__main__": + # mp.set_start_method('spawn', force=True) + # mp.set_start_method('spawn',force=True) + # dill.settings['recurse'] = True + # 替换 multiprocessing 的序列化机制为 dill + # mp.get_context().reduction.dump = dill.dump + parser = argparse.ArgumentParser() + + # dataset & dataloader + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--data", type=str, required='') + parser.add_argument("--sample_rate", type=int, default=1) + parser.add_argument("--train_fps", type=int, default=24) + parser.add_argument("--drop_short_ratio", type=float, default=1.0) + parser.add_argument("--speed_factor", type=float, default=1.0) + parser.add_argument("--num_frames", type=int, default=65) + parser.add_argument("--max_height", type=int, default=320) + parser.add_argument("--max_width", type=int, default=240) + parser.add_argument("--use_img_from_vid", action="store_true") + parser.add_argument("--use_image_num", type=int, default=0) + parser.add_argument("--model_max_length", type=int, default=512) + parser.add_argument('--cfg', type=float, default=0.1) + parser.add_argument("--dataloader_num_workers", type=int, default=10, 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("--group_data", action="store_true") + parser.add_argument("--hw_stride", type=int, default=32) + parser.add_argument("--skip_low_resolution", action="store_true") + parser.add_argument("--force_resolution", action="store_true") + parser.add_argument("--trained_data_global_step", type=int, default=None) + parser.add_argument("--use_decord", action="store_true") + + # text encoder & vae & diffusion model + parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122") + parser.add_argument('--enable_8bit_t5', action='store_true') + parser.add_argument('--tile_overlap_factor', type=float, default=0.125) + parser.add_argument('--enable_tiling', action='store_true') + parser.add_argument("--compress_kv", action="store_true") + parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers") + parser.add_argument('--use_rope', action='store_true') + parser.add_argument('--compress_kv_factor', type=int, default=1) + parser.add_argument('--interpolation_scale_h', type=float, default=1.0) + parser.add_argument('--interpolation_scale_w', type=float, default=1.0) + parser.add_argument('--interpolation_scale_t', type=float, default=1.0) + parser.add_argument("--downsampler", type=str, default=None) + parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") + parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") + 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("--pretrained", type=str, default=None) + parser.add_argument('--enable_stable_fp32', action='store_true') + parser.add_argument('--sparse1d', action='store_true') + parser.add_argument('--sparse2d', action='store_true') + parser.add_argument('--sparse_n', type=int, default=2) + parser.add_argument('--tile_sample_min_size', type=int, default=512) + parser.add_argument('--tile_sample_min_size_t', type=int, default=33) + parser.add_argument('--adapt_vae', action='store_true') + parser.add_argument('--use_motion', action='store_true') + parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") + + # diffusion setting + parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.") + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") + parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--ema_start_step", type=int, default=0) + parser.add_argument("--noise_offset", type=float, default=0.02, help="The scale of noise offset.") + 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.") + parser.add_argument('--rescale_betas_zero_snr', action='store_true') + + # validation & logs + parser.add_argument("--num_sampling_steps", type=int, default=50) + parser.add_argument('--guidance_scale', type=float, default=2.5) + parser.add_argument("--enable_tracker", action="store_true") + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + 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("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store.")) + parser.add_argument("--checkpointing_steps", type=int, default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument("--resume_from_checkpoint", type=str, default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + 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***." + ), + ) + parser.add_argument("--report_to", type=str, default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + # optimizer & scheduler + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") + parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]') + parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") + parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.") + parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.") + parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW") + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params") + parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.") + parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW") + parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--prodigy_beta3", type=float, default=None, + help="coefficients for computing the Prodidy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--lr_scheduler", type=str, default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument("--allow_tf32", action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel") + parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training") + + # inpaint + parser.add_argument("--t2v_ratio", type=float, default=0.0) # for inpainting mode + parser.add_argument("--i2v_ratio", type=float, default=0.0) # for inpainting mode + parser.add_argument("--transition_ratio", type=float, default=0.05) # for inpainting mode + parser.add_argument("--v2v_ratio", type=float, default=0.05) # for inpainting mode + parser.add_argument("--clear_video_ratio", type=float, default=0.0) # for inpainting mode + parser.add_argument("--Semantic_ratio", type=float, default=0.2) # for inpainting mode + parser.add_argument("--bbox_ratio", type=float, default=0.2) # for inpainting mode + parser.add_argument("--background_ratio", type=float, default=0.1) # for inpainting mode + parser.add_argument("--fixed_ratio", type=float, default=0.1) # for inpainting mode + parser.add_argument("--Semantic_expansion_ratio", type=float, default=0.1) # for inpainting mode + parser.add_argument("--fixed_bg_ratio", type=float, default=0.1) # for inpainting mode + parser.add_argument("--min_clear_ratio", type=float, default=0.1) # for inpainting mode + parser.add_argument("--default_text_ratio", type=float, default=0.5) # for inpainting mode + parser.add_argument("--pretrained_transformer_model_path", type=str, default=None) + parser.add_argument("--yolomodel_pathorname",type=str,default="/home/image_data/hxy/Open-Sora-Plan/opensora/dataset/yolov9c-seg.pt") + + #visualize + parser.add_argument("--max_sequence_length", type=int, default=512) + + + + args = parser.parse_args() + main(args) diff --git a/opensora/train/train_t2v.py b/opensora/train/train_t2v.py deleted file mode 100644 index b18828f6b..000000000 --- a/opensora/train/train_t2v.py +++ /dev/null @@ -1,846 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -""" -A minimal training script for DiT using PyTorch DDP. -""" -import argparse -import logging -import math -import os -import shutil -from pathlib import Path -from typing import Optional - -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 -import transformers -from accelerate import Accelerator -from accelerate.logging import get_logger -from accelerate.utils import ProjectConfiguration, set_seed -from huggingface_hub import create_repo -from packaging import version -from tqdm.auto import tqdm -from transformers import HfArgumentParser, TrainingArguments, AutoTokenizer - -import diffusers -from diffusers import DDPMScheduler, PNDMScheduler -from diffusers.optimization import get_scheduler -from diffusers.training_utils import EMAModel, compute_snr -from diffusers.utils import check_min_version, is_wandb_available - -from opensora.dataset import getdataset, ae_denorm -from opensora.models.ae import getae, getae_wrapper -from opensora.models.ae.videobase import CausalVQVAEModelWrapper, CausalVAEModelWrapper -from opensora.models.diffusion.diffusion import create_diffusion_T as create_diffusion -from opensora.models.diffusion.latte.modeling_latte import LatteT2V -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 - -# 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." - ) - 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 - - -################################################################################# -# Training Loop # -################################################################################# - -def main(args): - logging_dir = Path(args.output_dir, args.logging_dir) - - accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) - - accelerator = Accelerator( - gradient_accumulation_steps=args.gradient_accumulation_steps, - mixed_precision=args.mixed_precision, - log_with=args.report_to, - project_config=accelerator_project_config, - ) - - 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( - format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", - datefmt="%m/%d/%Y %H:%M:%S", - level=logging.INFO, - ) - logger.info(accelerator.state, main_process_only=False) - if accelerator.is_local_main_process: - transformers.utils.logging.set_verbosity_warning() - diffusers.utils.logging.set_verbosity_info() - else: - transformers.utils.logging.set_verbosity_error() - diffusers.utils.logging.set_verbosity_error() - - # If passed along, set the training seed now. - if args.seed is not None: - set_seed(args.seed) - - # Handle the repository creation - if accelerator.is_main_process: - if args.output_dir is not None: - os.makedirs(args.output_dir, exist_ok=True) - - # if args.push_to_hub: - # repo_id = create_repo( - # repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token - # ).repo_id - - # For mixed precision training we cast all non-trainable weigths to half-precision - # as these weights are only used for inference, keeping weights in full precision is not required. - weight_dtype = torch.float32 - if accelerator.mixed_precision == "fp16": - weight_dtype = torch.float16 - elif accelerator.mixed_precision == "bf16": - weight_dtype = torch.bfloat16 - - # Create model: - diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule - kwargs = {} - ae = getae_wrapper(args.ae)(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval() - if args.enable_tiling: - ae.vae.enable_tiling() - ae.vae.tile_overlap_factor = args.tile_overlap_factor - - kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True} - text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval() - - ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] - assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" - args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w - args.ae_stride = args.ae_stride_h - patch_size = args.model[-3:] - patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) - args.patch_size = patch_size_h - args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w - assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" - # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." - assert args.max_image_size % ae_stride_h == 0, f"Image size must be divisible by ae_stride_h, but found max_image_size ({args.max_image_size}), ae_stride_h ({ae_stride_h})." - - 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) - - if getae_wrapper(args.ae) == CausalVQVAEModelWrapper or getae_wrapper(args.ae) == CausalVAEModelWrapper: - args.video_length = video_length = args.num_frames // ae_stride_t + 1 - else: - video_length = args.num_frames // ae_stride_t - model = Diffusion_models[args.model]( - in_channels=ae_channel_config[args.ae], - out_channels=ae_channel_config[args.ae] * 2, - # caption_channels=4096, - # cross_attention_dim=1152, - attention_bias=True, - sample_size=latent_size, - num_vector_embeds=None, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - use_linear_projection=False, - only_cross_attention=False, - double_self_attention=False, - upcast_attention=False, - # norm_type="ada_norm_single", - norm_elementwise_affine=False, - norm_eps=1e-6, - attention_type='default', - video_length=video_length, - attention_mode=args.attention_mode, - compress_kv_factor=args.compress_kv_factor, - use_rope=args.use_rope, - model_max_length=args.model_max_length, - ) - model.gradient_checkpointing = args.gradient_checkpointing - - # # use pretrained model? - if args.pretrained: - if 'safetensors' in args.pretrained: - from safetensors.torch import load_file as safe_load - checkpoint = safe_load(args.pretrained, device="cpu") - else: - checkpoint = torch.load(args.pretrained, map_location='cpu')['model'] - model_state_dict = model.state_dict() - missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) - logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') - logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!') - - # Freeze vae and text encoders. - ae.requires_grad_(False) - text_enc.requires_grad_(False) - # Set model as trainable. - model.train() - - - # Move unet, vae and text_encoder to device and cast to weight_dtype - # The VAE is in float32 to avoid NaN losses. - # ae.to(accelerator.device, dtype=torch.float32) - ae.to(accelerator.device, dtype=weight_dtype) - # ae.to(accelerator.device) - text_enc.to(accelerator.device, dtype=weight_dtype) - # text_enc.to(accelerator.device) - - # Create EMA for the unet. - if args.use_ema: - ema_model = deepcopy(model) - ema_model = EMAModel(ema_model.parameters(), model_cls=LatteT2V, model_config=ema_model.config) - - # `accelerate` 0.16.0 will have better support for customized saving - if version.parse(accelerate.__version__) >= version.parse("0.16.0"): - # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format - def save_model_hook(models, weights, output_dir): - if accelerator.is_main_process: - if args.use_ema: - ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) - - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "model")) - if weights: # Don't pop if empty - # make sure to pop weight so that corresponding model is not saved again - weights.pop() - - def load_model_hook(models, input_dir): - if args.use_ema: - load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), LatteT2V) - ema_model.load_state_dict(load_model.state_dict()) - ema_model.to(accelerator.device) - del load_model - - for i in range(len(models)): - # pop models so that they are not loaded again - model = models.pop() - - # load diffusers style into model - load_model = LatteT2V.from_pretrained(input_dir, subfolder="model") - model.register_to_config(**load_model.config) - - model.load_state_dict(load_model.state_dict()) - del load_model - - accelerator.register_save_state_pre_hook(save_model_hook) - accelerator.register_load_state_pre_hook(load_model_hook) - - # Enable TF32 for faster training on Ampere GPUs, - # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices - if args.allow_tf32: - torch.backends.cuda.matmul.allow_tf32 = True - - if args.scale_lr: - args.learning_rate = ( - args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes - ) - - # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs - if args.use_8bit_adam: - try: - import bitsandbytes as bnb - except ImportError: - raise ImportError( - "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." - ) - - optimizer_class = bnb.optim.AdamW8bit - else: - optimizer_class = torch.optim.AdamW - - # Optimizer creation - params_to_optimize = model.parameters() - optimizer = optimizer_class( - params_to_optimize, - lr=args.learning_rate, - betas=(args.adam_beta1, args.adam_beta2), - weight_decay=args.adam_weight_decay, - eps=args.adam_epsilon, - ) - - # Setup data: - train_dataset = getdataset(args) - train_dataloader = torch.utils.data.DataLoader( - train_dataset, - shuffle=True, - collate_fn=Collate(args), - batch_size=args.train_batch_size, - num_workers=args.dataloader_num_workers, - ) - - # Scheduler and math around the number of training steps. - overrode_max_train_steps = False - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if args.max_train_steps is None: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - overrode_max_train_steps = True - - lr_scheduler = get_scheduler( - args.lr_scheduler, - optimizer=optimizer, - num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, - num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, - ) - - # Prepare everything with our `accelerator`. - model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, lr_scheduler - ) - - # We need to recalculate our total training steps as the size of the training dataloader may have changed. - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if overrode_max_train_steps: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - # Afterwards we recalculate our number of training epochs - args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) - - # We need to initialize the trackers we use, and also store our configuration. - # The trackers initializes automatically on the main process. - if accelerator.is_main_process: - accelerator.init_trackers(args.output_dir, config=vars(args)) - - # Train! - total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps - - logger.info("***** Running training *****") - logger.info(f" Num examples = {len(train_dataset)}") - logger.info(f" Num Epochs = {args.num_train_epochs}") - logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") - logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") - logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") - logger.info(f" Total optimization steps = {args.max_train_steps}") - global_step = 0 - first_epoch = 0 - - # Potentially load in the weights and states from a previous save - if args.resume_from_checkpoint: - if args.resume_from_checkpoint != "latest": - path = os.path.basename(args.resume_from_checkpoint) - else: - # Get the most recent checkpoint - dirs = os.listdir(args.output_dir) - dirs = [d for d in dirs if d.startswith("checkpoint")] - dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) - path = dirs[-1] if len(dirs) > 0 else None - - if path is None: - accelerator.print( - f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." - ) - args.resume_from_checkpoint = None - initial_global_step = 0 - else: - accelerator.print(f"Resuming from checkpoint {path}") - accelerator.load_state(os.path.join(args.output_dir, path)) - global_step = int(path.split("-")[1]) - - initial_global_step = global_step - first_epoch = global_step // num_update_steps_per_epoch - - else: - initial_global_step = 0 - - progress_bar = tqdm( - range(0, args.max_train_steps), - initial=initial_global_step, - desc="Steps", - # Only show the progress bar once on each machine. - disable=not accelerator.is_local_main_process, - ) - - for epoch in range(first_epoch, args.num_train_epochs): - train_loss = 0.0 - for step, (x, attn_mask, input_ids, cond_mask) in enumerate(train_dataloader): - with accelerator.accumulate(model): - # Sample noise that we'll add to the latents - - - x = x.to(accelerator.device, dtype=weight_dtype) # B C T+num_images H W, 16 + 4 - attn_mask = attn_mask.to(accelerator.device) # B L or B 1+num_images L - input_ids = input_ids.to(accelerator.device) # B L or B 1+num_images L - cond_mask = cond_mask.to(accelerator.device) # B L or B 1+num_images L - print('x.shape, attn_mask.shape, input_ids.shape, cond_mask.shape', x.shape, attn_mask.shape, input_ids.shape, cond_mask.shape) - - with torch.no_grad(): - # use for loop to avoid OOM, because T5 is too huge... - B, _, _ = input_ids.shape # B T+num_images L b 1+4, L - cond = torch.stack([text_enc(input_ids[i], cond_mask[i]) for i in range(B)]) # B 1+num_images L D - - # Map input images to latent space + normalize latents - if args.use_image_num == 0: - x = ae.encode(x) # B C T H W - cond = text_enc(input_ids, cond_mask) # B L -> B L D - else: - videos, images = x[:, :, :-args.use_image_num], x[:, :, -args.use_image_num:] - videos = ae.encode(videos) # B C T H W - - - def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None: - from examples.rec_imvi_vae import array_to_video - x = x.detach().cpu() - x = torch.clamp(x, -1, 1) - x = (x + 1) / 2 - x = x.permute(1, 2, 3, 0).numpy() - x = (255*x).astype(np.uint8) - array_to_video(x, fps=fps, output_file=output_file) - return - - # videos = ae.decode(videos.to(dtype=weight_dtype))[0] - # videos = videos.transpose(0, 1) - # custom_to_video(videos.to(torch.float32), fps=24, output_file='tmp.mp4') - # sys.exit() - - images = rearrange(images, 'b c t h w -> (b t) c 1 h w') - images = ae.encode(images) - - # import ipdb;ipdb.set_trace() - # images = ae.decode(images.to(dtype=weight_dtype)) - # for idx in range(args.use_image_num): - # x = images[idx, 0, :, :, :].to(torch.float32) - # x = x.squeeze() - # x = x.detach().cpu().numpy() - # x = np.clip(x, -1, 1) - # x = (x + 1) / 2 - # x = (255 * x).astype(np.uint8) - # x = x.transpose(1, 2, 0) - # from PIL import Image - # image = Image.fromarray(x) - # image.save(f'tmp{idx}.jpg') - # import sys - # sys.exit() - - - images = rearrange(images, '(b t) c 1 h w -> b c t h w', t=args.use_image_num) - x = torch.cat([videos, images], dim=2) # b c 17+4, h, w - - - - # print('(x.shape, attn_mask.shape, cond.shape, cond_mask.shape', x.shape, attn_mask.shape, cond.shape, cond_mask.shape) - model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, - encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) - t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=accelerator.device) - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() - - # Gather the losses across all processes for logging (if we use distributed training). - avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() - train_loss += avg_loss.item() / args.gradient_accumulation_steps - - # Backpropagate - accelerator.backward(loss) - if accelerator.sync_gradients: - params_to_clip = model.parameters() - accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - # Checks if the accelerator has performed an optimization step behind the scenes - if accelerator.sync_gradients: - progress_bar.update(1) - global_step += 1 - accelerator.log({"train_loss": train_loss}, step=global_step) - train_loss = 0.0 - - if args.use_deepspeed or accelerator.is_main_process: - if global_step % args.checkpointing_steps == 0: - # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` - if args.checkpoints_total_limit is not None: - checkpoints = os.listdir(args.output_dir) - checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] - checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) - - # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints - if len(checkpoints) >= args.checkpoints_total_limit: - num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 - removing_checkpoints = checkpoints[0:num_to_remove] - - logger.info( - f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" - ) - logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") - - for removing_checkpoint in removing_checkpoints: - removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) - shutil.rmtree(removing_checkpoint) - - save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") - accelerator.save_state(save_path) - logger.info(f"Saved state to {save_path}") - - logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} - progress_bar.set_postfix(**logs) - - if global_step >= args.max_train_steps: - break - - if accelerator.is_main_process: - if global_step % args.checkpointing_steps == 0: - if args.use_ema: - # Store the UNet parameters temporarily and load the EMA parameters to perform inference. - ema_model.store(model.parameters()) - 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() - - accelerator.wait_for_everyone() - accelerator.end_training() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--dataset", type=str, required=True) - parser.add_argument("--video_data_path", type=str, required=True) - parser.add_argument("--video_folder", type=str, default='') - parser.add_argument("--image_data_path", type=str, default='') - parser.add_argument("--image_folder", type=str, default='') - parser.add_argument("--sample_rate", type=int, default=1) - parser.add_argument("--num_frames", type=int, default=17) - parser.add_argument("--max_image_size", type=int, default=512) - parser.add_argument("--use_img_from_vid", action="store_true") - parser.add_argument("--use_image_num", type=int, default=0) - parser.add_argument("--model_max_length", type=int, default=300) - - parser.add_argument('--enable_8bit_t5', action='store_true') - parser.add_argument('--tile_overlap_factor', type=float, default=0.25) - parser.add_argument('--enable_tiling', action='store_true') - parser.add_argument("--compress_kv", action="store_true") - parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers") - parser.add_argument('--use_rope', action='store_true') - parser.add_argument('--compress_kv_factor', type=int, default=1) - - parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122") - parser.add_argument("--pretrained", type=str, default=None) - parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") - 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("--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, - default=None, - help="The output directory where the model predictions and checkpoints will be written.", - ) - parser.add_argument( - "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." - ) - parser.add_argument("--num_train_epochs", type=int, default=100) - parser.add_argument( - "--max_train_steps", - type=int, - default=None, - help="Total number of training steps to perform. If provided, overrides num_train_epochs.", - ) - parser.add_argument( - "--checkpointing_steps", - type=int, - default=500, - help=( - "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" - " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" - " training using `--resume_from_checkpoint`." - ), - ) - parser.add_argument( - "--checkpoints_total_limit", - type=int, - default=None, - help=("Max number of checkpoints to store."), - ) - parser.add_argument( - "--resume_from_checkpoint", - type=str, - default=None, - help=( - "Whether training should be resumed from a previous checkpoint. Use a path saved by" - ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' - ), - ) - parser.add_argument( - "--gradient_accumulation_steps", - type=int, - default=1, - help="Number of updates steps to accumulate before performing a backward/update pass.", - ) - parser.add_argument( - "--gradient_checkpointing", - action="store_true", - help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", - ) - parser.add_argument( - "--learning_rate", - type=float, - default=1e-4, - help="Initial learning rate (after the potential warmup period) to use.", - ) - parser.add_argument( - "--scale_lr", - action="store_true", - default=False, - help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", - ) - parser.add_argument( - "--lr_scheduler", - type=str, - default="constant", - help=( - 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' - ' "constant", "constant_with_warmup"]' - ), - ) - parser.add_argument( - "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." - ) - parser.add_argument( - "--timestep_bias_strategy", - type=str, - default="none", - choices=["earlier", "later", "range", "none"], - help=( - "The timestep bias strategy, which may help direct the model toward learning low or high frequency details." - " Choices: ['earlier', 'later', 'range', 'none']." - " The default is 'none', which means no bias is applied, and training proceeds normally." - " The value of 'later' will increase the frequency of the model's final training timesteps." - ), - ) - parser.add_argument( - "--timestep_bias_multiplier", - type=float, - default=1.0, - help=( - "The multiplier for the bias. Defaults to 1.0, which means no bias is applied." - " A value of 2.0 will double the weight of the bias, and a value of 0.5 will halve it." - ), - ) - parser.add_argument( - "--timestep_bias_begin", - type=int, - default=0, - help=( - "When using `--timestep_bias_strategy=range`, the beginning (inclusive) timestep to bias." - " Defaults to zero, which equates to having no specific bias." - ), - ) - parser.add_argument( - "--timestep_bias_end", - type=int, - default=1000, - help=( - "When using `--timestep_bias_strategy=range`, the final timestep (inclusive) to bias." - " Defaults to 1000, which is the number of timesteps that Stable Diffusion is trained on." - ), - ) - parser.add_argument( - "--timestep_bias_portion", - type=float, - default=0.25, - help=( - "The portion of timesteps to bias. Defaults to 0.25, which 25% of timesteps will be biased." - " A value of 0.5 will bias one half of the timesteps. The value provided for `--timestep_bias_strategy` determines" - " whether the biased portions are in the earlier or later timesteps." - ), - ) - parser.add_argument( - "--snr_gamma", - type=float, - default=None, - help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. " - "More details here: https://arxiv.org/abs/2303.09556.", - ) - parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") - parser.add_argument( - "--allow_tf32", - action="store_true", - help=( - "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" - " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" - ), - ) - parser.add_argument( - "--dataloader_num_workers", - type=int, - default=10, - help=( - "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." - ), - ) - parser.add_argument( - "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." - ) - parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") - parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") - parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") - parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") - parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") - parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") - parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") - 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.", - ) - parser.add_argument( - "--hub_model_id", - type=str, - default=None, - help="The name of the repository to keep in sync with the local `output_dir`.", - ) - 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***." - ), - ) - parser.add_argument( - "--report_to", - type=str, - default="tensorboard", - help=( - 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' - ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' - ), - ) - parser.add_argument( - "--mixed_precision", - type=str, - default=None, - choices=["no", "fp16", "bf16"], - help=( - "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" - " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" - " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." - ), - ) - parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") - parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") - - args = parser.parse_args() - main(args) diff --git a/opensora/train/train_t2v_diffusers.py b/opensora/train/train_t2v_diffusers.py new file mode 100644 index 000000000..f9c7c6f05 --- /dev/null +++ b/opensora/train/train_t2v_diffusers.py @@ -0,0 +1,992 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +A minimal training script for DiT using PyTorch DDP. +""" +import argparse +import logging +import math +import os +import shutil +from pathlib import Path +from typing import Optional +import gc +import numpy as np +from einops import rearrange +import torch.utils +import torch.utils.data +from tqdm import tqdm + +from opensora.adaptor.modules import replace_with_fp32_forwards + +try: + import torch_npu + from opensora.npu_config import npu_config + from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, \ + destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state + from opensora.acceleration.communications import prepare_parallel_data, broadcast +except: + torch_npu = None + npu_config = None + 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 + pass +import time +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 +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedType, ProjectConfiguration, set_seed +from packaging import version +from tqdm.auto import tqdm + +import diffusers +from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import EMAModel, compute_snr +from diffusers.utils import check_min_version, is_wandb_available + +from opensora.models.causalvideovae import ae_stride_config, ae_channel_config +from opensora.models.causalvideovae import ae_norm, ae_denorm +from opensora.models import CausalVAEModelWrapper +from opensora.models.text_encoder import get_text_enc, get_text_warpper +from opensora.dataset import getdataset +from opensora.models import CausalVAEModelWrapper +from opensora.models.diffusion import Diffusion_models, Diffusion_models_class +from opensora.utils.dataset_utils import Collate, LengthGroupedSampler +from opensora.sample.pipeline_opensora import OpenSoraPipeline +from opensora.models.causalvideovae import ae_stride_config, ae_wrapper + +# 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__) +from torch.utils.data import _utils +_utils.MP_STATUS_CHECK_INTERVAL = 1800.0 # dataloader timeout (default is 5.0s), we increase it to 1800s. + +@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 cat wearing sunglasses and working as a lifeguard at pool.", + "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, + 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 = opensora_pipeline( + positive_prompt.format(prompt), + negative_prompt=negative_prompt, + num_frames=args.num_frames, + height=args.max_height, + width=args.max_width, + num_inference_steps=args.num_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() + 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": + 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) + + del opensora_pipeline + gc.collect() + torch.cuda.empty_cache() + + +class ProgressInfo: + def __init__(self, global_step, train_loss=0.0): + self.global_step = global_step + self.train_loss = train_loss + + +################################################################################# +# Training Loop # +################################################################################# + +def main(args): + logging_dir = Path(args.output_dir, args.logging_dir) + + # use LayerNorm, GeLu, SiLu always as fp32 mode + if args.enable_stable_fp32: + replace_with_fp32_forwards() + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(args) + npu_config.seed_everything(args.seed) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + if args.num_frames != 1 and args.use_image_num == 0: + initialize_sequence_parallel_state(args.sp_size) + + 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.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed, device_specific=True) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Create model: + kwargs = {} + ae = ae_wrapper[args.ae](args.ae_path, cache_dir=args.cache_dir, **kwargs).eval() + if args.enable_tiling: + ae.vae.enable_tiling() + ae.vae.tile_overlap_factor = args.tile_overlap_factor + + kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True} + text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval() + + ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] + ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w) + assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" + args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w + args.ae_stride = args.ae_stride_h + patch_size = args.model[-3:] + patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + args.patch_size = patch_size_h + args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w + assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" + # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." + assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})." + assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})." + + args.stride_t = ae_stride_t * patch_size_t + args.stride = ae_stride_h * patch_size_h + latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w) + ae.latent_size = latent_size + + if args.num_frames % 2 == 1: + args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1 + else: + latent_size_t = args.num_frames // ae_stride_t + model = Diffusion_models[args.model]( + in_channels=ae_channel_config[args.ae], + out_channels=ae_channel_config[args.ae], + # caption_channels=4096, + # cross_attention_dim=1152, + attention_bias=True, + sample_size=latent_size, + sample_size_t=latent_size_t, + num_vector_embeds=None, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + use_linear_projection=False, + only_cross_attention=False, + double_self_attention=False, + upcast_attention=False, + # norm_type="ada_norm_single", + norm_elementwise_affine=False, + norm_eps=1e-6, + attention_type='default', + attention_mode=args.attention_mode, + interpolation_scale_h=args.interpolation_scale_h, + interpolation_scale_w=args.interpolation_scale_w, + interpolation_scale_t=args.interpolation_scale_t, + downsampler=args.downsampler, + # compress_kv_factor=args.compress_kv_factor, + use_rope=args.use_rope, + # model_max_length=args.model_max_length, + use_stable_fp32=args.enable_stable_fp32, + sparse1d=args.sparse1d, + sparse2d=args.sparse2d, + sparse_n=args.sparse_n, + use_motion=args.use_motion + ) + model.gradient_checkpointing = args.gradient_checkpointing + + # # use pretrained model? + if args.pretrained: + model_state_dict = model.state_dict() + if 'safetensors' in args.pretrained: # pixart series + from safetensors.torch import load_file as safe_load + # import ipdb;ipdb.set_trace() + pretrained_checkpoint = safe_load(args.pretrained, device="cpu") + pretrained_keys = set(list(pretrained_checkpoint.keys())) + model_keys = set(list(model_state_dict.keys())) + common_keys = list(pretrained_keys & model_keys) + checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()} + # if checkpoint['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and checkpoint['pos_embed.proj.weight'].ndim == 4: + # logger.info(f"Resize pos_embed, {checkpoint['pos_embed.proj.weight'].shape} -> {model.pos_embed.proj.weight.shape}") + # repeat = model.pos_embed.proj.weight.shape[2] + # checkpoint['pos_embed.proj.weight'] = checkpoint['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat) + # del checkpoint['proj_out.weight'], checkpoint['proj_out.bias'] + else: # latest stage training weight + checkpoint = torch.load(args.pretrained, map_location='cpu') + if 'model' in checkpoint: + checkpoint = checkpoint['model'] + missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) + logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') + logger.info(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!') + + # Freeze vae and text encoders. + ae.vae.requires_grad_(False) + text_enc.requires_grad_(False) + # Set model as trainable. + model.train() + + ae.vae.tile_sample_min_size = args.tile_sample_min_size + ae.vae.tile_sample_min_size_t = args.tile_sample_min_size_t + + noise_scheduler = DDPMScheduler(rescale_betas_zero_snr=args.rescale_betas_zero_snr) + # Move unet, vae and text_encoder to device and cast to weight_dtype + # The VAE is in float32 to avoid NaN losses. + ae.vae.to(accelerator.device, dtype=torch.float32) + # ae.vae.to(accelerator.device, dtype=weight_dtype) + text_enc.to(accelerator.device, dtype=weight_dtype) + + # Create EMA for the unet. + if args.use_ema: + ema_model = deepcopy(model) + ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step, + model_cls=Diffusion_models_class[args.model], model_config=ema_model.config) + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + if args.use_ema: + ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) + + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "model")) + if weights: # Don't pop if empty + # make sure to pop weight so that corresponding model is not saved again + weights.pop() + + def load_model_hook(models, input_dir): + if args.use_ema: + load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model]) + ema_model.load_state_dict(load_model.state_dict()) + ema_model.to(accelerator.device) + del load_model + + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model") + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + params_to_optimize = model.parameters() + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + logger.info(f"optimizer: {optimizer}") + + # Setup data: + + if args.trained_data_global_step is not None: + initial_global_step_for_sampler = args.trained_data_global_step + else: + initial_global_step_for_sampler = 0 + # print('initial_global_step, initial_global_step_for_sampler, total_batch_size', + # initial_global_step, initial_global_step_for_sampler, total_batch_size) + train_dataset = getdataset(args) + sampler = LengthGroupedSampler( + args.train_batch_size, + world_size=accelerator.num_processes, + gradient_accumulation_size=args.gradient_accumulation_steps, + initial_global_step=initial_global_step_for_sampler, + lengths=train_dataset.lengths, + group_data=args.group_data, + ) + train_dataloader = DataLoader( + train_dataset, + shuffle=False, + # pin_memory=True, + collate_fn=Collate(args), + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + sampler=sampler, + drop_last=True, + prefetch_factor=4 + ) + logger.info(f'after train_dataloader') + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, + num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, + ) + + # Prepare everything with our `accelerator`. + # model.requires_grad_(False) + # model.pos_embed.requires_grad_(True) + if args.adapt_vae: + model.requires_grad_(False) + for name, param in model.named_parameters(): + if 'pos_embed' in name or 'proj_out' in name: + param.requires_grad = True + logger.info(f'before accelerator.prepare') + model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, lr_scheduler + ) + logger.info(f'after accelerator.prepare') + if args.use_ema: + ema_model.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args)) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size + + logger.info("***** Running training *****") + logger.info(f" Model = {model}") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + logger.info(f" Total optimization steps (num_update_steps_per_epoch) = {num_update_steps_per_epoch}") + logger.info(f" Total training parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B") + + global_step = 0 + first_epoch = 0 + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + progress_info = ProgressInfo(global_step, train_loss=0.0) + + def sync_gradients_info(loss): + # Checks if the accelerator has performed an optimization step behind the scenes + if args.use_ema: + ema_model.step(model.parameters()) + progress_bar.update(1) + progress_info.global_step += 1 + end_time = time.time() + one_step_duration = end_time - start_time + accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step) + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(f"Step: [{progress_info.global_step}], local_loss={loss.detach().item()}, " + f"train_loss={progress_info.train_loss}, time_cost={one_step_duration}", + rank=0) + progress_info.train_loss = 0.0 + + # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues. + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: + if progress_info.global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if accelerator.is_main_process and args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + def run(model_input, model_kwargs, prof): + global start_time + start_time = time.time() + + noise = torch.randn_like(model_input) + if args.noise_offset: + # https://www.crosslabs.org//blog/diffusion-with-offset-noise + noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1), + device=model_input.device) + + bsz = model_input.shape[0] + current_step_frame = model_input.shape[2] + # Sample a random timestep for each image without bias. + timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device) + # print('accelerator.process_index, timesteps', accelerator.process_index, timesteps) + if current_step_frame != 1 and get_sequence_parallel_state(): # image do not need sp + broadcast(timesteps) + motion_score = model_kwargs.pop('motion_score', None) + if motion_score is not None: + raise NotImplementedError + + # Add noise to the model input according to the noise magnitude at each timestep + # (this is the forward diffusion process) + + noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) + model_pred = model( + noisy_model_input, + timesteps, + **model_kwargs + )[0] + # Get the target for loss depending on the prediction type + if args.prediction_type is not None: + # set prediction_type of scheduler if defined + noise_scheduler.register_to_config(prediction_type=args.prediction_type) + + if noise_scheduler.config.prediction_type == "epsilon": + target = noise + elif noise_scheduler.config.prediction_type == "v_prediction": + target = noise_scheduler.get_velocity(model_input, noise, timesteps) + elif noise_scheduler.config.prediction_type == "sample": + # We set the target to latents here, but the model_pred will return the noise sample prediction. + target = model_input + # We will have to subtract the noise residual from the prediction to get the target sample. + model_pred = model_pred - noise + else: + raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") + + mask = model_kwargs.get('attention_mask', None) + if torch.all(mask.bool()): + mask = None + if get_sequence_parallel_state(): + assert mask is None + b, c, _, _, _ = model_pred.shape + if mask is not None: + mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float() # b t h w -> b c t h w + mask = mask.reshape(b, -1) + if args.snr_gamma is None: + # model_pred: b c t h w, attention_mask: b t h w + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.reshape(b, -1) + if mask is not None: + loss = (loss * mask).sum() / mask.sum() # mean loss on unpad patches + else: + loss = loss.mean() + else: + # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. + # Since we predict the noise instead of x_0, the original formulation is slightly changed. + # This is discussed in Section 4.2 of the same paper. + snr = compute_snr(noise_scheduler, timesteps) + mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min( + dim=1 + )[0] + if noise_scheduler.config.prediction_type == "epsilon": + mse_loss_weights = mse_loss_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + mse_loss_weights = mse_loss_weights / (snr + 1) + else: + raise NameError(f'{noise_scheduler.config.prediction_type}') + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.reshape(b, -1) + mse_loss_weights = mse_loss_weights.reshape(b, 1) + if mask is not None: + loss = (loss * mask * mse_loss_weights).sum() / mask.sum() # mean loss on unpad patches + else: + loss = (loss * mse_loss_weights).mean() + + # Gather the losses across all processes for logging (if we use distributed training). + avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() + progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps + + # Backpropagate + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = model.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if accelerator.sync_gradients: + sync_gradients_info(loss) + + if accelerator.is_main_process: + + if progress_info.global_step % args.checkpointing_steps == 0: + + if args.enable_tracker: + log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + weight_dtype, progress_info.global_step) + + if args.use_ema and npu_config is None: + # Store the UNet parameters temporarily and load the EMA parameters to perform inference. + ema_model.store(model.parameters()) + ema_model.copy_to(model.parameters()) + log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + weight_dtype, progress_info.global_step, ema=True) + # Switch back to the original UNet parameters. + ema_model.restore(model.parameters()) + + if prof is not None: + prof.step() + + + return loss + + def train_one_step(step_, data_item_, prof_=None): + train_loss = 0.0 + x, attn_mask, input_ids, cond_mask, motion_score = data_item_ + # print(f'step: {step_}, rank: {accelerator.process_index}, x: {x.shape}') + assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))' + x = x.to(accelerator.device, dtype=ae.vae.dtype) # B C T+num_images H W, 16 + 4 + + attn_mask = attn_mask.to(accelerator.device) # B T+num_images H W + input_ids = input_ids.to(accelerator.device) # B 1+num_images L + cond_mask = cond_mask.to(accelerator.device) # B 1+num_images L + motion_score = motion_score.to(accelerator.device) if motion_score is not None else motion_score # B 1 + # if accelerator.process_index == 0: + # logger.info(f'rank: {accelerator.process_index}, x: {x.shape}, attn_mask: {attn_mask.shape}') + + with torch.no_grad(): + B, N, L = input_ids.shape # B 1+num_images L + # use batch inference + input_ids_ = input_ids.reshape(-1, L) + cond_mask_ = cond_mask.reshape(-1, L) + cond = text_enc(input_ids_, cond_mask_) # B 1+num_images L D + cond = cond.reshape(B, N, L, -1) + # Map input images to latent space + normalize latents + x = ae.encode(x) # B C T H W + + # def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None: + # from examples.rec_video import array_to_video + # x = x.detach().cpu() + # x = torch.clamp(x, -1, 1) + # x = (x + 1) / 2 + # x = x.permute(1, 2, 3, 0).numpy() + # x = (255*x).astype(np.uint8) + # array_to_video(x, fps=fps, output_file=output_file) + # return + # videos = ae.decode(x)[0] + # videos = videos.transpose(0, 1) + # custom_to_video(videos.to(torch.float32), fps=24, output_file='tmp.mp4') + # import sys;sys.exit() + current_step_frame = x.shape[2] + current_step_sp_state = get_sequence_parallel_state() + if args.sp_size != 1: # enable sp + if current_step_frame == 1: # but image do not need sp + set_sequence_parallel_state(False) + else: + set_sequence_parallel_state(True) + if get_sequence_parallel_state(): + x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask, + args.use_image_num) + for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size): + with accelerator.accumulate(model): + st_idx = iter * args.train_sp_batch_size + ed_idx = (iter + 1) * args.train_sp_batch_size + model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx], + attention_mask=attn_mask[st_idx: ed_idx], + encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num) + run(x[st_idx: ed_idx], model_kwargs, prof_) + + else: + with accelerator.accumulate(model): + assert not torch.any(torch.isnan(x)), 'after vae' + x = x.to(weight_dtype) + # print(f'rank: {accelerator.process_index}, x: {x.shape}') + model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, motion_score=motion_score, + encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) + run(x, model_kwargs, prof_) + + set_sequence_parallel_state(current_step_sp_state) # in case the next step use sp, which need broadcast(timesteps) + + if progress_info.global_step >= args.max_train_steps: + return True + + return False + + def train_one_epoch(prof_=None): + # for epoch in range(first_epoch, args.num_train_epochs): + progress_info.train_loss = 0.0 + if progress_info.global_step >= args.max_train_steps: + return True + + for step, data_item in enumerate(train_dataloader): + if train_one_step(step, data_item, prof_): + break + + if step >= 2 and torch_npu is not None and npu_config is not None: + npu_config.free_mm() + + if npu_config is not None and npu_config.on_npu and npu_config.profiling: + experimental_config = torch_npu.profiler._ExperimentalConfig( + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization + ) + profile_output_path = f"/home/image_data/npu_profiling_t2v/{os.getenv('PROJECT_NAME', 'local')}" + os.makedirs(profile_output_path, exist_ok=True) + + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU], + with_stack=True, + record_shapes=True, + profile_memory=True, + experimental_config=experimental_config, + schedule=torch_npu.profiler.schedule(wait=npu_config.profiling_step, warmup=0, active=1, repeat=1, + skip_first=0), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/") + ) as prof: + train_one_epoch(prof) + else: + train_one_epoch() + accelerator.wait_for_everyone() + accelerator.end_training() + if get_sequence_parallel_state(): + destroy_sequence_parallel_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # dataset & dataloader + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--data", type=str, required='') + parser.add_argument("--sample_rate", type=int, default=1) + parser.add_argument("--train_fps", type=int, default=24) + parser.add_argument("--drop_short_ratio", type=float, default=1.0) + parser.add_argument("--speed_factor", type=float, default=1.0) + parser.add_argument("--num_frames", type=int, default=65) + parser.add_argument("--max_height", type=int, default=320) + parser.add_argument("--max_width", type=int, default=240) + parser.add_argument("--use_img_from_vid", action="store_true") + parser.add_argument("--use_image_num", type=int, default=0) + parser.add_argument("--model_max_length", type=int, default=512) + parser.add_argument('--cfg', type=float, default=0.1) + parser.add_argument("--dataloader_num_workers", type=int, default=10, 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("--group_data", action="store_true") + parser.add_argument("--hw_stride", type=int, default=32) + parser.add_argument("--skip_low_resolution", action="store_true") + parser.add_argument("--force_resolution", action="store_true") + parser.add_argument("--trained_data_global_step", type=int, default=None) + parser.add_argument("--use_decord", action="store_true") + + # text encoder & vae & diffusion model + parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122") + parser.add_argument('--enable_8bit_t5', action='store_true') + parser.add_argument('--tile_overlap_factor', type=float, default=0.125) + parser.add_argument('--enable_tiling', action='store_true') + parser.add_argument("--compress_kv", action="store_true") + parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers") + parser.add_argument('--use_rope', action='store_true') + parser.add_argument('--compress_kv_factor', type=int, default=1) + parser.add_argument('--interpolation_scale_h', type=float, default=1.0) + parser.add_argument('--interpolation_scale_w', type=float, default=1.0) + parser.add_argument('--interpolation_scale_t', type=float, default=1.0) + parser.add_argument("--downsampler", type=str, default=None) + parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") + parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") + 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("--pretrained", type=str, default=None) + parser.add_argument('--enable_stable_fp32', action='store_true') + parser.add_argument('--sparse1d', action='store_true') + parser.add_argument('--sparse2d', action='store_true') + parser.add_argument('--sparse_n', type=int, default=2) + parser.add_argument('--tile_sample_min_size', type=int, default=512) + parser.add_argument('--tile_sample_min_size_t', type=int, default=33) + parser.add_argument('--adapt_vae', action='store_true') + parser.add_argument('--use_motion', action='store_true') + parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") + + # diffusion setting + parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.") + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") + parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--ema_start_step", type=int, default=0) + parser.add_argument("--noise_offset", type=float, default=0.0, help="The scale of noise offset.") + 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.") + parser.add_argument('--rescale_betas_zero_snr', action='store_true') + + # validation & logs + parser.add_argument("--num_sampling_steps", type=int, default=20) + 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.") + 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("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store.")) + parser.add_argument("--checkpointing_steps", type=int, default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument("--resume_from_checkpoint", type=str, default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + 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***." + ), + ) + parser.add_argument("--report_to", type=str, default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + # optimizer & scheduler + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") + parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]') + parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") + parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.") + parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.") + parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW") + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params") + parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.") + parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW") + parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--prodigy_beta3", type=float, default=None, + help="coefficients for computing the Prodidy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--lr_scheduler", type=str, default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument("--allow_tf32", action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel") + parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training") + + args = parser.parse_args() + main(args) diff --git a/opensora/train/train_t2v_diffusers_lora copy.py b/opensora/train/train_t2v_diffusers_lora copy.py new file mode 100644 index 000000000..7b1098a75 --- /dev/null +++ b/opensora/train/train_t2v_diffusers_lora copy.py @@ -0,0 +1,960 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +A minimal training script for DiT using PyTorch DDP. +""" +import argparse +import logging +import math +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 + +try: + import torch_npu + from opensora.npu_config import npu_config +except: + torch_npu = None + npu_config = None + 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 + pass +import time +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 +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedType, ProjectConfiguration, set_seed +from huggingface_hub import create_repo +from packaging import version +from tqdm.auto import tqdm +from transformers import HfArgumentParser, TrainingArguments, AutoTokenizer + +from peft import LoraConfig, PeftModel, get_peft_model + +import diffusers +from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import EMAModel, compute_snr +from diffusers.utils import check_min_version, is_wandb_available + +from opensora.dataset import getdataset, ae_denorm +from opensora.models.ae import getae, getae_wrapper +from opensora.models.ae.videobase import CausalVQVAEModelWrapper, CausalVAEModelWrapper +from opensora.models.diffusion.latte.modeling_latte import LatteT2V +from opensora.models.text_encoder import get_text_enc, get_text_warpper +from opensora.utils.dataset_utils import Collate +from opensora.utils.lora_utils import EMAModel_LoRA, maybe_zero_3, get_peft_state_maybe_zero_3 +from opensora.models.ae import ae_stride_config, ae_channel_config +from opensora.models.diffusion import Diffusion_models, Diffusion_models_class +from opensora.sample.pipeline_opensora import OpenSoraPipeline + + +# 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__) + + +@torch.inference_mode() +def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, ema=False): + validation_prompt = [ + "a cat wearing sunglasses and working as a lifeguard at pool.", + "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, + 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 = opensora_pipeline(prompt, + num_frames=args.num_frames, + # num_frames=1, + height=args.max_height, + width=args.max_width, + num_inference_steps=args.num_sampling_steps, + guidance_scale=args.guidance_scale, + enable_temporal_attentions=True, + num_images_per_prompt=1, + mask_feature=True, + max_sequence_length=150, + ).images + 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": + 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=30) + if tracker.name == "wandb": + import wandb + if videos.shape[1] == 1: + # assert args.num_frames == 1 + images = rearrange(videos, 'b 1 c h w -> (b 1) h w c') + # import ipdb;ipdb.set_trace() + 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=30) + for i, (video, prompt) in enumerate(zip(videos, validation_prompt)) + ] + } + # import ipdb;ipdb.set_trace() + if hasattr(model.pos_embed, 'temp_embed_gate'): + logs.update({'temp_embed_gate (tanh)': float(model.pos_embed.temp_embed_gate.tanh().item())}) + tracker.log(logs, step=global_step) + + del opensora_pipeline + gc.collect() + torch.cuda.empty_cache() + + +class ProgressInfo: + def __init__(self, global_step, train_loss=0.0): + self.global_step = global_step + self.train_loss = train_loss + + +################################################################################# +# Training Loop # +################################################################################# + +def main(args): + logging_dir = Path(args.output_dir, args.logging_dir) + + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(args) + npu_config.seed_everything() + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + 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.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Create model: + kwargs = {} + ae = getae_wrapper(args.ae)(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval() + if args.enable_tiling: + ae.vae.enable_tiling() + ae.vae.tile_overlap_factor = args.tile_overlap_factor + + kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True} + text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval() + + ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] + ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w) + assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" + args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w + args.ae_stride = args.ae_stride_h + patch_size = args.model[-3:] + patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + args.patch_size = patch_size_h + args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w + assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" + # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." + assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})." + assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})." + + args.stride_t = ae_stride_t * patch_size_t + args.stride = ae_stride_h * patch_size_h + latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w) + ae.latent_size = latent_size + + if args.num_frames % 2 == 1: + args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1 + else: + latent_size_t = args.num_frames // ae_stride_t + model = Diffusion_models[args.model]( + in_channels=ae_channel_config[args.ae], + out_channels=ae_channel_config[args.ae], + # caption_channels=4096, + # cross_attention_dim=1152, + attention_bias=True, + sample_size=latent_size, + sample_size_t=latent_size_t, + num_vector_embeds=None, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + use_linear_projection=False, + only_cross_attention=False, + double_self_attention=False, + upcast_attention=False, + # norm_type="ada_norm_single", + norm_elementwise_affine=False, + norm_eps=1e-6, + attention_type='default', + attention_mode=args.attention_mode, + interpolation_scale_h=args.interpolation_scale_h, + interpolation_scale_w=args.interpolation_scale_w, + interpolation_scale_t=args.interpolation_scale_t, + downsampler=args.downsampler, + # compress_kv_factor=args.compress_kv_factor, + use_rope=args.use_rope, + # model_max_length=args.model_max_length, + use_stable_fp32=args.enable_stable_fp32, + ) + model.gradient_checkpointing = args.gradient_checkpointing + + # NOTE replace some functions in model + hacked_model(model, model_type=ModelType.INPAINT_ONLY, model_cls=Diffusion_models_class[args.model]) + + # # use pretrained model? + if args.pretrained: + model_state_dict = model.state_dict() + if 'safetensors' in args.pretrained: # pixart series + from safetensors.torch import load_file as safe_load + # import ipdb;ipdb.set_trace() + pretrained_checkpoint = safe_load(args.pretrained, device="cpu") + pretrained_keys = set(list(pretrained_checkpoint.keys())) + model_keys = set(list(model_state_dict.keys())) + common_keys = list(pretrained_keys & model_keys) + checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()} + # if checkpoint['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and checkpoint['pos_embed.proj.weight'].ndim == 4: + # logger.info(f"Resize pos_embed, {checkpoint['pos_embed.proj.weight'].shape} -> {model.pos_embed.proj.weight.shape}") + # repeat = model.pos_embed.proj.weight.shape[2] + # checkpoint['pos_embed.proj.weight'] = checkpoint['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat) + # del checkpoint['proj_out.weight'], checkpoint['proj_out.bias'] + else: # latest stage training weight + checkpoint = torch.load(args.pretrained, map_location='cpu') + if 'model' in checkpoint: + checkpoint = checkpoint['model'] + missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) + logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') + logger.info(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!') + + # Freeze vae and text encoders. + ae.vae.requires_grad_(False) + text_enc.requires_grad_(False) + model.requires_grad_(False) + # # Set model as trainable. + # model.train() + + noise_scheduler = DDPMScheduler() + # Move unet, vae and text_encoder to device and cast to weight_dtype + # The VAE is in float32 to avoid NaN losses. + ae.vae.to(accelerator.device, dtype=torch.float32) + # ae.vae.to(accelerator.device, dtype=weight_dtype) + text_enc.to(accelerator.device, dtype=weight_dtype) + + + + + # now we will add new LoRA weights to the attention layers + # Set correct lora layers + if args.enable_lora: + lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.rank, + init_lora_weights="gaussian", + target_modules=["to_k", "to_q", "to_v", "to_out.0"], + ) + model = get_peft_model(model, lora_config) + # Create EMA for the unet. + if args.use_ema: + ema_model = deepcopy(model) + if args.enable_lora: # ema the whole lora_model + ema_model = EMAModel_LoRA(lora_config, parameters=ema_model.parameters(), update_after_step=args.ema_start_step, + model_cls=Diffusion_models_class[args.model], model_config=ema_model.config) + else: + ema_model = EMAModel(ema_model.parameters(), update_after_step=args.ema_start_step, + model_cls=Diffusion_models_class[args.model], model_config=ema_model.config) + + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + if args.use_ema: + if args.enable_lora: # only save lora weight + ema_model.save_pretrained(os.path.join(output_dir, "model_ema_lora")) + else: + ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) + + for i, model in enumerate(models): + if args.enable_lora: # only save lora weight + model.save_pretrained(os.path.join(output_dir, "model_lora")) + else: + model.save_pretrained(os.path.join(output_dir, "model")) + if weights: # Don't pop if empty + # make sure to pop weight so that corresponding model is not saved again + weights.pop() + + def load_model_hook(models, input_dir): + if args.use_ema: + if args.enable_lora: + load_model = EMAModel_LoRA.from_pretrained(os.path.join(input_dir, "model_ema_lora"), Diffusion_models_class[args.model], + lora_config, os.path.splitext(args.pretrained)) + ema_model.load_state_dict(load_model.state_dict()) + ema_model.to(accelerator.device) + del load_model + + else: + load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model]) + ema_model.load_state_dict(load_model.state_dict()) + ema_model.to(accelerator.device) + del load_model + + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + if args.enable_lora: + model = PeftModel.from_pretrained(model, os.path.join(input_dir, "model_lora")) + else: + # load diffusers style into model + load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model") + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + params_to_optimize = model.parameters() + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + # Setup data: + train_dataset = getdataset(args) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + shuffle=True, + # pin_memory=True, + collate_fn=Collate(args), + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + prefetch_factor=4 + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, + num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, + ) + + # Prepare everything with our `accelerator`. + model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, lr_scheduler + ) + if args.use_ema: + ema_model.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + # NOTE wandb + if accelerator.is_main_process: + accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args)) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size + logger.info("***** Running training *****") + logger.info(f" Model = {model}") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + logger.info(f" Total parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + if npu_config is not None: + train_dataset.n_used_elements = global_step * args.train_batch_size + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + progress_info = ProgressInfo(global_step, train_loss=0.0) + + def sync_gradients_info(loss): + # Checks if the accelerator has performed an optimization step behind the scenes + if args.use_ema: + ema_model.step(model.parameters()) + progress_bar.update(1) + progress_info.global_step += 1 + end_time = time.time() + one_step_duration = end_time - start_time + accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step) + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(f"Step: [{progress_info.global_step}], local_loss={loss.detach().item()}, " + f"train_loss={progress_info.train_loss}, time_cost={one_step_duration}", + rank=0) + progress_info.train_loss = 0.0 + + # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues. + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: + if progress_info.global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if accelerator.is_main_process and args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + def run(model_input, model_kwargs, prof): + global start_time + start_time = time.time() + + noise = torch.randn_like(model_input) + if args.noise_offset: + # https://www.crosslabs.org//blog/diffusion-with-offset-noise + noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1), + device=model_input.device) + + bsz = model_input.shape[0] + # Sample a random timestep for each image without bias. + timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device) + + # Add noise to the model input according to the noise magnitude at each timestep + # (this is the forward diffusion process) + + noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) + + model_pred = model( + noisy_model_input, + timesteps, + **model_kwargs + )[0] + # Get the target for loss depending on the prediction type + if args.prediction_type is not None: + # set prediction_type of scheduler if defined + noise_scheduler.register_to_config(prediction_type=args.prediction_type) + + if noise_scheduler.config.prediction_type == "epsilon": + target = noise + elif noise_scheduler.config.prediction_type == "v_prediction": + target = noise_scheduler.get_velocity(model_input, noise, timesteps) + elif noise_scheduler.config.prediction_type == "sample": + # We set the target to latents here, but the model_pred will return the noise sample prediction. + target = model_input + # We will have to subtract the noise residual from the prediction to get the target sample. + model_pred = model_pred - noise + else: + raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") + + if args.snr_gamma is None: + loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") + else: + # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. + # Since we predict the noise instead of x_0, the original formulation is slightly changed. + # This is discussed in Section 4.2 of the same paper. + snr = compute_snr(noise_scheduler, timesteps) + mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min( + dim=1 + )[0] + if noise_scheduler.config.prediction_type == "epsilon": + mse_loss_weights = mse_loss_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + mse_loss_weights = mse_loss_weights / (snr + 1) + + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights + loss = loss.mean() + + # Gather the losses across all processes for logging (if we use distributed training). + avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() + progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps + + # Backpropagate + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = model.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if accelerator.sync_gradients: + sync_gradients_info(loss) + + if accelerator.is_main_process: + for tracker in accelerator.trackers: + if tracker.name == "wandb": + if progress_info.global_step % args.checkpointing_steps != 0: + if hasattr(model, 'module') and hasattr(model.module.pos_embed, 'temp_embed_gate'): + tracker.log( + {'temp_embed_gate (tanh)': float(model.module.pos_embed.temp_embed_gate.tanh().item())}) + elif hasattr(model, 'pos_embed') and hasattr(model.pos_embed, 'temp_embed_gate'): + tracker.log( + {'temp_embed_gate (tanh)': float(model.pos_embed.temp_embed_gate.tanh().item())}) + + if progress_info.global_step % args.checkpointing_steps == 0: + + if args.enable_tracker: + log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + weight_dtype, progress_info.global_step) + + if args.use_ema: + # Store the UNet parameters temporarily and load the EMA parameters to perform inference. + ema_model.store(model.parameters()) + ema_model.copy_to(model.parameters()) + if npu_config is None: + log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + weight_dtype, progress_info.global_step, ema=True) + # Switch back to the original UNet parameters. + ema_model.restore(model.parameters()) + + if prof is not None: + prof.step() + + + return loss + + def train_one_step(step_, data_item_, prof_=None): + train_loss = 0.0 + x, attn_mask, input_ids, cond_mask = data_item_ + # Sample noise that we'll add to the latents + + if not args.multi_scale: + assert torch.all(attn_mask) + assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))' + x = x.to(accelerator.device, dtype=ae.vae.dtype) # B C T+num_images H W, 16 + 4 + + attn_mask = attn_mask.to(accelerator.device) # B T+num_images H W + input_ids = input_ids.to(accelerator.device) # B 1+num_images L + cond_mask = cond_mask.to(accelerator.device) # B 1+num_images L + # print('x.shape, attn_mask.shape, input_ids.shape, cond_mask.shape', x.shape, attn_mask.shape, input_ids.shape, cond_mask.shape) + + with torch.no_grad(): + # import ipdb;ipdb.set_trace() + # use for loop to avoid OOM, because T5 is too huge... + B, N, L = input_ids.shape # B 1+num_images L + # cond_ = torch.stack([text_enc(input_ids[i], cond_mask[i]) for i in range(B)]) # B 1+num_images L D + + # use batch inference + input_ids_ = input_ids.reshape(-1, L) + cond_mask_ = cond_mask.reshape(-1, L) + cond = text_enc(input_ids_, cond_mask_) # B 1+num_images L D + cond = cond.reshape(B, N, L, -1) + + # Map input images to latent space + normalize latents + if args.use_image_num == 0: + x = ae.encode(x) # B C T H W + else: + videos, images = x[:, :, :-args.use_image_num], x[:, :, -args.use_image_num:] + videos = ae.encode(videos) # B C T H W + images = rearrange(images, 'b c t h w -> (b t) c 1 h w') + images = ae.encode(images) + images = rearrange(images, '(b t) c 1 h w -> b c t h w', t=args.use_image_num) + x = torch.cat([videos, images], dim=2) # b c 17+4, h, w + + with accelerator.accumulate(model): + assert not torch.any(torch.isnan(x)), 'after vae' + x = x.to(weight_dtype) + model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, + encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) + run(x, model_kwargs, prof_) + + set_sequence_parallel_state(current_step_sp_state) # in case the next step use sp, which need broadcast(timesteps) + + if progress_info.global_step >= args.max_train_steps: + return True + + return False + + def train_all_epoch(prof_=None): + for epoch in range(first_epoch, args.num_train_epochs): + progress_info.train_loss = 0.0 + if progress_info.global_step >= args.max_train_steps: + return True + + for step, data_item in enumerate(train_dataloader): + if train_one_step(step, data_item, prof_): + break + + if step >= 2 and torch_npu is not None and npu_config is not None: + npu_config.free_mm() + + if npu_config is not None and npu_config.on_npu and npu_config.profiling: + experimental_config = torch_npu.profiler._ExperimentalConfig( + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization + ) + profile_output_path = f"/home/image_data/npu_profiling_t2v/{os.getenv('PROJECT_NAME', 'local')}" + os.makedirs(profile_output_path, exist_ok=True) + + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU], + with_stack=True, + record_shapes=True, + profile_memory=True, + experimental_config=experimental_config, + schedule=torch_npu.profiler.schedule(wait=npu_config.profiling_step, warmup=0, active=1, repeat=1, + skip_first=0), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/") + ) as prof: + train_all_epoch(prof) + else: + train_all_epoch() + accelerator.wait_for_everyone() + accelerator.end_training() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + + # dataset & dataloader + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--data", type=str, required='') + parser.add_argument("--sample_rate", type=int, default=1) + parser.add_argument("--num_frames", type=int, default=65) + parser.add_argument("--max_height", type=int, default=320) + parser.add_argument("--max_width", type=int, default=240) + parser.add_argument("--use_img_from_vid", action="store_true") + parser.add_argument("--use_image_num", type=int, default=0) + parser.add_argument("--model_max_length", type=int, default=512) + parser.add_argument("--multi_scale", action="store_true") + parser.add_argument('--cfg', type=float, default=0.1) + parser.add_argument("--dataloader_num_workers", type=int, default=10, 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.") + + # text encoder & vae & diffusion model + parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122") + parser.add_argument('--enable_8bit_t5', action='store_true') + parser.add_argument('--tile_overlap_factor', type=float, default=0.25) + parser.add_argument('--enable_tiling', action='store_true') + parser.add_argument("--compress_kv", action="store_true") + parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers") + parser.add_argument('--use_rope', action='store_true') + parser.add_argument('--compress_kv_factor', type=int, default=1) + parser.add_argument('--interpolation_scale_h', type=float, default=1.0) + parser.add_argument('--interpolation_scale_w', type=float, default=1.0) + parser.add_argument('--interpolation_scale_t', type=float, default=1.0) + parser.add_argument("--downsampler", type=str, default=None) + parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") + parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") + 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('--enable_stable_fp32', action='store_true') + parser.add_argument("--enable_lora", action="store_true") + parser.add_argument('--rank', type=int, default=64) + parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") + + # diffusion setting + parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.") + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") + parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--ema_start_step", type=int, default=0) + parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") + 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=50) + parser.add_argument('--guidance_scale', type=float, default=5.0) + parser.add_argument("--enable_tracker", action="store_true") + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + 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("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store.")) + parser.add_argument("--checkpointing_steps", type=int, default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument("--resume_from_checkpoint", type=str, default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + 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***." + ), + ) + parser.add_argument("--report_to", type=str, default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + # optimizer & scheduler + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") + parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]') + parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") + parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.") + parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.") + parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW") + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params") + parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.") + parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW") + parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--prodigy_beta3", type=float, default=None, + help="coefficients for computing the Prodidy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--lr_scheduler", type=str, default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument("--allow_tf32", action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") +<<<<<<< HEAD:opensora/train/train_inpaint.py + parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel") + parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training") + + parser.add_argument("--model_type", type=str, default='inpaint_only', choices=['inpaint_only', 'vip_only', 'vip_inpaint']) + parser.add_argument("--train_vip", action="store_true") + parser.add_argument("--need_validation", action="store_true") + # inpaint + parser.add_argument("--i2v_ratio", type=float, default=0.5) # for inpainting mode + parser.add_argument("--transition_ratio", type=float, default=0.4) # for inpainting mode + parser.add_argument("--v2v_ratio", type=float, default=0.1) # for inpainting mode + parser.add_argument("--clear_video_ratio", type=float, default=0.0) + parser.add_argument("--default_text_ratio", type=float, default=0.1) + parser.add_argument("--validation_dir", type=str, default=None, help="Path to the validation dataset.") + parser.add_argument("--image_encoder_name", type=str, default='laion/CLIP-ViT-H-14-laion2B-s32B-b79K') + parser.add_argument("--image_encoder_path", type=str, default=None) + parser.add_argument("--use_clip_mask", action="store_true") + parser.add_argument("--clip_loss_lambda", type=float, default=0.9) + parser.add_argument("--pretrained_transformer_model_path", type=str, default=None) + parser.add_argument("--pretrained_vip_adapter_path", type=str, default=None) + parser.add_argument("--vip_num_attention_heads", type=int, default=8) + parser.add_argument("--use_vae_preprocessed_mask", action="store_true") +======= +>>>>>>> anyres:opensora/train/train_t2v_diffusers_lora copy.py + + + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/opensora/train/train_t2v_diffusers_lora.py b/opensora/train/train_t2v_diffusers_lora.py new file mode 100644 index 000000000..be8e08183 --- /dev/null +++ b/opensora/train/train_t2v_diffusers_lora.py @@ -0,0 +1,1023 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +A minimal training script for DiT using PyTorch DDP. +""" +import argparse +import logging +import math +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 opensora.adaptor.modules import replace_with_fp32_forwards + +try: + import torch_npu + from opensora.npu_config import npu_config + from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, \ + destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state + from opensora.acceleration.communications import prepare_parallel_data, broadcast +except: + torch_npu = None + npu_config = None + 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 + pass +import time +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 +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedType, ProjectConfiguration, set_seed +from packaging import version +from tqdm.auto import tqdm + +import diffusers +from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import compute_snr +from diffusers.utils import check_min_version, is_wandb_available + +from opensora.models.causalvideovae import ae_stride_config, ae_channel_config +from opensora.models.causalvideovae import ae_norm, ae_denorm +from opensora.models import CausalVAEModelWrapper +from opensora.models.text_encoder import get_text_enc, get_text_warpper +from opensora.dataset import getdataset +from opensora.models import CausalVAEModelWrapper +from opensora.models.diffusion import Diffusion_models, Diffusion_models_class +from opensora.utils.dataset_utils import Collate, LengthGroupedSampler +from opensora.utils.ema_utils import EMAModel +from opensora.sample.pipeline_opensora import OpenSoraPipeline +from peft import LoraConfig, PeftModel, get_peft_model + +# 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__) + + +@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 cat wearing sunglasses and working as a lifeguard at pool.", + "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, + 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 = opensora_pipeline( + positive_prompt.format(prompt), + negative_prompt=negative_prompt, + num_frames=args.num_frames, + height=args.max_height, + width=args.max_width, + num_inference_steps=args.num_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() + 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": + 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) + + del opensora_pipeline + gc.collect() + torch.cuda.empty_cache() + + +class ProgressInfo: + def __init__(self, global_step, train_loss=0.0): + self.global_step = global_step + self.train_loss = train_loss + + +################################################################################# +# Training Loop # +################################################################################# + +def main(args): + logging_dir = Path(args.output_dir, args.logging_dir) + + # use LayerNorm, GeLu, SiLu always as fp32 mode + if args.enable_stable_fp32: + replace_with_fp32_forwards() + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(args) + npu_config.seed_everything(args.seed) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + if args.num_frames != 1 and args.use_image_num == 0: + initialize_sequence_parallel_state(args.sp_size) + + 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.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Create model: + kwargs = {} + ae = CausalVAEModelWrapper(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval() + if args.enable_tiling: + ae.vae.enable_tiling() + ae.vae.tile_overlap_factor = args.tile_overlap_factor + + kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True} + text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval() + + ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] + ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w) + assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" + args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w + args.ae_stride = args.ae_stride_h + patch_size = args.model[-3:] + patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + args.patch_size = patch_size_h + args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w + assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" + # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." + assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})." + assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})." + + args.stride_t = ae_stride_t * patch_size_t + args.stride = ae_stride_h * patch_size_h + latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w) + ae.latent_size = latent_size + + if args.num_frames % 2 == 1: + args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1 + else: + latent_size_t = args.num_frames // ae_stride_t + model = Diffusion_models[args.model]( + in_channels=ae_channel_config[args.ae], + out_channels=ae_channel_config[args.ae], + # caption_channels=4096, + # cross_attention_dim=1152, + attention_bias=True, + sample_size=latent_size, + sample_size_t=latent_size_t, + num_vector_embeds=None, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + use_linear_projection=False, + only_cross_attention=False, + double_self_attention=False, + upcast_attention=False, + # norm_type="ada_norm_single", + norm_elementwise_affine=False, + norm_eps=1e-6, + attention_type='default', + attention_mode=args.attention_mode, + interpolation_scale_h=args.interpolation_scale_h, + interpolation_scale_w=args.interpolation_scale_w, + interpolation_scale_t=args.interpolation_scale_t, + downsampler=args.downsampler, + # compress_kv_factor=args.compress_kv_factor, + use_rope=args.use_rope, + # model_max_length=args.model_max_length, + use_stable_fp32=args.enable_stable_fp32, + sparse1d=args.sparse1d, + sparse2d=args.sparse2d, + sparse_n=args.sparse_n, + ) + model.gradient_checkpointing = args.gradient_checkpointing + + # # use pretrained model? + if args.model_base: + model_state_dict = model.state_dict() + if '.safetensors' in args.model_base or any(['diffusion_pytorch_model.safetensors' in i for i in os.listdir(args.model_base)]): + from safetensors.torch import load_file as safe_load + # import ipdb;ipdb.set_trace() + if '.safetensors' in args.model_base: + model_base = args.model_base + else: + model_base = os.path.join(args.model_base, 'diffusion_pytorch_model.safetensors') + pretrained_checkpoint = safe_load(model_base, device="cpu") + pretrained_keys = set(list(pretrained_checkpoint.keys())) + model_keys = set(list(model_state_dict.keys())) + common_keys = list(pretrained_keys & model_keys) + checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()} + else: # latest stage training weight + checkpoint = torch.load(args.model_base, map_location='cpu') + if 'model' in checkpoint: + checkpoint = checkpoint['model'] + missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) + logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') + logger.info(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.model_base}!') + + # Freeze vae and text encoders. + ae.vae.requires_grad_(False) + text_enc.requires_grad_(False) + # Set model as trainable. + model.train() + + noise_scheduler = DDPMScheduler() + # Move unet, vae and text_encoder to device and cast to weight_dtype + # The VAE is in float32 to avoid NaN losses. + ae.vae.to(accelerator.device, dtype=torch.float32) + # ae.vae.to(accelerator.device, dtype=weight_dtype) + text_enc.to(accelerator.device, dtype=weight_dtype) + + + # now we will add new LoRA weights to the attention layers + # Set correct lora layers + if args.enable_lora: + lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.rank, + init_lora_weights="gaussian", + target_modules=["to_k", "to_q", "to_v", "to_out.0"], + ) + model = get_peft_model(model, lora_config) + logger.info(f'{model.print_trainable_parameters()}') + else: + lora_config = None + + # Create EMA for the unet. + if args.use_ema: + ema_model = deepcopy(model) + ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step, + model_cls=Diffusion_models_class[args.model], model_config=ema_model.config, lora_config=lora_config) + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + if args.use_ema: + ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) + + # import ipdb;ipdb.set_trace() + # if True: + # load_ema_model = EMAModel.from_pretrained(os.path.join(output_dir, "model_ema"), + # Diffusion_models_class[args.model], lora_config, args.model_base) + # res = 0 + # for p, p_ in zip(ema_model.shadow_params, load_ema_model.shadow_params): + # res += int(torch.allclose(p.detach().float().cpu(), p_.detach().float().cpu(), atol=1e-6)) + # # [int(torch.allclose(p.detach().cpu().float(), p_.detach().cpu().float(), atol=1e-6)) for p, p_ in zip(ema_model.shadow_params, load_ema_model.shadow_params)] + # print(f'total {len(ema_model.shadow_params)}, allclose {res}') + # import ipdb;ipdb.set_trace() + + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "model")) + if args.enable_lora: + unwrap_model = accelerator.unwrap_model(model) + unwrap_model_bak = deepcopy(unwrap_model) + + # import ipdb;ipdb.set_trace() + # if True: + # base_model = Diffusion_models_class[args.model].from_pretrained(args.model_base) + # model_lora = PeftModel.from_pretrained(base_model, os.path.join(output_dir, "model")) # 基模型+lora + # res = 0 + # for (n, p), (n_, p_) in zip(unwrap_model_bak.named_parameters(), model_lora.named_parameters()): + # res += int(torch.allclose(p.detach().float().cpu(), p_.detach().float().cpu(), atol=1e-6)) + # if torch.allclose(p.detach().float().cpu(), p_.detach().float().cpu(), atol=1e-6): + # print(n, n_) + # # [int(torch.allclose(p.detach().cpu().float(), p_.detach().cpu().float(), atol=1e-6)) for p, p_ in zip(unwrap_model_bak.parameters(), model_lora.parameters())] + # print(f'total {len(list(unwrap_model_bak.parameters()))}, allclose {res}') + # import ipdb;ipdb.set_trace() + + model_merge = unwrap_model_bak.merge_and_unload() + model_merge.save_pretrained(os.path.join(output_dir, "model")) + + del model_merge + del unwrap_model_bak + + gc.collect() + torch.cuda.empty_cache() + + if weights: # Don't pop if empty + # make sure to pop weight so that corresponding model is not saved again + weights.pop() + + def load_model_hook(models, input_dir): + if args.use_ema: + load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model], + lora_config, args.model_base) + ema_model.load_state_dict(load_model.state_dict()) + ema_model.to(accelerator.device) + del load_model + + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + if args.enable_lora: + # base_model = Diffusion_models_class[args.model].from_pretrained(args.model_base) + model = PeftModel.from_pretrained(model, os.path.join(input_dir, "model")) + else: + # load diffusers style into model + load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model") + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + params_to_optimize = model.parameters() + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + logger.info(f"optimizer: {optimizer}") + + # Setup data: + train_dataset = getdataset(args) + sampler = LengthGroupedSampler( + args.train_batch_size, + world_size=accelerator.num_processes, + lengths=train_dataset.lengths, + group_data=args.group_data, + ) if args.group_data and args.train_batch_size != 1 else None + train_dataloader = DataLoader( + train_dataset, + shuffle=sampler is None, + # pin_memory=True, + collate_fn=Collate(args), + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + sampler=sampler if args.group_data and args.train_batch_size != 1 else None, + drop_last=True, + # prefetch_factor=4 + ) + logger.info(f'after train_dataloader') + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, + num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, + ) + + # Prepare everything with our `accelerator`. + # model.requires_grad_(False) + # model.pos_embed.requires_grad_(True) + if args.adapt_vae: + model.requires_grad_(False) + for name, param in model.named_parameters(): + if 'pos_embed' in name or 'proj_out' in name: + param.requires_grad = True + logger.info(f'before accelerator.prepare') + model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, lr_scheduler + ) + logger.info(f'after accelerator.prepare') + if args.use_ema: + ema_model.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args)) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size + logger.info("***** Running training *****") + logger.info(f" Model = {model}") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + logger.info(f" Total training parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + if npu_config is not None: + train_dataset.n_used_elements = global_step * args.train_batch_size + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + progress_info = ProgressInfo(global_step, train_loss=0.0) + + def sync_gradients_info(loss): + # Checks if the accelerator has performed an optimization step behind the scenes + if args.use_ema: + ema_model.step(model.parameters()) + progress_bar.update(1) + progress_info.global_step += 1 + end_time = time.time() + one_step_duration = end_time - start_time + accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step) + if torch_npu is not None and npu_config is not None: + npu_config.print_msg(f"Step: [{progress_info.global_step}], local_loss={loss.detach().item()}, " + f"train_loss={progress_info.train_loss}, time_cost={one_step_duration}", + rank=0) + progress_info.train_loss = 0.0 + + # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues. + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: + if progress_info.global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if accelerator.is_main_process and args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + def run(model_input, model_kwargs, prof): + global start_time + start_time = time.time() + + noise = torch.randn_like(model_input) + if args.noise_offset: + # https://www.crosslabs.org//blog/diffusion-with-offset-noise + noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1), + device=model_input.device) + + bsz = model_input.shape[0] + current_step_frame = model_input.shape[2] + # Sample a random timestep for each image without bias. + timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device) + if current_step_frame != 1 and get_sequence_parallel_state(): # image do not need sp + broadcast(timesteps) + + # Add noise to the model input according to the noise magnitude at each timestep + # (this is the forward diffusion process) + + noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) + model_pred = model( + noisy_model_input, + timesteps, + **model_kwargs + )[0] + # Get the target for loss depending on the prediction type + if args.prediction_type is not None: + # set prediction_type of scheduler if defined + noise_scheduler.register_to_config(prediction_type=args.prediction_type) + + if noise_scheduler.config.prediction_type == "epsilon": + target = noise + elif noise_scheduler.config.prediction_type == "v_prediction": + target = noise_scheduler.get_velocity(model_input, noise, timesteps) + elif noise_scheduler.config.prediction_type == "sample": + # We set the target to latents here, but the model_pred will return the noise sample prediction. + target = model_input + # We will have to subtract the noise residual from the prediction to get the target sample. + model_pred = model_pred - noise + else: + raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") + + mask = model_kwargs.get('attention_mask', None) + if torch.all(mask.bool()): + mask = None + if get_sequence_parallel_state(): + assert mask is None + b, c, _, _, _ = model_pred.shape + if mask is not None: + mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float() # b t h w -> b c t h w + mask = mask.reshape(b, -1) + if args.snr_gamma is None: + # model_pred: b c t h w, attention_mask: b t h w + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.reshape(b, -1) + if mask is not None: + loss = (loss * mask).sum() / mask.sum() # mean loss on unpad patches + else: + loss = loss.mean() + else: + # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. + # Since we predict the noise instead of x_0, the original formulation is slightly changed. + # This is discussed in Section 4.2 of the same paper. + snr = compute_snr(noise_scheduler, timesteps) + mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min( + dim=1 + )[0] + if noise_scheduler.config.prediction_type == "epsilon": + mse_loss_weights = mse_loss_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + mse_loss_weights = mse_loss_weights / (snr + 1) + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.reshape(b, -1) + mse_loss_weights = mse_loss_weights.reshape(b, 1) + if mask is not None: + loss = (loss * mask * mse_loss_weights).sum() / mask.sum() # mean loss on unpad patches + else: + loss = (loss * mse_loss_weights).mean() + + # Gather the losses across all processes for logging (if we use distributed training). + avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() + progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps + + # Backpropagate + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = model.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if accelerator.sync_gradients: + sync_gradients_info(loss) + + if accelerator.is_main_process: + + if progress_info.global_step % args.checkpointing_steps == 0: + + if args.enable_tracker: + log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + weight_dtype, progress_info.global_step) + + if args.use_ema and npu_config is None: + # Store the UNet parameters temporarily and load the EMA parameters to perform inference. + ema_model.store(model.parameters()) + ema_model.copy_to(model.parameters()) + log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, + weight_dtype, progress_info.global_step, ema=True) + # Switch back to the original UNet parameters. + ema_model.restore(model.parameters()) + + gc.collect() + torch.cuda.empty_cache() + + if prof is not None: + prof.step() + + + return loss + + def train_one_step(step_, data_item_, prof_=None): + train_loss = 0.0 + x, attn_mask, input_ids, cond_mask = data_item_ + assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))' + x = x.to(accelerator.device, dtype=ae.vae.dtype) # B C T+num_images H W, 16 + 4 + + attn_mask = attn_mask.to(accelerator.device) # B T+num_images H W + input_ids = input_ids.to(accelerator.device) # B 1+num_images L + cond_mask = cond_mask.to(accelerator.device) # B 1+num_images L + # if accelerator.process_index == 0: + # logger.info(f'rank: {accelerator.process_index}, x: {x.shape}, attn_mask: {attn_mask.shape}') + + with torch.no_grad(): + B, N, L = input_ids.shape # B 1+num_images L + # use batch inference + input_ids_ = input_ids.reshape(-1, L) + cond_mask_ = cond_mask.reshape(-1, L) + cond = text_enc(input_ids_, cond_mask_) # B 1+num_images L D + cond = cond.reshape(B, N, L, -1) + # Map input images to latent space + normalize latents + x = ae.encode(x) # B C T H W + + # def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None: + # from examples.rec_video import array_to_video + # x = x.detach().cpu() + # x = torch.clamp(x, -1, 1) + # x = (x + 1) / 2 + # x = x.permute(1, 2, 3, 0).numpy() + # x = (255*x).astype(np.uint8) + # array_to_video(x, fps=fps, output_file=output_file) + # return + # videos = ae.decode(x)[0] + # videos = videos.transpose(0, 1) + # custom_to_video(videos.to(torch.float32), fps=24, output_file='tmp.mp4') + # import sys;sys.exit() + current_step_frame = x.shape[2] + current_step_sp_state = get_sequence_parallel_state() + if args.sp_size != 1: # enable sp + if current_step_frame == 1: # but image do not need sp + set_sequence_parallel_state(False) + else: + set_sequence_parallel_state(True) + if get_sequence_parallel_state(): + x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask, + args.use_image_num) + for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size): + with accelerator.accumulate(model): + st_idx = iter * args.train_sp_batch_size + ed_idx = (iter + 1) * args.train_sp_batch_size + model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx], + attention_mask=attn_mask[st_idx: ed_idx], + encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num) + run(x[st_idx: ed_idx], model_kwargs, prof_) + + else: + with accelerator.accumulate(model): + assert not torch.any(torch.isnan(x)), 'after vae' + x = x.to(weight_dtype) + model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, + encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) + run(x, model_kwargs, prof_) + + set_sequence_parallel_state(current_step_sp_state) # in case the next step use sp, which need broadcast(timesteps) + + if progress_info.global_step >= args.max_train_steps: + return True + + return False + + def train_all_epoch(prof_=None): + for epoch in range(first_epoch, args.num_train_epochs): + progress_info.train_loss = 0.0 + if progress_info.global_step >= args.max_train_steps: + return True + + for step, data_item in enumerate(train_dataloader): + if train_one_step(step, data_item, prof_): + break + + if step >= 2 and torch_npu is not None and npu_config is not None: + npu_config.free_mm() + + if npu_config is not None and npu_config.on_npu and npu_config.profiling: + experimental_config = torch_npu.profiler._ExperimentalConfig( + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization + ) + profile_output_path = f"/home/image_data/npu_profiling_t2v/{os.getenv('PROJECT_NAME', 'local')}" + os.makedirs(profile_output_path, exist_ok=True) + + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU], + with_stack=True, + record_shapes=True, + profile_memory=True, + experimental_config=experimental_config, + schedule=torch_npu.profiler.schedule(wait=npu_config.profiling_step, warmup=0, active=1, repeat=1, + skip_first=0), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/") + ) as prof: + train_all_epoch(prof) + else: + train_all_epoch() + accelerator.wait_for_everyone() + accelerator.end_training() + if get_sequence_parallel_state(): + destroy_sequence_parallel_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # dataset & dataloader + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--data", type=str, required='') + parser.add_argument("--sample_rate", type=int, default=1) + parser.add_argument("--train_fps", type=int, default=24) + parser.add_argument("--drop_short_ratio", type=float, default=1.0) + parser.add_argument("--speed_factor", type=float, default=1.0) + parser.add_argument("--num_frames", type=int, default=65) + parser.add_argument("--max_height", type=int, default=320) + parser.add_argument("--max_width", type=int, default=240) + parser.add_argument("--use_img_from_vid", action="store_true") + parser.add_argument("--use_image_num", type=int, default=0) + parser.add_argument("--model_max_length", type=int, default=512) + parser.add_argument('--cfg', type=float, default=0.1) + parser.add_argument("--dataloader_num_workers", type=int, default=10, 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("--group_data", action="store_true") + parser.add_argument("--hw_stride", type=int, default=32) + parser.add_argument("--skip_low_resolution", action="store_true") + parser.add_argument("--force_resolution", action="store_true") + + # text encoder & vae & diffusion model + parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122") + parser.add_argument('--enable_8bit_t5', action='store_true') + parser.add_argument('--tile_overlap_factor', type=float, default=0.125) + parser.add_argument('--enable_tiling', action='store_true') + parser.add_argument("--compress_kv", action="store_true") + parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers") + parser.add_argument('--use_rope', action='store_true') + parser.add_argument('--compress_kv_factor', type=int, default=1) + parser.add_argument('--interpolation_scale_h', type=float, default=1.0) + parser.add_argument('--interpolation_scale_w', type=float, default=1.0) + parser.add_argument('--interpolation_scale_t', type=float, default=1.0) + parser.add_argument("--downsampler", type=str, default=None) + parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") + parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") + 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("--model_base", type=str, default='') + parser.add_argument('--enable_stable_fp32', action='store_true') + parser.add_argument("--enable_lora", action="store_true") + parser.add_argument('--rank', type=int, default=64) + parser.add_argument('--sparse1d', action='store_true') + parser.add_argument('--sparse2d', action='store_true') + parser.add_argument('--sparse_n', type=int, default=2) + parser.add_argument('--adapt_vae', action='store_true') + parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") + + # diffusion setting + parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.") + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") + parser.add_argument("--ema_decay", type=float, default=0.9999) + parser.add_argument("--ema_start_step", type=int, default=0) + parser.add_argument("--noise_offset", type=float, default=0.02, help="The scale of noise offset.") + 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('--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.") + 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("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store.")) + parser.add_argument("--checkpointing_steps", type=int, default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument("--resume_from_checkpoint", type=str, default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + 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***." + ), + ) + parser.add_argument("--report_to", type=str, default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + # optimizer & scheduler + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") + parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]') + parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") + parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.") + parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.") + parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW") + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.") + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params") + parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.") + parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW") + parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--prodigy_beta3", type=float, default=None, + help="coefficients for computing the Prodidy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--lr_scheduler", type=str, default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument("--allow_tf32", action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel") + parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training") + + args = parser.parse_args() + main(args) diff --git a/opensora/train/train_t2v_feature.py b/opensora/train/train_t2v_feature.py deleted file mode 100644 index fbd2ea4c3..000000000 --- a/opensora/train/train_t2v_feature.py +++ /dev/null @@ -1,787 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -""" -A minimal training script for DiT using PyTorch DDP. -""" -import argparse -import logging -import math -import os -import shutil -from pathlib import Path -from typing import Optional - -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 -import transformers -from accelerate import Accelerator -from accelerate.logging import get_logger -from accelerate.utils import ProjectConfiguration, set_seed -from huggingface_hub import create_repo -from packaging import version -from tqdm.auto import tqdm -from transformers import HfArgumentParser, TrainingArguments, AutoTokenizer - -import diffusers -from diffusers import DDPMScheduler, PNDMScheduler -from diffusers.optimization import get_scheduler -from diffusers.training_utils import EMAModel, compute_snr -from diffusers.utils import check_min_version, is_wandb_available - -from opensora.dataset import getdataset, ae_denorm -from opensora.models.ae import getae, getae_wrapper -from opensora.models.ae.videobase import CausalVQVAEModelWrapper, CausalVAEModelWrapper -from opensora.models.diffusion.diffusion import create_diffusion_T as create_diffusion -from opensora.models.diffusion.latte.modeling_latte import LatteT2V -from opensora.models.text_encoder import get_text_enc -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 - -# 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." - ) - 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 - - -################################################################################# -# Training Loop # -################################################################################# - -def main(args): - logging_dir = Path(args.output_dir, args.logging_dir) - - accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) - - accelerator = Accelerator( - gradient_accumulation_steps=args.gradient_accumulation_steps, - mixed_precision=args.mixed_precision, - log_with=args.report_to, - project_config=accelerator_project_config, - ) - - 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( - format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", - datefmt="%m/%d/%Y %H:%M:%S", - level=logging.INFO, - ) - logger.info(accelerator.state, main_process_only=False) - if accelerator.is_local_main_process: - transformers.utils.logging.set_verbosity_warning() - diffusers.utils.logging.set_verbosity_info() - else: - transformers.utils.logging.set_verbosity_error() - diffusers.utils.logging.set_verbosity_error() - - # If passed along, set the training seed now. - if args.seed is not None: - set_seed(args.seed) - - # Handle the repository creation - if accelerator.is_main_process: - if args.output_dir is not None: - os.makedirs(args.output_dir, exist_ok=True) - - # if args.push_to_hub: - # repo_id = create_repo( - # repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token - # ).repo_id - - # Create model: - - diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule - # ae = getae(args).eval() - # text_enc = get_text_enc(args).eval() - - ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] - args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w - args.ae_stride = args.ae_stride_h - patch_size = args.model[-3:] - patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) - args.patch_size = patch_size_h - args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w - assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" - assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" - # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." - assert args.max_image_size % ae_stride_h == 0, f"Image size must be divisible by ae_stride_h, but found max_image_size ({args.max_image_size}), ae_stride_h ({ae_stride_h})." - - latent_size = (args.max_image_size // ae_stride_h, args.max_image_size // ae_stride_w) - - if getae_wrapper(args.ae) == CausalVQVAEModelWrapper or getae_wrapper(args.ae) == CausalVAEModelWrapper: - args.video_length = video_length = args.num_frames // ae_stride_t + 1 - else: - args.video_length = video_length = args.num_frames // ae_stride_t - model = Diffusion_models[args.model]( - in_channels=ae_channel_config[args.ae], - out_channels=ae_channel_config[args.ae] * 2, - # caption_channels=4096, - # cross_attention_dim=1152, - attention_bias=True, - sample_size=latent_size, - num_vector_embeds=None, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - use_linear_projection=False, - only_cross_attention=False, - double_self_attention=False, - upcast_attention=False, - # norm_type="ada_norm_single", - norm_elementwise_affine=False, - norm_eps=1e-6, - attention_type='default', - video_length=video_length, - attention_mode=args.attention_mode, - # compress_kv=args.compress_kv - ) - model.gradient_checkpointing = args.gradient_checkpointing - - # # use pretrained model? - if args.pretrained: - if 'safetensors' in args.pretrained: - from safetensors.torch import load_file as safe_load - checkpoint = safe_load(args.pretrained, device="cpu") - else: - checkpoint = torch.load(args.pretrained, map_location='cpu')['model'] - model_state_dict = model.state_dict() - missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) - logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') - logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!') - # load from pixart-alpha - # pixelart_alpha = torch.load(args.pretrained, map_location='cpu')['state_dict'] - # checkpoint = {} - # for k, v in pixelart_alpha.items(): - # if 'x_embedder' in k or 't_embedder' in k or 'y_embedder' in k: - # checkpoint[k] = v - # if k.startswith('blocks'): - # k_spilt = k.split('.') - # blk_id = str(int(k_spilt[1]) * 2) - # k_spilt[1] = blk_id - # new_k = '.'.join(k_spilt) - # checkpoint[new_k] = v - # missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) - # logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)} keys from {args.pretrained}!') - - # Freeze vae and text encoders. - # ae.requires_grad_(False) - # text_enc.requires_grad_(False) - # Set model as trainable. - model.train() - - # For mixed precision training we cast all non-trainable weigths to half-precision - # as these weights are only used for inference, keeping weights in full precision is not required. - weight_dtype = torch.float32 - if accelerator.mixed_precision == "fp16": - weight_dtype = torch.float16 - elif accelerator.mixed_precision == "bf16": - weight_dtype = torch.bfloat16 - - # Move unet, vae and text_encoder to device and cast to weight_dtype - # The VAE is in float32 to avoid NaN losses. - # ae.to(accelerator.device, dtype=torch.float32) - # text_enc.to(accelerator.device, dtype=weight_dtype) - - # Create EMA for the unet. - if args.use_ema: - ema_model = deepcopy(model) - ema_model = EMAModel(ema_model.parameters(), model_cls=LatteT2V, model_config=ema_model.config) - - # `accelerate` 0.16.0 will have better support for customized saving - if version.parse(accelerate.__version__) >= version.parse("0.16.0"): - # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format - def save_model_hook(models, weights, output_dir): - if accelerator.is_main_process: - if args.use_ema: - ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) - - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "model")) - if weights: # Don't pop if empty - # make sure to pop weight so that corresponding model is not saved again - weights.pop() - - def load_model_hook(models, input_dir): - if args.use_ema: - load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), LatteT2V) - ema_model.load_state_dict(load_model.state_dict()) - ema_model.to(accelerator.device) - del load_model - - for i in range(len(models)): - # pop models so that they are not loaded again - model = models.pop() - - # load diffusers style into model - load_model = LatteT2V.from_pretrained(input_dir, subfolder="model") - model.register_to_config(**load_model.config) - - model.load_state_dict(load_model.state_dict()) - del load_model - - accelerator.register_save_state_pre_hook(save_model_hook) - accelerator.register_load_state_pre_hook(load_model_hook) - - # Enable TF32 for faster training on Ampere GPUs, - # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices - if args.allow_tf32: - torch.backends.cuda.matmul.allow_tf32 = True - - if args.scale_lr: - args.learning_rate = ( - args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes - ) - - # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs - if args.use_8bit_adam: - try: - import bitsandbytes as bnb - except ImportError: - raise ImportError( - "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." - ) - - optimizer_class = bnb.optim.AdamW8bit - else: - optimizer_class = torch.optim.AdamW - - # Optimizer creation - params_to_optimize = model.parameters() - optimizer = optimizer_class( - params_to_optimize, - lr=args.learning_rate, - betas=(args.adam_beta1, args.adam_beta2), - weight_decay=args.adam_weight_decay, - eps=args.adam_epsilon, - ) - - # Setup data: - train_dataset = getdataset(args) - train_dataloader = torch.utils.data.DataLoader( - train_dataset, - shuffle=True, - # collate_fn=Collate(args), # TODO: do not enable dynamic mask in this point - batch_size=args.train_batch_size, - num_workers=args.dataloader_num_workers, - ) - - # Scheduler and math around the number of training steps. - overrode_max_train_steps = False - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if args.max_train_steps is None: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - overrode_max_train_steps = True - - lr_scheduler = get_scheduler( - args.lr_scheduler, - optimizer=optimizer, - num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, - num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, - ) - - # Prepare everything with our `accelerator`. - model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, lr_scheduler - ) - - # We need to recalculate our total training steps as the size of the training dataloader may have changed. - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if overrode_max_train_steps: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - # Afterwards we recalculate our number of training epochs - args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) - - # We need to initialize the trackers we use, and also store our configuration. - # The trackers initializes automatically on the main process. - if accelerator.is_main_process: - accelerator.init_trackers(args.output_dir, config=vars(args)) - - # Train! - total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps - - logger.info("***** Running training *****") - logger.info(f" Num examples = {len(train_dataset)}") - logger.info(f" Num Epochs = {args.num_train_epochs}") - logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") - logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") - logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") - logger.info(f" Total optimization steps = {args.max_train_steps}") - global_step = 0 - first_epoch = 0 - - # Potentially load in the weights and states from a previous save - if args.resume_from_checkpoint: - if args.resume_from_checkpoint != "latest": - path = os.path.basename(args.resume_from_checkpoint) - else: - # Get the most recent checkpoint - dirs = os.listdir(args.output_dir) - dirs = [d for d in dirs if d.startswith("checkpoint")] - dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) - path = dirs[-1] if len(dirs) > 0 else None - - if path is None: - accelerator.print( - f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." - ) - args.resume_from_checkpoint = None - initial_global_step = 0 - else: - accelerator.print(f"Resuming from checkpoint {path}") - accelerator.load_state(os.path.join(args.output_dir, path)) - global_step = int(path.split("-")[1]) - - initial_global_step = global_step - first_epoch = global_step // num_update_steps_per_epoch - - else: - initial_global_step = 0 - - progress_bar = tqdm( - range(0, args.max_train_steps), - initial=initial_global_step, - desc="Steps", - # Only show the progress bar once on each machine. - disable=not accelerator.is_local_main_process, - ) - - for epoch in range(first_epoch, args.num_train_epochs): - train_loss = 0.0 - for step, (x, cond, cond_mask) in enumerate(train_dataloader): - with accelerator.accumulate(model): - # Sample noise that we'll add to the latents - x = x.to(accelerator.device) # B C T H W - # attn_mask = attn_mask.to(device) # B T H W - # assert torch.all(attn_mask.bool()), 'do not enable dynamic input' - attn_mask = None - cond = cond.to(accelerator.device) # B L or B 1+num_images L - cond_mask = cond_mask.to(accelerator.device) # B L or B 1+num_images L - # print(args.use_image_num, x.shape, cond.shape, cond_mask.shape, cond_mask) - model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, - encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) - t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=accelerator.device) - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() - - # Gather the losses across all processes for logging (if we use distributed training). - avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() - train_loss += avg_loss.item() / args.gradient_accumulation_steps - - # Backpropagate - accelerator.backward(loss) - if accelerator.sync_gradients: - params_to_clip = model.parameters() - accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - # Checks if the accelerator has performed an optimization step behind the scenes - if accelerator.sync_gradients: - progress_bar.update(1) - global_step += 1 - accelerator.log({"train_loss": train_loss}, step=global_step) - train_loss = 0.0 - - if args.use_deepspeed or accelerator.is_main_process: - if global_step % args.checkpointing_steps == 0: - # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` - if args.checkpoints_total_limit is not None: - checkpoints = os.listdir(args.output_dir) - checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] - checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) - - # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints - if len(checkpoints) >= args.checkpoints_total_limit: - num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 - removing_checkpoints = checkpoints[0:num_to_remove] - - logger.info( - f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" - ) - logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") - - for removing_checkpoint in removing_checkpoints: - removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) - shutil.rmtree(removing_checkpoint) - - save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") - accelerator.save_state(save_path) - logger.info(f"Saved state to {save_path}") - - logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} - progress_bar.set_postfix(**logs) - - if global_step >= args.max_train_steps: - break - - if accelerator.is_main_process: - validation_prompt = "The majestic beauty of a waterfall cascading down a cliff into a serene lake. The camera angle provides a bird's eye view of the waterfall." - if global_step % args.checkpointing_steps == 0: - logger.info(f"Running validation... \n" - f"Generating {args.num_validation_videos} videos with prompt: {validation_prompt}") - if args.use_ema: - # Store the UNet parameters temporarily and load the EMA parameters to perform inference. - ema_model.store(model.parameters()) - ema_model.copy_to(model.parameters()) - - if args.enable_tracker: - 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() - - accelerator.wait_for_everyone() - accelerator.end_training() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--dataset", type=str, required=True) - parser.add_argument("--data_path", type=str, required=True) - parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122") - parser.add_argument("--num_classes", type=int, default=1000) - parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--sample_rate", type=int, default=4) - parser.add_argument("--num_frames", type=int, default=16) - parser.add_argument("--max_image_size", type=int, default=128) - parser.add_argument("--dynamic_frames", action="store_true") - parser.add_argument("--compress_kv", action="store_true") - parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math") - parser.add_argument("--pretrained", type=str, default=None) - - parser.add_argument('--tile_overlap_factor', type=float, default=0.25) - parser.add_argument('--enable_tiling', action='store_true') - - parser.add_argument("--video_folder", type=str, default='') - parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl') - parser.add_argument("--model_max_length", type=int, default=120) - - parser.add_argument("--enable_tracker", action="store_true") - parser.add_argument("--use_image_num", type=int, default=0) - parser.add_argument("--use_img_from_vid", 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, - default=None, - help="The output directory where the model predictions and checkpoints will be written.", - ) - parser.add_argument( - "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." - ) - parser.add_argument("--num_train_epochs", type=int, default=100) - parser.add_argument( - "--max_train_steps", - type=int, - default=None, - help="Total number of training steps to perform. If provided, overrides num_train_epochs.", - ) - parser.add_argument( - "--checkpointing_steps", - type=int, - default=500, - help=( - "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" - " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" - " training using `--resume_from_checkpoint`." - ), - ) - parser.add_argument( - "--checkpoints_total_limit", - type=int, - default=None, - help=("Max number of checkpoints to store."), - ) - parser.add_argument( - "--resume_from_checkpoint", - type=str, - default=None, - help=( - "Whether training should be resumed from a previous checkpoint. Use a path saved by" - ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' - ), - ) - parser.add_argument( - "--gradient_accumulation_steps", - type=int, - default=1, - help="Number of updates steps to accumulate before performing a backward/update pass.", - ) - parser.add_argument( - "--gradient_checkpointing", - action="store_true", - help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", - ) - parser.add_argument( - "--learning_rate", - type=float, - default=1e-4, - help="Initial learning rate (after the potential warmup period) to use.", - ) - parser.add_argument( - "--scale_lr", - action="store_true", - default=False, - help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", - ) - parser.add_argument( - "--lr_scheduler", - type=str, - default="constant", - help=( - 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' - ' "constant", "constant_with_warmup"]' - ), - ) - parser.add_argument( - "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." - ) - parser.add_argument( - "--timestep_bias_strategy", - type=str, - default="none", - choices=["earlier", "later", "range", "none"], - help=( - "The timestep bias strategy, which may help direct the model toward learning low or high frequency details." - " Choices: ['earlier', 'later', 'range', 'none']." - " The default is 'none', which means no bias is applied, and training proceeds normally." - " The value of 'later' will increase the frequency of the model's final training timesteps." - ), - ) - parser.add_argument( - "--timestep_bias_multiplier", - type=float, - default=1.0, - help=( - "The multiplier for the bias. Defaults to 1.0, which means no bias is applied." - " A value of 2.0 will double the weight of the bias, and a value of 0.5 will halve it." - ), - ) - parser.add_argument( - "--timestep_bias_begin", - type=int, - default=0, - help=( - "When using `--timestep_bias_strategy=range`, the beginning (inclusive) timestep to bias." - " Defaults to zero, which equates to having no specific bias." - ), - ) - parser.add_argument( - "--timestep_bias_end", - type=int, - default=1000, - help=( - "When using `--timestep_bias_strategy=range`, the final timestep (inclusive) to bias." - " Defaults to 1000, which is the number of timesteps that Stable Diffusion is trained on." - ), - ) - parser.add_argument( - "--timestep_bias_portion", - type=float, - default=0.25, - help=( - "The portion of timesteps to bias. Defaults to 0.25, which 25% of timesteps will be biased." - " A value of 0.5 will bias one half of the timesteps. The value provided for `--timestep_bias_strategy` determines" - " whether the biased portions are in the earlier or later timesteps." - ), - ) - parser.add_argument( - "--snr_gamma", - type=float, - default=None, - help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. " - "More details here: https://arxiv.org/abs/2303.09556.", - ) - parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") - parser.add_argument( - "--allow_tf32", - action="store_true", - help=( - "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" - " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" - ), - ) - parser.add_argument( - "--dataloader_num_workers", - type=int, - default=10, - help=( - "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." - ), - ) - parser.add_argument( - "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." - ) - parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") - parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") - parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") - parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") - parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") - parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") - parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") - 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.", - ) - parser.add_argument( - "--hub_model_id", - type=str, - default=None, - help="The name of the repository to keep in sync with the local `output_dir`.", - ) - 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***." - ), - ) - parser.add_argument( - "--report_to", - type=str, - default="tensorboard", - help=( - 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' - ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' - ), - ) - parser.add_argument( - "--mixed_precision", - type=str, - default=None, - choices=["no", "fp16", "bf16"], - help=( - "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" - " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" - " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." - ), - ) - parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") - parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") - - args = parser.parse_args() - main(args) diff --git a/opensora/train/train_t2v_t5_feature.py b/opensora/train/train_t2v_t5_feature.py deleted file mode 100644 index a8f3543e0..000000000 --- a/opensora/train/train_t2v_t5_feature.py +++ /dev/null @@ -1,825 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -""" -A minimal training script for DiT using PyTorch DDP. -""" -import argparse -import logging -import math -import os -import shutil -import sys -from pathlib import Path -from typing import Optional - -import numpy as np -from PIL import Image -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 -import transformers -from accelerate import Accelerator -from accelerate.logging import get_logger -from accelerate.utils import ProjectConfiguration, set_seed -from huggingface_hub import create_repo -from packaging import version -from tqdm.auto import tqdm -from transformers import HfArgumentParser, TrainingArguments, AutoTokenizer - -import diffusers -from diffusers import DDPMScheduler, PNDMScheduler -from diffusers.optimization import get_scheduler -from diffusers.training_utils import EMAModel, compute_snr -from diffusers.utils import check_min_version, is_wandb_available - -from examples.rec_imvi_vae import custom_to_video -from opensora.dataset import getdataset, ae_denorm -from opensora.models.ae import getae, getae_wrapper -from opensora.models.ae.videobase import CausalVQVAEModelWrapper, CausalVAEModelWrapper -from opensora.models.diffusion.diffusion import create_diffusion_T as create_diffusion -from opensora.models.diffusion.latte.modeling_latte import LatteT2V -from opensora.models.text_encoder import get_text_enc -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 - -# 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." - ) - 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 - - -################################################################################# -# Training Loop # -################################################################################# - -def main(args): - logging_dir = Path(args.output_dir, args.logging_dir) - - accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) - - accelerator = Accelerator( - gradient_accumulation_steps=args.gradient_accumulation_steps, - mixed_precision=args.mixed_precision, - log_with=args.report_to, - project_config=accelerator_project_config, - ) - - 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( - format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", - datefmt="%m/%d/%Y %H:%M:%S", - level=logging.INFO, - ) - logger.info(accelerator.state, main_process_only=False) - if accelerator.is_local_main_process: - transformers.utils.logging.set_verbosity_warning() - diffusers.utils.logging.set_verbosity_info() - else: - transformers.utils.logging.set_verbosity_error() - diffusers.utils.logging.set_verbosity_error() - - # If passed along, set the training seed now. - if args.seed is not None: - set_seed(args.seed) - - # Handle the repository creation - if accelerator.is_main_process: - if args.output_dir is not None: - os.makedirs(args.output_dir, exist_ok=True) - - # if args.push_to_hub: - # repo_id = create_repo( - # repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token - # ).repo_id - - # Create model: - - diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule - ae = getae_wrapper(args.ae)(args.ae_path).eval() - if args.enable_tiling: - ae.vae.enable_tiling() - ae.vae.tile_overlap_factor = args.tile_overlap_factor - # text_enc = get_text_enc(args).eval() - - ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] - args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w - args.ae_stride = args.ae_stride_h - patch_size = args.model[-3:] - patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) - args.patch_size = patch_size_h - args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w - assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" - assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" - # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." - assert args.max_image_size % ae_stride_h == 0, f"Image size must be divisible by ae_stride_h, but found max_image_size ({args.max_image_size}), ae_stride_h ({ae_stride_h})." - - latent_size = (args.max_image_size // ae_stride_h, args.max_image_size // ae_stride_w) - - if getae_wrapper(args.ae) == CausalVQVAEModelWrapper or getae_wrapper(args.ae) == CausalVAEModelWrapper: - args.video_length = video_length = args.num_frames // ae_stride_t + 1 - else: - args.video_length = video_length = args.num_frames // ae_stride_t - model = Diffusion_models[args.model]( - in_channels=ae_channel_config[args.ae], - out_channels=ae_channel_config[args.ae] * 2, - # caption_channels=4096, - # cross_attention_dim=1152, - attention_bias=True, - sample_size=latent_size, - num_vector_embeds=None, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - use_linear_projection=False, - only_cross_attention=False, - double_self_attention=False, - upcast_attention=False, - # norm_type="ada_norm_single", - norm_elementwise_affine=False, - norm_eps=1e-6, - attention_type='default', - video_length=video_length, - attention_mode=args.attention_mode, - # compress_kv=args.compress_kv - ) - model.gradient_checkpointing = args.gradient_checkpointing - - # # use pretrained model? - if args.pretrained: - if 'safetensors' in args.pretrained: - from safetensors.torch import load_file as safe_load - checkpoint = safe_load(args.pretrained, device="cpu") - else: - checkpoint = torch.load(args.pretrained, map_location='cpu')['model'] - model_state_dict = model.state_dict() - missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) - logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') - logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!') - # load from pixart-alpha - # pixelart_alpha = torch.load(args.pretrained, map_location='cpu')['state_dict'] - # checkpoint = {} - # for k, v in pixelart_alpha.items(): - # if 'x_embedder' in k or 't_embedder' in k or 'y_embedder' in k: - # checkpoint[k] = v - # if k.startswith('blocks'): - # k_spilt = k.split('.') - # blk_id = str(int(k_spilt[1]) * 2) - # k_spilt[1] = blk_id - # new_k = '.'.join(k_spilt) - # checkpoint[new_k] = v - # missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) - # logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)} keys from {args.pretrained}!') - - # Freeze vae and text encoders. - ae.requires_grad_(False) - # text_enc.requires_grad_(False) - # Set model as trainable. - model.train() - - # For mixed precision training we cast all non-trainable weigths to half-precision - # as these weights are only used for inference, keeping weights in full precision is not required. - weight_dtype = torch.float32 - if accelerator.mixed_precision == "fp16": - weight_dtype = torch.float16 - elif accelerator.mixed_precision == "bf16": - weight_dtype = torch.bfloat16 - - # Move unet, vae and text_encoder to device and cast to weight_dtype - # The VAE is in float32 to avoid NaN losses. - # ae.to(accelerator.device, dtype=torch.float32) - ae.to(accelerator.device, dtype=weight_dtype) - model.to(accelerator.device, dtype=weight_dtype) - # text_enc.to(accelerator.device, dtype=weight_dtype) - - # Create EMA for the unet. - if args.use_ema: - ema_model = deepcopy(model) - ema_model = EMAModel(ema_model.parameters(), model_cls=LatteT2V, model_config=ema_model.config) - - # `accelerate` 0.16.0 will have better support for customized saving - if version.parse(accelerate.__version__) >= version.parse("0.16.0"): - # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format - def save_model_hook(models, weights, output_dir): - if accelerator.is_main_process: - if args.use_ema: - ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) - - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "model")) - if weights: # Don't pop if empty - # make sure to pop weight so that corresponding model is not saved again - weights.pop() - - def load_model_hook(models, input_dir): - if args.use_ema: - load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), LatteT2V) - ema_model.load_state_dict(load_model.state_dict()) - ema_model.to(accelerator.device) - del load_model - - for i in range(len(models)): - # pop models so that they are not loaded again - model = models.pop() - - # load diffusers style into model - load_model = LatteT2V.from_pretrained(input_dir, subfolder="model") - model.register_to_config(**load_model.config) - - model.load_state_dict(load_model.state_dict()) - del load_model - - accelerator.register_save_state_pre_hook(save_model_hook) - accelerator.register_load_state_pre_hook(load_model_hook) - - # Enable TF32 for faster training on Ampere GPUs, - # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices - if args.allow_tf32: - torch.backends.cuda.matmul.allow_tf32 = True - - if args.scale_lr: - args.learning_rate = ( - args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes - ) - - # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs - if args.use_8bit_adam: - try: - import bitsandbytes as bnb - except ImportError: - raise ImportError( - "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." - ) - - optimizer_class = bnb.optim.AdamW8bit - else: - optimizer_class = torch.optim.AdamW - - # Optimizer creation - params_to_optimize = model.parameters() - optimizer = optimizer_class( - params_to_optimize, - lr=args.learning_rate, - betas=(args.adam_beta1, args.adam_beta2), - weight_decay=args.adam_weight_decay, - eps=args.adam_epsilon, - ) - - # Setup data: - train_dataset = getdataset(args) - train_dataloader = torch.utils.data.DataLoader( - train_dataset, - shuffle=True, - # collate_fn=Collate(args), # TODO: do not enable dynamic mask in this point - batch_size=args.train_batch_size, - num_workers=args.dataloader_num_workers, - ) - - # Scheduler and math around the number of training steps. - overrode_max_train_steps = False - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if args.max_train_steps is None: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - overrode_max_train_steps = True - - lr_scheduler = get_scheduler( - args.lr_scheduler, - optimizer=optimizer, - num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, - num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, - ) - - # Prepare everything with our `accelerator`. - model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, lr_scheduler - ) - - # We need to recalculate our total training steps as the size of the training dataloader may have changed. - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if overrode_max_train_steps: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - # Afterwards we recalculate our number of training epochs - args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) - - # We need to initialize the trackers we use, and also store our configuration. - # The trackers initializes automatically on the main process. - if accelerator.is_main_process: - accelerator.init_trackers(args.output_dir, config=vars(args)) - - # Train! - total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps - - logger.info("***** Running training *****") - logger.info(f" Num examples = {len(train_dataset)}") - logger.info(f" Num Epochs = {args.num_train_epochs}") - logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") - logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") - logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") - logger.info(f" Total optimization steps = {args.max_train_steps}") - global_step = 0 - first_epoch = 0 - - # Potentially load in the weights and states from a previous save - if args.resume_from_checkpoint: - if args.resume_from_checkpoint != "latest": - path = os.path.basename(args.resume_from_checkpoint) - else: - # Get the most recent checkpoint - dirs = os.listdir(args.output_dir) - dirs = [d for d in dirs if d.startswith("checkpoint")] - dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) - path = dirs[-1] if len(dirs) > 0 else None - - if path is None: - accelerator.print( - f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." - ) - args.resume_from_checkpoint = None - initial_global_step = 0 - else: - accelerator.print(f"Resuming from checkpoint {path}") - accelerator.load_state(os.path.join(args.output_dir, path)) - global_step = int(path.split("-")[1]) - - initial_global_step = global_step - first_epoch = global_step // num_update_steps_per_epoch - - else: - initial_global_step = 0 - - progress_bar = tqdm( - range(0, args.max_train_steps), - initial=initial_global_step, - desc="Steps", - # Only show the progress bar once on each machine. - disable=not accelerator.is_local_main_process, - ) - - for epoch in range(first_epoch, args.num_train_epochs): - train_loss = 0.0 - for step, (x, cond, cond_mask) in enumerate(train_dataloader): - with accelerator.accumulate(model): - # Sample noise that we'll add to the latents - x = x.to(accelerator.device) # B C T H W - # print(x.dtype) - # attn_mask = attn_mask.to(device) # B T H W - # assert torch.all(attn_mask.bool()), 'do not enable dynamic input' - attn_mask = None - cond = cond.to(accelerator.device, dtype=weight_dtype) # B L or B 1+num_images L - cond_mask = cond_mask.to(accelerator.device) # B L or B 1+num_images L - - with torch.no_grad(): - # Map input images to latent space + normalize latents - if args.use_image_num == 0: - x = ae.encode(x.to(dtype=weight_dtype)) # B C T H W - else: - videos, images = x[:, :, :-args.use_image_num], x[:, :, -args.use_image_num:] - videos = ae.encode(videos.to(dtype=weight_dtype)) # B C T H W - - # videos = ae.decode(videos.to(dtype=weight_dtype))[0] - # videos = videos.transpose(0, 1) - # custom_to_video(videos.to(torch.float32), fps=24, output_file='tmp.mp4') - - images = rearrange(images, 'b c t h w -> (b t) c 1 h w') - images = ae.encode(images.to(dtype=weight_dtype)) - - # images = ae.decode(images.to(dtype=weight_dtype)) - # x = images[0, 0, :, :, :].to(torch.float32) - # x = x.squeeze() - # x = x.detach().cpu().numpy() - # x = np.clip(x, -1, 1) - # x = (x + 1) / 2 - # x = (255 * x).astype(np.uint8) - # x = x.transpose(1, 2, 0) - # image = Image.fromarray(x) - # image.save('tmp.jpg') - # sys.exit() - - images = rearrange(images, '(b t) c 1 h w -> b c t h w', t=args.use_image_num) - x = torch.cat([videos, images], dim=2) - - # print(args.use_image_num, x.shape, cond.shape, cond_mask.shape, cond_mask) - model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, - encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) - t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=accelerator.device) - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() - - # Gather the losses across all processes for logging (if we use distributed training). - avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() - train_loss += avg_loss.item() / args.gradient_accumulation_steps - - # Backpropagate - accelerator.backward(loss) - if accelerator.sync_gradients: - params_to_clip = model.parameters() - accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - # Checks if the accelerator has performed an optimization step behind the scenes - if accelerator.sync_gradients: - progress_bar.update(1) - global_step += 1 - accelerator.log({"train_loss": train_loss}, step=global_step) - train_loss = 0.0 - - if args.use_deepspeed or accelerator.is_main_process: - if global_step % args.checkpointing_steps == 0: - # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` - if args.checkpoints_total_limit is not None: - checkpoints = os.listdir(args.output_dir) - checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] - checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) - - # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints - if len(checkpoints) >= args.checkpoints_total_limit: - num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 - removing_checkpoints = checkpoints[0:num_to_remove] - - logger.info( - f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" - ) - logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") - - for removing_checkpoint in removing_checkpoints: - removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) - shutil.rmtree(removing_checkpoint) - - save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") - accelerator.save_state(save_path) - logger.info(f"Saved state to {save_path}") - - logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} - progress_bar.set_postfix(**logs) - - if global_step >= args.max_train_steps: - break - - if accelerator.is_main_process: - validation_prompt = "The majestic beauty of a waterfall cascading down a cliff into a serene lake. The camera angle provides a bird's eye view of the waterfall." - if global_step % args.checkpointing_steps == 0: - logger.info(f"Running validation... \n" - f"Generating {args.num_validation_videos} videos with prompt: {validation_prompt}") - if args.use_ema: - # Store the UNet parameters temporarily and load the EMA parameters to perform inference. - ema_model.store(model.parameters()) - ema_model.copy_to(model.parameters()) - if args.enable_tracker: - 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(500)) - 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=24) - if tracker.name == "wandb": - tracker.log( - { - "validation": [ - wandb.Video(video, caption=f"{i}: {validation_prompt}", fps=24) - for i, video in enumerate(videos) - ] - } - ) - - del ae_, text_enc_, model_, diffusion_, tokenizer_ - # del ae_, model_, diffusion_, tokenizer_ - torch.cuda.empty_cache() - - accelerator.wait_for_everyone() - accelerator.end_training() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--dataset", type=str, required=True) - parser.add_argument("--data_path", type=str, required=True) - parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122") - parser.add_argument("--num_classes", type=int, default=1000) - parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--sample_rate", type=int, default=4) - parser.add_argument("--num_frames", type=int, default=16) - parser.add_argument("--max_image_size", type=int, default=128) - parser.add_argument("--dynamic_frames", action="store_true") - parser.add_argument("--compress_kv", action="store_true") - parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math") - parser.add_argument("--pretrained", type=str, default=None) - - parser.add_argument('--tile_overlap_factor', type=float, default=0.25) - parser.add_argument('--enable_tiling', action='store_true') - - parser.add_argument("--video_folder", type=str, default='') - parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl') - parser.add_argument("--model_max_length", type=int, default=120) - - parser.add_argument("--use_image_num", type=int, default=0) - parser.add_argument("--use_img_from_vid", action="store_true") - 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, - default=None, - help="The output directory where the model predictions and checkpoints will be written.", - ) - parser.add_argument( - "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." - ) - parser.add_argument("--num_train_epochs", type=int, default=100) - parser.add_argument( - "--max_train_steps", - type=int, - default=None, - help="Total number of training steps to perform. If provided, overrides num_train_epochs.", - ) - parser.add_argument( - "--checkpointing_steps", - type=int, - default=500, - help=( - "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" - " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" - " training using `--resume_from_checkpoint`." - ), - ) - parser.add_argument( - "--checkpoints_total_limit", - type=int, - default=None, - help=("Max number of checkpoints to store."), - ) - parser.add_argument( - "--resume_from_checkpoint", - type=str, - default=None, - help=( - "Whether training should be resumed from a previous checkpoint. Use a path saved by" - ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' - ), - ) - parser.add_argument( - "--gradient_accumulation_steps", - type=int, - default=1, - help="Number of updates steps to accumulate before performing a backward/update pass.", - ) - parser.add_argument( - "--gradient_checkpointing", - action="store_true", - help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", - ) - parser.add_argument( - "--learning_rate", - type=float, - default=1e-4, - help="Initial learning rate (after the potential warmup period) to use.", - ) - parser.add_argument( - "--scale_lr", - action="store_true", - default=False, - help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", - ) - parser.add_argument( - "--lr_scheduler", - type=str, - default="constant", - help=( - 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' - ' "constant", "constant_with_warmup"]' - ), - ) - parser.add_argument( - "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." - ) - parser.add_argument( - "--timestep_bias_strategy", - type=str, - default="none", - choices=["earlier", "later", "range", "none"], - help=( - "The timestep bias strategy, which may help direct the model toward learning low or high frequency details." - " Choices: ['earlier', 'later', 'range', 'none']." - " The default is 'none', which means no bias is applied, and training proceeds normally." - " The value of 'later' will increase the frequency of the model's final training timesteps." - ), - ) - parser.add_argument( - "--timestep_bias_multiplier", - type=float, - default=1.0, - help=( - "The multiplier for the bias. Defaults to 1.0, which means no bias is applied." - " A value of 2.0 will double the weight of the bias, and a value of 0.5 will halve it." - ), - ) - parser.add_argument( - "--timestep_bias_begin", - type=int, - default=0, - help=( - "When using `--timestep_bias_strategy=range`, the beginning (inclusive) timestep to bias." - " Defaults to zero, which equates to having no specific bias." - ), - ) - parser.add_argument( - "--timestep_bias_end", - type=int, - default=1000, - help=( - "When using `--timestep_bias_strategy=range`, the final timestep (inclusive) to bias." - " Defaults to 1000, which is the number of timesteps that Stable Diffusion is trained on." - ), - ) - parser.add_argument( - "--timestep_bias_portion", - type=float, - default=0.25, - help=( - "The portion of timesteps to bias. Defaults to 0.25, which 25% of timesteps will be biased." - " A value of 0.5 will bias one half of the timesteps. The value provided for `--timestep_bias_strategy` determines" - " whether the biased portions are in the earlier or later timesteps." - ), - ) - parser.add_argument( - "--snr_gamma", - type=float, - default=None, - help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. " - "More details here: https://arxiv.org/abs/2303.09556.", - ) - parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") - parser.add_argument( - "--allow_tf32", - action="store_true", - help=( - "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" - " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" - ), - ) - parser.add_argument( - "--dataloader_num_workers", - type=int, - default=10, - help=( - "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." - ), - ) - parser.add_argument( - "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." - ) - parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") - parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") - parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") - parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") - parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") - parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") - parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") - 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.", - ) - parser.add_argument( - "--hub_model_id", - type=str, - default=None, - help="The name of the repository to keep in sync with the local `output_dir`.", - ) - 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***." - ), - ) - parser.add_argument( - "--report_to", - type=str, - default="tensorboard", - help=( - 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' - ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' - ), - ) - parser.add_argument( - "--mixed_precision", - type=str, - default=None, - choices=["no", "fp16", "bf16"], - help=( - "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" - " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" - " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." - ), - ) - parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") - parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") - - args = parser.parse_args() - main(args) \ No newline at end of file diff --git a/opensora/train/train_videogpt.py b/opensora/train/train_videogpt.py deleted file mode 100644 index e953af872..000000000 --- a/opensora/train/train_videogpt.py +++ /dev/null @@ -1,72 +0,0 @@ -import sys - -sys.path.append(".") - -from opensora.models.ae.videobase.dataset_videobase import VideoDataset -from opensora.models.ae.videobase import ( - VQVAEModel, - VQVAEConfiguration, - VQVAETrainer, -) -import argparse -from typing import Optional -from accelerate.utils import set_seed -from transformers import HfArgumentParser, TrainingArguments -from dataclasses import dataclass, field, asdict - - -@dataclass -class VQVAEArgument: - embedding_dim: int = (field(default=256),) - n_codes: int = (field(default=2048),) - n_hiddens: int = (field(default=240),) - n_res_layers: int = (field(default=4),) - resolution: int = (field(default=128),) - sequence_length: int = (field(default=16),) - downsample: str = (field(default="4,4,4"),) - no_pos_embd: bool = (True,) - data_path: str = field(default=None, metadata={"help": "data path"}) - - -@dataclass -class VQVAETrainingArgument(TrainingArguments): - remove_unused_columns: Optional[bool] = field( - default=False, - metadata={ - "help": "Remove columns not required by the model when using an nlp.Dataset." - }, - ) - - -def train(args, vqvae_args: VQVAEArgument, training_args: VQVAETrainingArgument): - # Load Config - config = VQVAEConfiguration( - embedding_dim=vqvae_args.embedding_dim, - n_codes=vqvae_args.n_codes, - n_hiddens=vqvae_args.n_hiddens, - n_res_layers=vqvae_args.n_res_layers, - resolution=vqvae_args.resolution, - sequence_length=vqvae_args.sequence_length, - downsample=vqvae_args.downsample, - no_pos_embd=vqvae_args.no_pos_embd, - ) - # Load Model - model = VQVAEModel(config) - # Load Dataset - dataset = VideoDataset( - args.data_path, - sequence_length=args.sequence_length, - resolution=config.resolution, - ) - # Load Trainer - trainer = VQVAETrainer(model, training_args, train_dataset=dataset) - trainer.train() - - -if __name__ == "__main__": - parser = HfArgumentParser((VQVAEArgument, VQVAETrainingArgument)) - vqvae_args, training_args = parser.parse_args_into_dataclasses() - args = argparse.Namespace(**vars(vqvae_args), **vars(training_args)) - set_seed(args.seed) - - train(args, vqvae_args, training_args) diff --git a/opensora/utils/communications.py b/opensora/utils/communications.py new file mode 100644 index 000000000..c7a79e3a6 --- /dev/null +++ b/opensora/utils/communications.py @@ -0,0 +1,122 @@ +import torch +import torch.distributed as dist +from einops import rearrange +from opensora.utils.parallel_states import nccl_info + +def broadcast(input_: torch.Tensor): + sp_size = nccl_info.world_size + src = nccl_info.rank // sp_size * sp_size + dist.broadcast(input_, src=src, group=nccl_info.group) + +_COUNT = 0 +def _all_to_all( + input_: torch.Tensor, + scatter_dim: int, + gather_dim: int, +): + group = nccl_info.group + sp_size = nccl_info.world_size + input_list = [t.contiguous() for t in torch.tensor_split(input_, sp_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + +def _single_all_to_all( + input_: torch.Tensor, + scatter_dim: int, + gather_dim: int, + enable_HCCL=False, +): + + sp_size = nccl_info.world_size + inp_shape = list(input_.shape) + inp_shape[scatter_dim] = inp_shape[scatter_dim] // sp_size + if scatter_dim < 1: + input_t = input_.reshape( + [sp_size, inp_shape[scatter_dim]] + \ + inp_shape[scatter_dim + 1:] + ) + else: + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + input_t = input_.reshape( + [-1, sp_size, inp_shape[scatter_dim]] + \ + inp_shape[scatter_dim + 1:] + ).transpose(0, 1).contiguous() + + output = torch.empty_like(input_t) + dist.all_to_all_single(output, input_t, group=nccl_info.group) + # if scattering the seq-dim, transpose the heads back to the original dimension + if scatter_dim < 1: + output = output.transpose(0, 1).contiguous() + + return output.reshape( + inp_shape[: gather_dim] + [inp_shape[gather_dim] * sp_size, ] + inp_shape[gather_dim + 1:]) + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, scatter_dim, gather_dim, all_to_all_func): + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.all_to_all = all_to_all_func + output = ctx.all_to_all(input_, scatter_dim, gather_dim) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = ctx.all_to_all( + grad_output, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + +def all_to_all_SBH( + input_: torch.Tensor, + scatter_dim: int = 1, + gather_dim: int = 0, +): + return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all) + +def all_to_all_BSND( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all) + + +def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num): + def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): + hidden_states = _single_all_to_all(hidden_states, scatter_dim=2, gather_dim=0, enable_HCCL=True) + encoder_hidden_states = _single_all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0, enable_HCCL=True) + attention_mask = _single_all_to_all(attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True) + encoder_attention_mask = _single_all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True) + return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask + + sp_size = nccl_info.world_size + frame = hidden_states.shape[2] + assert frame % sp_size == 0, "frame should be a multiple of sp_size" + + encoder_hidden_states = rearrange(encoder_hidden_states, 'b 1 (n x) h -> b n x h', + n=sp_size, x=encoder_hidden_states.shape[2]//sp_size).contiguous() + hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask = all_to_all(hidden_states, + encoder_hidden_states, + attention_mask.repeat(1, sp_size, 1, 1), + encoder_attention_mask.repeat(1, sp_size, 1)) + + return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num \ No newline at end of file diff --git a/opensora/utils/dataset_utils.py b/opensora/utils/dataset_utils.py index 8ab656446..2f9169b50 100644 --- a/opensora/utils/dataset_utils.py +++ b/opensora/utils/dataset_utils.py @@ -3,6 +3,16 @@ import decord from torch.nn import functional as F import torch +from typing import Optional +import torch.utils +import torch.utils.data +import torch +from torch.utils.data import Sampler +from typing import List +from collections import Counter, defaultdict +import random + +from opensora.utils.mask_utils import MaskProcessor IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG'] @@ -43,99 +53,315 @@ def pad_to_multiple(number, ds_stride): return number + padding class Collate: - def __init__(self, args): - self.max_image_size = args.max_image_size + def __init__(self, args,YOLOmodel): + self.batch_size = args.train_batch_size + self.group_data = args.group_data + self.force_resolution = args.force_resolution + + self.max_height = args.max_height + self.max_width = args.max_width self.ae_stride = args.ae_stride + self.ae_stride_t = args.ae_stride_t self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride) - self.ae_stride_1hw = (1, self.ae_stride, self.ae_stride) self.patch_size = args.patch_size self.patch_size_t = args.patch_size_t - self.patch_size_thw = (self.patch_size_t, self.patch_size, self.patch_size) - self.patch_size_1hw = (1, self.patch_size, self.patch_size) self.num_frames = args.num_frames self.use_image_num = args.use_image_num - self.max_thw = (self.num_frames, self.max_image_size, self.max_image_size) - self.max_1hw = (1, self.max_image_size, self.max_image_size) + self.max_thw = (self.num_frames, self.max_height, self.max_width) + + self.mask_processor = MaskProcessor(args,YOLOmodel) + def package(self, batch): - batch_tubes_vid = [i['video_data']['video'] for i in batch] # b [c t h w] - input_ids_vid = torch.stack([i['video_data']['input_ids'] for i in batch]) # b 1 l - cond_mask_vid = torch.stack([i['video_data']['cond_mask'] for i in batch]) # b 1 l - batch_tubes_img, input_ids_img, cond_mask_img = None, None, None - if self.use_image_num != 0: - batch_tubes_img = [j for i in batch for j in i['image_data']['image']] # b*num_img [c 1 h w] - input_ids_img = torch.stack([i['image_data']['input_ids'] for i in batch]) # b image_num l - cond_mask_img = torch.stack([i['image_data']['cond_mask'] for i in batch]) # b image_num l - return batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img + batch_tubes = [i['pixel_values'] for i in batch] # b [c t h w] + input_ids = [i['input_ids'] for i in batch] # b [1 l] + cond_mask = [i['cond_mask'] for i in batch] # b [1 l] + motion_score = [i['motion_score'] for i in batch] # List[float] + assert all([i is None for i in motion_score]) or all([i is not None for i in motion_score]) + if all([i is None for i in motion_score]): + motion_score = None + return batch_tubes, input_ids, cond_mask, motion_score def __call__(self, batch): - batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img = self.package(batch) + batch_tubes, input_ids, cond_mask, motion_score = self.package(batch) + # b 2c+1 * t * h * w + masked_batch_tubes = [] + + for pixed_values in batch_tubes: + masked_video,video,mask = self.mask_processor(pixed_values) + masked_batch_tube = torch.cat((video,masked_video,mask),dim=1) + masked_batch_tubes.append(masked_batch_tube) ds_stride = self.ae_stride * self.patch_size t_ds_stride = self.ae_stride_t * self.patch_size_t - if self.use_image_num == 0: - pad_batch_tubes, attention_mask = self.process(batch_tubes_vid, t_ds_stride, ds_stride, - self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True) - # attention_mask: b t h w - input_ids, cond_mask = input_ids_vid.squeeze(1), cond_mask_vid.squeeze(1) # b 1 l -> b l - else: - pad_batch_tubes_vid, attention_mask_vid = self.process(batch_tubes_vid, t_ds_stride, ds_stride, - self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True) - # attention_mask_vid: b t h w - pad_batch_tubes_img, attention_mask_img = self.process(batch_tubes_img, 1, ds_stride, - self.max_1hw, self.ae_stride_1hw, self.patch_size_1hw, extra_1=False) - pad_batch_tubes_img = rearrange(pad_batch_tubes_img, '(b i) c 1 h w -> b c i h w', i=self.use_image_num) - attention_mask_img = rearrange(attention_mask_img, '(b i) 1 h w -> b i h w', i=self.use_image_num) - pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=2) # concat at temporal, video first - # attention_mask_img: b num_img h w - attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1) # b t+num_img h w - input_ids = torch.cat([input_ids_vid, input_ids_img], dim=1) # b 1+num_img hw - cond_mask = torch.cat([cond_mask_vid, cond_mask_img], dim=1) # b 1+num_img hw - return pad_batch_tubes, attention_mask, input_ids, cond_mask - - def process(self, batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw, extra_1): + + pad_batch_tubes, attention_mask, input_ids, cond_mask, motion_score = self.process(masked_batch_tubes, input_ids, + cond_mask, motion_score, + t_ds_stride, ds_stride, + self.max_thw, self.ae_stride_thw) + + # pad_batch_tubes, attention_mask, input_ids, cond_mask, motion_score = self.process(batch_tubes, input_ids, + # cond_mask, motion_score, + # t_ds_stride, ds_stride, + # self.max_thw, self.ae_stride_thw) + assert not torch.any(torch.isnan(pad_batch_tubes)), 'after pad_batch_tubes' + return pad_batch_tubes, attention_mask, input_ids, cond_mask, motion_score + + def process(self, batch_tubes, input_ids, cond_mask, motion_score, t_ds_stride, ds_stride, max_thw, ae_stride_thw): # pad to max multiple of ds_stride batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] - max_t, max_h, max_w = max_thw - pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1 if extra_1 else max_t, t_ds_stride), \ + assert len(batch_input_size) == self.batch_size + if self.group_data or self.batch_size == 1: # + len_each_batch = batch_input_size + idx_length_dict = dict([*zip(list(range(self.batch_size)), len_each_batch)]) + count_dict = Counter(len_each_batch) + if len(count_dict) != 1: + sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1]) + # import ipdb;ipdb.set_trace() + # print(batch, idx_length_dict, count_dict, sorted_by_value) + pick_length = sorted_by_value[-1][0] # the highest frequency + candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length] + random_select_batch = [random.choice(candidate_batch) for _ in range(len(len_each_batch) - len(candidate_batch))] + print(batch_input_size, idx_length_dict, count_dict, sorted_by_value, pick_length, candidate_batch, random_select_batch) + pick_idx = candidate_batch + random_select_batch + + batch_tubes = [batch_tubes[i] for i in pick_idx] + batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] + input_ids = [input_ids[i] for i in pick_idx] # b [1, l] + cond_mask = [cond_mask[i] for i in pick_idx] # b [1, l] + if motion_score is not None: + motion_score = [motion_score[i] for i in pick_idx] # b [1, l] + + for i in range(1, self.batch_size): + assert batch_input_size[0] == batch_input_size[i] + max_t = max([i[1] for i in batch_input_size]) + max_h = max([i[2] for i in batch_input_size]) + max_w = max([i[3] for i in batch_input_size]) + else: + max_t, max_h, max_w = max_thw + pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1+self.ae_stride_t, t_ds_stride), \ pad_to_multiple(max_h, ds_stride), \ pad_to_multiple(max_w, ds_stride) - pad_max_t = pad_max_t + 1 if extra_1 else pad_max_t - each_pad_t_h_w = [[pad_max_t - i.shape[1], - pad_max_h - i.shape[2], - pad_max_w - i.shape[3]] for i in batch_tubes] - pad_batch_tubes = [F.pad(im, - (0, pad_w, - 0, pad_h, - 0, pad_t), value=0) for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)] + pad_max_t = pad_max_t + 1 - self.ae_stride_t + each_pad_t_h_w = [ + [ + pad_max_t - i.shape[1], + pad_max_h - i.shape[2], + pad_max_w - i.shape[3] + ] for i in batch_tubes + ] + pad_batch_tubes = [ + F.pad(im, (0, pad_w, 0, pad_h, 0, pad_t), value=0) + for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes) + ] pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0) - # make attention_mask - first_channel_first_frame, first_channel_other_frame = pad_batch_tubes[:, :1, :1], pad_batch_tubes[:, :1, 1:] # first channel to make attention_mask - attention_mask_first_frame = F.max_pool3d(first_channel_first_frame, kernel_size=(1, *ae_stride_thw[1:]), stride=(1, *ae_stride_thw[1:])) - if first_channel_other_frame.numel() != 0: - attention_mask_other_frame = F.max_pool3d(first_channel_other_frame, kernel_size=ae_stride_thw, stride=ae_stride_thw) - attention_mask = torch.cat([attention_mask_first_frame, attention_mask_other_frame], dim=2) - else: - attention_mask = attention_mask_first_frame - attention_mask = attention_mask[:, 0].bool().float() # b t h w, do not channel - # max_tube_size = [pad_max_t, pad_max_h, pad_max_w] - # max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]), - # max_tube_size[1] // ae_stride_thw[1], - # max_tube_size[2] // ae_stride_thw[2]] - # max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]), - # max_latent_size[1] // patch_size_thw[1], - # max_latent_size[2] // patch_size_thw[2]] - # valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)), - # int(math.ceil(i[2] / ds_stride)), - # int(math.ceil(i[3] / ds_stride))] for i in batch_input_size] - # attention_mask = [F.pad(torch.ones(i), - # (0, max_patchify_latent_size[2] - i[2], - # 0, max_patchify_latent_size[1] - i[1], - # 0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size] - # attention_mask = torch.stack(attention_mask) # b t h w - - return pad_batch_tubes, attention_mask + + max_tube_size = [pad_max_t, pad_max_h, pad_max_w] + max_latent_size = [ + ((max_tube_size[0]-1) // ae_stride_thw[0] + 1), + max_tube_size[1] // ae_stride_thw[1], + max_tube_size[2] // ae_stride_thw[2] + ] + valid_latent_size = [ + [ + int(math.ceil((i[1]-1) / ae_stride_thw[0])) + 1, + int(math.ceil(i[2] / ae_stride_thw[1])), + int(math.ceil(i[3] / ae_stride_thw[2])) + ] for i in batch_input_size] + attention_mask = [ + F.pad(torch.ones(i, dtype=pad_batch_tubes.dtype), (0, max_latent_size[2] - i[2], + 0, max_latent_size[1] - i[1], + 0, max_latent_size[0] - i[0]), value=0) for i in valid_latent_size] + attention_mask = torch.stack(attention_mask) # b t h w + if self.batch_size == 1 or self.group_data: + if not torch.all(attention_mask.bool()): + print(batch_input_size, (max_t, max_h, max_w), (pad_max_t, pad_max_h, pad_max_w), each_pad_t_h_w, max_latent_size, valid_latent_size) + assert torch.all(attention_mask.bool()) + + input_ids = torch.stack(input_ids) # b 1 l + cond_mask = torch.stack(cond_mask) # b 1 l + motion_score = torch.tensor(motion_score) if motion_score is not None else motion_score # b + + return pad_batch_tubes, attention_mask, input_ids, cond_mask, motion_score + + + + +def group_data_fun(lengths, generator=None): + counter = Counter(lengths) + grouped_indices = defaultdict(list) + for idx, item in enumerate(lengths): + grouped_indices[counter[item]].append(idx) + grouped_indices = dict(grouped_indices) + sorted_indices = [grouped_indices[count] for count in sorted(grouped_indices, reverse=True)] + + shuffle_sorted_indices = [] + for indice in sorted_indices: + shuffle_idx = torch.randperm(len(indice), generator=generator).tolist() + shuffle_sorted_indices.extend([indice[idx] for idx in shuffle_idx]) + return shuffle_sorted_indices + +def last_group_data_fun(shuffled_megabatches, lengths): + re_shuffled_megabatches = [] + # print('shuffled_megabatches', len(shuffled_megabatches)) + for i_megabatch, megabatch in enumerate(shuffled_megabatches): + re_megabatch = [] + for i_batch, batch in enumerate(megabatch): + assert len(batch) != 0 + + len_each_batch = [lengths[i] for i in batch] + idx_length_dict = dict([*zip(batch, len_each_batch)]) + count_dict = Counter(len_each_batch) + if len(count_dict) != 1: + sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1]) + # import ipdb;ipdb.set_trace() + # print(batch, idx_length_dict, count_dict, sorted_by_value) + pick_length = sorted_by_value[-1][0] # the highest frequency + candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length] + random_select_batch = [random.choice(candidate_batch) for i in range(len(len_each_batch) - len(candidate_batch))] + # print(batch, idx_length_dict, count_dict, sorted_by_value, pick_length, candidate_batch, random_select_batch) + batch = candidate_batch + random_select_batch + # print(batch) + + for i in range(1, len(batch)-1): + # if not lengths[batch[0]] == lengths[batch[i]]: + # print(batch, [lengths[i] for i in batch]) + # import ipdb;ipdb.set_trace() + assert lengths[batch[0]] == lengths[batch[i]] + re_megabatch.append(batch) + re_shuffled_megabatches.append(re_megabatch) + + + # for megabatch, re_megabatch in zip(shuffled_megabatches, re_shuffled_megabatches): + # for batch, re_batch in zip(megabatch, re_megabatch): + # for i, re_i in zip(batch, re_batch): + # if i != re_i: + # print(i, re_i) + return re_shuffled_megabatches + +def split_to_even_chunks(indices, lengths, num_chunks, batch_size): + """ + Split a list of indices into `chunks` chunks of roughly equal lengths. + """ + + if len(indices) % num_chunks != 0: + chunks = [indices[i::num_chunks] for i in range(num_chunks)] + else: + num_indices_per_chunk = len(indices) // num_chunks + + chunks = [[] for _ in range(num_chunks)] + cur_chunk = 0 + for index in indices: + chunks[cur_chunk].append(index) + if len(chunks[cur_chunk]) == num_indices_per_chunk: + cur_chunk += 1 + pad_chunks = [] + for idx, chunk in enumerate(chunks): + if batch_size != len(chunk): + assert batch_size > len(chunk) + if len(chunk) != 0: + chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))] + else: + chunk = random.choice(pad_chunks) + print(chunks[idx], '->', chunk) + pad_chunks.append(chunk) + return pad_chunks + +def get_length_grouped_indices(lengths, batch_size, world_size, gradient_accumulation_size, initial_global_step, generator=None, group_data=False, seed=42): + # We need to use torch for the random part as a distributed sampler will set the random seed for torch. + if generator is None: + generator = torch.Generator().manual_seed(seed) # every rank will generate a fixed order but random index + # print('lengths', lengths) + + if group_data: + indices = group_data_fun(lengths, generator) + else: + indices = torch.randperm(len(lengths), generator=generator).tolist() + # print('indices', len(indices)) + + # print('sort indices', len(indices)) + # print('sort indices', indices) + # print('sort lengths', [lengths[i] for i in indices]) + + megabatch_size = world_size * batch_size + megabatches = [indices[i: i + megabatch_size] for i in range(0, len(lengths), megabatch_size)] + # print('megabatches', len(megabatches)) + # print('\nmegabatches', megabatches) + megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches] + # print('sort megabatches', len(megabatches)) + # megabatches_len = [[lengths[i] for i in megabatch] for megabatch in megabatches] + # print('\nsorted megabatches', megabatches) + # print('\nsorted megabatches_len', megabatches_len) + # import ipdb;ipdb.set_trace() + megabatches = [split_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches] + # print('nsplit_to_even_chunks megabatches', len(megabatches)) + # print('\nsplit_to_even_chunks megabatches', megabatches) + # print('\nsplit_to_even_chunks len', [lengths[i] for megabatch in megabatches for batch in megabatch for i in batch]) + # return [i for megabatch in megabatches for batch in megabatch for i in batch] + + indices_mega = torch.randperm(len(megabatches), generator=generator).tolist() + + shuffled_megabatches = [megabatches[i] for i in indices_mega] + # print('shuffled_megabatches', len(shuffled_megabatches)) + if group_data: + shuffled_megabatches = last_group_data_fun(shuffled_megabatches, lengths) + + initial_global_step = initial_global_step * gradient_accumulation_size + print('shuffled_megabatches', len(shuffled_megabatches)) + print('have been trained idx:', len(shuffled_megabatches[:initial_global_step])) + # print('shuffled_megabatches[:10]', shuffled_megabatches[:10]) + # print('have been trained idx:', shuffled_megabatches[:initial_global_step]) + shuffled_megabatches = shuffled_megabatches[initial_global_step:] + print('after shuffled_megabatches', len(shuffled_megabatches)) + # print('after shuffled_megabatches[:10]', shuffled_megabatches[:10]) + + # print('\nshuffled_megabatches', shuffled_megabatches) + # import ipdb;ipdb.set_trace() + # print('\nshuffled_megabatches len', [[i, lengths[i]] for megabatch in shuffled_megabatches for batch in megabatch for i in batch]) + + return [i for megabatch in shuffled_megabatches for batch in megabatch for i in batch] + + +class LengthGroupedSampler(Sampler): + r""" + Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while + keeping a bit of randomness. + """ + + def __init__( + self, + batch_size: int, + world_size: int, + gradient_accumulation_size: int, + initial_global_step: int, + lengths: Optional[List[int]] = None, + group_data=False, + generator=None, + ): + if lengths is None: + raise ValueError("Lengths must be provided.") + + self.batch_size = batch_size + self.world_size = world_size + self.initial_global_step = initial_global_step + self.gradient_accumulation_size = gradient_accumulation_size + self.lengths = lengths + self.group_data = group_data + self.generator = generator + # print('self.lengths, self.initial_global_step, self.batch_size, self.world_size, self.gradient_accumulation_size', + # len(self.lengths), self.initial_global_step, self.batch_size, self.world_size, self.gradient_accumulation_size) + + def __len__(self): + return len(self.lengths) - self.initial_global_step * self.batch_size * self.world_size * self.gradient_accumulation_size + + def __iter__(self): + indices = get_length_grouped_indices(self.lengths, self.batch_size, self.world_size, + self.gradient_accumulation_size, self.initial_global_step, + group_data=self.group_data, generator=self.generator) + # print(len(indices), indices[23640:23690]) + # import sys;sys.exit() + return iter(indices) diff --git a/opensora/utils/ema.py b/opensora/utils/ema.py new file mode 100644 index 000000000..a906efee9 --- /dev/null +++ b/opensora/utils/ema.py @@ -0,0 +1,328 @@ +import contextlib +import copy +import random +from typing import Any, Dict, Iterable, List, Optional, Union + +from diffusers.utils import ( + deprecate, + is_torchvision_available, + is_transformers_available, +) + +if is_transformers_available(): + import transformers + +if is_torchvision_available(): + from torchvision import transforms + +import numpy as np +import torch + + +# Adapted from diffusers-style ema https://github.com/huggingface/diffusers/blob/main/src/diffusers/training_utils.py#L263 +class EMAModel: + """ + Exponential Moving Average of models weights + """ + + def __init__( + self, + parameters: Iterable[torch.nn.Parameter], + decay: float = 0.9999, + min_decay: float = 0.0, + update_after_step: int = 0, + use_ema_warmup: bool = False, + inv_gamma: Union[float, int] = 1.0, + power: Union[float, int] = 2 / 3, + model_cls: Optional[Any] = None, + model_config: Dict[str, Any] = None, + **kwargs, + ): + """ + Args: + parameters (Iterable[torch.nn.Parameter]): The parameters to track. + decay (float): The decay factor for the exponential moving average. + min_decay (float): The minimum decay factor for the exponential moving average. + update_after_step (int): The number of steps to wait before starting to update the EMA weights. + use_ema_warmup (bool): Whether to use EMA warmup. + inv_gamma (float): + Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True. + power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True. + device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA + weights will be stored on CPU. + + @crowsonkb's notes on EMA Warmup: + If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan + to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps), + gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 + at 215.4k steps). + """ + + if isinstance(parameters, torch.nn.Module): + deprecation_message = ( + "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. " + "Please pass the parameters of the module instead." + ) + deprecate( + "passing a `torch.nn.Module` to `ExponentialMovingAverage`", + "1.0.0", + deprecation_message, + standard_warn=False, + ) + parameters = parameters.parameters() + + # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility + use_ema_warmup = True + + if kwargs.get("max_value", None) is not None: + deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead." + deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False) + decay = kwargs["max_value"] + + if kwargs.get("min_value", None) is not None: + deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead." + deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False) + min_decay = kwargs["min_value"] + + parameters = list(parameters) + self.shadow_params = [p.clone().detach() for p in parameters] + + if kwargs.get("device", None) is not None: + deprecation_message = "The `device` argument is deprecated. Please use `to` instead." + deprecate("device", "1.0.0", deprecation_message, standard_warn=False) + self.to(device=kwargs["device"]) + + self.temp_stored_params = None + + self.decay = decay + self.min_decay = min_decay + self.update_after_step = update_after_step + self.use_ema_warmup = use_ema_warmup + self.inv_gamma = inv_gamma + self.power = power + self.optimization_step = 0 + self.cur_decay_value = None # set in `step()` + + self.model_cls = model_cls + self.model_config = model_config + + @classmethod + def extract_ema_kwargs(cls, kwargs): + """ + Extracts the EMA kwargs from the kwargs of a class method. + """ + ema_kwargs = {} + for key in [ + "decay", + "min_decay", + "optimization_step", + "update_after_step", + "use_ema_warmup", + "inv_gamma", + "power", + ]: + if kwargs.get(key, None) is not None: + ema_kwargs[key] = kwargs.pop(key) + return ema_kwargs + + @classmethod + def from_pretrained(cls, path, model_cls) -> "EMAModel": + config = model_cls.load_config(path) + ema_kwargs = cls.extract_ema_kwargs(config) + model = model_cls.from_pretrained(path) + + ema_model = cls(model.parameters(), model_cls=model_cls, model_config=config) + + ema_model.load_state_dict(ema_kwargs) + return ema_model + + def save_pretrained(self, path): + if self.model_cls is None: + raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.") + + if self.model_config is None: + raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.") + + model = self.model_cls.from_config(self.model_config) + state_dict = self.state_dict() + state_dict.pop("shadow_params", None) + + model.register_to_config(**state_dict) + self.copy_to(model.parameters()) + model.save_pretrained(path) + + def get_decay(self, optimization_step: int) -> float: + """ + Compute the decay factor for the exponential moving average. + """ + step = max(0, optimization_step - self.update_after_step - 1) + + if step <= 0: + return 0.0 + + if self.use_ema_warmup: + cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power + else: + cur_decay_value = (1 + step) / (10 + step) + + cur_decay_value = min(cur_decay_value, self.decay) + # make sure decay is not smaller than min_decay + cur_decay_value = max(cur_decay_value, self.min_decay) + return cur_decay_value + + @torch.no_grad() + def step(self, parameters: Iterable[torch.nn.Parameter]): + if isinstance(parameters, torch.nn.Module): + deprecation_message = ( + "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. " + "Please pass the parameters of the module instead." + ) + deprecate( + "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`", + "1.0.0", + deprecation_message, + standard_warn=False, + ) + parameters = parameters.parameters() + + parameters = list(parameters) + + self.optimization_step += 1 + + # Compute the decay factor for the exponential moving average. + decay = self.get_decay(self.optimization_step) + self.cur_decay_value = decay + one_minus_decay = 1 - decay + + context_manager = contextlib.nullcontext + if is_transformers_available() and transformers.deepspeed.is_deepspeed_zero3_enabled(): + import deepspeed + + for s_param, param in zip(self.shadow_params, parameters): + if is_transformers_available() and transformers.deepspeed.is_deepspeed_zero3_enabled(): + context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None) + + with context_manager(): + if param.requires_grad: + s_param.sub_(one_minus_decay * (s_param - param)) + else: + s_param.copy_(param) + + def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None: + """ + Copy current averaged parameters into given collection of parameters. + + Args: + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + updated with the stored moving averages. If `None`, the parameters with which this + `ExponentialMovingAverage` was initialized will be used. + """ + parameters = list(parameters) + for s_param, param in zip(self.shadow_params, parameters): + param.data.copy_(s_param.to(param.device).data) + + + def to(self, device=None, dtype=None) -> None: + r"""Move internal buffers of the ExponentialMovingAverage to `device`. + + Args: + device: like `device` argument to `torch.Tensor.to` + """ + # .to() on the tensors handles None correctly + self.shadow_params = [ + p.to(device=device, dtype=dtype) if p.is_floating_point() else p.to(device=device) + for p in self.shadow_params + ] + + def state_dict(self) -> dict: + r""" + Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during + checkpointing to save the ema state dict. + """ + # Following PyTorch conventions, references to tensors are returned: + # "returns a reference to the state and not its copy!" - + # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict + return { + "decay": self.decay, + "min_decay": self.min_decay, + "optimization_step": self.optimization_step, + "update_after_step": self.update_after_step, + "use_ema_warmup": self.use_ema_warmup, + "inv_gamma": self.inv_gamma, + "power": self.power, + "shadow_params": self.shadow_params, + } + + def store(self, parameters: Iterable[torch.nn.Parameter]) -> None: + r""" + Args: + Save the current parameters for restoring later. + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + temporarily stored. + """ + self.temp_stored_params = [param.detach().cpu().clone() for param in parameters] + + def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None: + r""" + Args: + Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without: + affecting the original optimization process. Store the parameters before the `copy_to()` method. After + validation (or model saving), use this to restore the former parameters. + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + updated with the stored parameters. If `None`, the parameters with which this + `ExponentialMovingAverage` was initialized will be used. + """ + if self.temp_stored_params is None: + raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`") + for c_param, param in zip(self.temp_stored_params, parameters): + param.data.copy_(c_param.data) + + # Better memory-wise. + self.temp_stored_params = None + + def load_state_dict(self, state_dict: dict) -> None: + r""" + Args: + Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the + ema state dict. + state_dict (dict): EMA state. Should be an object returned + from a call to :meth:`state_dict`. + """ + # deepcopy, to be consistent with module API + state_dict = copy.deepcopy(state_dict) + + self.decay = state_dict.get("decay", self.decay) + if self.decay < 0.0 or self.decay > 1.0: + raise ValueError("Decay must be between 0 and 1") + + self.min_decay = state_dict.get("min_decay", self.min_decay) + if not isinstance(self.min_decay, float): + raise ValueError("Invalid min_decay") + + self.optimization_step = state_dict.get("optimization_step", self.optimization_step) + if not isinstance(self.optimization_step, int): + raise ValueError("Invalid optimization_step") + + self.update_after_step = state_dict.get("update_after_step", self.update_after_step) + if not isinstance(self.update_after_step, int): + raise ValueError("Invalid update_after_step") + + self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup) + if not isinstance(self.use_ema_warmup, bool): + raise ValueError("Invalid use_ema_warmup") + + self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma) + if not isinstance(self.inv_gamma, (float, int)): + raise ValueError("Invalid inv_gamma") + + self.power = state_dict.get("power", self.power) + if not isinstance(self.power, (float, int)): + raise ValueError("Invalid power") + + shadow_params = state_dict.get("shadow_params", None) + if shadow_params is not None: + self.shadow_params = shadow_params + if not isinstance(self.shadow_params, list): + raise ValueError("shadow_params must be a list") + if not all(isinstance(p, torch.Tensor) for p in self.shadow_params): + raise ValueError("shadow_params must all be Tensors") diff --git a/opensora/utils/ema_utils.py b/opensora/utils/ema_utils.py new file mode 100644 index 000000000..7e8d2649c --- /dev/null +++ b/opensora/utils/ema_utils.py @@ -0,0 +1,60 @@ + +from peft import get_peft_model, PeftModel +import os +from copy import deepcopy +import torch +import json +from diffusers.training_utils import EMAModel as diffuser_EMAModel + + + +class EMAModel(diffuser_EMAModel): + def __init__(self, parameters, **kwargs): + self.lora_config = kwargs.pop('lora_config', None) + super().__init__(parameters, **kwargs) + + @classmethod + def from_pretrained(cls, path, model_cls, lora_config, model_base) -> "EMAModel": + # 1. load model + if lora_config is not None: + # 1.1 load origin model + model_base = model_cls.from_pretrained(model_base) # model_base + config = model_base.config + # 1.2 convert to lora model automatically and load lora weight + model = PeftModel.from_pretrained(model_base, path) # lora_origin_model + else: + model = model_cls.from_pretrained(path) + config = model.config + # 3. ema the whole model + ema_model = cls(model.parameters(), model_cls=model_cls, model_config=config, lora_config=lora_config) + # 4. load ema_config, e.g decay... + with open(os.path.join(path, 'ema_config.json'), 'r') as f: + state_dict = json.load(f) + ema_model.load_state_dict(state_dict) + return ema_model + + def save_pretrained(self, path): + if self.model_cls is None: + raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.") + + if self.model_config is None: + raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.") + # 1. init a base model randomly + model = self.model_cls.from_config(self.model_config) + # 1.1 convert lora_model + if self.lora_config is not None: + model = get_peft_model(model, self.lora_config) + # 2. ema_model copy to model + self.copy_to(model.parameters()) + # 3. save weight + if self.lora_config is not None: + model.save_pretrained(path) # only lora weight + merge_model = model.merge_and_unload() + merge_model.save_pretrained(path) # merge_model weight + else: + merge_model.save_pretrained(path) # model weight + # 4. save ema_config, e.g decay... + state_dict = self.state_dict() # lora_model weight + state_dict.pop("shadow_params", None) + with open(os.path.join(path, 'ema_config.json'), 'w') as f: + json.dump(state_dict, f, indent=2) \ No newline at end of file diff --git a/opensora/utils/freeinit_utils.py b/opensora/utils/freeinit_utils.py new file mode 100644 index 000000000..a55cc2940 --- /dev/null +++ b/opensora/utils/freeinit_utils.py @@ -0,0 +1,140 @@ +import torch +import torch.fft as fft +import math + + +def freq_mix_3d(x, noise, LPF): + """ + Noise reinitialization. + + Args: + x: diffused latent + noise: randomly sampled noise + LPF: low pass filter + """ + # FFT + x_freq = fft.fftn(x, dim=(-3, -2, -1)) + x_freq = fft.fftshift(x_freq, dim=(-3, -2, -1)) + noise_freq = fft.fftn(noise, dim=(-3, -2, -1)) + noise_freq = fft.fftshift(noise_freq, dim=(-3, -2, -1)) + + # frequency mix + HPF = 1 - LPF + x_freq_low = x_freq * LPF + noise_freq_high = noise_freq * HPF + x_freq_mixed = x_freq_low + noise_freq_high # mix in freq domain + + # IFFT + x_freq_mixed = fft.ifftshift(x_freq_mixed, dim=(-3, -2, -1)) + x_mixed = fft.ifftn(x_freq_mixed, dim=(-3, -2, -1)).real + + return x_mixed + + +def get_freq_filter(shape, device, filter_type, n, d_s, d_t): + """ + Form the frequency filter for noise reinitialization. + + Args: + shape: shape of latent (B, C, T, H, W) + filter_type: type of the freq filter + n: (only for butterworth) order of the filter, larger n ~ ideal, smaller n ~ gaussian + d_s: normalized stop frequency for spatial dimensions (0.0-1.0) + d_t: normalized stop frequency for temporal dimension (0.0-1.0) + """ + if filter_type == "gaussian": + return gaussian_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device) + elif filter_type == "ideal": + return ideal_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device) + elif filter_type == "box": + return box_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device) + elif filter_type == "butterworth": + return butterworth_low_pass_filter(shape=shape, n=n, d_s=d_s, d_t=d_t).to(device) + else: + raise NotImplementedError + +def gaussian_low_pass_filter(shape, d_s=0.25, d_t=0.25): + """ + Compute the gaussian low pass filter mask. + + Args: + shape: shape of the filter (volume) + d_s: normalized stop frequency for spatial dimensions (0.0-1.0) + d_t: normalized stop frequency for temporal dimension (0.0-1.0) + """ + T, H, W = shape[-3], shape[-2], shape[-1] + mask = torch.zeros(shape) + if d_s==0 or d_t==0: + return mask + for t in range(T): + for h in range(H): + for w in range(W): + d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2) + mask[..., t,h,w] = math.exp(-1/(2*d_s**2) * d_square) + return mask + + +def butterworth_low_pass_filter(shape, n=4, d_s=0.25, d_t=0.25): + """ + Compute the butterworth low pass filter mask. + + Args: + shape: shape of the filter (volume) + n: order of the filter, larger n ~ ideal, smaller n ~ gaussian + d_s: normalized stop frequency for spatial dimensions (0.0-1.0) + d_t: normalized stop frequency for temporal dimension (0.0-1.0) + """ + T, H, W = shape[-3], shape[-2], shape[-1] + mask = torch.zeros(shape) + if d_s==0 or d_t==0: + return mask + for t in range(T): + for h in range(H): + for w in range(W): + d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2) + mask[..., t,h,w] = 1 / (1 + (d_square / d_s**2)**n) + return mask + + +def ideal_low_pass_filter(shape, d_s=0.25, d_t=0.25): + """ + Compute the ideal low pass filter mask. + + Args: + shape: shape of the filter (volume) + d_s: normalized stop frequency for spatial dimensions (0.0-1.0) + d_t: normalized stop frequency for temporal dimension (0.0-1.0) + """ + T, H, W = shape[-3], shape[-2], shape[-1] + mask = torch.zeros(shape) + if d_s==0 or d_t==0: + return mask + for t in range(T): + for h in range(H): + for w in range(W): + d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2) + mask[..., t,h,w] = 1 if d_square <= d_s*2 else 0 + return mask + + +def box_low_pass_filter(shape, d_s=0.25, d_t=0.25): + """ + Compute the ideal low pass filter mask (approximated version). + + Args: + shape: shape of the filter (volume) + d_s: normalized stop frequency for spatial dimensions (0.0-1.0) + d_t: normalized stop frequency for temporal dimension (0.0-1.0) + """ + T, H, W = shape[-3], shape[-2], shape[-1] + mask = torch.zeros(shape) + if d_s==0 or d_t==0: + return mask + + threshold_s = round(int(H // 2) * d_s) + threshold_t = round(T // 2 * d_t) + + cframe, crow, ccol = T // 2, H // 2, W //2 + mask[..., cframe - threshold_t:cframe + threshold_t, crow - threshold_s:crow + threshold_s, ccol - threshold_s:ccol + threshold_s] = 1.0 + + return mask \ No newline at end of file diff --git a/opensora/utils/gpu_mem_track.py b/opensora/utils/gpu_mem_track.py new file mode 100644 index 000000000..7b7fb04de --- /dev/null +++ b/opensora/utils/gpu_mem_track.py @@ -0,0 +1,113 @@ +import gc +import datetime +import inspect + +import torch +import numpy as np + +dtype_memory_size_dict = { + torch.float64: 64/8, + torch.double: 64/8, + torch.float32: 32/8, + torch.float: 32/8, + torch.float16: 16/8, + torch.half: 16/8, + torch.int64: 64/8, + torch.long: 64/8, + torch.int32: 32/8, + torch.int: 32/8, + torch.int16: 16/8, + torch.short: 16/6, + torch.uint8: 8/8, + torch.int8: 8/8, +} +# compatibility of torch1.0 +if getattr(torch, "bfloat16", None) is not None: + dtype_memory_size_dict[torch.bfloat16] = 16/8 +if getattr(torch, "bool", None) is not None: + dtype_memory_size_dict[torch.bool] = 8/8 # pytorch use 1 byte for a bool, see https://github.com/pytorch/pytorch/issues/41571 + +def get_mem_space(x): + try: + ret = dtype_memory_size_dict[x] + except KeyError: + print(f"dtype {x} is not supported!") + return ret + +class MemTracker(object): + """ + Class used to track pytorch memory usage + Arguments: + detail(bool, default True): whether the function shows the detail gpu memory usage + path(str): where to save log file + verbose(bool, default False): whether show the trivial exception + device(int): GPU number, default is 0 + """ + def __init__(self, detail=True, path='', verbose=False, device=0): + self.print_detail = detail + self.last_tensor_sizes = set() + self.gpu_profile_fn = path + f'{datetime.datetime.now():%d-%b-%y-%H:%M:%S}-gpu_mem_track.txt' + self.verbose = verbose + self.begin = True + self.device = device + + def get_tensors(self): + for obj in gc.get_objects(): + try: + if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)): + tensor = obj + else: + continue + if tensor.is_cuda: + yield tensor + except Exception as e: + if self.verbose: + print('A trivial exception occured: {}'.format(e)) + + def get_tensor_usage(self): + sizes = [np.prod(np.array(tensor.size())) * get_mem_space(tensor.dtype) for tensor in self.get_tensors()] + return np.sum(sizes) / 1024**2 + + def get_allocate_usage(self): + return torch.cuda.memory_allocated() / 1024**2 + + def clear_cache(self): + gc.collect() + torch.cuda.empty_cache() + + def print_all_gpu_tensor(self, file=None): + for x in self.get_tensors(): + print(x.size(), x.dtype, np.prod(np.array(x.size()))*get_mem_space(x.dtype)/1024**2, file=file) + + def track(self): + """ + Track the GPU memory usage + """ + frameinfo = inspect.stack()[1] + where_str = frameinfo.filename + ' line ' + str(frameinfo.lineno) + ': ' + frameinfo.function + + with open(self.gpu_profile_fn, 'a+') as f: + + if self.begin: + f.write(f"GPU Memory Track | {datetime.datetime.now():%d-%b-%y-%H:%M:%S} |" + f" Total Tensor Used Memory:{self.get_tensor_usage():<7.1f}Mb" + f" Total Allocated Memory:{self.get_allocate_usage():<7.1f}Mb\n\n") + self.begin = False + + if self.print_detail is True: + ts_list = [(tensor.size(), tensor.dtype) for tensor in self.get_tensors()] + new_tensor_sizes = {(type(x), + tuple(x.size()), + ts_list.count((x.size(), x.dtype)), + np.prod(np.array(x.size()))*get_mem_space(x.dtype)/1024**2, + x.dtype) for x in self.get_tensors()} + for t, s, n, m, data_type in new_tensor_sizes - self.last_tensor_sizes: + f.write(f'+ | {str(n)} * Size:{str(s):<20} | Memory: {str(m*n)[:6]} M | {str(t):<20} | {data_type}\n') + for t, s, n, m, data_type in self.last_tensor_sizes - new_tensor_sizes: + f.write(f'- | {str(n)} * Size:{str(s):<20} | Memory: {str(m*n)[:6]} M | {str(t):<20} | {data_type}\n') + + self.last_tensor_sizes = new_tensor_sizes + + f.write(f"\nAt {where_str:<50}" + f" Total Tensor Used Memory:{self.get_tensor_usage():<7.1f}Mb" + f" Total Allocated Memory:{self.get_allocate_usage():<7.1f}Mb\n\n") \ No newline at end of file diff --git a/opensora/utils/lora_utils.py b/opensora/utils/lora_utils.py new file mode 100644 index 000000000..e4159e9fd --- /dev/null +++ b/opensora/utils/lora_utils.py @@ -0,0 +1,44 @@ + +from peft import get_peft_model, PeftModel +import os +from copy import deepcopy +import torch +import json + +def maybe_zero_3(param, ignore_status=False, name=None): + from deepspeed import zero + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + if hasattr(param, "ds_id"): + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if not ignore_status: + logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}") + with zero.GatheredParameters([param]): + param = param.data.detach().cpu().clone() + else: + param = param.detach().cpu().clone() + return param + +# Borrowed from peft.utils.get_peft_model_state_dict +def get_peft_state_maybe_zero_3(named_params, bias): + if bias == "none": + to_return = {k: t for k, t in named_params if "lora_" in k} + elif bias == "all": + to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k} + elif bias == "lora_only": + to_return = {} + maybe_lora_bias = {} + lora_bias_names = set() + for k, t in named_params: + if "lora_" in k: + to_return[k] = t + bias_name = k.split("lora_")[0] + "bias" + lora_bias_names.add(bias_name) + elif "bias" in k: + maybe_lora_bias[k] = t + for k, t in maybe_lora_bias: + if bias_name in lora_bias_names: + to_return[bias_name] = t + else: + raise NotImplementedError + to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()} + return to_return diff --git a/opensora/utils/mask_utils.py b/opensora/utils/mask_utils.py new file mode 100644 index 000000000..e61467ed1 --- /dev/null +++ b/opensora/utils/mask_utils.py @@ -0,0 +1,548 @@ +import random +from opensora.dataset.inpaint_utils import get_mask_tensor,MaskType +import torch +from opensora.dataset.transform import ToTensorVideo +from opensora.models.causalvideovae import ae_norm +from torchvision import transforms +import os +from PIL import Image +import numpy as np +import cv2 +from enum import Enum, auto +from ultralytics import YOLO + + +import torch_npu +from torch_npu.contrib import transfer_to_npu + + +os.environ['YOLO_VERBOSS'] = 'False' + +class MaskType(Enum): + Semantic_mask = auto() + bbox_mask = auto() + background_mask = auto() + fixed_mask = auto() + Semantic_expansion_mask = auto() + fixed_bg_mask = auto() + t2iv_mask = auto() + i2v_mask = auto() + transition_mask = auto() + v2v_mask = auto() + clear_mask = auto() + random_mask = auto() + + + + +class single_info: + def __init__(self, id, label, shape) -> None: + self.id = id + self.label = label + self.shape = shape + self.frame_indexes = [] + self.infos = [] + def update(self,frame_index,box,conf,mask): + self.frame_indexes.append(frame_index) + info = dict( + box=box, + conf=conf, + mask=mask, + ) + self.infos.append(info) + def return_dict(self,): + return dict( + id=self.id, + label=self.label, + frame_size=self.shape, + frame_index_list = self.frame_indexes, + infos_list = self.infos + ) + +def save_videos_from_pil(pil_images, path, fps=24): + """ + pil_images: list[Image,...] + """ + os.makedirs(os.path.dirname(path), exist_ok=True) + + image = pil_images[0] + + image = ndarray_to_pil(pil_images[0]) + width, height = image.size + + + codec = "libx264" + container = av.open(path, "w") + stream = container.add_stream(codec, rate=fps) + + stream.width = width + stream.height = height + + for pil_image in pil_images: + # pil_image = Image.fromarray(image_arr).convert("RGB") + pil_image = ndarray_to_pil(pil_image) + av_frame = av.VideoFrame.from_image(pil_image) + container.mux(stream.encode(av_frame)) + container.mux(stream.encode()) + container.close() + +def read_frames(video_tensor) -> list: + """ + 读取视频，返回一个元素类型为ndarray的列表 + """ + # container = av.open(video_path) + T = video_tensor.shape[0] + frames = [] + for t in range(T): + frame_tensor = video_tensor[t] + frame_tensor = frame_tensor.cpu().numpy() + frame_tensor = np.transpose(frame_tensor, (1, 2, 0)) + frames.append(frame_tensor) + return frames + + +def get_masked_image(image: np.ndarray, mask: np.ndarray) -> np.ndarray: + mask = mask.astype(bool) + if len(mask.shape) == 2: + mask = np.expand_dims(mask, axis=2) + masked_img = image * (1-mask) + return masked_img # shape: [H,W,C]; range: [0, 255] + +def get_bbox_image(image: np.ndarray,bbox,obj_id): + # cv2.rectangle(image, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2) + bbox_image = image.copy() + bbox_image[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])] = 0 + # cv2.putText(image, f'ID: {obj_id}', (int(bbox[0]), int(bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 0, 0), 2) + return bbox_image + + + +def select_bg_from_video(bg_masks, video): + new_container = [] + for index, frame in enumerate(video): + + mask = bg_masks[index] + masked_frame = get_masked_image(frame, mask) + new_container.append(masked_frame) + return new_container + +def get_random_box(image_tensor): + + H, W,C = image_tensor.shape + + box_min_size = min(H,W)/2 + box_max_size = max(H,W)/2 + + # 随机确定 box 的宽高 + box_width = random.randint(box_min_size, min(box_max_size, W)) + box_height = random.randint(box_min_size, min(box_max_size, H)) + + # 随机确定 box 的左上角坐标 + x_start = random.randint(0, W - box_width) + y_start = random.randint(0, H - box_height) + + box = (x_start, y_start, x_start + box_width, y_start + box_height) + + return box + +def combine_masks_and_get_background(masks): + """ + 合并所有 mask 并取反得到背景 mask + """ + combined_mask = np.any(masks, axis=0) + background_mask = np.logical_not(combined_mask) + return background_mask + +def parser_results_for_ids(results, frame_size=None): + id_record = [] + single_info_ins = {} + background_masks = [] + for frame_index, result in enumerate(results): + result = result[0] + if frame_index == 0 and frame_size is None: + frame_size = result.boxes.orig_shape + id = result.boxes.id + + # 如果没有检测到物体 + if id is None: + background_masks.append(np.ones((frame_size)) * 255) + continue + + id = id.tolist() + cls = result.boxes.cls.tolist() #每个id对应的label + conf = result.boxes.conf.tolist() #每个id对应的预测置信度 + box_n = result.boxes.xyxy.tolist() #每个id对应的box + mask = result.masks.data.cpu().detach().numpy() #每个id对应的mask + background_masks.append(combine_masks_and_get_background(mask)) + + for i, iden in enumerate(id): + if iden not in id_record: + id_record.append(iden) + single_info_ins[iden] = single_info(iden, cls[i], frame_size) + single_info_ins[iden].update(frame_index, box_n[i], conf[i],mask[i]) + return_list = [] + for _, value in single_info_ins.items(): + return_list.append(value.return_dict()) + return return_list, background_masks + + +def get_mask(video_tensor,mask_type,yole_model): + + video = read_frames(video_tensor=video_tensor) + + # video_tensor_batch = video_tensor.unsqueeze(1) + T,C,H,W = video_tensor.shape + + # video = video_tensor + + tracker = yole_model + + results = [] + + + for t in range(T): + frame_tensor = video_tensor[t] # 获取当前帧, (C, H, W) + frame_tensor = frame_tensor.data.cpu().numpy() # 转为numpy + frame_tensor = np.transpose(frame_tensor, (1, 2, 0)) + + # 进行推理 + result = tracker.track(frame_tensor,save=False, retina_masks=True, agnostic_nms=True,half=True,verbose=False,nms=False) + + # 保存结果 + results.append(result) + + + parser_res, background_masks = parser_results_for_ids(results) + + select_index = -1 + object_info = [] + frame_indexes = [] + infos = [] + + + #随机选择一个被追踪物体 + if len(parser_res) is not 0: + select_index = random.randint(0, len(parser_res)-1) + object_info = parser_res[select_index] + frame_indexes = object_info['frame_index_list'] + infos = object_info['infos_list'] + # print("infos size",len(infos)) + # print("frame_indexed",len(frame_indexes)) + else: + mask_type = MaskType.fixed_mask + + + + # print("frame_indexed:",frame_indexes) + # print("infos:",infos) + + # mask_type = MaskType.Semantic_mask + + if mask_type == MaskType.Semantic_mask or mask_type == MaskType.Semantic_expansion_mask: + Semantic_masks = [] + mask_container = [] + info_index = 0 + for index, frame in enumerate(video): + if index in frame_indexes: + + mask = infos[info_index]['mask'] + info_index = info_index + 1 + + if mask_type == MaskType.Semantic_expansion_mask: + kernel = np.ones((5, 5), np.uint8) + # 进行膨胀操作 + mask = cv2.dilate(mask, kernel, iterations=1) + + # 计算掩码中前景像素的数量 + foreground_pixels = np.sum(mask) + + # 计算图像的总像素数 + total_pixels = mask.size # 或者使用 image.shape[0] * image.shape[1] + + # 计算比例 + ratio = foreground_pixels / total_pixels + + if ratio < 0.2: + if random.random() < 0.5: + mask_type = MaskType.fixed_mask + break + + masked_frame = get_masked_image(frame, mask) + mask_container.append(masked_frame) + Semantic_masks.append(mask) + else: + mask_container.append(np.zeros_like(frame)) + Semantic_masks.append(np.zeros_like(frame)[:,:,0]) + if mask_type == MaskType.Semantic_mask or mask_type == MaskType.Semantic_expansion_mask: + return mask_container, Semantic_masks + + if mask_type == MaskType.bbox_mask: + boxes_masks = [] + box_container = [] + + info_index = 0 + + for index, frame in enumerate(video): + if index in frame_indexes: + bbox = infos[info_index]['box'] + info_index = info_index + 1 + + boxed_frame = get_bbox_image(frame, bbox, object_info['id']) + box_container.append(boxed_frame) + boxmask = np.zeros_like(frame)[:,:,0] + boxmask[int(bbox[1]): int(bbox[3]), int(bbox[0]): int(bbox[2])] = 1 + boxes_masks.append(boxmask) + else: + box_container.append(frame) + boxes_masks.append(np.zeros_like(frame)[:,:,0]) + + return box_container, boxes_masks + + if mask_type == MaskType.background_mask: + bg_container = select_bg_from_video(background_masks, video) + return bg_container, background_masks + + if mask_type == MaskType.fixed_mask or mask_type == MaskType.fixed_bg_mask: + fixed_mask_container = [] + fixed_masks = [] + + box = get_random_box(video[0]) + for index , frame in enumerate(video): + if mask_type == MaskType.fixed_mask: + boxed_frame = frame.copy() + boxed_frame[int(box[1]): int(box[3]), int(box[0]): int(box[2])] = 0 + fixed_mask_container.append(boxed_frame) + + fixed_mask = np.zeros_like(frame)[:,:,0] + fixed_mask[int(box[1]): int(box[3]), int(box[0]): int(box[2])] = 1 + fixed_masks.append(fixed_mask) + if mask_type == MaskType.fixed_bg_mask: + boxed_frame = frame.copy() + + fixed_mask = np.zeros_like(frame)[:,:,0] + fixed_mask[int(box[1]): int(box[3]), int(box[0]): int(box[2])] = 1 + fixed_mask = 1 - fixed_mask + fixed_masks.append(fixed_mask) + + boxed_bg_frame = get_masked_image(boxed_frame, fixed_mask) + fixed_mask_container.append(boxed_bg_frame) + + return fixed_mask_container, fixed_masks + + + +def video_to_tensor(video_path): + # 打开视频文件 + cap = cv2.VideoCapture(video_path) + + frames = [] + + while True: + ret, frame = cap.read() + if not ret: + break + + # 将 BGR 转换为 RGB + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # 转换为张量并添加到帧列表中 + frame_tensor = torch.from_numpy(frame).permute(2, 0, 1) # (H, W, C) -> (C, H, W) + frames.append(frame_tensor) + + cap.release() + + # 将所有帧组合成一个四维张量 + video_tensor = torch.stack(frames) # (T, C, H, W) + + return video_tensor + + +def ndarray_to_pil(image: np.ndarray) -> Image: + if np.max(image) <= 1.1: + image = image * 255 + image = image.astype(np.uint8) + return Image.fromarray(image) + +def get_random_type(): + # 数字列表 + mask_type = [MaskType.Semantic_mask, MaskType.bbox_mask, MaskType.background_mask, MaskType.fixed_mask, MaskType.Semantic_expansion_mask, MaskType.fixed_bg_mask] + + # 概率权重列表（总和应为 1 或任意正数比例） + weights = [0.3, 0.2, 0.1, 0.1, 0.2, 0.1] # 例如，第一个数字1的概率是0.1 + + # 从1-6中按设定概率随机选择一个数字 + chosen_number = random.choices(mask_type, weights=weights)[0] + + return chosen_number + +def get_mask_tensor(video_tensor,mask_type,yolomodel): + + # return video_tensor,video_tensor + + masked_video_container,masks_container = get_mask(video_tensor,mask_type,yolomodel) + + masked_frames = [torch.from_numpy(frame.transpose(2,0,1)) for frame in masked_video_container] + masked_video = torch.stack(masked_frames) + + masks = [torch.from_numpy(mask.reshape(1,mask.shape[0],mask.shape[1])) for mask in masks_container] + # import ipdb; ipdb.set_trace() + mask = torch.stack(masks) + + return masked_video,mask + + + +class MaskProcessor: + def __init__(self,args,YOLOmodel): + # ratio + # transform + self.num_frames = args.num_frames + if self.num_frames != 1: + # inpaint + self.t2v_ratio = args.t2v_ratio + self.i2v_ratio = args.i2v_ratio + self.transition_ratio = args.transition_ratio + self.v2v_ratio = args.v2v_ratio + self.clear_video_ratio = args.clear_video_ratio + self.Semantic_ratio = args.Semantic_ratio + self.bbox_ratio = args.bbox_ratio + self.background_ratio = args.background_ratio + self.fixed_ratio = args.fixed_ratio + self.Semantic_expansion_ratio = args.Semantic_expansion_ratio + self.fixed_bg_ratio = args.fixed_bg_ratio + assert self.t2v_ratio + self.i2v_ratio + self.transition_ratio + self.v2v_ratio + self.clear_video_ratio + self.Semantic_ratio + self.bbox_ratio + self.background_ratio + self.fixed_ratio + self.fixed_bg_ratio + self.Semantic_expansion_ratio < 1, 'The sum of t2v_ratio, i2v_ratio, transition_ratio, v2v_ratio and clear video ratio should be less than 1.' + + self.min_clear_ratio = 0.0 if args.min_clear_ratio is None else args.min_clear_ratio + assert self.min_clear_ratio >= 0 and self.min_clear_ratio <= 1, 'min_clear_ratio should be in the range of [0, 1].' + + + self.transform = transforms.Compose([ + ToTensorVideo(), + ae_norm[args.ae] + ]) + + self.init_mask_func() + self.init_ratio() + + self.default_text_ratio = args.default_text_ratio + + self.yolomodel = YOLOmodel + + def init_mask_func(self): + # mask: ones_like (t 1 h w) + def t2iv(mask): + mask[:] = 1 + return mask + + def i2v(mask): + mask[0] = 0 + return mask + + def transition(mask): + mask[0] = 0 + mask[-1] = 0 + return mask + + def v2v(mask): + end_idx = random.randint(int(mask.shape[0] * self.min_clear_ratio), mask.shape[0]) + mask[:end_idx] = 0 + return mask + + def clear(mask): + mask[:] = 0 + return mask + + def random_mask(mask): + num_to_select = random.randint(int(mask.shape[0] * self.min_clear_ratio), mask.shape[0]) + selected_indices = random.sample(range(mask.shape[0]), num_to_select) + mask[selected_indices] = 0 + return mask + + def Semantic_mask(video_tensor): + return get_mask_tensor(video_tensor,MaskType.Semantic_mask,self.yolomodel) + + def bbox_mask(video_tensor): + return get_mask_tensor(video_tensor,MaskType.bbox_mask,self.yolomodel) + + def background_mask(video_tensor): + return get_mask_tensor(video_tensor,MaskType.background_mask,self.yolomodel) + + def fixed_mask(video_tensor): + return get_mask_tensor(video_tensor,MaskType.fixed_mask,self.yolomodel) + + def Semantic_expansion_mask(video_tensor): + return get_mask_tensor(video_tensor,MaskType.Semantic_expansion_mask,self.yolomodel) + + def fixed_bg_mask(video_tensor): + return get_mask_tensor(video_tensor,MaskType.fixed_bg_mask,self.yolomodel) + + + + self.mask_functions = { + MaskType.t2iv_mask: t2iv, + MaskType.i2v_mask: i2v, + MaskType.transition_mask: transition, + MaskType.v2v_mask: v2v, + MaskType.clear_mask: clear, + MaskType.random_mask: random_mask, + MaskType.Semantic_mask:Semantic_mask, + MaskType.bbox_mask:bbox_mask, + MaskType.background_mask:background_mask, + MaskType.fixed_mask:fixed_mask, + MaskType.Semantic_expansion_mask:Semantic_expansion_mask, + MaskType.fixed_bg_mask:fixed_bg_mask + } + + def init_ratio(self): + + self.mask_func_weights_video = { + MaskType.t2iv_mask: self.t2v_ratio, + MaskType.i2v_mask: self.i2v_ratio, + MaskType.transition_mask: self.transition_ratio, + MaskType.v2v_mask: self.v2v_ratio, + MaskType.clear_mask: self.clear_video_ratio, + MaskType.Semantic_mask:self.Semantic_ratio, + MaskType.bbox_mask:self.bbox_ratio, + MaskType.background_mask:self.background_ratio, + MaskType.fixed_mask:self.fixed_ratio, + MaskType.Semantic_expansion_mask:self.Semantic_expansion_ratio, + MaskType.fixed_bg_mask:self.fixed_bg_ratio, + MaskType.random_mask: 1 - self.t2v_ratio - self.i2v_ratio - self.transition_ratio - self.v2v_ratio - self.clear_video_ratio - self.Semantic_ratio - self.bbox_ratio - self.background_ratio - self.fixed_ratio - self.Semantic_expansion_ratio - self.fixed_bg_ratio + + } + + self.mask_func_weights_image = { + 't2iv': 0.9, + 'clear': 0.1 + } + + # t c h w + def __call__(self,pixel_values): + # pixel_values shape (T, C, H, W) + # 1 means masked, 0 means not masked + t, c, h, w = pixel_values.shape + mask = torch.ones([t, 1, h, w], device=pixel_values.device, dtype=pixel_values.dtype) + + mask_func_name = random.choices(list(self.mask_func_weights_video.keys()), list(self.mask_func_weights_video.values()))[0] + frame_mask_list = [MaskType.t2iv_mask,MaskType.i2v_mask,MaskType.transition_mask,MaskType.v2v_mask,MaskType.clear_mask,MaskType.random_mask] + pos_mask_list = [MaskType.Semantic_mask,MaskType.bbox_mask,MaskType.background_mask,MaskType.fixed_mask,MaskType.Semantic_expansion_mask,MaskType.fixed_bg_mask] + + + if mask_func_name in frame_mask_list: + mask = self.mask_functions[mask_func_name](mask) + masked_pixel_values = pixel_values * (mask < 0.5) + + if mask_func_name in pos_mask_list: + masked_pixel_values,mask = self.mask_functions[mask_func_name](pixel_values) + # save_video(masked_pixel_values.permute(0, 2, 3, 1).cpu().numpy(), 'masked_video.mp4') + + # import ipdb; ipdb.set_trace() + + pixel_values = self.transform(pixel_values) + masked_pixel_values = self.transform(masked_pixel_values.to(torch.uint8)) + + + + return masked_pixel_values, pixel_values, mask + + diff --git a/opensora/utils/parallel_states.py b/opensora/utils/parallel_states.py new file mode 100644 index 000000000..def564f71 --- /dev/null +++ b/opensora/utils/parallel_states.py @@ -0,0 +1,44 @@ +import torch +import torch.distributed as dist +import os + +class COMM_INFO: + def __init__(self): + self.group = None + self.world_size = 0 + self.rank = -1 + +nccl_info = COMM_INFO() +_SEQUENCE_PARALLEL_STATE = False +def initialize_sequence_parallel_state(sequence_parallel_size): + global _SEQUENCE_PARALLEL_STATE + if sequence_parallel_size > 1: + _SEQUENCE_PARALLEL_STATE = True + initialize_sequence_parallel_group(sequence_parallel_size) + +def set_sequence_parallel_state(state): + global _SEQUENCE_PARALLEL_STATE + _SEQUENCE_PARALLEL_STATE = state + +def get_sequence_parallel_state(): + return _SEQUENCE_PARALLEL_STATE + +def initialize_sequence_parallel_group(sequence_parallel_size): + """Initialize the sequence parallel group.""" + rank = int(os.getenv('RANK', '0')) + world_size = int(os.getenv("WORLD_SIZE", '1')) + assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size" + # hccl + nccl_info.world_size = sequence_parallel_size + nccl_info.rank = rank + num_sequence_parallel_groups: int = world_size // sequence_parallel_size + for i in range(num_sequence_parallel_groups): + ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + nccl_info.group = group + + +def destroy_sequence_parallel_group(): + """Destroy the sequence parallel group.""" + dist.destroy_process_group() diff --git a/opensora/utils/utils.py b/opensora/utils/utils.py index ba94b6aa1..601b79d7c 100644 --- a/opensora/utils/utils.py +++ b/opensora/utils/utils.py @@ -15,6 +15,7 @@ from torch import inf from PIL import Image from typing import Union, Iterable +import collections from collections import OrderedDict from torch.utils.tensorboard import SummaryWriter @@ -32,6 +33,11 @@ _tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]] +def to_2tuple(x): + if isinstance(x, collections.abc.Iterable): + return x + return (x, x) + def find_model(model_name): """ Finds a pre-trained Latte model, downloading it if necessary. Alternatively, loads a model from a local path. @@ -297,8 +303,9 @@ def save_video_grid(video, nrow=None): video_grid = torch.zeros((t, (padding + h) * nrow + padding, (padding + w) * ncol + padding, c), dtype=torch.uint8) - print(video_grid.shape) + print('save_video_grid video_grid.shape', video_grid.shape) for i in range(b): + print('i', i) r = i // ncol c = i % ncol start_r = (padding + h) * r @@ -336,10 +343,9 @@ def collect_env(): bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}') # noqa -def text_preprocessing(text): +def text_preprocessing(text, support_Chinese=True): # The exact text cleaning as was in the training stage: - text = clean_caption(text) - text = clean_caption(text) + text = clean_caption(text, support_Chinese=support_Chinese) return text def basic_clean(text): @@ -347,7 +353,7 @@ def basic_clean(text): text = html.unescape(html.unescape(text)) return text.strip() -def clean_caption(caption): +def clean_caption(caption, support_Chinese=True): caption = str(caption) caption = ul.unquote_plus(caption) caption = caption.strip().lower() @@ -378,7 +384,8 @@ def clean_caption(caption): caption = re.sub(r'[\u3300-\u33ff]+', '', caption) caption = re.sub(r'[\u3400-\u4dbf]+', '', caption) caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption) - caption = re.sub(r'[\u4e00-\u9fff]+', '', caption) + if not support_Chinese: + caption = re.sub(r'[\u4e00-\u9fff]+', '', caption) # Chinese ####################################################### # все виды тире / all types of dash --> "-" @@ -455,5 +462,10 @@ def clean_caption(caption): return caption.strip() - +if __name__ == '__main__': + + # caption = re.sub(r'[\u4e00-\u9fff]+', '', caption) + a = "امرأة مسنة بشعر أبيض ووجه مليء بالتجاعيد تجلس داخل سيارة قديمة الطراز، تنظر من خلال النافذة الجانبية بتعبير تأملي أو حزين قليلاً." + print(a) + print(text_preprocessing(a)) diff --git a/pyproject.toml b/pyproject.toml index 5a3624fd0..431e0c32e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "opensora" -version = "1.0.0" +version = "1.2.0" description = "Reproduce OpenAI's Sora." readme = "README.md" requires-python = ">=3.8" @@ -14,14 +14,14 @@ classifiers = [ ] dependencies = [ "torch==2.1.0", "torchvision==0.16.0", - "transformers==4.39.1", "accelerate==0.28.0", + "transformers==4.40.1", "accelerate==0.29.3", "tokenizers==0.19.1", "diffusers==0.28.0", "albumentations==1.4.0", "av==11.0.0", "decord==0.6.0", "einops==0.7.0", "fastapi==0.110.0", "gdown==5.1.0", "h5py==3.10.0", "idna==3.6", 'imageio==2.34.0', "matplotlib==3.7.5", "numpy==1.24.4", "omegaconf==2.1.1", "opencv-python==4.9.0.80", "opencv-python-headless==4.9.0.80", "pandas==2.0.3", "pillow==10.2.0", "pydub==0.25.1", "pytorch-lightning==2.2.1", "pytorchvideo==0.1.5", "PyYAML==6.0.1", "regex==2023.12.25", "requests==2.31.0", "scikit-learn==1.3.2", "scipy==1.10.1", "six==1.16.0", "test-tube==0.7.5", "timm==0.9.16", "torchdiffeq==0.2.3", "torchmetrics==1.3.2", "tqdm==4.66.2", "urllib3==2.2.1", "uvicorn==0.27.1", - "diffusers==0.27.2", "scikit-video==1.1.11", "imageio-ffmpeg==0.4.9", "sentencepiece==0.1.99", "beautifulsoup4==4.12.3", + "scikit-video==1.1.11", "imageio-ffmpeg==0.4.9", "sentencepiece==0.1.99", "beautifulsoup4==4.12.3", "ftfy==6.1.3", "moviepy==1.0.3", "wandb==0.16.3", "tensorboard==2.14.0", "pydantic==2.6.4", "gradio==4.0.0", "xformers==0.0.22.post7" ] diff --git a/scripts/accelerate_configs/deepspeed_zero2_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_config.yaml index 43ec1e80c..b15390113 100644 --- a/scripts/accelerate_configs/deepspeed_zero2_config.yaml +++ b/scripts/accelerate_configs/deepspeed_zero2_config.yaml @@ -5,9 +5,9 @@ deepspeed_config: fsdp_config: {} machine_rank: 0 main_process_ip: null -main_process_port: 29501 +main_process_port: 29503 main_training_function: main num_machines: 1 num_processes: 8 gpu_ids: 0,1,2,3,4,5,6,7 -use_cpu: false \ No newline at end of file +use_cpu: false diff --git a/scripts/accelerate_configs/hostfile b/scripts/accelerate_configs/hostfile index a22693099..901cae636 100644 --- a/scripts/accelerate_configs/hostfile +++ b/scripts/accelerate_configs/hostfile @@ -1,2 +1,4 @@ -gpu55 slots=8 # your server name and GPU in total -gpu117 slots=8 +node032 slots=8 +node034 slots=8 +node035 slots=8 +node037 slots=8 diff --git a/scripts/accelerate_configs/hostfile1 b/scripts/accelerate_configs/hostfile1 new file mode 100644 index 000000000..2940421cf --- /dev/null +++ b/scripts/accelerate_configs/hostfile1 @@ -0,0 +1,2 @@ +node014 slots=8 +node026 slots=8 diff --git a/scripts/accelerate_configs/hostfile2 b/scripts/accelerate_configs/hostfile2 new file mode 100644 index 000000000..6af909735 --- /dev/null +++ b/scripts/accelerate_configs/hostfile2 @@ -0,0 +1,2 @@ +node109 slots=8 +node114 slots=8 diff --git a/scripts/accelerate_configs/hostfile40 b/scripts/accelerate_configs/hostfile40 new file mode 100644 index 000000000..d098b9db3 --- /dev/null +++ b/scripts/accelerate_configs/hostfile40 @@ -0,0 +1,40 @@ +node30 slots=8 +node004 slots=8 +node005 slots=8 +node006 slots=8 +node007 slots=8 +node008 slots=8 +node010 slots=8 +node011 slots=8 +node012 slots=8 +node013 slots=8 +node019 slots=8 +node026 slots=8 +node027 slots=8 +node028 slots=8 +node029 slots=8 +node031 slots=8 +node032 slots=8 +node033 slots=8 +node034 slots=8 +node035 slots=8 +node037 slots=8 +node056 slots=8 +node059 slots=8 +node060 slots=8 +node061 slots=8 +node062 slots=8 +node064 slots=8 +node065 slots=8 +node072 slots=8 +node074 slots=8 +node076 slots=8 +node077 slots=8 +node068 slots=8 +node063 slots=8 +node075 slots=8 +node003 slots=8 +node058 slots=8 +node109 slots=8 +node066 slots=8 +node114 slots=8 \ No newline at end of file diff --git a/scripts/accelerate_configs/multi_node_example.yaml b/scripts/accelerate_configs/multi_node_example.yaml index 4e3cac79c..a8a65f012 100644 --- a/scripts/accelerate_configs/multi_node_example.yaml +++ b/scripts/accelerate_configs/multi_node_example.yaml @@ -2,14 +2,14 @@ compute_environment: LOCAL_MACHINE distributed_type: DEEPSPEED deepspeed_config: deepspeed_config_file: scripts/accelerate_configs/zero2.json - deepspeed_hostfile: /remote-home1/yeyang/Open-Sora-Plan/scripts/accelerate_configs/hostfile + deepspeed_hostfile: scripts/accelerate_configs/hostfile fsdp_config: {} machine_rank: 0 -main_process_ip: 10.10.10.55 -main_process_port: 29501 +main_process_ip: 100.64.24.32 +main_process_port: 29522 main_training_function: main -num_machines: 2 -num_processes: 16 +num_machines: 32 +num_processes: 256 rdzv_backend: static same_network: true tpu_env: [] diff --git a/scripts/accelerate_configs/multi_node_example1.yaml b/scripts/accelerate_configs/multi_node_example1.yaml new file mode 100644 index 000000000..2288c4f30 --- /dev/null +++ b/scripts/accelerate_configs/multi_node_example1.yaml @@ -0,0 +1,18 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero2.json + deepspeed_hostfile: scripts/accelerate_configs/hostfile1 +fsdp_config: {} +machine_rank: 0 +main_process_ip: 100.64.24.14 +main_process_port: 29522 +main_training_function: main +num_machines: 2 +num_processes: 16 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/scripts/accelerate_configs/multi_node_example2.yaml b/scripts/accelerate_configs/multi_node_example2.yaml new file mode 100644 index 000000000..4dd53c8ea --- /dev/null +++ b/scripts/accelerate_configs/multi_node_example2.yaml @@ -0,0 +1,18 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero2.json + deepspeed_hostfile: scripts/accelerate_configs/hostfile2 +fsdp_config: {} +machine_rank: 0 +main_process_ip: 100.64.24.109 +main_process_port: 29502 +main_training_function: main +num_machines: 4 +num_processes: 32 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/scripts/accelerate_configs/multi_node_example_by_ddp.yaml b/scripts/accelerate_configs/multi_node_example_by_ddp.yaml new file mode 100644 index 000000000..fa16c29a6 --- /dev/null +++ b/scripts/accelerate_configs/multi_node_example_by_ddp.yaml @@ -0,0 +1,13 @@ +compute_environment: LOCAL_MACHINE +distributed_type: MULTI_GPU +fsdp_config: {} +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false \ No newline at end of file diff --git a/scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml b/scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml new file mode 100644 index 000000000..9c26fd8c5 --- /dev/null +++ b/scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml @@ -0,0 +1,16 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero2_npu.json + deepspeed_multinode_launcher: standard +fsdp_config: {} +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false \ No newline at end of file diff --git a/scripts/accelerate_configs/zero2.json b/scripts/accelerate_configs/zero2.json index bc26a0a5f..8dcdd154d 100644 --- a/scripts/accelerate_configs/zero2.json +++ b/scripts/accelerate_configs/zero2.json @@ -1,6 +1,6 @@ { "fp16": { - "enabled": "auto", + "enabled": false, "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, @@ -10,6 +10,8 @@ "bf16": { "enabled": "auto" }, + "communication_data_type": "fp32", + "gradient_clipping": 1.0, "train_micro_batch_size_per_gpu": "auto", "train_batch_size": "auto", "gradient_accumulation_steps": "auto", diff --git a/scripts/accelerate_configs/zero2_npu.json b/scripts/accelerate_configs/zero2_npu.json new file mode 100644 index 000000000..6f1ecdf2c --- /dev/null +++ b/scripts/accelerate_configs/zero2_npu.json @@ -0,0 +1,25 @@ +{ + "fp16": { + "enabled": false, + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "communication_data_type": "fp32", + "gradient_clipping": 1.0, + "train_micro_batch_size_per_gpu": "auto", + "train_batch_size": "auto", + "gradient_accumulation_steps": "auto", + "zero_optimization": { + "stage": 2, + "overlap_comm": true, + "allgather_bucket_size": 536870912, + "contiguous_gradients": true, + "reduce_bucket_size": 536870912 + } +} \ No newline at end of file diff --git a/scripts/causalvae/eval.sh b/scripts/causalvae/eval.sh index ed8bd3eda..0b7b909cc 100644 --- a/scripts/causalvae/eval.sh +++ b/scripts/causalvae/eval.sh @@ -1,12 +1,21 @@ -python opensora/eval/eval_common_metric.py \ - --batch_size 2 \ - --real_video_dir ..//test_eval/release/origin \ - --generated_video_dir ../test_eval/release \ - --device cuda \ - --sample_fps 10 \ - --crop_size 256 \ - --resolution 256 \ - --num_frames 17 \ - --sample_rate 1 \ - --subset_size 100 \ - --metric ssim \ No newline at end of file +# REAL_DATASET_DIR=/remote-home1/dataset/OpenMMLab___Kinetics-400/raw/Kinetics-400/videos_val/ +REAL_DATASET_DIR=../dataset/webvid/videos +EXP_NAME=decoder +SAMPLE_RATE=3 +NUM_FRAMES=33 +RESOLUTION=256 +SUBSET_SIZE=50 +METRIC=ssim + +python opensora/models/causalvideovae/eval/eval_common_metric.py \ + --batch_size 1 \ + --real_video_dir ${REAL_DATASET_DIR} \ + --generated_video_dir /remote-home1/lzj/dataset/gen/${EXP_NAME}_sr${SAMPLE_RATE}_nf${NUM_FRAMES}_res${RESOLUTION}_subset${SUBSET_SIZE} \ + --device cuda:0 \ + --sample_fps 3 \ + --sample_rate ${SAMPLE_RATE} \ + --num_frames ${NUM_FRAMES} \ + --resolution ${RESOLUTION} \ + --subset_size ${SUBSET_SIZE} \ + --crop_size ${RESOLUTION} \ + --metric ${METRIC} \ No newline at end of file diff --git a/scripts/causalvae/gen_video.sh b/scripts/causalvae/gen_video.sh deleted file mode 100644 index 2f340d62e..000000000 --- a/scripts/causalvae/gen_video.sh +++ /dev/null @@ -1,13 +0,0 @@ -python examples/rec_video_vae.py \ - --batch_size 1 \ - --real_video_dir ../test_eval/eyes_test \ - --generated_video_dir ../test_eval/eyes_gen \ - --device cuda \ - --sample_fps 10 \ - --sample_rate 1 \ - --num_frames 17 \ - --resolution 512 \ - --crop_size 512 \ - --num_workers 8 \ - --ckpt results/pretrained_488 \ - --enable_tiling \ No newline at end of file diff --git a/scripts/causalvae/prepare_eval.sh b/scripts/causalvae/prepare_eval.sh new file mode 100644 index 000000000..8ce0b5ba8 --- /dev/null +++ b/scripts/causalvae/prepare_eval.sh @@ -0,0 +1,24 @@ +export CUDA_VISIBLE_DEVICES=0 +REAL_DATASET_DIR=valid/ +EXP_NAME=decoder +SAMPLE_RATE=1 +NUM_FRAMES=33 +RESOLUTION=512 +SUBSET_SIZE=1 +CKPT=/storage/lcm/Causal-Video-VAE/results/488dim8 + +python opensora/models/causalvideovae/sample/rec_video_vae.py \ + --batch_size 1 \ + --real_video_dir ${REAL_DATASET_DIR} \ + --generated_video_dir valid_gen/${EXP_NAME}_sr${SAMPLE_RATE}_nf${NUM_FRAMES}_res${RESOLUTION}_subset${SUBSET_SIZE} \ + --device cuda \ + --sample_fps 24 \ + --sample_rate ${SAMPLE_RATE} \ + --num_frames ${NUM_FRAMES} \ + --resolution ${RESOLUTION} \ + --subset_size ${SUBSET_SIZE} \ + --crop_size ${RESOLUTION} \ + --num_workers 8 \ + --ckpt ${CKPT} \ + --output_origin \ + --enable_tiling \ No newline at end of file diff --git a/scripts/causalvae/rec_image.sh b/scripts/causalvae/rec_image.sh new file mode 100644 index 000000000..3c65c54b3 --- /dev/null +++ b/scripts/causalvae/rec_image.sh @@ -0,0 +1,8 @@ +CUDA_VISIBLE_DEVICES=0 python examples/rec_image.py \ + --ae_path "/storage/dataset/test140k" \ + --image_path /storage/dataset/image/anytext3m/ocr_data/Art/images/gt_5544.jpg \ + --rec_path rec.jpg \ + --device cuda \ + --short_size 512 \ + --ae CausalVAEModel_4x8x8 \ + --enable_tiling \ No newline at end of file diff --git a/scripts/causalvae/rec_video.sh b/scripts/causalvae/rec_video.sh new file mode 100644 index 000000000..a317fc3bb --- /dev/null +++ b/scripts/causalvae/rec_video.sh @@ -0,0 +1,12 @@ +CUDA_VISIBLE_DEVICES=1 python examples/rec_video.py \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --video_path /storage/dataset/mixkit-train-passing-the-rails-4462_resize1080p.mp4 \ + --rec_path rec_tile.mp4 \ + --device cpu \ + --sample_rate 1 \ + --num_frames 17 \ + --height 256 \ + --width 256 \ + --fps 30 \ + --enable_tiling \ No newline at end of file diff --git a/scripts/causalvae/reconstruction.sh b/scripts/causalvae/reconstruction.sh deleted file mode 100644 index 3e89d2c28..000000000 --- a/scripts/causalvae/reconstruction.sh +++ /dev/null @@ -1,11 +0,0 @@ -CUDA_VISIBLE_DEVICES=0 python examples/rec_imvi_vae.py \ - --ae_path "../../Open-Sora-Plan/CausalVAEModel_4x8x8/" \ - --video_path origin_tiger.mp4 \ - --rec_path rec.mp4 \ - --device cuda \ - --sample_rate 1 \ - --num_frames 65 \ - --resolution 512 \ - --crop_size 512 \ - --ae CausalVAEModel_4x8x8 \ - --enable_tiling \ No newline at end of file diff --git a/scripts/causalvae/release.json b/scripts/causalvae/release.json deleted file mode 100644 index f572f3207..000000000 --- a/scripts/causalvae/release.json +++ /dev/null @@ -1,71 +0,0 @@ -{ - "_class_name": "CausalVAEModel", - "_diffusers_version": "0.27.2", - "attn_resolutions": [], - "decoder_attention": "AttnBlock3D", - "decoder_conv_in": "CausalConv3d", - "decoder_conv_out": "CausalConv3d", - "decoder_mid_resnet": "ResnetBlock3D", - "decoder_resnet_blocks": [ - "ResnetBlock3D", - "ResnetBlock3D", - "ResnetBlock3D", - "ResnetBlock3D" - ], - "decoder_spatial_upsample": [ - "", - "SpatialUpsample2x", - "SpatialUpsample2x", - "SpatialUpsample2x" - ], - "decoder_temporal_upsample": [ - "", - "", - "TimeUpsample2x", - "TimeUpsample2x" - ], - "double_z": true, - "dropout": 0.0, - "embed_dim": 4, - "encoder_attention": "AttnBlock3D", - "encoder_conv_in": "CausalConv3d", - "encoder_conv_out": "CausalConv3d", - "encoder_mid_resnet": "ResnetBlock3D", - "encoder_resnet_blocks": [ - "ResnetBlock3D", - "ResnetBlock3D", - "ResnetBlock3D", - "ResnetBlock3D" - ], - "encoder_spatial_downsample": [ - "SpatialDownsample2x", - "SpatialDownsample2x", - "SpatialDownsample2x", - "" - ], - "encoder_temporal_downsample": [ - "TimeDownsample2x", - "TimeDownsample2x", - "", - "" - ], - "hidden_size": 128, - "hidden_size_mult": [ - 1, - 2, - 4, - 4 - ], - "loss_params": { - "disc_start": 2001, - "disc_weight": 0.5, - "kl_weight": 1e-06, - "logvar_init": 0.0 - }, - "loss_type": "opensora.models.ae.videobase.losses.LPIPSWithDiscriminator", - "lr": 1e-05, - "num_res_blocks": 2, - "q_conv": "CausalConv3d", - "resolution": 256, - "z_channels": 4 -} diff --git a/scripts/causalvae/train.sh b/scripts/causalvae/train.sh index 1a49bf50f..d97ba141f 100644 --- a/scripts/causalvae/train.sh +++ b/scripts/causalvae/train.sh @@ -1,15 +1,45 @@ -python opensora/train/train_causalvae.py \ - --exp_name "exp_name" \ +# export https_proxy=http://127.0.0.1:8998 +# export http_proxy=http://127.0.0.1:8998 +unset https_proxy +unset http_proxy +export WANDB_PROJECT=causalvideovae_2.0 +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 + +EXP_NAME=latent4_resume_2 + +torchrun \ + --nnodes=1 --nproc_per_node=8 \ + --master_addr=localhost \ + --master_port=29600 \ + scripts/causalvae/train_causalvae.py \ + --exp_name ${EXP_NAME} \ + --model_config scripts/config.json \ + --video_path /storage/dataset/pexels/ \ + --eval_video_path /storage/dataset/pixabay_v2/folder_01 \ + --resolution 320 \ + --epochs 1000 \ + --num_frames 25 \ --batch_size 1 \ - --precision bf16 \ - --max_steps 40000 \ - --save_steps 100 \ - --output_dir results/causalvae_ \ - --video_path /remote-home1/dataset/data_split_tt \ - --video_num_frames 17 \ - --resolution 256 \ - --sample_rate 1 \ - --n_nodes 1 \ - --devices 1 \ - --num_workers 8 \ - --load_from_checkpoint ./results/pretrained_488/ \ No newline at end of file + --disc_start 2000 \ + --save_ckpt_step 2000 \ + --eval_steps 500 \ + --eval_batch_size 1 \ + --eval_num_frames 33 \ + --eval_sample_rate 3 \ + --eval_subset_size 50 \ + --eval_lpips \ + --ema \ + --ema_decay 0.999 \ + --perceptual_weight 1.0 \ + --loss_type l1 \ + --disc_cls causalvideovae.model.losses.LPIPSWithDiscriminator3D \ + --not_resume_training_process \ + --resume_from_checkpoint /storage/lcm/Causal-Video-VAE/results/latent8_3d-lr1.00e-05-bs1-rs320-sr2-fr25/checkpoint-14000.ckpt \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_image.sh b/scripts/text_condition/gpu/sample_image.sh new file mode 100644 index 000000000..22ff081f4 --- /dev/null +++ b/scripts/text_condition/gpu/sample_image.sh @@ -0,0 +1,21 @@ +export TASK_QUEUE_ENABLE=0 + +CUDA_VISIBLE_DEVICES=1 python opensora/sample/sample_t2v.py \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x8_vae8_anyx320x320_lr5e-5_snr5_noioff0.02_ema9999_sparse1d4_newdit_l_122_rope_mt5xxl_mj/checkpoint-218000/model_ema \ + --version 65x512x512 \ + --num_frames 1 \ + --height 320 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt examples/prompt_list_2.txt \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "WFVAE_DISTILL_FORMAL" \ + --save_img_path "sample_image_vae8_newdit_218k_320x320_test_sam" \ + --fps 24 \ + --guidance_scale 4.5 \ + --num_sampling_steps 28 \ + --enable_tiling \ + --max_sequence_length 512 \ + --sample_method DPMSolverMultistep \ + --model_type sparsedit \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_image_sparse.sh b/scripts/text_condition/gpu/sample_image_sparse.sh new file mode 100644 index 000000000..9c2bb39b8 --- /dev/null +++ b/scripts/text_condition/gpu/sample_image_sparse.sh @@ -0,0 +1,19 @@ +CUDA_VISIBLE_DEVICES=3 python opensora/sample/sample_t2v.py \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs1x8x16_1x720p_lr1e-4_snr5_noioff0.02_ema9999_sparse1d16_dit_l_122_rope_mt5xxl_mj_fromhw480p/checkpoint-4000/model_ema \ + --version 65x512x512 \ + --num_frames 1 \ + --height 720 \ + --width 1280 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt examples/prompt_list_1.txt \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --save_img_path "sample_images_sparse1d16_dit_9k" \ + --fps 24 \ + --guidance_scale 4.5 \ + --num_sampling_steps 20 \ + --enable_tiling \ + --max_sequence_length 512 \ + --sample_method DPMSolverMultistep \ + --model_type dit \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_inpaint_ddp.sh b/scripts/text_condition/gpu/sample_inpaint_ddp.sh new file mode 100644 index 000000000..c8aaadb1d --- /dev/null +++ b/scripts/text_condition/gpu/sample_inpaint_ddp.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --nnodes=1 --nproc_per_node 2 --master_port 29502 \ + -m opensora.sample.sample_inpaint_ddp \ + --model_path /storage/gyy/hw/Open-Sora-Plan/test_sparse_inpaint/checkpoint-84000/model \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/gyy/hw/Open-Sora-Plan/test_prompt.txt \ + --conditional_images_path /storage/gyy/hw/Open-Sora-Plan/test_cond_imgs_path.txt \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "./sample_test_inpaint_sparse" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 1234 \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_t2v.sh b/scripts/text_condition/gpu/sample_t2v.sh new file mode 100644 index 000000000..47b425897 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v.sh @@ -0,0 +1,21 @@ +CUDA_VISIBLE_DEVICES=5 python opensora/sample/sample_t2v.py \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr5e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-20000/model \ + --version 65x512x512 \ + --num_frames 1 \ + --height 320 \ + --width 160 \ + --cache_dir "cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt examples/prompt_list_0.txt \ + --ae CausalVAEModel_D8_4x8x8 \ + --ae_path "/storage/dataset/new488dim8/last" \ + --save_img_path "./sample_video_dit_vae8_newmodel_anyx93x320x160_img" \ + --fps 24 \ + --guidance_scale 5.0 \ + --num_sampling_steps 50 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 1.0 \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_t2v_ddp.sh b/scripts/text_condition/gpu/sample_t2v_ddp.sh new file mode 100644 index 000000000..7d1fac080 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp.sh @@ -0,0 +1,26 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=3,4,6,7 torchrun --nnodes=1 --nproc_per_node 4 --master_port 29502 \ + -m opensora.sample.sample_t2v_ddp \ + --model_path bs32x8x2_anyx93x320x320_fps16_lr2e-6_snr5_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m_vpred_zerosnr/checkpoint-49000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/new/7.19anyres/Open-Sora-Plan/examples/sora_refine.txt \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "./sample_video_dit_vae8_newmodel_anyx93x160x320_sora_m0.9_49k_s100_wf_xformer_refine_vpred_zerosnr_trailing_lowlr" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 1234 \ + --num_samples_per_prompt 1 \ + --prediction_type "v_prediction" \ + --rescale_betas_zero_snr \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_t2v_ddp1.sh b/scripts/text_condition/gpu/sample_t2v_ddp1.sh new file mode 100644 index 000000000..2ebff1f17 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp1.sh @@ -0,0 +1,26 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --nnodes=1 --nproc_per_node 8 --master_port 29502 \ + -m opensora.sample.sample_t2v_ddp \ + --model_path bs32x8x2_anyx93x320x320_fps16_lr1e-5_snr5_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m_vpred_zerosnr/checkpoint-3000/model \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/new/7.19anyres/Open-Sora-Plan/examples/sora_refine.txt \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "./sample_video_dit_vae8_newmodel_anyx93x160x320_sora_m0.9_3k_s100_wf_xformer_refine_vpred_zerosnr_trailing" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 1234 \ + --num_samples_per_prompt 1 \ + --prediction_type "v_prediction" \ + --rescale_betas_zero_snr \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy0.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy0.sh new file mode 100644 index 000000000..e820621d4 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy0.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=0 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29502 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/overall_consistency.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/overall_consistency_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy1.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy1.sh new file mode 100644 index 000000000..b15f935a2 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy1.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29503 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/multiple_objects.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/multiple_objects_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy2.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy2.sh new file mode 100644 index 000000000..7fc4e499c --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy2.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=2 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29504 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/object_class.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/object_class_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy3.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy3.sh new file mode 100644 index 000000000..f1f5c7975 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy3.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=3 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29505 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/human_action.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/human_action_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy4.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy4.sh new file mode 100644 index 000000000..868da3592 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy4.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=4 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29506 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/appearance_style.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/appearance_style_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy5.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy5.sh new file mode 100644 index 000000000..eecb258b1 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy5.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=5 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29507 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/spatial_relationship.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/spatial_relationship_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy6.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy6.sh new file mode 100644 index 000000000..e0096dcda --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy6.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=6 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29508 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/temporal_style.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/temporal_style_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_ddp_yy7.sh b/scripts/text_condition/gpu/sample_t2v_ddp_yy7.sh new file mode 100644 index 000000000..b986b3fc3 --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_ddp_yy7.sh @@ -0,0 +1,23 @@ +# --save_img_path "./sample_video_26500ema_61x480p_k333_s122_cfg5.0_step50" \ +# CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=7 torchrun --nnodes=1 --nproc_per_node 1 --master_port 29509 \ + -m opensora.sample.sample_t2v_ddp_vbench_gpt \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-415000/model_ema \ + --version 65x512x512 \ + --num_frames 93 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt /storage/ongoing/refine_model/prompts_per_dimension/csv_gpt/scene.csv \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/node-user/scene_gpt" \ + --fps 18 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "sparsedit" \ + --motion_score 0.9 \ + --seed 42 43 44 45 46 diff --git a/scripts/text_condition/gpu/sample_t2v_sp.sh b/scripts/text_condition/gpu/sample_t2v_sp.sh new file mode 100644 index 000000000..2ffc0ee5c --- /dev/null +++ b/scripts/text_condition/gpu/sample_t2v_sp.sh @@ -0,0 +1,22 @@ + +torchrun --nnodes=1 --nproc_per_node 8 --master_port 29503 \ + -m opensora.sample.sample_t2v_sp \ + --model_path /storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs8x8x1_29x480p_lr1e-5_snr5_noioff0.02_ema999_sparse1d4_dit_l_122_rope_t5xxl_movie_panda \ + --version 65x512x512 \ + --num_frames 29 \ + --height 480 \ + --width 640 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt examples/prompt_list_2.txt \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --save_img_path "./sample_video_test" \ + --fps 24 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "dit" \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_v2v.sh b/scripts/text_condition/gpu/sample_v2v.sh new file mode 100644 index 000000000..1a5d37a21 --- /dev/null +++ b/scripts/text_condition/gpu/sample_v2v.sh @@ -0,0 +1,30 @@ +export MASTER_ADDR='localhost' +export TOKENIZERS_PARALLELISM=false +MODEL_PATH=/storage/gyy/hw/Open-Sora-Plan/runs/inpaint_only_480p_f93_bs4x8x1_lr1e-5_snrgamma5_0_noiseoffset0_02_ema0_999/checkpoint-14400/model_ema +# export HF_DATASETS_OFFLINE=1 +# export TRANSFORMERS_OFFLINE=1 + +torchrun --nproc_per_node=8 --master_port=29501 opensora/sample/sample_v2v.py \ + --model_path $MODEL_PATH \ + --model_type 'inpaint_only' \ + --num_frames 93 \ + --height 480 \ + --width 640 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "/storage/dataset/test140k" \ + --save_img_path "./samples/md" \ + --fps 24 \ + --guidance_scale 5.0 \ + --num_sampling_steps 50 \ + --enable_tiling \ + --max_sequence_length 512 \ + --sample_method PNDM \ + --validation_dir "./md_validation" \ + --pretrained_transformer_model_path /storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/bs16x8x1_93x480p_lr1e-4_snr5_ema999_opensora122_rope_mt5xxl_high_pandamovie_speed1.0/checkpoint-3500/model_ema \ + --pretrained_vipnet_path /storage/gyy/hw/Open-Sora-Plan/runs/videoip_3d_480p_f29_bs2x16_lr1e-5_snrgamma5_0_noiseoffset0_02_dino518_ema0_999/checkpoint-14000/model \ + --image_encoder_name vit_giant_patch14_reg4_dinov2.lvd142m \ + --image_encoder_path /storage/cache_dir/hub/models--timm--vit_giant_patch14_reg4_dinov2.lvd142m/snapshots/a2208b21b069f6b2e45999870fcce4b7e43d1a2c/model.safetensors \ + --max_sample_num 24 \ + --seed 42 \ No newline at end of file diff --git a/scripts/text_condition/gpu/sample_video_ddp.sh b/scripts/text_condition/gpu/sample_video_ddp.sh new file mode 100644 index 000000000..a6d82f217 --- /dev/null +++ b/scripts/text_condition/gpu/sample_video_ddp.sh @@ -0,0 +1,21 @@ + +torchrun --nnodes=1 --nproc_per_node 8 --master_port 29503 \ + -m opensora.sample.sample_t2v_ddp \ + --model_path /path/to/checkpoint-xxx/model_ema \ + --num_frames 93 \ + --height 720 \ + --width 1280 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt examples/prompt_list_0.txt \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/path/to/causalvideovae" \ + --save_img_path "./sample_93x480p_cfg7.5_step100_eulera" \ + --fps 24 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --max_sequence_length 512 \ + --sample_method EulerAncestralDiscrete \ + --model_type "dit" \ No newline at end of file diff --git a/scripts/text_condition/gpu/train_inpaint_all_in_one.sh b/scripts/text_condition/gpu/train_inpaint_all_in_one.sh new file mode 100644 index 000000000..b13bb626d --- /dev/null +++ b/scripts/text_condition/gpu/train_inpaint_all_in_one.sh @@ -0,0 +1,91 @@ +# PROJECT="video_test" +PROJECT="inpaint_3d_720p_f93_bs16x8x1_lr1e-5_snrgamma5_0_noiseoffset0_02_dino518_ema0_999" +export WANDB_API_KEY="720d886d8c437c2142c88056a1eab8ef78d64a1f" +export WANDB_MODE="online" +export ENTITY="yunyangge" +export PROJECT=$PROJECT +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export TOKENIZERS_PARALLELISM=false +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 + +export PDSH_RCMD_TYPE=ssh + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_inpaint_all_in_one.py \ + --model OpenSoraInpaint-ROPE-L/122 \ + --text_encoder_name google/mt5-xxl \ + --image_encoder_name vit_giant_patch14_reg4_dinov2.lvd142m \ + --image_encoder_path /storage/cache_dir/hub/models--timm--vit_giant_patch14_reg4_dinov2.lvd142m/snapshots/a2208b21b069f6b2e45999870fcce4b7e43d1a2c/model.safetensors \ + --cache_dir "/storage/cache_dir" \ + --dataset inpaint \ + --model_type inpaint_only \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "/storage/dataset/test140k" \ + --data "scripts/train_data/video_data.txt" \ + --sample_rate 1 \ + --num_frames 93 \ + --use_image_num 0 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --gradient_checkpointing \ + --max_train_steps=500000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --enable_tracker \ + --checkpointing_steps=200 \ + --output_dir runs/$PROJECT \ + --allow_tf32 \ + --model_max_length 512 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --validation_dir "validation_dir" \ + --guidance_scale 5.0 \ + --num_sampling_steps 50 \ + --ema_start_step 0 \ + --use_ema \ + --cfg 0.1 \ + --i2v_ratio 0.5 \ + --transition_ratio 0.4 \ + --v2v_ratio 0.1 \ + --clear_video_ratio 0.0 \ + --default_text_ratio 0.5 \ + --seed 42 \ + --snr_gamma 5.0 \ + --noise_offset 0.02 \ + --vip_num_attention_heads 16 \ + --ema_decay 0.999 \ + --use_rope \ + --speed_factor 1.5 \ + --group_frame \ + --pretrained_transformer_model_path "/storage/ongoing/new/Open-Sora-Plan/bs32x8x1_93x720p_lr2e-5_snr5_ema999_opensora122_rope_fp32_mt5xxl_sucai_aes5.5_speed1.5/checkpoint-400/model_ema" \ + --pretrained_vip_adapter_path "/storage/gyy/hw/Open-Sora-Plan/runs/videoip_3d_480p_f93_bs1x16_lr1e-5_snrgamma5_0_noiseoffset0_02_dino518_ema0_999/checkpoint-3000/model" \ + # --sp_size 8 \ + # --train_sp_batch_size 2 \ + # --train_vip \ + # --need_validation \ + # --resume_from_checkpoint "latest" \ + # --zero_terminal_snr \ + # 基模型权重没有参与训练所以一定要加载 \ No newline at end of file diff --git a/scripts/text_condition/gpu/train_inpaint_only.sh b/scripts/text_condition/gpu/train_inpaint_only.sh new file mode 100644 index 000000000..4dfe8ca58 --- /dev/null +++ b/scripts/text_condition/gpu/train_inpaint_only.sh @@ -0,0 +1,88 @@ +# PROJECT="video_test" +PROJECT="inpaint_only_480p_f93_bs4x8x1_lr1e-5_snrgamma5_0_noiseoffset0_02_ema0_999" +export WANDB_API_KEY="720d886d8c437c2142c88056a1eab8ef78d64a1f" +export WANDB_MODE="online" +export ENTITY="yunyangge" +export PROJECT=$PROJECT +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export TOKENIZERS_PARALLELISM=false +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 + +export PDSH_RCMD_TYPE=ssh + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_inpaint_all_in_one.py \ + --model OpenSoraInpaint-ROPE-L/122 \ + --text_encoder_name google/mt5-xxl \ + --image_encoder_name vit_giant_patch14_reg4_dinov2.lvd142m \ + --image_encoder_path /storage/cache_dir/hub/models--timm--vit_giant_patch14_reg4_dinov2.lvd142m/snapshots/a2208b21b069f6b2e45999870fcce4b7e43d1a2c/model.safetensors \ + --cache_dir "/storage/cache_dir" \ + --dataset inpaint \ + --model_type inpaint_only \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "/storage/dataset/test140k" \ + --data "scripts/train_data/video_data_sucai.txt" \ + --sample_rate 1 \ + --num_frames 93 \ + --use_image_num 0 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --gradient_checkpointing \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --enable_tracker \ + --checkpointing_steps=200 \ + --output_dir runs/$PROJECT \ + --allow_tf32 \ + --model_max_length 512 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --validation_dir "validation_dir" \ + --guidance_scale 5.0 \ + --num_sampling_steps 50 \ + --ema_start_step 0 \ + --use_ema \ + --cfg 0.05 \ + --i2v_ratio 0.4 \ + --transition_ratio 0.4 \ + --v2v_ratio 0.1 \ + --clear_video_ratio 0.0 \ + --default_text_ratio 0.5 \ + --seed 42 \ + --snr_gamma 5.0 \ + --noise_offset 0.02 \ + --ema_decay 0.999 \ + --use_rope \ + --group_frame \ + --resume_from_checkpoint "latest" \ + # --pretrained_transformer_model_path "/storage/gyy/hw/Open-Sora-Plan/runs/inpaint_only_480p_f93_bs4x8x1_lr1e-5_snrgamma5_0_noiseoffset0_02_ema0_999_old_script/checkpoint-25800/model_ema" \ + # --pretrained_vip_adapter_path "/storage/gyy/hw/Open-Sora-Plan/pretrained_models/pretrained_vip_9000.pth" + # --speed_factor 1.5 \ + # --vip_num_attention_heads 16 \ + # --train_vip \ + # --zero_terminal_snr \ + # 基模型权重没有参与训练所以一定要加载 \ No newline at end of file diff --git a/scripts/text_condition/gpu/train_inpaint_sparse1d_newmodel_motion.sh b/scripts/text_condition/gpu/train_inpaint_sparse1d_newmodel_motion.sh new file mode 100644 index 000000000..9a650d736 --- /dev/null +++ b/scripts/text_condition/gpu/train_inpaint_sparse1d_newmodel_motion.sh @@ -0,0 +1,80 @@ +export WANDB_KEY="720d886d8c437c2142c88056a1eab8ef78d64a1f" +export WANDB_MODE="online" +export ENTITY="yunyang" +export PROJECT="test_sparse_inpaint" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=25 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export NCCL_IB_RETRY_CNT=32 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_inpaint.py \ + --model OpenSoraInpaint-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset inpaint \ + --data "scripts/train_data/merge_data_debug.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 320 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --use_motion \ + --train_fps 16 \ + --seed 1234 \ + --group_data \ + --t2v_ratio 0.1 \ + --i2v_ratio 0.0 \ + --transition_ratio 0.0 \ + --v2v_ratio 0.0 \ + --clear_video_ratio 0.0 \ + --min_clear_ratio 0.5 \ + --default_text_ratio 0.5 \ + --pretrained_transformer_model_path "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-526000/model_ema" \ + --output_dir="test_sparse_inpaint" > training_log_new.txt + # --resume_from_checkpoint="latest" \ diff --git a/scripts/text_condition/gpu/train_t2v.sh b/scripts/text_condition/gpu/train_t2v.sh new file mode 100644 index 000000000..c719c9eb7 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v.sh @@ -0,0 +1,70 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_panda_movie_mj.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 29 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 1.0 \ + --force_resolution \ + --pretrained "/storage/dataset/hw29/image/model_ema/diffusion_pytorch_model.safetensors" \ + --output_dir="bs16x8x1_img_29x480p_lr2e-5_snr5_noioff0.02_ema9999_dit_l_122_rope_mt5xxl_movie_panda_mj" diff --git a/scripts/text_condition/gpu/train_t2v_img.sh b/scripts/text_condition/gpu/train_t2v_img.sh new file mode 100644 index 000000000..880f95148 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_img.sh @@ -0,0 +1,70 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D8_4x8x8 \ + --ae_path "/storage/dataset/new488dim8/last" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 320 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=8 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --enable_tracker \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs16x8x8_vae8_any320x320_lr1e-4_snr5_noioff0.02_ema9999_dit_l_122_rope_mt5xxl_mj/checkpoint-150000/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --output_dir="bs16x8x8_vae8_any320x320_lr2e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_122_rope_mt5xxl_mj" diff --git a/scripts/text_condition/gpu/train_t2v_img_lora.sh b/scripts/text_condition/gpu/train_t2v_img_lora.sh new file mode 100644 index 000000000..ccfb4d717 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_img_lora.sh @@ -0,0 +1,70 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_t2v_diffusers_lora.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 720 \ + --max_width 1280 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.5 \ + --interpolation_scale_w 2.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --enable_tracker \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --force_resolution \ + --model_base "/storage/dataset/hw29/image/model_ema" \ + --output_dir="bs1x8x16_1x720p_lr1e-5_snr5_noioff0.02_ema9999_lora64_dit_l_122_rope_mt5xxl_mj_fromhw480p" \ + --enable_lora --rank 64 diff --git a/scripts/text_condition/gpu/train_t2v_img_sparse1d.sh b/scripts/text_condition/gpu/train_t2v_img_sparse1d.sh new file mode 100644 index 000000000..ad6b7a6e2 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_img_sparse1d.sh @@ -0,0 +1,71 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example2.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D8_4x8x8 \ + --ae_path "/storage/dataset/new488dim8/last" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --enable_tracker \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs16x8x8_vae8_any320x320_lr2e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_122_rope_mt5xxl_mj/checkpoint-84000/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --force_resolution \ + --sparse1d --sparse_n 4 \ + --output_dir="bs4x8x16_vae8_1x480p_lr2e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_122_rope_mt5xxl_mj" diff --git a/scripts/text_condition/gpu/train_t2v_img_sparse1d_newmodel.sh b/scripts/text_condition/gpu/train_t2v_img_sparse1d_newmodel.sh new file mode 100644 index 000000000..7f5a48abf --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_img_sparse1d_newmodel.sh @@ -0,0 +1,70 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V1-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/image_data.txt" \ + --ae CausalVAEModel_D8_4x8x8 \ + --ae_path "/storage/dataset/new488dim8/last" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 320 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=4 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=5e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --enable_tracker \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs16x8x8_vae8_anyx320x320_lr2e-5_snr5_noioff0.02_ema9999_sparse1d4_newdit_l_122_rope_mt5xxl_mj/checkpoint-57000/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --output_dir="bs32x8x8_vae8_anyx320x320_lr5e-5_snr5_noioff0.02_ema9999_sparse1d4_newdit_l_122_rope_mt5xxl_mj" diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d.sh b/scripts/text_condition/gpu/train_t2v_sparse1d.sh new file mode 100644 index 000000000..091d32ccf --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d.sh @@ -0,0 +1,74 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example1.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_panda_movie_mj.txt" \ + --ae CausalVAEModel_D8_4x8x8 \ + --ae_path "/storage/dataset/new488dim8/last" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 176 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=500 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --force_resolution \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs16x8x8_vae8_any320x320_lr1e-4_snr5_noioff0.02_ema9999_dit_l_122_rope_mt5xxl_mj/checkpoint-150000/model_ema/diffusion_pytorch_model.safetensors" \ + --output_dir="bs8x8x1_93x176x320_fps16_lr2e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_sucaiaes5_fromimgsparse1d4" \ + --tile_sample_min_size 512 \ + --tile_sample_min_size_t 125 \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 --train_fps 16 diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d16_image.sh b/scripts/text_condition/gpu/train_t2v_sparse1d16_image.sh new file mode 100644 index 000000000..19f6e77dd --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d16_image.sh @@ -0,0 +1,73 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 720 \ + --max_width 1280 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.5 \ + --interpolation_scale_w 2.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --sparse1d \ + --sparse_n 16 \ + --force_resolution \ + --enable_tracker \ + --pretrained "/storage/dataset/hw29/image/model_ema/diffusion_pytorch_model.safetensors" \ + --output_dir="bs1x8x16_1x720p_lr1e-4_snr5_noioff0.02_ema9999_sparse1d16_dit_l_122_rope_mt5xxl_mj_fromhw480p" diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d_debug.sh b/scripts/text_condition/gpu/train_t2v_sparse1d_debug.sh new file mode 100644 index 000000000..1bb41112c --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d_debug.sh @@ -0,0 +1,74 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V1-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/video_data_debug.txt" \ + --ae CausalVAEModel_D8_4x8x8 \ + --ae_path "/storage/dataset/new488dim8/last" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 720 \ + --max_width 1280 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=500 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 1.0 \ + --force_resolution \ + --output_dir="debug" \ + --tile_sample_min_size 512 \ + --tile_sample_min_size_t 33 \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 --train_fps 16 \ + --train_sp_batch_size 4 --sp_size 8 diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d_image.sh b/scripts/text_condition/gpu/train_t2v_sparse1d_image.sh new file mode 100644 index 000000000..63a84c467 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d_image.sh @@ -0,0 +1,73 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example1.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=8 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --sparse1d \ + --sparse_n 4 \ + --force_resolution \ + --enable_tracker \ + --pretrained "/storage/dataset/hw29/image/model_ema/diffusion_pytorch_model.safetensors" \ + --output_dir="bs8x8x8_1x480p_lr2e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_122_rope_mt5xxl_mj_fromhw480p" diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion.sh b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion.sh new file mode 100644 index 000000000..ca5883318 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion.sh @@ -0,0 +1,75 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=25 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export NCCL_IB_RETRY_CNT=32 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V2-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 320 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr5e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-118000/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --use_motion \ + --train_fps 16 \ + --seed 1234 \ + --trained_data_global_step 42000 \ + --group_data \ + --use_decord \ + --output_dir="bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m" > training_log_new.txt diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_.sh b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_.sh new file mode 100644 index 000000000..b3639f5fe --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_.sh @@ -0,0 +1,74 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=25 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export NCCL_IB_RETRY_CNT=32 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V2-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_debug.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.5 \ + --interpolation_scale_w 2.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --force_resolution \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr5e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-118000/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --use_motion \ + --train_fps 16 \ + --seed 1234 \ + --trained_data_global_step 0 \ + --output_dir="debug" diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_480p.sh b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_480p.sh new file mode 100644 index 000000000..b1fc43c84 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_480p.sh @@ -0,0 +1,76 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=25 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export NCCL_IB_RETRY_CNT=32 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V2-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 640 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=500 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/checkpoint-526000/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --use_motion \ + --train_fps 16 \ + --seed 1234 \ + --trained_data_global_step 0 \ + --group_data \ + --use_decord \ + --force_resolution \ + --output_dir="bs32x8x1_anyx93x640x640_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_aesimg18m" > training_log_new.txt diff --git a/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_v_pred.sh b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_v_pred.sh new file mode 100644 index 000000000..f2679686d --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse1d_newmodel_motion_v_pred.sh @@ -0,0 +1,77 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=25 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export NCCL_IB_RETRY_CNT=32 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V2-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/storage/lcm/Causal-Video-VAE/results/WFVAE_DISTILL_FORMAL" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 320 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 8 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.0 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --pretrained "/storage/ongoing/new/7.19anyres/Open-Sora-Plan/bs32x8x2_anyx93x320x320_fps16_lr2e-6_snr5_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m_vpred_zerosnr/checkpoint-45100/model_ema/diffusion_pytorch_model.safetensors" \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --use_motion \ + --train_fps 16 \ + --seed 1234 \ + --trained_data_global_step 0 \ + --group_data \ + --use_decord \ + --prediction_type "v_prediction" \ + --rescale_betas_zero_snr \ + --output_dir="bs32x8x1_anyx93x320x320_fps16_lr2e-6_snr5_ema9999_sparse1d4_dit_l_mt5xxl_40m_vpred_zerosnr" diff --git a/scripts/text_condition/gpu/train_t2v_sparse2d.sh b/scripts/text_condition/gpu/train_t2v_sparse2d.sh new file mode 100644 index 000000000..be229b422 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse2d.sh @@ -0,0 +1,73 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example1.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/video_data_sucai_aes5.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 720 \ + --max_width 1280 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=2e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=250 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 1.0 \ + --force_resolution \ + --pretrained "/storage/dataset/Open-Sora-Plan-v1.2.0/93x720p/diffusion_pytorch_model.safetensors" \ + --output_dir="bs8x8x1_93x720p_lr2e-5_snr5_noioff0.02_ema9999_sparse2d4_dit_l_122_rope_mt5xxl_sucaiaes5" \ + --tile_sample_min_size 512 \ + --tile_sample_min_size_t 33 \ + --sparse2d --sparse_n 4 diff --git a/scripts/text_condition/gpu/train_t2v_sparse2d_image.sh b/scripts/text_condition/gpu/train_t2v_sparse2d_image.sh new file mode 100644 index 000000000..e8a2d6d10 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparse2d_image.sh @@ -0,0 +1,73 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 720 \ + --max_width 1280 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.5 \ + --interpolation_scale_w 2.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --sparse2d \ + --sparse_n 4 \ + --force_resolution \ + --enable_tracker \ + --pretrained "/storage/dataset/hw29/image/model_ema/diffusion_pytorch_model.safetensors" \ + --output_dir="bs1x8x16_1x720p_lr1e-4_snr5_noioff0.02_ema9999_sparse2d4_dit_l_122_rope_mt5xxl_mj_fromhw480p" diff --git a/scripts/text_condition/gpu/train_t2v_sparsedebug_image.sh b/scripts/text_condition/gpu/train_t2v_sparsedebug_image.sh new file mode 100644 index 000000000..fed165ba4 --- /dev/null +++ b/scripts/text_condition/gpu/train_t2v_sparsedebug_image.sh @@ -0,0 +1,71 @@ +export WANDB_KEY="953e958793b218efb850fa194e85843e2c3bd88b" +# export WANDB_MODE="offline" +export ENTITY="linbin" +export PROJECT="bs32x8x2_61x480p_lr1e-4_snr5_noioff0.02_opensora122_rope_mt5xxl_pandamovie_aes_mo_sucai_mo_speed1.2" +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export PDSH_RCMD_TYPE=ssh +# NCCL setting +export GLOO_SOCKET_IFNAME=bond0 +export NCCL_SOCKET_IFNAME=bond0 +export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 +export NCCL_IB_GID_INDEX=3 +export NCCL_IB_TC=162 +export NCCL_IB_TIMEOUT=22 +export NCCL_PXN_DISABLE=0 +export NCCL_IB_QPS_PER_CONNECTION=4 +export NCCL_ALGO=Ring +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +# export NCCL_ALGO=Tree + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset t2v \ + --data "scripts/train_data/merge_data_mj.txt" \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/storage/dataset/488dim4_plus" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 720 \ + --max_width 1280 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.5 \ + --interpolation_scale_w 2.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --tile_overlap_factor 0.125 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --use_rope \ + --resume_from_checkpoint="latest" \ + --group_data \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --force_resolution \ + --enable_tracker \ + --pretrained "/storage/dataset/hw29/image/model_ema/diffusion_pytorch_model.safetensors" \ + --output_dir="debug" diff --git a/scripts/text_condition/npu/sample_image.sh b/scripts/text_condition/npu/sample_image.sh new file mode 100644 index 000000000..e44aeaf53 --- /dev/null +++ b/scripts/text_condition/npu/sample_image.sh @@ -0,0 +1,18 @@ +export TASK_QUEUE_ENABLE=0 +torchrun --nproc_per_node=1 opensora/sample/sample_t2v_on_npu.py \ + --model_path bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/model_ema \ + --num_frames 1 \ + --height 320 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name /home/image_data/mt5-xxl \ + --text_prompt examples/prompt_list_0.txt \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "WFVAE_DISTILL_FORMAL" \ + --save_img_path "./test_image" \ + --fps 24 \ + --guidance_scale 5.0 \ + --num_sampling_steps 24 \ + --sample_method PNDM \ + --model_type "sparsedit" \ + --motion_score 0.9 \ \ No newline at end of file diff --git a/scripts/text_condition/npu/sample_inpaint.sh b/scripts/text_condition/npu/sample_inpaint.sh new file mode 100644 index 000000000..1037f512d --- /dev/null +++ b/scripts/text_condition/npu/sample_inpaint.sh @@ -0,0 +1,23 @@ +export TASK_QUEUE_ENABLE=0 +torchrun --nnodes=1 --nproc_per_node=1 --master_port 29502 \ + -m opensora.sample.sample_inpaint_ddp \ + --model_path /home/save_dir/runs/allinpaint_stage1/checkpoint-18000/model \ + --num_frames 93 \ + --height 320 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name /home/image_data/mt5-xxl \ + --text_prompt /home/image_data/test_prompt.txt \ + --conditional_images_path "/home/image_data/test_cond_imgs_path.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/home/image_data/lb/Open-Sora-Plan/WFVAE_DISTILL_FORMAL" \ + --save_img_path "/home/image_data/hxy/data/video" \ + --fps 24 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --sample_method EulerAncestralDiscrete \ + --motion_score 0.9 \ + --video_path "/home/image_data/hxy/data/video/000184_cut.mp4" \ + --mask_path "/home/image_data/hxy/data/video/000001_bbox_cut.mp4" \ + --prediction_type "v_prediction" \ + --rescale_betas_zero_snr \ No newline at end of file diff --git a/scripts/text_condition/npu/sample_inpaint_bak.sh b/scripts/text_condition/npu/sample_inpaint_bak.sh new file mode 100644 index 000000000..005be4739 --- /dev/null +++ b/scripts/text_condition/npu/sample_inpaint_bak.sh @@ -0,0 +1,34 @@ +WEIGHT_PATH="/home/opensora/pre_weights/" + +export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$(pwd)" +export MASTER_PORT=12359 + +if [ -z "$SAMPLE_SAVE_PATH" ]; then + export SAMPLE_SAVE_PATH="/home/image_data/sample_videos" +fi + +if [ -z "$SAMPLE_HEIGHT" ]; then + echo "You should set both envs of SAMPLE_HEIGHT and SAMPLE_WIDTH" + return +fi + +torchrun --nproc_per_node=8 opensora/sample/sample_inpaint_on_npu.py \ + --model_path /home/image_data/checkpoints/${PROJECT_NAME} \ + --num_frames ${NUM_FRAME} \ + --height $SAMPLE_HEIGHT \ + --width $SAMPLE_WIDTH \ + --cache_dir "../cache_dir" \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --text_prompt "/home/image_data/checkpoints/i2v_validation_dir/prompt.txt" \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --save_img_path "${SAMPLE_SAVE_PATH}/${PROJECT_NAME}" \ + --fps 24 \ + --max_sequence_length 512 \ + --guidance_scale 5.0 \ + --num_sampling_steps 50 \ + --sample_method PNDM \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --model_3d \ + --validation_dir "/home/image_data/checkpoints/i2v_validation_dir" \ No newline at end of file diff --git a/scripts/text_condition/npu/sample_video.sh b/scripts/text_condition/npu/sample_video.sh new file mode 100644 index 000000000..24f03cd25 --- /dev/null +++ b/scripts/text_condition/npu/sample_video.sh @@ -0,0 +1,18 @@ +export TASK_QUEUE_ENABLE=0 +torchrun --nproc_per_node=1 opensora/sample/sample_t2v_on_npu.py \ + --model_path bs32x8x1_anyx93x320x320_fps16_lr1e-5_snr5_noioff0.02_ema9999_sparse1d4_dit_l_mt5xxl_alldata100m/model_ema \ + --num_frames 29 \ + --height 160 \ + --width 320 \ + --cache_dir "../cache_dir" \ + --text_encoder_name /home/image_data/mt5-xxl \ + --text_prompt examples/prompt_list_0.txt \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "WFVAE_DISTILL_FORMAL" \ + --save_img_path "./test_video" \ + --fps 24 \ + --guidance_scale 5.0 \ + --num_sampling_steps 24 \ + --sample_method PNDM \ + --model_type "sparsedit" \ + --motion_score 0.9 \ \ No newline at end of file diff --git a/scripts/text_condition/npu/train_image3d_240p.sh b/scripts/text_condition/npu/train_image3d_240p.sh new file mode 100644 index 000000000..a85907a82 --- /dev/null +++ b/scripts/text_condition/npu/train_image3d_240p.sh @@ -0,0 +1,50 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE + + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 240 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --seed=10 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_image3d_240p_rope.sh b/scripts/text_condition/npu/train_image3d_240p_rope.sh new file mode 100644 index 000000000..18e938642 --- /dev/null +++ b/scripts/text_condition/npu/train_image3d_240p_rope.sh @@ -0,0 +1,52 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE + + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 240 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/image3d_rope_256p_zp_umt5/checkpoint-146000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_image3d_256p.sh b/scripts/text_condition/npu/train_image3d_256p.sh new file mode 100644 index 000000000..375bfd201 --- /dev/null +++ b/scripts/text_condition/npu/train_image3d_256p.sh @@ -0,0 +1,50 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE + + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 256 \ + --max_width 256 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --seed=10 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_image3d_256p_rope.sh b/scripts/text_condition/npu/train_image3d_256p_rope.sh new file mode 100644 index 000000000..34792dcfc --- /dev/null +++ b/scripts/text_condition/npu/train_image3d_256p_rope.sh @@ -0,0 +1,51 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE + + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 256 \ + --max_width 256 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=16 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --seed=10 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --use_rope \ + --noise_offset 0.02 \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_image3d_480p.sh b/scripts/text_condition/npu/train_image3d_480p.sh new file mode 100644 index 000000000..67e51a1dc --- /dev/null +++ b/scripts/text_condition/npu/train_image3d_480p.sh @@ -0,0 +1,54 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=8 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/image3d_256p_zp_umt5_from_initial_layer_40_head_16/checkpoint-176000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" + diff --git a/scripts/text_condition/npu/train_image3d_480p_rope.sh b/scripts/text_condition/npu/train_image3d_480p_rope.sh new file mode 100644 index 000000000..138d6d011 --- /dev/null +++ b/scripts/text_condition/npu/train_image3d_480p_rope.sh @@ -0,0 +1,54 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=8 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/image3d_rope_256p_zp_umt5/checkpoint-146000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_imageudit_480p.sh b/scripts/text_condition/npu/train_imageudit_480p.sh new file mode 100644 index 000000000..398c49b0f --- /dev/null +++ b/scripts/text_condition/npu/train_imageudit_480p.sh @@ -0,0 +1,51 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model UDiTT2V-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/DeepFloyd/t5-v1_1-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=4 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=250 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --downsampler "k33_s22" \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_imageuditultra_480p_new.sh b/scripts/text_condition/npu/train_imageuditultra_480p_new.sh new file mode 100644 index 000000000..26c2bd087 --- /dev/null +++ b/scripts/text_condition/npu/train_imageuditultra_480p_new.sh @@ -0,0 +1,52 @@ +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model UDiTUltraT2V-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames 1 \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 20 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="constant" \ + --seed=10 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=250 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --enable_tracker \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --downsampler "k33_s22" \ + --resume_from_checkpoint="latest" \ + --checkpoints_total_limit 3 \ No newline at end of file diff --git a/scripts/text_condition/npu/train_inpaint_sparse1d_newmodel_motion.sh b/scripts/text_condition/npu/train_inpaint_sparse1d_newmodel_motion.sh new file mode 100644 index 000000000..1d45767f4 --- /dev/null +++ b/scripts/text_condition/npu/train_inpaint_sparse1d_newmodel_motion.sh @@ -0,0 +1,83 @@ +export WANDB_KEY="c54943d667ed1abb58ed994e739462e66bc1aee2" +export WANDB_MODE="online" +export ENTITY="hexianyi" +export PROJECT=$PROJECT_NAME +# export PROJECT='test' +export HF_DATASETS_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 + +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +# export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_inpaint.py \ + --model OpenSoraInpaint-L/122 \ + --text_encoder_name google/mt5-xxl \ + --cache_dir "../../cache_dir/" \ + --dataset inpaint \ + --data "scripts/train_data/video_data_debug.txt" \ + --ae WFVAEModel_D8_4x8x8 \ + --ae_path "/home/image_data/lb/Open-Sora-Plan/WFVAE_DISTILL_FORMAL" \ + --sample_rate 1 \ + --num_frames 93 \ + --max_height 320 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 0 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.0 \ + --use_rope \ + --skip_low_resolution \ + --speed_factor 1.0 \ + --ema_decay 0.9999 \ + --drop_short_ratio 0.0 \ + --hw_stride 32 \ + --sparse1d --sparse_n 4 \ + --use_motion \ + --train_fps 16 \ + --seed 1234 \ + --trained_data_global_step 0 \ + --group_data \ + --use_decord \ + --prediction_type "v_prediction" \ + --rescale_betas_zero_snr \ + --t2v_ratio 0.0 \ + --i2v_ratio 0.0 \ + --transition_ratio 0.0 \ + --v2v_ratio 0.0 \ + --Semantic_ratio 0.2\ + --bbox_ratio 0.2\ + --background_ratio 0.2\ + --fixed_ratio 0.1\ + --Semantic_expansion_ratio 0.1\ + --fixed_bg_ratio 0.1\ + --clear_video_ratio 0.0 \ + --min_clear_ratio 0.25 \ + --default_text_ratio 0.0 \ + --output_dir /home/save_dir/runs/$PROJECT \ + --pretrained_transformer_model_path "/home/image_data/captions/vpre_latest_134k/model_ema" \ + --yolomodel_pathorname "/home/image_data/hxy/Open-Sora-Plan/opensora/dataset/yolov9c-seg.pt"\ + # --resume_from_checkpoint="/home/save_dir/runs/allinpaint_stage1/checkpoint-13000" + # 切part是resume，不是pretrained + # --snr_gamma 5.0 \ diff --git a/scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh b/scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh new file mode 100644 index 000000000..b7f751312 --- /dev/null +++ b/scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh @@ -0,0 +1,61 @@ +export WANDB_KEY="720d886d8c437c2142c88056a1eab8ef78d64a1f" +# # export WANDB_MODE="offline" +export ENTITY="yunyangge" +WEIGHT_PATH="/home/opensora/pre_weights/" +env +# export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_inpaint.py \ + --model OpenSoraInpaint-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset i2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 8 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-5 \ + --lr_scheduler="constant" \ + --seed=42 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=200 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/gyy_pretrained_f125/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" \ + --i2v_ratio 0.5 \ + --transition_ratio 0.4 \ + --default_text_ratio 0.5 \ + --group_frame \ + --speed_factor 1.1 \ diff --git a/scripts/text_condition/npu/train_video21d_nx240p.sh b/scripts/text_condition/npu/train_video21d_nx240p.sh new file mode 100644 index 000000000..6e14288ab --- /dev/null +++ b/scripts/text_condition/npu/train_video21d_nx240p.sh @@ -0,0 +1,51 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model LatteT2V-S/122 \ + --text_encoder_name ${WEIGHT_PATH}/DeepFloyd/t5-v1_1-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 240 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 15 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --seed=10 \ + --lr_scheduler="cosine" \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --seed=10 \ + --checkpointing_steps=1000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --pretrained "${WEIGHT_PATH}/t2v.pt" \ + --cfg 0.1 \ + --resume_from_checkpoint="latest" \ No newline at end of file diff --git a/scripts/text_condition/npu/train_video3d_nx240p.sh b/scripts/text_condition/npu/train_video3d_nx240p.sh new file mode 100644 index 000000000..7a0244e9e --- /dev/null +++ b/scripts/text_condition/npu/train_video3d_nx240p.sh @@ -0,0 +1,55 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 240 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 8 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="constant" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/yc_video3d_rope_240p/checkpoint-98000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" \ + --enable_stable_fp32 diff --git a/scripts/text_condition/npu/train_video3d_nx480p.sh b/scripts/text_condition/npu/train_video3d_nx480p.sh new file mode 100644 index 000000000..a691fff7c --- /dev/null +++ b/scripts/text_condition/npu/train_video3d_nx480p.sh @@ -0,0 +1,54 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 8 \ + --gradient_accumulation_steps=8 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/image3d_rope_480p_from_pretrain/checkpoint-14000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/npu/train_video3d_sp_nx240p.sh b/scripts/text_condition/npu/train_video3d_sp_nx240p.sh new file mode 100644 index 000000000..3511f2d60 --- /dev/null +++ b/scripts/text_condition/npu/train_video3d_sp_nx240p.sh @@ -0,0 +1,55 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 240 \ + --max_width 320 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 0.5 \ + --interpolation_scale_w 0.5 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 8 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="constant" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/yc_image3d_rope_240p_pretrain_seed_data/checkpoint-18000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" \ + --sp_size 8 \ + --train_sp_batch_size 4 \ + --enable_stable_fp32 diff --git a/scripts/text_condition/npu/train_video3d_sp_nx480p.sh b/scripts/text_condition/npu/train_video3d_sp_nx480p.sh new file mode 100644 index 000000000..f5b7c9777 --- /dev/null +++ b/scripts/text_condition/npu/train_video3d_sp_nx480p.sh @@ -0,0 +1,56 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model OpenSoraT2V-ROPE-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 8 \ + --gradient_accumulation_steps=8 \ + --max_train_steps=1000000 \ + --learning_rate=4e-5 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=500 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=2000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --enable_tiling \ + --tile_overlap_factor 0.125 \ + --use_rope \ + --noise_offset 0.02 \ + --pretrained "/home/image_data/checkpoints/image3d_rope_480p_from_pretrain/checkpoint-14000/model_ema/diffusion_pytorch_model.safetensors" \ + --resume_from_checkpoint="latest" \ + --sp_size 8 \ + --train_sp_batch_size 1 diff --git a/scripts/text_condition/npu/train_videoudit_nx480p.sh b/scripts/text_condition/npu/train_videoudit_nx480p.sh new file mode 100644 index 000000000..167f3cb45 --- /dev/null +++ b/scripts/text_condition/npu/train_videoudit_nx480p.sh @@ -0,0 +1,53 @@ +export PROJECT=$PROJECT_NAME +WEIGHT_PATH="/home/opensora/pre_weights/" +env +export WANDB_MODE='offline' +export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE +export HCCL_ALGO="level0:NA;level1:H-D_R" + +accelerate launch \ + --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \ + --machine_rank=${MACHINE_RANK} \ + --main_process_ip=${MAIN_PROCESS_IP_VALUE} \ + opensora/train/train_t2v_diffusers.py \ + --model UDiTT2V-L/122 \ + --text_encoder_name ${WEIGHT_PATH}/DeepFloyd/t5-v1_1-xxl \ + --cache_dir "../cache_dir" \ + --dataset t2v \ + --ae CausalVAEModel_4x8x8 \ + --ae_path "${WEIGHT_PATH}/test140k/" \ + --video_data "./scripts/train_data/video_data_on_npu.txt" \ + --image_data "./scripts/train_data/image_data_on_npu.txt" \ + --sample_rate 1 \ + --num_frames ${NUM_FRAME} \ + --max_height 480 \ + --max_width 640 \ + --interpolation_scale_t 1.0 \ + --interpolation_scale_h 1.0 \ + --interpolation_scale_w 1.0 \ + --attention_mode xformers \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --dataloader_num_workers 10 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=1000000 \ + --learning_rate=1e-4 \ + --lr_scheduler="cosine" \ + --seed=10 \ + --lr_warmup_steps=0 \ + --mixed_precision="bf16" \ + --report_to="wandb" \ + --checkpointing_steps=1000 \ + --output_dir="/home/image_data/checkpoints/${PROJECT}/" \ + --allow_tf32 \ + --model_max_length 512 \ + --use_image_num 0 \ + --enable_tiling \ + --snr_gamma 5.0 \ + --use_ema \ + --ema_start_step 0 \ + --cfg 0.1 \ + --noise_offset 0.02 \ + --downsampler "k333_s222" \ + + --resume_from_checkpoint="latest" diff --git a/scripts/text_condition/sample_image.sh b/scripts/text_condition/sample_image.sh deleted file mode 100644 index 303dac6d7..000000000 --- a/scripts/text_condition/sample_image.sh +++ /dev/null @@ -1,12 +0,0 @@ -CUDA_VISIBLE_DEVICES=0 python opensora/sample/sample_t2v.py \ - --model_path LanguageBind/Open-Sora-Plan-v1.0.0 \ - --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --text_prompt examples/prompt_list_0.txt \ - --ae CausalVAEModel_4x8x8 \ - --version 65x512x512 \ - --save_img_path "./sample_images/prompt_list_0" \ - --fps 24 \ - --guidance_scale 7.5 \ - --num_sampling_steps 250 \ - --enable_tiling \ - --force_images diff --git a/scripts/text_condition/sample_video.sh b/scripts/text_condition/sample_video.sh deleted file mode 100644 index fb6b7bd10..000000000 --- a/scripts/text_condition/sample_video.sh +++ /dev/null @@ -1,14 +0,0 @@ -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 \ - --version 65x512x512 \ - --image_size 512 \ - --cache_dir "/remote-home1/yeyang/sora4.8/Open-Sora-Plan/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/" \ - --save_img_path "./sample_videos" \ - --fps 24 \ - --guidance_scale 10.0 \ - --num_sampling_steps 50 \ - --enable_tiling diff --git a/scripts/text_condition/train_imageae.sh b/scripts/text_condition/train_imageae.sh deleted file mode 100644 index fff9c8afe..000000000 --- a/scripts/text_condition/train_imageae.sh +++ /dev/null @@ -1,34 +0,0 @@ -export WANDB_KEY="" -export ENTITY="" -export PROJECT="t2v-f16s3-img4-128-imgvae188-bf16-gc-xformers" -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_t2v.py \ - --model LatteT2V-XL/122 \ - --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --dataset t2v \ - --ae stabilityai/sd-vae-ft-mse \ - --data_path /remote-home1/dataset/sharegpt4v_path_cap_.json \ - --video_folder /remote-home1/dataset/data_split \ - --sample_rate 1 \ - --num_frames 17 \ - --max_image_size 256 \ - --gradient_checkpointing \ - --attention_mode xformers \ - --train_batch_size=4 \ - --dataloader_num_workers 10 \ - --gradient_accumulation_steps=1 \ - --max_train_steps=1000000 \ - --learning_rate=2e-05 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --mixed_precision="bf16" \ - --report_to="wandb" \ - --checkpointing_steps=500 \ - --output_dir="t2v-f17-256-img4-imagevae488-bf16-ckpt-xformers-bs4-lr2e-5-t5" \ - --allow_tf32 \ - --pretrained t2v.pt \ - --use_deepspeed \ - --model_max_length 300 \ - --use_image_num 4 \ - --use_img_from_vid diff --git a/scripts/text_condition/train_videoae_65x1024x1024.sh b/scripts/text_condition/train_videoae_65x1024x1024.sh deleted file mode 100644 index 6b5d4fb09..000000000 --- a/scripts/text_condition/train_videoae_65x1024x1024.sh +++ /dev/null @@ -1,39 +0,0 @@ -export WANDB_KEY="" -export ENTITY="" -export PROJECT="1024" -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_t2v.py \ - --model LatteT2V-XL/122 \ - --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --cache_dir "../../Open-Sora-Plan/cache_dir" \ - --dataset t2v \ - --ae CausalVAEModel_4x8x8 \ - --ae_path "../../Open-Sora-Plan/CausalVAEModel_4x8x8/" \ - --video_data_path "../../Open-Sora-Plan/sharegpt4v_path_cap_64x512x512_mixkit.json" \ - --video_folder /remote-home1/dataset/data_split_tt \ - --image_data_path "../../../dataset/image_114054.json" \ - --image_folder "../../../dataset/picture" \ - --sample_rate 1 \ - --num_frames 65 \ - --max_image_size 1024 \ - --gradient_checkpointing \ - --attention_mode xformers \ - --train_batch_size=2 \ - --dataloader_num_workers 10 \ - --gradient_accumulation_steps=1 \ - --max_train_steps=1000000 \ - --learning_rate=2e-05 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --mixed_precision="bf16" \ - --report_to="tensorboard" \ - --checkpointing_steps=100 \ - --output_dir="1024" \ - --allow_tf32 \ - --pretrained ../../65x512x512/diffusion_pytorch_model.safetensors \ - --use_deepspeed \ - --model_max_length 300 \ - --use_image_num 4 \ - --use_img_from_vid \ - --enable_tiling diff --git a/scripts/text_condition/train_videoae_65x512x512.sh b/scripts/text_condition/train_videoae_65x512x512.sh deleted file mode 100644 index 30414c65e..000000000 --- a/scripts/text_condition/train_videoae_65x512x512.sh +++ /dev/null @@ -1,38 +0,0 @@ -export WANDB_KEY="" -export ENTITY="" -export PROJECT="512_useimg" -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_t2v.py \ - --model LatteT2V-XL/122 \ - --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --cache_dir "../../Open-Sora-Plan/cache_dir" \ - --dataset t2v \ - --ae CausalVAEModel_4x8x8 \ - --ae_path "../../Open-Sora-Plan/CausalVAEModel_4x8x8/" \ - --video_data_path "../../Open-Sora-Plan/sharegpt4v_path_cap_64x512x512_mixkit.json" \ - --video_folder /remote-home1/dataset/data_split_tt \ - --image_data_path "../../../dataset/image_114054.json" \ - --image_folder "../../../dataset/picture" \ - --sample_rate 1 \ - --num_frames 65 \ - --max_image_size 512 \ - --gradient_checkpointing \ - --attention_mode xformers \ - --train_batch_size=4 \ - --dataloader_num_workers 10 \ - --gradient_accumulation_steps=1 \ - --max_train_steps=1000000 \ - --learning_rate=2e-05 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --mixed_precision="bf16" \ - --report_to="tensorboard" \ - --checkpointing_steps=500 \ - --output_dir="512_useimg" \ - --allow_tf32 \ - --pretrained t2v.pt \ - --use_deepspeed \ - --model_max_length 300 \ - --use_image_num 4 \ - --enable_tiling diff --git a/scripts/text_condition/train_videoae_65x512x512_d64.sh b/scripts/text_condition/train_videoae_65x512x512_d64.sh deleted file mode 100644 index d93ea8773..000000000 --- a/scripts/text_condition/train_videoae_65x512x512_d64.sh +++ /dev/null @@ -1,39 +0,0 @@ -export WANDB_KEY="" -export ENTITY="" -export PROJECT="512_d64" -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_t2v.py \ - --model LatteT2V-D64-XL/122 \ - --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --cache_dir "../../Open-Sora-Plan/cache_dir" \ - --dataset t2v \ - --ae CausalVAEModel_4x8x8 \ - --ae_path "../../Open-Sora-Plan/CausalVAEModel_4x8x8/" \ - --video_data_path "../../Open-Sora-Plan/sharegpt4v_path_cap_64x512x512_mixkit.json" \ - --video_folder /remote-home1/dataset/data_split_tt \ - --image_data_path "../../../dataset/image_114054.json" \ - --image_folder "../../../dataset/picture" \ - --sample_rate 1 \ - --num_frames 65 \ - --max_image_size 512 \ - --gradient_checkpointing \ - --attention_mode xformers \ - --train_batch_size=4 \ - --dataloader_num_workers 10 \ - --gradient_accumulation_steps=1 \ - --max_train_steps=1000000 \ - --learning_rate=2e-05 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --mixed_precision="bf16" \ - --report_to="tensorboard" \ - --checkpointing_steps=200 \ - --output_dir="512_d64" \ - --allow_tf32 \ - --pretrained t2v.pt \ - --use_deepspeed \ - --model_max_length 300 \ - --use_image_num 4 \ - --use_img_from_vid \ - --enable_tiling diff --git a/scripts/text_condition/train_videoae_65x512x512_rope.sh b/scripts/text_condition/train_videoae_65x512x512_rope.sh deleted file mode 100644 index dcd9a87b2..000000000 --- a/scripts/text_condition/train_videoae_65x512x512_rope.sh +++ /dev/null @@ -1,40 +0,0 @@ -export WANDB_KEY="" -export ENTITY="" -export PROJECT="512_rope_abs" -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_t2v.py \ - --model LatteT2V-XL/122 \ - --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --cache_dir "../../Open-Sora-Plan/cache_dir" \ - --dataset t2v \ - --ae CausalVAEModel_4x8x8 \ - --ae_path "../../Open-Sora-Plan/CausalVAEModel_4x8x8/" \ - --video_data_path "../../Open-Sora-Plan/sharegpt4v_path_cap_64x512x512_mixkit.json" \ - --video_folder /remote-home1/dataset/data_split_tt \ - --image_data_path "../../../dataset/image_114054.json" \ - --image_folder "../../../dataset/picture" \ - --sample_rate 1 \ - --num_frames 65 \ - --max_image_size 512 \ - --gradient_checkpointing \ - --attention_mode xformers \ - --train_batch_size=4 \ - --dataloader_num_workers 10 \ - --gradient_accumulation_steps=1 \ - --max_train_steps=1000000 \ - --learning_rate=2e-05 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --mixed_precision="bf16" \ - --report_to="tensorboard" \ - --checkpointing_steps=200 \ - --output_dir="512_rope_abs" \ - --allow_tf32 \ - --pretrained 512_useimg/checkpoint-5500/model/diffusion_pytorch_model.safetensors \ - --use_deepspeed \ - --model_max_length 300 \ - --use_image_num 4 \ - --use_img_from_vid \ - --enable_tiling \ - --use_rope diff --git a/scripts/train_data/current_data_on_npu.txt b/scripts/train_data/current_data_on_npu.txt new file mode 100644 index 000000000..a8a14d937 --- /dev/null +++ b/scripts/train_data/current_data_on_npu.txt @@ -0,0 +1 @@ +/home/obs_data,/home/image_data/captions/2_16_4764771.json diff --git a/scripts/train_data/image_data.txt b/scripts/train_data/image_data.txt new file mode 100644 index 000000000..b19a5ff40 --- /dev/null +++ b/scripts/train_data/image_data.txt @@ -0,0 +1,6 @@ +/storage/dataset/image/sam,/storage/anno_jsons/sam_image_11185255_resolution.json +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_1712571_resolution.json +/storage/dataset/image/human_images,/storage/anno_jsons/human_images_162094_resolution.json +/storage/dataset/image/anytext3m,/storage/anno_jsons/anytext_3509994_resolution.json +/storage/dataset/civitai/Images_civitai_v1,/storage/anno_jsons/civitai_v1_1940032_resolution.json +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.json diff --git a/scripts/train_data/image_data_debug.txt b/scripts/train_data/image_data_debug.txt new file mode 100644 index 000000000..0e85c866c --- /dev/null +++ b/scripts/train_data/image_data_debug.txt @@ -0,0 +1 @@ +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.json \ No newline at end of file diff --git a/scripts/train_data/image_data_notext.txt b/scripts/train_data/image_data_notext.txt new file mode 100644 index 000000000..4e3a08d17 --- /dev/null +++ b/scripts/train_data/image_data_notext.txt @@ -0,0 +1,3 @@ +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_4615265.json +/storage/dataset/image/human_images,/storage/anno_jsons/human_images_162094.json +/storage/dataset/image/sam,/storage/anno_jsons/sam_image_11185255.json \ No newline at end of file diff --git a/scripts/train_data/image_data_on_npu.txt b/scripts/train_data/image_data_on_npu.txt new file mode 100644 index 000000000..a65548d77 --- /dev/null +++ b/scripts/train_data/image_data_on_npu.txt @@ -0,0 +1,4 @@ +/home/local_dataset_6t/image_data_obs/mj_unzip_files,/home/opensora/captions/linbin_captions/tuzhan_mj_4615265.json +/home/local_dataset_6t/image_data_obs/mj_unzip_files,/home/opensora/captions/tuzhan_mj_4615265_cn.json +/home/local_dataset_6t/image_data_obs/sa_unzip_files,/home/opensora/captions/linbin_captions/sam_image_11185255.json +/home/local_dataset_6t/image_data_obs/images/,/home/opensora/captions/linbin_captions/human_images_162094.json \ No newline at end of file diff --git a/scripts/train_data/merge_data.txt b/scripts/train_data/merge_data.txt new file mode 100644 index 000000000..05244e1ab --- /dev/null +++ b/scripts/train_data/merge_data.txt @@ -0,0 +1,5 @@ +/storage,/storage/anno_jsons/stage2_pandamovie34m_aes4.5_sucai5m_sam3m_c1m_vidal3m_1_16_3045439_shuffle.json +/storage,/storage/anno_jsons/stage2_pandamovie34m_aes4.5_sucai5m_sam3m_c1m_vidal3m_2_16_3045439_shuffle.json +/storage,/storage/anno_jsons/stage2_pandamovie34m_aes4.5_sucai5m_sam3m_c1m_vidal3m_3_16_3045439_shuffle.json +/storage,/storage/anno_jsons/stage2_pandamovie34m_aes4.5_sucai5m_sam3m_c1m_vidal3m_4_16_3045439_shuffle.json +/storage,/storage/anno_jsons/stage2_pandamovie34m_aes4.5_sucai5m_sam3m_c1m_vidal3m_5_16_3045439_shuffle.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_debug.txt b/scripts/train_data/merge_data_debug.txt new file mode 100644 index 000000000..3450a2574 --- /dev/null +++ b/scripts/train_data/merge_data_debug.txt @@ -0,0 +1 @@ +/storage/dataset/movie,/storage/dataset/movie/TV01_clips_final_478625_llavanext_217405.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_jsonl.txt b/scripts/train_data/merge_data_jsonl.txt new file mode 100644 index 000000000..5b7c2a780 --- /dev/null +++ b/scripts/train_data/merge_data_jsonl.txt @@ -0,0 +1 @@ +/storage,/storage/anno_jsons/stage1_nomj_sam5m_1_4_22673408.jsonl \ No newline at end of file diff --git a/scripts/train_data/merge_data_mj.txt b/scripts/train_data/merge_data_mj.txt new file mode 100644 index 000000000..53e0412b9 --- /dev/null +++ b/scripts/train_data/merge_data_mj.txt @@ -0,0 +1,4 @@ +/storage/dataset/civitai/Images_civitai_v1,/storage/anno_jsons/civitai_v1_1940032_resolution.json +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.json +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_4615265_resolution.json +/storage/dataset/image/human_images,/storage/anno_jsons/human_images_162094_resolution.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_movie.txt b/scripts/train_data/merge_data_movie.txt new file mode 100644 index 000000000..fc0c33ddf --- /dev/null +++ b/scripts/train_data/merge_data_movie.txt @@ -0,0 +1,11 @@ + +/storage/dataset/movie,/storage/dataset/movie/bbc01_clips_final_1610998_llavanext_246875.json +/storage/dataset/movie,/storage/dataset/movie/bbc02_clips_final_1441415_llavanext_289761.json +/storage/dataset/movie,/storage/dataset/movie/bbc03_clips_final_1905074_llavanext_519162.json +/storage/dataset/movie,/storage/dataset/movie/bbc04_clips_final_1718543_llavanext_249487.json +/storage/dataset/movie,/storage/dataset/movie/bbc05_clips_final_2684676_llavanext_416525.json +/storage/dataset/movie,/storage/dataset/movie/TV01_clips_final_478625_llavanext_217405.json +/storage/dataset/image/sam,/storage/anno_jsons/sam_image_11185255_resolution.json +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_1712571_resolution.json +/storage/dataset/image/human_images,/storage/anno_jsons/human_images_162094_resolution.json +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_on_npu.txt b/scripts/train_data/merge_data_on_npu.txt new file mode 100644 index 000000000..46cfd4222 --- /dev/null +++ b/scripts/train_data/merge_data_on_npu.txt @@ -0,0 +1,16 @@ +/home/image_data/captions/1_16_4441840.json +/home/image_data/captions/2_16_4441840.json +/home/image_data/captions/3_16_4441840.json +/home/image_data/captions/4_16_4441840.json +/home/image_data/captions/5_16_4441840.json +/home/image_data/captions/6_16_4441840.json +/home/image_data/captions/7_16_4441840.json +/home/image_data/captions/8_16_4441840.json +/home/image_data/captions/9_16_4441840.json +/home/image_data/captions/10_16_4441840.json +/home/image_data/captions/11_16_4441839.json +/home/image_data/captions/12_16_4441839.json +/home/image_data/captions/13_16_4441839.json +/home/image_data/captions/14_16_4441839.json +/home/image_data/captions/15_16_4441839.json +/home/image_data/captions/16_16_4441839.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_panda_movie.txt b/scripts/train_data/merge_data_panda_movie.txt new file mode 100644 index 000000000..bbe859653 --- /dev/null +++ b/scripts/train_data/merge_data_panda_movie.txt @@ -0,0 +1,7 @@ +/storage/dataset/panda70m,/storage/dataset/filter_aes/filter_panda70m_json/panda_53473351_18923406_4.75plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc01_clips_final_1610998_566004_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc02_clips_final_1441415_663717_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc03_clips_final_1905075_1178482_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc04_clips_final_1718543_639395_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc05_clips_final_2684681_1060985_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/TV01_clips_final_478625_364795_4.5plus.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_panda_movie_mj.txt b/scripts/train_data/merge_data_panda_movie_mj.txt new file mode 100644 index 000000000..8a52a3a27 --- /dev/null +++ b/scripts/train_data/merge_data_panda_movie_mj.txt @@ -0,0 +1,10 @@ +/storage/dataset/panda70m,/storage/dataset/filter_aes/filter_panda70m_json/panda_53473351_18923406_4.75plus.json +/storage/dataset/civitai/Images_civitai_v1,/storage/anno_jsons/civitai_v1_1940032_resolution.json +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.json +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_4615265_resolution.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc01_clips_final_1610998_566004_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc02_clips_final_1441415_663717_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc03_clips_final_1905075_1178482_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc04_clips_final_1718543_639395_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/bbc05_clips_final_2684681_1060985_4.5plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_all_4.5_json/TV01_clips_final_478625_364795_4.5plus.json \ No newline at end of file diff --git a/scripts/train_data/merge_data_parquet.txt b/scripts/train_data/merge_data_parquet.txt new file mode 100644 index 000000000..23ad123f4 --- /dev/null +++ b/scripts/train_data/merge_data_parquet.txt @@ -0,0 +1,13 @@ +/storage/dataset,/storage/zhubin/liuyihang/add_motion/output/Webvid_output_rela_path.parquet +/storage/dataset,/storage/zhubin/liuyihang/add_motion/output/VIDAL_output_new_rela_path.parquet +/storage/dataset,/storage/zhubin/liuyihang/add_aes/output/sucai_aes.parquet +/storage/dataset/movie,/storage/dataset/movie/bbc01_clips_final_1610998_llavanext_246875.parquet +/storage/dataset/movie,/storage/dataset/movie/bbc02_clips_final_1441415_llavanext_289761.parquet +/storage/dataset/movie,/storage/dataset/movie/bbc03_clips_final_1905074_llavanext_519162.parquet +/storage/dataset/movie,/storage/dataset/movie/bbc04_clips_final_1718543_llavanext_249487.parquet +/storage/dataset/movie,/storage/dataset/movie/bbc05_clips_final_2684676_llavanext_416525.parquet +/storage/dataset/movie,/storage/dataset/movie/TV01_clips_final_478625_llavanext_217405.parquet +/storage/dataset/image/sam,/storage/anno_jsons/sam_image_11185255_resolution.parquet +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_1712571_resolution.parquet +/storage/dataset/image/human_images,/storage/anno_jsons/human_images_162094_resolution.parquet +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.parquet \ No newline at end of file diff --git a/scripts/train_data/merge_data_zb.txt b/scripts/train_data/merge_data_zb.txt new file mode 100644 index 000000000..ffe6dea43 --- /dev/null +++ b/scripts/train_data/merge_data_zb.txt @@ -0,0 +1,16 @@ +/storage/dataset,/storage/zhubin/liuyihang/add_aes/output/panda_aes.json +/storage/dataset,/storage/zhubin/liuyihang/add_motion/output/Webvid_output_rela_path.json +/storage/dataset,/storage/zhubin/liuyihang/add_motion/output/VIDAL_output_new_rela_path.json +/storage/dataset,/storage/zhubin/liuyihang/add_aes/output/sucai_aes.json + +/storage/dataset/movie,/storage/dataset/movie/bbc01_clips_final_1610998_llavanext_246875.json +/storage/dataset/movie,/storage/dataset/movie/bbc02_clips_final_1441415_llavanext_289761.json +/storage/dataset/movie,/storage/dataset/movie/bbc03_clips_final_1905074_llavanext_519162.json +/storage/dataset/movie,/storage/dataset/movie/bbc04_clips_final_1718543_llavanext_249487.json +/storage/dataset/movie,/storage/dataset/movie/bbc05_clips_final_2684676_llavanext_416525.json +/storage/dataset/movie,/storage/dataset/movie/TV01_clips_final_478625_llavanext_217405.json +/storage/dataset/image/sam,/storage/anno_jsons/sam_image_11185255_resolution.json +/storage/dataset/image/tuzhan_mj,/storage/anno_jsons/tuzhan_mj_1712571_resolution.json +/storage/dataset/image/human_images,/storage/anno_jsons/human_images_162094_resolution.json +/storage/dataset/civitai/Images_civitai_v1,/storage/anno_jsons/civitai_v1_1940032_resolution.json +/storage/dataset/ideogram/Images_ideogram_v1,/storage/anno_jsons/ideogram_v1_71637_resolution.json \ No newline at end of file diff --git a/scripts/train_data/video_data.txt b/scripts/train_data/video_data.txt new file mode 100644 index 000000000..e0bffc577 --- /dev/null +++ b/scripts/train_data/video_data.txt @@ -0,0 +1,7 @@ +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_json/bbc01_250508_117513_5.0plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_json/bbc02_289778_128670_5.0plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_json/bbc03_519184_281100_5.0plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_json/bbc04_249497_102768_5.0plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_json/bbc05_416548_207570_5.0plus.json +/storage/dataset/movie,/storage/dataset/filter_aes/filter_movie_json/TV01_217419_141366_5.0plus.json +/storage/dataset/panda70m,/storage/dataset/filter_aes/filter_panda70m_json/panda70m_last_6268414_flowValue_940432_5.4plus.json \ No newline at end of file diff --git a/scripts/train_data/video_data_513.txt b/scripts/train_data/video_data_513.txt new file mode 100644 index 000000000..796c7d1f3 --- /dev/null +++ b/scripts/train_data/video_data_513.txt @@ -0,0 +1 @@ +/storage/dataset/mixkit,/storage/anno_jsons/video_mixkit_65f_54735.json \ No newline at end of file diff --git a/scripts/train_data/video_data_aesmovie_panda.txt b/scripts/train_data/video_data_aesmovie_panda.txt new file mode 100644 index 000000000..b26ad56a8 --- /dev/null +++ b/scripts/train_data/video_data_aesmovie_panda.txt @@ -0,0 +1,7 @@ +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc01_250508.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc02_289778.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc03_519184.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc04_249497.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc05_416548.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/TV01_217419.json +/storage/dataset/panda70m,/storage/dataset/panda70m/panda70m_last_6268414_flowValue.json \ No newline at end of file diff --git a/scripts/train_data/video_data_aesmovie_sucai_panda.txt b/scripts/train_data/video_data_aesmovie_sucai_panda.txt new file mode 100644 index 000000000..f583df7f6 --- /dev/null +++ b/scripts/train_data/video_data_aesmovie_sucai_panda.txt @@ -0,0 +1,16 @@ +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc01_250508.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc02_289778.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc03_519184.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc04_249497.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/bbc05_416548.json +/storage/dataset/movie,/storage/dataset/movie/merge_cap_lb_7_11/TV01_217419.json +/storage/dataset/panda70m,/storage/dataset/panda70m/panda70m_last_6268414_flowValue.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_canva_final_95441.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_clipchamp_final_452264.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_coverr_final_3002.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_istock_final_815070.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_kapwing_final_68473.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_mixkit_final_4490.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_pexels_final_267395.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_pixabay_v2_final_21608.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_storyblocks_final_1270947.json \ No newline at end of file diff --git a/scripts/train_data/video_data_debug.txt b/scripts/train_data/video_data_debug.txt new file mode 100644 index 000000000..4c9bd7bd1 --- /dev/null +++ b/scripts/train_data/video_data_debug.txt @@ -0,0 +1 @@ +/home/obs_data/20240426/20240704-02-bbc-01,/home/image_data/captions/TV01_clips_final_478625_llavanext_217405_aes478625.json \ No newline at end of file diff --git a/scripts/train_data/video_data_on_npu.txt b/scripts/train_data/video_data_on_npu.txt new file mode 100644 index 000000000..eac812cee --- /dev/null +++ b/scripts/train_data/video_data_on_npu.txt @@ -0,0 +1,16 @@ +/home/obs_data,/home/image_data/captions/1_16_4764771.json +/home/obs_data,/home/image_data/captions/2_16_4764771.json +/home/obs_data,/home/image_data/captions/3_16_4764771.json +/home/obs_data,/home/image_data/captions/4_16_4764771.json +/home/obs_data,/home/image_data/captions/5_16_4764770.json +/home/obs_data,/home/image_data/captions/6_16_4764770.json +/home/obs_data,/home/image_data/captions/7_16_4764770.json +/home/obs_data,/home/image_data/captions/8_16_4764770.json +/home/obs_data,/home/image_data/captions/9_16_4764770.json +/home/obs_data,/home/image_data/captions/10_16_4764770.json +/home/obs_data,/home/image_data/captions/11_16_4764770.json +/home/obs_data,/home/image_data/captions/12_16_4764770.json +/home/obs_data,/home/image_data/captions/13_16_4764770.json +/home/obs_data,/home/image_data/captions/14_16_4764770.json +/home/obs_data,/home/image_data/captions/15_16_4764770.json +/home/obs_data,/home/image_data/captions/16_16_4764770.json diff --git a/scripts/train_data/video_data_sucai.txt b/scripts/train_data/video_data_sucai.txt new file mode 100644 index 000000000..31537b50c --- /dev/null +++ b/scripts/train_data/video_data_sucai.txt @@ -0,0 +1,9 @@ +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_canva_final_95441.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_clipchamp_final_452264.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_coverr_final_3002.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_istock_final_815070.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_kapwing_final_68473.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_mixkit_final_4490.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_pexels_final_267395.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_pixabay_v2_final_21608.json +/storage/dataset,/storage/zhubin/video_statistics_data/task1.5/Final_format_dataset_data_v2/step1.5_storyblocks_final_1270947.json \ No newline at end of file diff --git a/scripts/train_data/video_data_sucai_aes5.txt b/scripts/train_data/video_data_sucai_aes5.txt new file mode 100644 index 000000000..fc4e028c0 --- /dev/null +++ b/scripts/train_data/video_data_sucai_aes5.txt @@ -0,0 +1,9 @@ +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_canva_final_95441_65422_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_clipchamp_final_452264_231178_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_coverr_final_3002_1274_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_istock_final_815070_405120_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_kapwing_final_68473_27757_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_mixkit_final_4490_2634_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_pexels_final_267395_148323_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_pixabay_v2_final_21608_11400_5.0plus.json +/storage/dataset,/storage/dataset/filter_aes/filter_sucai_json/step1.5_storyblocks_final_1270947_825948_5.0plus.json \ No newline at end of file diff --git a/scripts/videogpt/train_videogpt.sh b/scripts/videogpt/train_videogpt.sh deleted file mode 100644 index 6c06e0a64..000000000 --- a/scripts/videogpt/train_videogpt.sh +++ /dev/null @@ -1,30 +0,0 @@ - -accelerate launch \ - --config_file scripts/accelerate_configs/ddp_config.yaml \ - opensora/train/train_videogpt.py \ - --do_train \ - --seed 1234 \ - --data_path "/remote-home/yeyang/UCF-101/" \ - --per_device_train_batch_size 1 \ - --gradient_accumulation_steps 1 \ - --learning_rate 7e-4 \ - --weight_decay 0. \ - --max_steps 20000 \ - --lr_scheduler_type cosine \ - --max_grad_norm 1.0 \ - --save_strategy steps \ - --save_total_limit 5 \ - --logging_steps 5 \ - --save_steps 1000 \ - --n_codes 2048 \ - --n_hiddens 240 \ - --embedding_dim 4 \ - --n_res_layers 4 \ - --downsample "4,4,4" \ - --resolution 240 \ - --sequence_length 16 \ - --output_dir results/videogpt_488_256_16 \ - --bf16 True \ - --fp16 False \ - --report_to tensorboard \ - --dataloader_num_workers 10 diff --git a/scripts/videogpt/train_videogpt_dsz2.sh b/scripts/videogpt/train_videogpt_dsz2.sh deleted file mode 100644 index 7875333eb..000000000 --- a/scripts/videogpt/train_videogpt_dsz2.sh +++ /dev/null @@ -1,30 +0,0 @@ -export ACCELERATE_GRADIENT_ACCUMULATION_STEPS=1 - -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_videogpt.py \ - --do_train \ - --seed 1234 \ - --data_path "datasets/UCF-101/" \ - --per_device_train_batch_size 32 \ - --gradient_accumulation_steps $ACCELERATE_GRADIENT_ACCUMULATION_STEPS \ - --learning_rate 7e-4 \ - --weight_decay 0. \ - --num_train_epochs 2 \ - --lr_scheduler_type cosine \ - --max_grad_norm 1.0 \ - --save_strategy steps \ - --save_total_limit 5 \ - --logging_steps 5 \ - --save_steps 10000 \ - --n_codes 1024 \ - --n_hiddens 240 \ - --embedding_dim 4 \ - --n_res_layers 4 \ - --downsample "4,4,4" \ - --resolution 128 \ - --sequence_length 16 \ - --output_dir results/videogpt_444_128 \ - --bf16 True \ - --fp16 False \ - --report_to tensorboard diff --git a/scripts/videogpt/train_videogpt_dsz3.sh b/scripts/videogpt/train_videogpt_dsz3.sh deleted file mode 100644 index 484084236..000000000 --- a/scripts/videogpt/train_videogpt_dsz3.sh +++ /dev/null @@ -1,30 +0,0 @@ -export ACCELERATE_GRADIENT_ACCUMULATION_STEPS=1 - -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero3_config.yaml \ - opensora/train/train_videogpt.py \ - --do_train \ - --seed 1234 \ - --data_path "datasets/UCF-101/" \ - --per_device_train_batch_size 32 \ - --gradient_accumulation_steps $ACCELERATE_GRADIENT_ACCUMULATION_STEPS \ - --learning_rate 7e-4 \ - --weight_decay 0. \ - --num_train_epochs 2 \ - --lr_scheduler_type cosine \ - --max_grad_norm 1.0 \ - --save_strategy steps \ - --save_total_limit 5 \ - --logging_steps 5 \ - --save_steps 10000 \ - --n_codes 1024 \ - --n_hiddens 240 \ - --embedding_dim 4 \ - --n_res_layers 4 \ - --downsample "4,4,4" \ - --resolution 128 \ - --sequence_length 16 \ - --output_dir results/videogpt_444_128 \ - --bf16 True \ - --fp16 False \ - --report_to tensorboard diff --git a/seed.txt b/seed.txt new file mode 100644 index 000000000..71abb8505 --- /dev/null +++ b/seed.txt @@ -0,0 +1,200 @@ +54772192 +1818559 +96453589 +31562712 +57322344 +69472081 +111012 +21408319 +92640719 +61294860 +21327469 +53263210 +40582092 +28394045 +71159969 +2157594 +58200604 +23406796 +76254691 +21916831 +62388905 +63738670 +26204979 +7771881 +59943270 +80583057 +95793202 +52859755 +2075209 +99525631 +33755453 +35407358 +27398776 +18898781 +74552050 +61167570 +35845404 +33774409 +79535361 +89984239 +57416385 +83451617 +93069601 +7626451 +90670347 +12456687 +88103643 +36406985 +53298385 +79583560 +31475157 +69403517 +10853392 +79612930 +57941600 +64597087 +3837970 +55989364 +30518141 +30458584 +90178174 +78102020 +85241869 +56602460 +4287298 +69650713 +84917682 +22260866 +86060151 +50690452 +79466300 +15142491 +16189759 +48108189 +40868115 +26720024 +57424386 +92772187 +66605653 +82892525 +43330140 +47838312 +54852636 +91021440 +38201699 +88828952 +22709517 +40070762 +13020663 +83479581 +12206776 +88705935 +51064471 +70264569 +62293291 +67887552 +91338327 +56759295 +42384628 +19062215 +89549704 +36762545 +65980816 +80073026 +49304259 +84459031 +46163305 +92297335 +91470376 +4243674 +54249696 +94095222 +79354938 +27136429 +38474371 +43628631 +72990601 +48425792 +71070649 +62709171 +20040550 +46032149 +40764114 +50477866 +79805183 +35062693 +15657913 +72197366 +10639508 +53191828 +74447311 +71914086 +59133527 +98084299 +91062387 +83689168 +41848719 +21283911 +15325792 +81328486 +77584362 +52674907 +9253255 +89806702 +67348822 +68776016 +42791133 +93708934 +47683245 +43970228 +21994327 +57504058 +20830613 +31554532 +76251166 +34808236 +59340023 +47416276 +13518519 +80865077 +82164806 +95435324 +41236744 +64021281 +78588359 +5529302 +63957306 +60796342 +91025348 +57426714 +77145225 +22724514 +33351724 +28468671 +68301472 +25527807 +27943666 +52461789 +38671842 +81392557 +25561570 +46263168 +89589577 +77038019 +71423560 +67019718 +56009735 +68932922 +98988290 +26565463 +60536315 +13782449 +32549829 +5696904 +56946718 +17171986 +6197126 +1463826 +84204677 +12099676 \ No newline at end of file diff --git a/seed_list.py b/seed_list.py new file mode 100644 index 000000000..de4bee0f1 --- /dev/null +++ b/seed_list.py @@ -0,0 +1,8 @@ +import random +import numpy as np + + +randnum = np.random.randint(0, 1e8, size=200) + +with open('seed.txt', 'w') as f: + f.write('\n'.join([str(i) for i in randnum])) \ No newline at end of file diff --git a/tools/get_img_info.py b/tools/get_img_info.py new file mode 100644 index 000000000..f09375f8e --- /dev/null +++ b/tools/get_img_info.py @@ -0,0 +1,76 @@ +# import cv2 +# from tqdm import tqdm +# from glob import glob +# import json +# import os + +# def get_image_size(image_path): +# """ +# Given an image path, return its width and height. +# """ +# try: +# image = cv2.imread(image_path) +# height, width = image.shape[:2] +# return height, width +# except Exception as e: +# return None, None + +# image_root = '/storage/dataset/image/human_images/' +# save_root = '/storage/dataset/image' +# os.makedirs(save_root, exist_ok=True) +# save_name = 'human_images_{}_resolution.json' +# all_paths = glob(os.path.join(image_root, '**', f'*.jpg'), recursive=True) +# items = [] +# for i in tqdm(all_paths): +# height, width = get_image_size(i) +# path = i.replace(image_root if image_root.endswith('/') else image_root + '/', '') +# item = dict(path=path, resolution=dict(height=height, width=width)) +# items.append(item) +# with open(os.path.join(save_root, save_name.format(len(items))), 'w') as f: +# json.dump(items, f, indent=2) + + + + +import cv2 +from tqdm import tqdm +from glob import glob +import json +import os +from multiprocessing import Pool + +def get_image_size(image_path): + """ + Given an image path, return its width and height. + """ + try: + image = cv2.imread(image_path) + height, width = image.shape[:2] + return image_path, height, width + except Exception as e: + return image_path, None, None + +def process_image_paths(image_paths): + items = [] + for image_path, height, width in image_paths: + path = image_path.replace(image_root if image_root.endswith('/') else image_root + '/', '') + item = dict(path=path, resolution=dict(height=height, width=width)) + items.append(item) + return items + +if __name__ == '__main__': + image_root = '/storage/dataset/image/tuzhan_mj' + save_root = '/storage/dataset/image' + os.makedirs(save_root, exist_ok=True) + save_name = 'tuzhan_mj_{}_resolution.json' + all_paths = glob(os.path.join(image_root, '**', '*.jpg'), recursive=True) + + num_processes = os.cpu_count() # Use the number of CPU cores + num_processes = 128 # Use the number of CPU cores + with Pool(num_processes) as pool: + results = list(tqdm(pool.imap(get_image_size, all_paths), total=len(all_paths))) + + items = process_image_paths(results) + + with open(os.path.join(save_root, save_name.format(len(items))), 'w') as f: + json.dump(items, f, indent=2) diff --git a/tools/merge_imginfo_to_anno.py b/tools/merge_imginfo_to_anno.py new file mode 100644 index 000000000..067b7b9ee --- /dev/null +++ b/tools/merge_imginfo_to_anno.py @@ -0,0 +1,37 @@ + +from tqdm import tqdm +from glob import glob +import json +import os + +# anno_path = '/storage/anno_jsons/human_images_162094.json' +# img_info_path = '/storage/dataset/image/human_images_162094_resolution.json' +# save_root = '/storage/anno_jsons' +# save_name = 'human_images_{}_resolution.json' + + +# anno_path = '/storage/anno_jsons/tuzhan_mj_4615265.json' +# img_info_path = '/storage/dataset/image/tuzhan_mj_4615530_resolution.json' +# save_root = '/storage/anno_jsons' +# save_name = 'tuzhan_mj_{}_resolution.json' + + +anno_path = '/storage/anno_jsons/sam_image_11185255.json' +img_info_path = '/storage/dataset/image/sam_image_11185362_resolution.json' +save_root = '/storage/anno_jsons' +save_name = 'sam_image_{}_resolution.json' + +with open(anno_path, 'r') as f: + anno = json.load(f) +with open(img_info_path, 'r') as f: + img_info = json.load(f) +img_info = {i['path']: i['resolution'] for i in img_info} + +items = [] +cnt = 0 +for i in tqdm(anno): + resolution = img_info[i['path']] + i['resolution'] = resolution + items.append(i) +with open(os.path.join(save_root, save_name.format(len(items))), 'w') as f: + json.dump(items, f, indent=2) \ No newline at end of file diff --git a/validation_dir/prompt.txt b/validation_dir/prompt.txt new file mode 100644 index 000000000..592bd5a4d --- /dev/null +++ b/validation_dir/prompt.txt @@ -0,0 +1,27 @@ +A rocket ascends slowly into the sky. +A coffee cup with "anytext" foam floating on it. +A beautiful girl is blinking her eyes. +A flower and bird painting.This video has a silky transition effect. +The scenery of snow capped mountains.This video has a silky transition effect. +The night view of the city.This video has a silky transition effect. +A rocket ascends slowly into the sky. +Along the coast, variously sized boats float on the lake. +The landscape at sunset is profound and expansive. +A beautiful girl is blinking her eyes. +A beautiful girl is blinking her eyes. +A coffee cup with "anytext" foam floating on it. +A person is eating pizza on a plate. +A group of people are cheering and celebrating. +A green dinosaur is playing tennis. +A female elf is smiling. +The girl with red hair is turning her head and smiling. +A racing car is speeding on the road. Environmental perspective. +A family photo of four people smiling and playing around with each other. +The old lady hugged her granddaughter and they laughed happily. +Drones fly from snow capped mountains to cliffs, with waterfalls flowing. This video has a silky transition effect. +There is no transition video from the cliff to the bottom, and the waterfall is flowing. +The scene from the seabed to the cliff.This video has a silky transition effect. +The boy in the video abandoned his girlfriend and instead embraced the girl in red clothes. +A beautiful girl is stroking her hair and blinking her eyes. +A fashionable female model strolls on the streets of Tokyo. +A panda playing guitar on the grass. \ No newline at end of file