overfit SP

Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com>
Co-authored-by: Will Lin <wlsaidhi@gmail.com>

.github/ISSUE_TEMPLATE/1-bug-report.yml

@@ -8,7 +8,7 @@ body:
   attributes:
     label: Environment
     description: |
-      Please share your environment with us. You can run the command **python fastvideo/utils/env_utils.py** and copy-paste its output below.
+      Please share your environment with us. You can run the command **python fastvideo/utils/collect_env.py** and copy-paste its output below.
     placeholder: FastVideo version, platform, python version, cuda version...
   validations:
     required: true

.github/workflows/pr-test.yml

@@ -77,6 +77,8 @@ jobs:
               - 'fastvideo/v1/models/dits/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/transformers/**'
+              - 'fastvideo/v1/layers/**'
+              - 'fastvideo/v1/attention/**'
 
   encoder-test:
     needs: change-filter
@@ -141,8 +143,8 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [
-          {version: "3.10", tag: "latest"},
-          {version: "3.11", tag: "py3.11-latest"},
+          # {version: "3.10", tag: "latest"},
+          # {version: "3.11", tag: "py3.11-latest"},
           {version: "3.12", tag: "py3.12-latest"}
         ]
     uses: ./.github/workflows/runpod-test.yml

csrc/sliding_tile_attention/test/test_sta.py

@@ -2,7 +2,6 @@ import torch
 from flex_sta_ref import get_sliding_tile_attention_mask
 from st_attn import sliding_tile_attention
 from torch.nn.attention.flex_attention import flex_attention
-# from flash_attn_interface import flash_attn_func
 from tqdm import tqdm
 
 flex_attention = torch.compile(flex_attention, dynamic=False)
@@ -23,7 +22,7 @@ def h100_fwd_kernel_test(Q, K, V, kernel_size):
 def generate_tensor(shape, mean, std, dtype, device):
     tensor = torch.randn(shape, dtype=dtype, device=device)
 
-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    magnitude = torch.linalg.norm(tensor, dim=-1, keepdim=True)
     scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
 
     return scaled_tensor.contiguous()

distill_v1.sh

@@ -1,46 +0,0 @@
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-
-DATA_DIR=./data
-# IP=[MASTER NODE IP]
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-    # --gradient_checkpointing\
-    # --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \
-    # --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
-torchrun --nnodes 1 --nproc_per_node 4\
-    fastvideo/v1/pipelines/training_pipeline.py\
-    --inference_mode False\
-    --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
-    --cache_dir "/home/ray/.cache"\
-    --data_json_path "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/videos2caption.json"\
-    --validation_prompt_dir "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/validation"\
-    --train_batch_size=1\
-    --num_latent_t 1 \
-    --sp_size 4 \
-    --tp_size 4 \
-    --train_sp_batch_size 1\
-    --dataloader_num_workers 4\
-    --gradient_accumulation_steps=1\
-    --max_train_steps=320\
-    --learning_rate=1e-6\
-    --mixed_precision="bf16"\
-    --checkpointing_steps=64\
-    --validation_steps 64\
-    --validation_sampling_steps "2,4,8" \
-    --checkpoints_total_limit 3\
-    --allow_tf32\
-    --ema_start_step 0\
-    --cfg 0.0\
-    --log_validation\
-    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\
-    --tracker_project_name Hunyuan_Distill \
-    --num_height 720 \
-    --num_width 1280 \
-    --num_frames  125 \
-    --shift 17 \
-    --validation_guidance_scale "1.0" \
-    --num_euler_timesteps 50 \
-    --multi_phased_distill_schedule "4000-1" \
-    --not_apply_cfg_solver \
-    --master_weight_type "bf16"

fastvideo/__init__.py

@@ -1,5 +1,6 @@
 from fastvideo.v1.configs.pipelines import PipelineConfig
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.entrypoints.video_generator import VideoGenerator
+from fastvideo.version import __version__
 
-__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam"]
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]

fastvideo/data_preprocess/preprocess.py

@@ -77,17 +77,6 @@ if __name__ == "__main__":
         default=8,
         help="Batch size (per device) for the training dataloader.",
     )
-    parser.add_argument(
-        "--samples_per_file",
-        type=int,
-        default=64
-    )
-    parser.add_argument(
-        "--flush_frequency",
-        type=int,
-        default=256,
-        help="how often to save to parquet files"
-    )
     parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
     parser.add_argument("--max_height", type=int, default=480)
     parser.add_argument("--max_width", type=int, default=848)

fastvideo/data_preprocess/preprocess_text_embeddings_v1.py

@@ -1,199 +0,0 @@
-import argparse
-import json
-import os
-
-import torch
-import torch.distributed as dist
-from accelerate.logging import get_logger
-from diffusers.utils import export_to_video
-from diffusers.video_processor import VideoProcessor
-from torch.utils.data import DataLoader, Dataset
-from torch.utils.data.distributed import DistributedSampler
-from tqdm import tqdm
-
-# from fastvideo.utils.load import load_text_encoder, load_vae
-
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
-from fastvideo.v1.utils import maybe_download_model
-from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
-from fastvideo.v1.configs.models.encoders.t5 import T5Config
-
-logger = get_logger(__name__)
-
-class T5dataset(Dataset):
-
-    def __init__(
-        self,
-        json_path,
-        vae_debug,
-    ):
-        self.json_path = json_path
-        self.vae_debug = vae_debug
-        with open(self.json_path, "r") as f:
-            train_dataset = json.load(f)
-            self.train_dataset = sorted(train_dataset, key=lambda x: x["latent_path"])
-
-    def __getitem__(self, idx):
-        caption = self.train_dataset[idx]["caption"]
-        filename = self.train_dataset[idx]["latent_path"].split(".")[0]
-        length = self.train_dataset[idx]["length"]
-        if self.vae_debug:
-            latents = torch.load(
-                os.path.join(args.output_dir, "latent", self.train_dataset[idx]["latent_path"]),
-                map_location="cpu",
-            )
-        else:
-            latents = []
-
-        return dict(caption=caption, latents=latents, filename=filename, length=length)
-
-    def __len__(self):
-        return len(self.train_dataset)
-
-
-def main(args):
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    rank = int(os.getenv("RANK", 0))
-    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
-    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
-    print("world_size", world_size, "local rank", local_rank)
-    device = torch.device(f"cuda:{local_rank}")
-    torch.cuda.set_device(device)
-    world_group = get_world_group()
-
-    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    # torch.cuda.set_device(local_rank)
-    # if not dist.is_initialized():
-    #     dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-
-    videoprocessor = VideoProcessor(vae_scale_factor=8)
-    os.makedirs(args.output_dir, exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
-
-    vae_precision = "fp16" 
-    text_encoder_precision = "fp32"
-    fastvideo_args = FastVideoArgs(model_path=args.model_path,
-                         use_cpu_offload=False,
-                         vae_precision=vae_precision,
-                         text_encoder_precisions=(text_encoder_precision,))
-    fastvideo_args.device = device
-    fastvideo_args.device_str = f"cuda:{local_rank}"
-
-    # fastvideo_args.dit_config = HunyuanVideoConfig()
-    fastvideo_args.vae_config = WanVAEConfig()
-    fastvideo_args.text_encoder_configs = (T5Config(),)
-
-    # vae_loader = VAELoader()
-    # vae = vae_loader.load_vae()
-    text_encoder_loader = TextEncoderLoader()
-    tokenizer_loader = TokenizerLoader()
-
-    model_path = args.model_path
-    path = maybe_download_model(model_path)
-    # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-    ENCODER_PATH = os.path.join(path, "text_encoder")
-    TOKENIZER_PATH = os.path.join(path, "tokenizer")
-    print(ENCODER_PATH)
-    text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
-    tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
-
-    latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
-    train_dataset = T5dataset(latents_json_path, args.vae_debug)
-    # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
-    # vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
-    # vae.enable_tiling()
-    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-    )
-
-    json_data = []
-    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
-        with torch.inference_mode():
-            # with torch.autocast("cuda", dtype=torch.float32):
-            print(data["caption"])
-            text_inputs = tokenizer(data["caption"], **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
-                fastvideo_args.device)
-            input_ids = text_inputs["input_ids"]
-            attention_mask = text_inputs["attention_mask"]
-            outputs = text_encoder(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                output_hidden_states=True,
-            )
-            from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
-            post_process_func = t5_postprocess_text
-            prompt_embeds = post_process_func(outputs)
-            prompt_attention_mask = attention_mask
-            if args.vae_debug:
-                latents = data["latents"]
-                video = vae.decode(latents.to(device), return_dict=False)[0]
-                video = videoprocessor.postprocess_video(video)
-            for idx, video_name in enumerate(data["filename"]):
-                prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
-                video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
-                prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask",
-                                                            video_name + ".pt")
-                # save latent
-                torch.save(prompt_embeds[idx], prompt_embed_path)
-                torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
-                print(f"sample {video_name} saved")
-                if args.vae_debug:
-                    export_to_video(video[idx], video_path, fps=16)
-                item = {}
-                item["length"] = int(data["length"][idx])
-                item["latent_path"] = video_name + ".pt"
-                item["prompt_embed_path"] = video_name + ".pt"
-                item["prompt_attention_mask"] = video_name + ".pt"
-                item["caption"] = data["caption"][idx]
-                json_data.append(item)
-    dist.barrier()
-    local_data = json_data
-    gathered_data = [None] * world_size
-    dist.all_gather_object(gathered_data, local_data)
-    if local_rank == 0:
-        # os.remove(latents_json_path)
-        all_json_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f:
-            json.dump(all_json_data, f, indent=4)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # dataset & dataloader
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    # parser.add_argument("--model_type", type=str, default="mochi")
-    # text encoder & vae & diffusion model
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=1,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--train_batch_size",
-        type=int,
-        default=1,
-        help="Batch size (per device) for the training dataloader.",
-    )
-    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    parser.add_argument("--vae_debug", action="store_true")
-    args = parser.parse_args()
-    main(args)

fastvideo/data_preprocess/preprocess_vae_latents_v1.py

@@ -1,151 +0,0 @@
-import argparse
-import json
-import os
-
-import torch
-# import torch.distributed as dist
-# from accelerate.logging import get_logger
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from tqdm import tqdm
-
-from fastvideo.dataset import getdataset
-# from fastvideo.utils.load import load_vae
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.loader.component_loader import VAELoader
-from fastvideo.v1.utils import maybe_download_model
-from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
-
-logger = init_logger(__name__)
-
-model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-path = maybe_download_model(model_path)
-# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-VAE_PATH = os.path.join(path, "vae")
-print(VAE_PATH)
-
-
-def main(args):
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    rank = int(os.getenv("RANK", 0))
-    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
-    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
-    print("world_size", world_size, "local rank", local_rank)
-    device = torch.device(f"cuda:{local_rank}")
-    torch.cuda.set_device(device)
-    world_group = get_world_group()
-
-    vae_precision = "fp16" 
-    fastvideo_args = FastVideoArgs(model_path=VAE_PATH,
-                         use_cpu_offload=False,
-                         vae_precision=vae_precision)
-    fastvideo_args.device = device
-    # fastvideo_args.dit_config = HunyuanVideoConfig()
-    fastvideo_args.vae_config = WanVAEConfig()
-
-    train_dataset = getdataset(args)
-    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-    )
-
-
-    # encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    # torch.cuda.set_device(local_rank)
-    # if not dist.is_initialized():
-    #     dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-    vae_loader = VAELoader()
-    vae = vae_loader.load(VAE_PATH, "vae", fastvideo_args)
-    # vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
-    # vae.enable_tiling()
-    os.makedirs(args.output_dir, exist_ok=True)
-    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
-
-    json_data = []
-    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
-        with torch.inference_mode():
-            with torch.autocast("cuda", dtype=torch.float16):
-                latents = vae.encode(data["pixel_values"].to(device)).sample()
-            for idx, video_path in enumerate(data["path"]):
-                video_name = os.path.basename(video_path).split(".")[0]
-                latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
-                torch.save(latents[idx].to(torch.bfloat16), latent_path)
-                item = {}
-                item["length"] = latents[idx].shape[1]
-                item["latent_path"] = video_name + ".pt"
-                item["caption"] = data["text"][idx]
-                json_data.append(item)
-                print(f"{video_name} processed")
-    world_group.barrier()
-    local_data = json_data
-    gathered_data = [None] * world_size
-    for i in range(world_size):
-        if local_rank == i:
-            world_group.broadcast_object(local_data, src=i)
-        else:
-            gathered_data[i] = world_group.broadcast_object(None, src=i)
-    gathered_data[local_rank] = json_data
-    print(gathered_data)
-    if local_rank == 0:
-        all_json_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f:
-            json.dump(all_json_data, f, indent=4)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # dataset & dataloader
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    # parser.add_argument("--model_type", type=str, default="mochi")
-    parser.add_argument("--data_merge_path", type=str, required=True)
-    parser.add_argument("--num_frames", type=int, default=163)
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=1,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--train_batch_size",
-        type=int,
-        default=16,
-        help="Batch size (per device) for the training dataloader.",
-    )
-    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--max_height", type=int, default=480)
-    parser.add_argument("--max_width", type=int, default=848)
-    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
-    parser.add_argument("--group_frame", action="store_true")  # TODO
-    parser.add_argument("--group_resolution", action="store_true")  # TODO
-    parser.add_argument("--dataset", default="t2v")
-    parser.add_argument("--train_fps", type=int, default=30)
-    parser.add_argument("--use_image_num", type=int, default=0)
-    parser.add_argument("--text_max_length", type=int, default=256)
-    parser.add_argument("--speed_factor", type=float, default=1.0)
-    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
-    # text encoder & vae & diffusion model
-    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument("--cfg", type=float, default=0.0)
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    parser.add_argument(
-        "--logging_dir",
-        type=str,
-        default="logs",
-        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
-    )
-
-    args = parser.parse_args()
-    main(args)

fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py

@@ -1,115 +0,0 @@
-import argparse
-import os
-
-import torch
-# import torch.distributed as dist
-from accelerate.logging import get_logger
-
-# from fastvideo.utils.load import load_text_encoder
-
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
-from fastvideo.v1.utils import maybe_download_model
-from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
-from fastvideo.v1.configs.models.encoders.t5 import T5Config
-
-logger = get_logger(__name__)
-
-
-def main(args):
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    rank = int(os.getenv("RANK", 0))
-    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
-    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
-    print("world_size", world_size, "local rank", local_rank)
-    device = torch.device(f"cuda:{local_rank}")
-    torch.cuda.set_device(device)
-    world_group = get_world_group()
-
-    vae_precision = "fp16" 
-    text_encoder_precision = "fp32"
-    fastvideo_args = FastVideoArgs(model_path=args.model_path,
-                         use_cpu_offload=False,
-                         vae_precision=vae_precision,
-                         text_encoder_precisions=(text_encoder_precision,))
-    fastvideo_args.device = device
-    fastvideo_args.device_str = f"cuda:{local_rank}"
-
-    # fastvideo_args.dit_config = HunyuanVideoConfig()
-    fastvideo_args.vae_config = WanVAEConfig()
-    fastvideo_args.text_encoder_configs = (T5Config(),)
-
-    # vae_loader = VAELoader()
-    # vae = vae_loader.load_vae()
-    text_encoder_loader = TextEncoderLoader()
-    tokenizer_loader = TokenizerLoader()
-
-    model_path = args.model_path
-    path = maybe_download_model(model_path)
-    # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-    ENCODER_PATH = os.path.join(path, "text_encoder")
-    TOKENIZER_PATH = os.path.join(path, "tokenizer")
-    print(ENCODER_PATH)
-    text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
-    tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
-
-    # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
-    # autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16
-    # output_dir/validation/prompt_attention_mask
-    # output_dir/validation/prompt_embed
-    os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True)
-    os.makedirs(
-        os.path.join(args.output_dir, "validation", "prompt_attention_mask"),
-        exist_ok=True,
-    )
-    os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True)
-
-    with open(args.validation_prompt_txt, "r", encoding="utf-8") as file:
-        lines = file.readlines()
-    prompts = [line.strip() for line in lines]
-    for prompt in prompts:
-        with torch.inference_mode():
-            # with torch.autocast("cuda", dtype=autocast_type):
-            text_inputs = tokenizer(prompt, **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
-                fastvideo_args.device)
-            input_ids = text_inputs["input_ids"]
-            attention_mask = text_inputs["attention_mask"]
-            outputs = text_encoder(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                output_hidden_states=True,
-            )
-            from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
-            post_process_func = t5_postprocess_text
-            prompt_embeds = post_process_func(outputs)
-            prompt_attention_mask = attention_mask
-
-            file_name = prompt.split(".")[0]
-            prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt")
-            prompt_attention_mask_path = os.path.join(
-                args.output_dir,
-                "validation",
-                "prompt_attention_mask",
-                f"{file_name}.pt",
-            )
-            torch.save(prompt_embeds[0], prompt_embed_path)
-            torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
-            print(f"sample {file_name} saved")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    # dataset & dataloader
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--validation_prompt_txt", type=str)
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default=None,
-        help="The output directory where the model predictions and checkpoints will be written.",
-    )
-    args = parser.parse_args()
-    main(args)

fastvideo/dataset/__init__.py

@@ -0,0 +1,97 @@
+from torchvision import transforms
+from torchvision.transforms import Lambda
+from transformers import AutoTokenizer
+
+from fastvideo.dataset.t2v_datasets import T2V_dataset
+from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
+
+
+def getdataset(args):
+    temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
+    norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
+    resize_topcrop = [
+        CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True),
+    ]
+    resize = [
+        CenterCropResizeVideo((args.max_height, args.max_width)),
+    ]
+    transform = transforms.Compose([
+        # Normalize255(),
+        *resize,
+    ])
+    transform_topcrop = transforms.Compose([
+        Normalize255(),
+        *resize_topcrop,
+        norm_fun,
+    ])
+    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+    if args.dataset == "t2v":
+        return T2V_dataset(
+            args,
+            transform=transform,
+            temporal_sample=temporal_sample,
+            tokenizer=tokenizer,
+            transform_topcrop=transform_topcrop,
+        )
+
+    raise NotImplementedError(args.dataset)
+
+
+if __name__ == "__main__":
+    import random
+
+    from accelerate import Accelerator
+    from tqdm import tqdm
+
+    from fastvideo.dataset.t2v_datasets import dataset_prog
+
+    args = type(
+        "args",
+        (),
+        {
+            "ae": "CausalVAEModel_4x8x8",
+            "dataset": "t2v",
+            "attention_mode": "xformers",
+            "use_rope": True,
+            "text_max_length": 300,
+            "max_height": 320,
+            "max_width": 240,
+            "num_frames": 1,
+            "use_image_num": 0,
+            "interpolation_scale_t": 1,
+            "interpolation_scale_h": 1,
+            "interpolation_scale_w": 1,
+            "cache_dir": "../cache_dir",
+            "image_data": "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
+            "video_data": "1",
+            "train_fps": 24,
+            "drop_short_ratio": 1.0,
+            "use_img_from_vid": False,
+            "speed_factor": 1.0,
+            "cfg": 0.1,
+            "text_encoder_name": "google/mt5-xxl",
+            "dataloader_num_workers": 10,
+        },
+    )
+    accelerator = Accelerator()
+    dataset = getdataset(args)
+    num = len(dataset_prog.img_cap_list)
+    zero = 0
+    for idx in tqdm(range(num)):
+        image_data = dataset_prog.img_cap_list[idx]
+        caps = [i["cap"] if isinstance(i["cap"], list) else [i["cap"]] for i in image_data]
+        try:
+            caps = [[random.choice(i)] for i in caps]
+        except Exception as e:
+            print(e)
+            # import ipdb;ipdb.set_trace()
+            print(image_data)
+            zero += 1
+            continue
+        assert caps[0] is not None and len(caps[0]) > 0
+    print(num, zero)
+    import ipdb
+
+    ipdb.set_trace()
+    print("end")

fastvideo/dataset/latent_datasets.py

@@ -0,0 +1,118 @@
+import json
+import os
+import random
+
+import torch
+from torch.utils.data import Dataset
+
+
+class LatentDataset(Dataset):
+
+    def __init__(
+        self,
+        json_path,
+        num_latent_t,
+        cfg_rate,
+    ):
+        # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
+        self.json_path = json_path
+        self.cfg_rate = cfg_rate
+        self.datase_dir_path = os.path.dirname(json_path)
+        self.video_dir = os.path.join(self.datase_dir_path, "video")
+        self.latent_dir = os.path.join(self.datase_dir_path, "latent")
+        self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed")
+        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
+        with open(self.json_path, "r") as f:
+            self.data_anno = json.load(f)
+        # json.load(f) already keeps the order
+        # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
+        self.num_latent_t = num_latent_t
+        # just zero embeddings [256, 4096]
+        self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
+        # 256 zeros
+        self.uncond_prompt_mask = torch.zeros(256).bool()
+        self.lengths = [data_item["length"] if "length" in data_item else 1 for data_item in self.data_anno]
+
+    def __getitem__(self, idx):
+        latent_file = self.data_anno[idx]["latent_path"]
+        prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
+        prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"]
+        # load
+        latent = torch.load(
+            os.path.join(self.latent_dir, latent_file),
+            map_location="cpu",
+            weights_only=True,
+        )
+        latent = latent.squeeze(0)[:, -self.num_latent_t:]
+        if random.random() < self.cfg_rate:
+            prompt_embed = self.uncond_prompt_embed
+            prompt_attention_mask = self.uncond_prompt_mask
+        else:
+            prompt_embed = torch.load(
+                os.path.join(self.prompt_embed_dir, prompt_embed_file),
+                map_location="cpu",
+                weights_only=True,
+            )
+            prompt_attention_mask = torch.load(
+                os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file),
+                map_location="cpu",
+                weights_only=True,
+            )
+        return latent, prompt_embed, prompt_attention_mask
+
+    def __len__(self):
+        return len(self.data_anno)
+
+
+def latent_collate_function(batch):
+    # return latent, prompt, latent_attn_mask, text_attn_mask
+    # latent_attn_mask: # b t h w
+    # text_attn_mask: b 1 l
+    # needs to check if the latent/prompt' size and apply padding & attn mask
+    latents, prompt_embeds, prompt_attention_masks = zip(*batch)
+    # calculate max shape
+    max_t = max([latent.shape[1] for latent in latents])
+    max_h = max([latent.shape[2] for latent in latents])
+    max_w = max([latent.shape[3] for latent in latents])
+
+    # padding
+    latents = [
+        torch.nn.functional.pad(
+            latent,
+            (
+                0,
+                max_t - latent.shape[1],
+                0,
+                max_h - latent.shape[2],
+                0,
+                max_w - latent.shape[3],
+            ),
+        ) for latent in latents
+    ]
+    # attn mask
+    latent_attn_mask = torch.ones(len(latents), max_t, max_h, max_w)
+    # set to 0 if padding
+    for i, latent in enumerate(latents):
+        latent_attn_mask[i, latent.shape[1]:, :, :] = 0
+        latent_attn_mask[i, :, latent.shape[2]:, :] = 0
+        latent_attn_mask[i, :, :, latent.shape[3]:] = 0
+
+    prompt_embeds = torch.stack(prompt_embeds, dim=0)
+    prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0)
+    latents = torch.stack(latents, dim=0)
+    return latents, prompt_embeds, latent_attn_mask, prompt_attention_masks
+
+
+if __name__ == "__main__":
+    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
+    for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
+        print(
+            latent.shape,
+            prompt_embed.shape,
+            latent_attn_mask.shape,
+            prompt_attention_mask.shape,
+        )
+        import pdb
+
+        pdb.set_trace()

fastvideo/dataset/t2v_datasets.py

@@ -0,0 +1,324 @@
+import json
+import math
+import os
+import random
+from collections import Counter
+from os.path import join as opj
+
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+from PIL import Image
+from torch.utils.data import Dataset
+
+from fastvideo.utils.dataset_utils import DecordInit
+from fastvideo.utils.logging_ import main_print
+
+
+class SingletonMeta(type):
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class DataSetProg(metaclass=SingletonMeta):
+
+    def __init__(self):
+        self.cap_list = []
+        self.elements = []
+        self.num_workers = 1
+        self.n_elements = 0
+        self.worker_elements = dict()
+        self.n_used_elements = dict()
+
+    def set_cap_list(self, num_workers, cap_list, n_elements):
+        self.num_workers = num_workers
+        self.cap_list = cap_list
+        self.n_elements = n_elements
+        self.elements = list(range(n_elements))
+        random.shuffle(self.elements)
+        print(f"n_elements: {len(self.elements)}", flush=True)
+
+        for i in range(self.num_workers):
+            self.n_used_elements[i] = 0
+            per_worker = int(math.ceil(len(self.elements) / float(self.num_workers)))
+            start = i * per_worker
+            end = min(start + per_worker, len(self.elements))
+            self.worker_elements[i] = self.elements[start:end]
+
+    def get_item(self, work_info):
+        if work_info is None:
+            worker_id = 0
+        else:
+            worker_id = work_info.id
+
+        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])]
+        self.n_used_elements[worker_id] += 1
+        return idx
+
+
+dataset_prog = DataSetProg()
+
+
+def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16):
+    if h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio:
+        return True
+    return False
+
+
+class T2V_dataset(Dataset):
+
+    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
+        self.data = args.data_merge_path
+        self.num_frames = args.num_frames
+        self.train_fps = args.train_fps
+        self.use_image_num = args.use_image_num
+        self.transform = transform
+        self.transform_topcrop = transform_topcrop
+        self.temporal_sample = temporal_sample
+        self.tokenizer = tokenizer
+        self.text_max_length = args.text_max_length
+        self.cfg = args.cfg
+        self.speed_factor = args.speed_factor
+        self.max_height = args.max_height
+        self.max_width = args.max_width
+        self.drop_short_ratio = args.drop_short_ratio
+        assert self.speed_factor >= 1
+        self.v_decoder = DecordInit()
+        self.video_length_tolerance_range = args.video_length_tolerance_range
+        self.support_Chinese = True
+        if "mt5" not in args.text_encoder_name:
+            self.support_Chinese = False
+
+        cap_list = self.get_cap_list()
+
+        assert len(cap_list) > 0
+        cap_list, self.sample_num_frames = self.define_frame_index(cap_list)
+        self.lengths = self.sample_num_frames
+
+        n_elements = len(cap_list)
+        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list, n_elements)
+
+        print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
+
+    def set_checkpoint(self, n_used_elements):
+        for i in range(len(dataset_prog.n_used_elements)):
+            dataset_prog.n_used_elements[i] = n_used_elements
+
+    def __len__(self):
+        return dataset_prog.n_elements
+
+    def __getitem__(self, idx):
+
+        data = self.get_data(idx)
+        return data
+
+    def get_data(self, idx):
+        path = dataset_prog.cap_list[idx]["path"]
+        if path.endswith(".mp4"):
+            return self.get_video(idx)
+        else:
+            return self.get_image(idx)
+
+    def get_video(self, idx):
+        video_path = dataset_prog.cap_list[idx]["path"]
+        assert os.path.exists(video_path), f"file {video_path} do not exist!"
+        frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
+        torchvision_video, _, metadata = torchvision.io.read_video(video_path, output_format="TCHW")
+        video = torchvision_video[frame_indices]
+        video = self.transform(video)
+        video = rearrange(video, "t c h w -> c t h w")
+        video = video.to(torch.uint8)
+        assert video.dtype == torch.uint8
+
+        h, w = video.shape[-2:]
+        assert (
+            h / w <= 17 / 16 and h / w >= 8 / 16
+        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
+
+        video = video.float() / 127.5 - 1.0
+
+        text = dataset_prog.cap_list[idx]["cap"]
+        if not isinstance(text, list):
+            text = [text]
+        text = [random.choice(text)]
+
+        text = text[0] if random.random() > self.cfg else ""
+        text_tokens_and_mask = self.tokenizer(
+            text,
+            max_length=self.text_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"]
+        cond_mask = text_tokens_and_mask["attention_mask"]
+        return dict(
+            pixel_values=video,
+            text=text,
+            input_ids=input_ids,
+            cond_mask=cond_mask,
+            path=video_path,
+        )
+
+    def get_image(self, idx):
+        image_data = dataset_prog.cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
+
+        image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
+        image = torch.from_numpy(np.array(image))  # [h, w, c]
+        image = rearrange(image, "h w c -> c h w").unsqueeze(0)  #  [1 c h w]
+        # for i in image:
+        #     h, w = i.shape[-2:]
+        #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
+
+        image = (self.transform_topcrop(image) if "human_images" in image_data["path"] else self.transform(image)
+                 )  #  [1 C H W] -> num_img [1 C H W]
+        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
+
+        image = image.float() / 127.5 - 1.0
+
+        caps = (image_data["cap"] if isinstance(image_data["cap"], list) else [image_data["cap"]])
+        caps = [random.choice(caps)]
+        text = caps
+        input_ids, cond_mask = [], []
+        text = text[0] if random.random() > self.cfg else ""
+        text_tokens_and_mask = self.tokenizer(
+            text,
+            max_length=self.text_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"]  # 1, l
+        cond_mask = text_tokens_and_mask["attention_mask"]  # 1, l
+        return dict(
+            pixel_values=image,
+            text=text,
+            input_ids=input_ids,
+            cond_mask=cond_mask,
+            path=image_data["path"],
+        )
+
+    def define_frame_index(self, cap_list):
+        new_cap_list = []
+        sample_num_frames = []
+        cnt_too_long = 0
+        cnt_too_short = 0
+        cnt_no_cap = 0
+        cnt_no_resolution = 0
+        cnt_resolution_mismatch = 0
+        cnt_movie = 0
+        cnt_img = 0
+        for i in cap_list:
+            path = i["path"]
+            cap = i.get("cap", None)
+            # ======no caption=====
+            if cap is None:
+                cnt_no_cap += 1
+                continue
+            if path.endswith(".mp4"):
+                # ======no fps and duration=====
+                duration = i.get("duration", None)
+                fps = i.get("fps", None)
+                if fps is None or duration is None:
+                    continue
+
+                # ======resolution mismatch=====
+                resolution = i.get("resolution", None)
+                if resolution is None:
+                    cnt_no_resolution += 1
+                    continue
+                else:
+                    if (resolution.get("height", None) is None or resolution.get("width", None) is None):
+                        cnt_no_resolution += 1
+                        continue
+                    height, width = i["resolution"]["height"], i["resolution"]["width"]
+                    aspect = self.max_height / self.max_width
+                    hw_aspect_thr = 1.5
+                    is_pick = filter_resolution(
+                        height,
+                        width,
+                        max_h_div_w_ratio=hw_aspect_thr * aspect,
+                        min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
+                    )
+                    if not is_pick:
+                        print("resolution mismatch")
+                        cnt_resolution_mismatch += 1
+                        continue
+
+                # import ipdb;ipdb.set_trace()
+                i["num_frames"] = math.ceil(fps * duration)
+                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
+                if i["num_frames"] / fps > self.video_length_tolerance_range * (
+                        self.num_frames / self.train_fps *
+                        self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
+                    cnt_too_long += 1
+                    continue
+
+                # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
+                frame_interval = fps / self.train_fps
+                start_frame_idx = 0
+                frame_indices = np.arange(start_frame_idx, i["num_frames"], frame_interval).astype(int)
+
+                # comment out it to enable dynamic frames training
+                if (len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio):
+                    cnt_too_short += 1
+                    continue
+
+                #  too long video will be temporal-crop randomly
+                if len(frame_indices) > self.num_frames:
+                    begin_index, end_index = self.temporal_sample(len(frame_indices))
+                    frame_indices = frame_indices[begin_index:end_index]
+                    # frame_indices = frame_indices[:self.num_frames]  # head crop
+                i["sample_frame_index"] = frame_indices.tolist()
+                new_cap_list.append(i)
+                i["sample_num_frames"] = len(i["sample_frame_index"])  # will use in dataloader(group sampler)
+                sample_num_frames.append(i["sample_num_frames"])
+            elif path.endswith(".jpg"):  # image
+                cnt_img += 1
+                new_cap_list.append(i)
+                i["sample_num_frames"] = 1
+                sample_num_frames.append(i["sample_num_frames"])
+            else:
+                raise NameError(
+                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image")
+        # import ipdb;ipdb.set_trace()
+        main_print(
+            f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
+            f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
+            f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
+            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}")
+        return new_cap_list, sample_num_frames
+
+    def decord_read(self, path, frame_indices):
+        decord_vr = self.v_decoder(path)
+        video_data = decord_vr.get_batch(frame_indices).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+    def read_jsons(self, data):
+        cap_lists = []
+        with open(data, "r") as f:
+            folder_anno = [i.strip().split(",") for i in f.readlines() if len(i.strip()) > 0]
+        print(folder_anno)
+        for folder, anno in folder_anno:
+            with open(anno, "r") as f:
+                sub_list = json.load(f)
+            for i in range(len(sub_list)):
+                sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
+            cap_lists += sub_list
+        return cap_lists
+
+    def get_cap_list(self):
+        cap_lists = self.read_jsons(self.data)
+        return cap_lists

fastvideo/dataset/transform.py

@@ -0,0 +1,608 @@
+import numbers
+import random
+
+import torch
+from PIL import Image
+
+
+def _is_tensor_video_clip(clip):
+    if not torch.is_tensor(clip):
+        raise TypeError("clip should be Tensor. Got %s" % type(clip))
+
+    if not clip.ndimension() == 4:
+        raise ValueError("clip should be 4D. Got %dD" % clip.dim())
+
+    return True
+
+
+def center_crop_arr(pil_image, image_size):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return Image.fromarray(arr[crop_y:crop_y + image_size, crop_x:crop_x + image_size])
+
+
+def crop(clip, i, j, h, w):
+    """
+    Args:
+        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+    """
+    if len(clip.size()) != 4:
+        raise ValueError("clip should be a 4D tensor")
+    return clip[..., i:i + h, j:j + w]
+
+
+def resize(clip, target_size, interpolation_mode):
+    if len(target_size) != 2:
+        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+    return torch.nn.functional.interpolate(
+        clip,
+        size=target_size,
+        mode=interpolation_mode,
+        align_corners=True,
+        antialias=True,
+    )
+
+
+def resize_scale(clip, target_size, interpolation_mode):
+    if len(target_size) != 2:
+        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+    H, W = clip.size(-2), clip.size(-1)
+    scale_ = target_size[0] / min(H, W)
+    return torch.nn.functional.interpolate(
+        clip,
+        scale_factor=scale_,
+        mode=interpolation_mode,
+        align_corners=True,
+        antialias=True,
+    )
+
+
+def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"):
+    """
+    Do spatial cropping and resizing to the video clip
+    Args:
+        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        i (int): i in (i,j) i.e coordinates of the upper left corner.
+        j (int): j in (i,j) i.e coordinates of the upper left corner.
+        h (int): Height of the cropped region.
+        w (int): Width of the cropped region.
+        size (tuple(int, int)): height and width of resized clip
+    Returns:
+        clip (torch.tensor): Resized and cropped clip. Size is (T, C, H, W)
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    clip = crop(clip, i, j, h, w)
+    clip = resize(clip, size, interpolation_mode)
+    return clip
+
+
+def center_crop(clip, crop_size):
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    h, w = clip.size(-2), clip.size(-1)
+    th, tw = crop_size
+    if h < th or w < tw:
+        raise ValueError("height and width must be no smaller than crop_size")
+
+    i = int(round((h - th) / 2.0))
+    j = int(round((w - tw) / 2.0))
+    return crop(clip, i, j, th, tw)
+
+
+def center_crop_using_short_edge(clip):
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    h, w = clip.size(-2), clip.size(-1)
+    if h < w:
+        th, tw = h, h
+        i = 0
+        j = int(round((w - tw) / 2.0))
+    else:
+        th, tw = w, w
+        i = int(round((h - th) / 2.0))
+        j = 0
+    return crop(clip, i, j, th, tw)
+
+
+def center_crop_th_tw(clip, th, tw, top_crop):
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+
+    # import ipdb;ipdb.set_trace()
+    h, w = clip.size(-2), clip.size(-1)
+    tr = th / tw
+    if h / w > tr:
+        new_h = int(w * tr)
+        new_w = w
+    else:
+        new_h = h
+        new_w = int(h / tr)
+
+    i = 0 if top_crop else int(round((h - new_h) / 2.0))
+    j = int(round((w - new_w) / 2.0))
+    return crop(clip, i, j, new_h, new_w)
+
+
+def random_shift_crop(clip):
+    """
+    Slide along the long edge, with the short edge as crop size
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    h, w = clip.size(-2), clip.size(-1)
+
+    if h <= w:
+        short_edge = h
+    else:
+        short_edge = w
+
+    th, tw = short_edge, short_edge
+
+    i = torch.randint(0, h - th + 1, size=(1, )).item()
+    j = torch.randint(0, w - tw + 1, size=(1, )).item()
+    return crop(clip, i, j, th, tw)
+
+
+def normalize_video(clip):
+    """
+    Convert tensor data type from uint8 to float, divide value by 255.0 and
+    permute the dimensions of clip tensor
+    Args:
+        clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
+    Return:
+        clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W)
+    """
+    _is_tensor_video_clip(clip)
+    if not clip.dtype == torch.uint8:
+        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
+    # return clip.float().permute(3, 0, 1, 2) / 255.0
+    return clip.float() / 255.0
+
+
+def normalize(clip, mean, std, inplace=False):
+    """
+    Args:
+        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
+        mean (tuple): pixel RGB mean. Size is (3)
+        std (tuple): pixel standard deviation. Size is (3)
+    Returns:
+        normalized clip (torch.tensor): Size is (T, C, H, W)
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    if not inplace:
+        clip = clip.clone()
+    mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device)
+    # print(mean)
+    std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device)
+    clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None])
+    return clip
+
+
+def hflip(clip):
+    """
+    Args:
+        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
+    Returns:
+        flipped clip (torch.tensor): Size is (T, C, H, W)
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    return clip.flip(-1)
+
+
+class RandomCropVideo:
+
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: randomly cropped video clip.
+                size is (T, C, OH, OW)
+        """
+        i, j, h, w = self.get_params(clip)
+        return crop(clip, i, j, h, w)
+
+    def get_params(self, clip):
+        h, w = clip.shape[-2:]
+        th, tw = self.size
+
+        if h < th or w < tw:
+            raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
+
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = torch.randint(0, h - th + 1, size=(1, )).item()
+        j = torch.randint(0, w - tw + 1, size=(1, )).item()
+
+        return i, j, th, tw
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+class SpatialStrideCropVideo:
+
+    def __init__(self, stride):
+        self.stride = stride
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: cropped video clip by stride.
+                size is (T, C, OH, OW)
+        """
+        i, j, h, w = self.get_params(clip)
+        return crop(clip, i, j, h, w)
+
+    def get_params(self, clip):
+        h, w = clip.shape[-2:]
+
+        th, tw = h // self.stride * self.stride, w // self.stride * self.stride
+
+        return 0, 0, th, tw  # from top-left
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+class LongSideResizeVideo:
+    """
+    First use the long side,
+    then resize to the specified size
+    """
+
+    def __init__(
+        self,
+        size,
+        skip_low_resolution=False,
+        interpolation_mode="bilinear",
+    ):
+        self.size = size
+        self.skip_low_resolution = skip_low_resolution
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: scale resized video clip.
+                size is (T, C, 512, *) or (T, C, *, 512)
+        """
+        _, _, h, w = clip.shape
+        if self.skip_low_resolution and max(h, w) <= self.size:
+            return clip
+        if h > w:
+            w = int(w * self.size / h)
+            h = self.size
+        else:
+            h = int(h * self.size / w)
+            w = self.size
+        resize_clip = resize(clip, target_size=(h, w), interpolation_mode=self.interpolation_mode)
+        return resize_clip
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class CenterCropResizeVideo:
+    """
+    First use the short side for cropping length,
+    center crop video, then resize to the specified size
+    """
+
+    def __init__(
+        self,
+        size,
+        top_crop=False,
+        interpolation_mode="bilinear",
+    ):
+        if len(size) != 2:
+            raise ValueError(f"size should be tuple (height, width), instead got {size}")
+        self.size = size
+        self.top_crop = top_crop
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: scale resized / center cropped video clip.
+                size is (T, C, crop_size, crop_size)
+        """
+        # clip_center_crop = center_crop_using_short_edge(clip)
+        clip_center_crop = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop)
+        # import ipdb;ipdb.set_trace()
+        clip_center_crop_resize = resize(
+            clip_center_crop,
+            target_size=self.size,
+            interpolation_mode=self.interpolation_mode,
+        )
+        return clip_center_crop_resize
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class UCFCenterCropVideo:
+    """
+    First scale to the specified size in equal proportion to the short edge,
+    then center cropping
+    """
+
+    def __init__(
+        self,
+        size,
+        interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: scale resized / center cropped video clip.
+                size is (T, C, crop_size, crop_size)
+        """
+        clip_resize = resize_scale(clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode)
+        clip_center_crop = center_crop(clip_resize, self.size)
+        return clip_center_crop
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class KineticsRandomCropResizeVideo:
+    """
+    Slide along the long edge, with the short edge as crop size. And resie to the desired size.
+    """
+
+    def __init__(
+        self,
+        size,
+        interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        clip_random_crop = random_shift_crop(clip)
+        clip_resize = resize(clip_random_crop, self.size, self.interpolation_mode)
+        return clip_resize
+
+
+class CenterCropVideo:
+
+    def __init__(
+        self,
+        size,
+        interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: center cropped video clip.
+                size is (T, C, crop_size, crop_size)
+        """
+        clip_center_crop = center_crop(clip, self.size)
+        return clip_center_crop
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class Normalize:
+    """
+    Normalize the video clip by mean subtraction and division by standard deviation
+    Args:
+        mean (3-tuple): pixel RGB mean
+        std (3-tuple): pixel RGB standard deviation
+        inplace (boolean): whether do in-place normalization
+    """
+
+    def __init__(self, mean, std, inplace=False):
+        self.mean = mean
+        self.std = std
+        self.inplace = inplace
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): video clip must be normalized. Size is (C, T, H, W)
+        """
+        return normalize(clip, self.mean, self.std, self.inplace)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})"
+
+
+class Normalize255:
+    """
+    Convert tensor data type from uint8 to float, divide value by 255.0 and
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
+        Return:
+            clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W)
+        """
+        return normalize_video(clip)
+
+    def __repr__(self) -> str:
+        return self.__class__.__name__
+
+
+class RandomHorizontalFlipVideo:
+    """
+    Flip the video clip along the horizontal direction with a given probability
+    Args:
+        p (float): probability of the clip being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Size is (T, C, H, W)
+        Return:
+            clip (torch.tensor): Size is (T, C, H, W)
+        """
+        if random.random() < self.p:
+            clip = hflip(clip)
+        return clip
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(p={self.p})"
+
+
+#  ------------------------------------------------------------
+#  ---------------------  Sampling  ---------------------------
+#  ------------------------------------------------------------
+class TemporalRandomCrop(object):
+    """Temporally crop the given frame indices at a random location.
+
+    Args:
+        size (int): Desired length of frames will be seen in the model.
+    """
+
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, total_frames):
+        rand_end = max(0, total_frames - self.size - 1)
+        begin_index = random.randint(0, rand_end)
+        end_index = min(begin_index + self.size, total_frames)
+        return begin_index, end_index
+
+
+class DynamicSampleDuration(object):
+    """Temporally crop the given frame indices at a random location.
+
+    Args:
+        size (int): Desired length of frames will be seen in the model.
+    """
+
+    def __init__(self, t_stride, extra_1):
+        self.t_stride = t_stride
+        self.extra_1 = extra_1
+
+    def __call__(self, t, h, w):
+        if self.extra_1:
+            t = t - 1
+        truncate_t_list = list(range(t + 1))[t // 2:][::self.t_stride]  # need half at least
+        truncate_t = random.choice(truncate_t_list)
+        if self.extra_1:
+            truncate_t = truncate_t + 1
+        return 0, truncate_t
+
+
+if __name__ == "__main__":
+    import os
+
+    import numpy as np
+    import torchvision.io as io
+    from torchvision import transforms
+    from torchvision.utils import save_image
+
+    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW")
+
+    trans = transforms.Compose([
+        Normalize255(),
+        RandomHorizontalFlipVideo(),
+        UCFCenterCropVideo(512),
+        # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+    ])
+
+    target_video_len = 32
+    frame_interval = 1
+    total_frames = len(vframes)
+    print(total_frames)
+
+    temporal_sample = TemporalRandomCrop(target_video_len * frame_interval)
+
+    # Sampling video frames
+    start_frame_ind, end_frame_ind = temporal_sample(total_frames)
+    # print(start_frame_ind)
+    # print(end_frame_ind)
+    assert end_frame_ind - start_frame_ind >= target_video_len
+    frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int)
+    print(frame_indice)
+
+    select_vframes = vframes[frame_indice]
+    print(select_vframes.shape)
+    print(select_vframes.dtype)
+
+    select_vframes_trans = trans(select_vframes)
+    print(select_vframes_trans.shape)
+    print(select_vframes_trans.dtype)
+
+    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(dtype=torch.uint8)
+    print(select_vframes_trans_int.dtype)
+    print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
+
+    io.write_video("./test.avi", select_vframes_trans_int.permute(0, 2, 3, 1), fps=8)
+
+    for i in range(target_video_len):
+        save_image(
+            select_vframes_trans[i],
+            os.path.join("./test000", "%04d.png" % i),
+            normalize=True,
+            value_range=(-1, 1),
+        )

fastvideo/distill_wan.py

@@ -1,870 +0,0 @@
-# !/bin/python3
-# isort: skip_file
-import argparse
-import math
-import os
-import time
-from collections import deque
-
-import torch
-import torch.distributed as dist
-import wandb
-from accelerate.utils import set_seed
-from diffusers.optimization import get_scheduler
-from diffusers.utils import check_min_version
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import ShardingStrategy
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from tqdm.auto import tqdm
-
-from fastvideo.dataset.latent_datasets import (LatentDataset,
-                                               latent_collate_function)
-from fastvideo.distill.solver import EulerSolver, extract_into_tensor
-from fastvideo.utils.checkpoint import (save_checkpoint, save_lora_checkpoint)
-from fastvideo.utils.communications import (broadcast,
-                                            sp_parallel_dataloader_wrapper)
-from fastvideo.utils.dataset_utils import LengthGroupedSampler
-from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group,
-                                             get_sequence_parallel_state)
-from fastvideo.utils.validation import log_validation
-
-from fastvideo.v1.utils import maybe_download_model
-from fastvideo.v1.configs.models.dits import WanVideoConfig
-from fastvideo.v1.models.loader.component_loader import TransformerLoader, SchedulerLoader
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
-from fastvideo.v1.forward_context import set_forward_context
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.31.0")
-
-logger = init_logger(__name__)
-
-BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join(
-                                      'data', BASE_MODEL_PATH))
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-SCHEDULER_PATH = os.path.join(MODEL_PATH, "scheduler")
-
-
-def reshard_fsdp(model):
-    for m in FSDP.fsdp_modules(model):
-        if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
-            torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
-
-
-def get_norm(model_pred, norms, gradient_accumulation_steps):
-    fro_norm = (
-        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
-        gradient_accumulation_steps)
-    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) /
-                              gradient_accumulation_steps)
-    absolute_mean = torch.mean(
-        torch.abs(model_pred)) / gradient_accumulation_steps
-    absolute_max = torch.max(
-        torch.abs(model_pred)) / gradient_accumulation_steps
-    dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
-    dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
-    dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
-    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
-    norms["largest singular value"] += torch.mean(largest_singular_value).item()
-    norms["absolute mean"] += absolute_mean.item()
-    norms["absolute max"] += absolute_max.item()
-
-
-def distill_one_step(
-    transformer,
-    model_type,
-    teacher_transformer,
-    ema_transformer,
-    optimizer,
-    lr_scheduler,
-    loader,
-    noise_scheduler,
-    solver,
-    noise_random_generator,
-    gradient_accumulation_steps,
-    sp_size,
-    max_grad_norm,
-    uncond_prompt_embed,
-    uncond_prompt_mask,
-    num_euler_timesteps,
-    multiphase,
-    not_apply_cfg_solver,
-    distill_cfg,
-    ema_decay,
-    pred_decay_weight,
-    pred_decay_type,
-    hunyuan_teacher_disable_cfg,
-):
-    total_loss = 0.0
-    optimizer.zero_grad()
-    model_pred_norm = {
-        "fro": 0.0,  # codespell:ignore
-        "largest singular value": 0.0,
-        "absolute mean": 0.0,
-        "absolute max": 0.0,
-    }
-    for _ in range(gradient_accumulation_steps):
-        (
-            latents,
-            encoder_hidden_states,
-            latents_attention_mask,
-            encoder_attention_mask,
-        ) = next(loader)
-        # model_input = normalize_dit_input(model_type, latents)
-        model_input = latents
-        noise = torch.randn_like(model_input)
-        bsz = model_input.shape[0]
-        index = torch.randint(0,
-                              num_euler_timesteps, (bsz, ),
-                              device=model_input.device).long()
-        if sp_size > 1:
-            broadcast(index)
-        # Add noise according to flow matching.
-        # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
-        sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
-        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index,
-                                          model_input.shape)
-
-        timesteps = (sigmas *
-                     noise_scheduler.config.num_train_timesteps).view(-1)
-        # if squeeze to [], unsqueeze to [1]
-
-        timesteps_prev = (sigmas_prev *
-                          noise_scheduler.config.num_train_timesteps).view(-1)
-        noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
-        noisy_model_input = noisy_model_input.to(torch.bfloat16)
-
-        forward_batch = ForwardBatch(data_type="video", enable_teacache=False)
-        # Predict the noise residual
-        with torch.autocast("cuda", dtype=torch.bfloat16):
-            teacher_kwargs = {
-                "hidden_states": noisy_model_input,
-                "encoder_hidden_states": encoder_hidden_states,
-                "timestep": timesteps,
-                "encoder_attention_mask": encoder_attention_mask,  # B, L
-                "return_dict": False,
-            }
-            if hunyuan_teacher_disable_cfg:
-                teacher_kwargs["guidance"] = torch.tensor(
-                    [1000.0],
-                    device=noisy_model_input.device,
-                    dtype=torch.bfloat16)
-            with set_forward_context(current_timestep=0,
-                                     attn_metadata=None,
-                                     forward_batch=forward_batch):
-                with torch.autograd.graph.save_on_cpu(pin_memory=True):
-                    model_pred = transformer(**teacher_kwargs)
-
-        # if accelerator.is_main_process:
-        model_pred, end_index = solver.euler_style_multiphase_pred(
-            noisy_model_input, model_pred, index, multiphase)
-        with torch.no_grad():
-            w = distill_cfg
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                with set_forward_context(current_timestep=0,
-                                         attn_metadata=None,
-                                         forward_batch=forward_batch):
-                    cond_teacher_output = teacher_transformer(
-                        noisy_model_input,
-                        encoder_hidden_states,
-                        timesteps,
-                        encoder_attention_mask,  # B, L
-                        return_dict=False,
-                    ).float()
-            if not_apply_cfg_solver:
-                uncond_teacher_output = cond_teacher_output
-            else:
-                # Get teacher model prediction on noisy_latents and unconditional embedding
-                with torch.autocast("cuda", dtype=torch.bfloat16):
-                    with set_forward_context(current_timestep=0,
-                                             attn_metadata=None,
-                                             forward_batch=forward_batch):
-                        uncond_teacher_output = teacher_transformer(
-                            noisy_model_input,
-                            uncond_prompt_embed.unsqueeze(0).expand(
-                                bsz, -1, -1),
-                            timesteps,
-                            uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
-                            return_dict=False,
-                        ).float()
-            teacher_output = uncond_teacher_output + w * (cond_teacher_output -
-                                                          uncond_teacher_output)
-            x_prev = solver.euler_step(noisy_model_input, teacher_output,
-                                       index).to(torch.bfloat16)
-
-        # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
-        with torch.no_grad():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                if ema_transformer is not None:
-                    target_pred = ema_transformer(
-                        x_prev.float(),
-                        encoder_hidden_states,
-                        timesteps_prev,
-                        encoder_attention_mask,  # B, L
-                        return_dict=False,
-                    )[0]
-                else:
-                    with set_forward_context(current_timestep=0,
-                                             attn_metadata=None,
-                                             forward_batch=forward_batch):
-                        with torch.autograd.graph.save_on_cpu(pin_memory=True):
-                            target_pred = transformer(
-                                x_prev,
-                                encoder_hidden_states,
-                                timesteps_prev,
-                                encoder_attention_mask,  # B, L
-                                return_dict=False,
-                            )
-
-            target, end_index = solver.euler_style_multiphase_pred(
-                x_prev, target_pred, index, multiphase, True)
-
-        huber_c = 0.001
-        # loss = loss.mean()
-        loss = (torch.mean(
-            torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) -
-            huber_c) / gradient_accumulation_steps)
-        if pred_decay_weight > 0:
-            if pred_decay_type == "l1":
-                pred_decay_loss = (
-                    torch.mean(torch.sqrt(model_pred.float()**2)) *
-                    pred_decay_weight / gradient_accumulation_steps)
-                loss += pred_decay_loss
-            elif pred_decay_type == "l2":
-                # essnetially k2?
-                pred_decay_loss = (torch.mean(model_pred.float()**2) *
-                                   pred_decay_weight /
-                                   gradient_accumulation_steps)
-                loss += pred_decay_loss
-            else:
-                assert NotImplementedError("pred_decay_type is not implemented")
-
-        # calculate model_pred norm and mean
-        get_norm(model_pred.detach().float(), model_pred_norm,
-                 gradient_accumulation_steps)
-        loss.backward()
-
-        avg_loss = loss.detach().clone()
-        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
-        total_loss += avg_loss.item()
-
-    # update ema
-    if ema_transformer is not None:
-        reshard_fsdp(ema_transformer)
-        for p_averaged, p_model in zip(ema_transformer.parameters(),
-                                       transformer.parameters()):
-            with torch.no_grad():
-                p_averaged.copy_(
-                    torch.lerp(p_averaged.detach(), p_model.detach(),
-                               1 - ema_decay))
-
-    # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
-    grad_norm = torch.nn.utils.clip_grad_norm_(transformer.parameters(),
-                                               max_norm=max_grad_norm)
-    optimizer.step()
-    lr_scheduler.step()
-
-    return total_loss, grad_norm.item(), model_pred_norm
-
-
-def main(args):
-    torch.backends.cuda.matmul.allow_tf32 = True
-
-    local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    rank = int(os.environ.get("RANK", 0))
-    world_size = int(os.environ.get("WORLD_SIZE", 1))
-    torch.cuda.set_device(rank)
-    init_distributed_environment(world_size=world_size,
-                                 rank=rank,
-                                 local_rank=local_rank)
-    initialize_model_parallel(tensor_model_parallel_size=args.sp_size,
-                              sequence_model_parallel_size=args.sp_size)
-    fastvideo_args = FastVideoArgs(
-        model_path=MODEL_PATH,
-        num_gpus=world_size,
-        use_cpu_offload=False,
-        precision=args.master_weight_type,
-        dit_config=WanVideoConfig(),
-        device_str="cuda",
-    )
-    fastvideo_args.check_fastvideo_args()
-
-    device_str = f"cuda:{rank}"
-    device = torch.device(device_str)
-    fastvideo_args.device = device
-
-    # If passed along, set the training seed now. On GPU...
-    if args.seed is not None:
-        # TODO: t within the same seq parallel group should be the same. Noise should be different.
-        set_seed(args.seed + rank)
-    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
-    noise_random_generator = None
-
-    # Handle the repository creation
-    if rank <= 0 and args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    # For mixed precision training we cast all non-trainable weights to half-precision
-    # as these weights are only used for inference, keeping weights in full precision is not required.
-
-    # Create model:
-
-    logger.info("--> loading model from %s", TRANSFORMER_PATH)
-
-    fastvideo_args.device = device
-    transformer_loader = TransformerLoader()
-    transformer = transformer_loader.load(TRANSFORMER_PATH, "", fastvideo_args)
-    transformer = transformer.train()
-    transformer.requires_grad_(True)
-
-    teacher_loader = TransformerLoader()
-    teacher_transformer = teacher_loader.load(TRANSFORMER_PATH, "",
-                                              fastvideo_args)
-    if args.use_ema:
-        ema_transformer = teacher_loader.load(TRANSFORMER_PATH, "",
-                                              fastvideo_args)
-    else:
-        ema_transformer = None
-
-    logger.info(
-        "  Total training parameters = %s M",
-        sum(p.numel()
-            for p in transformer.parameters() if p.requires_grad) / 1e6)
-    logger.info("--> model loaded")
-
-    teacher_transformer.requires_grad_(False)
-    if args.use_ema:
-        ema_transformer.requires_grad_(False)
-
-    # scheduler
-    noise_scheduler_loader = SchedulerLoader()
-    noise_scheduler = noise_scheduler_loader.load(SCHEDULER_PATH, "",
-                                                  fastvideo_args)
-    solver = EulerSolver(
-        noise_scheduler.sigmas.numpy()[::-1],
-        noise_scheduler.config.num_train_timesteps,
-        euler_timesteps=args.num_euler_timesteps,
-    )
-    solver.to(device)
-    params_to_optimize = transformer.parameters()
-    params_to_optimize = list(
-        filter(lambda p: p.requires_grad, params_to_optimize))
-
-    optimizer = torch.optim.AdamW(
-        params_to_optimize,
-        lr=args.learning_rate,
-        betas=(0.9, 0.999),
-        weight_decay=args.weight_decay,
-        eps=1e-8,
-    )
-
-    init_steps = 0
-    logger.info("optimizer: %s", optimizer)
-
-    # todo add lr scheduler
-    lr_scheduler = get_scheduler(
-        args.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * world_size,
-        num_training_steps=args.max_train_steps * world_size,
-        num_cycles=args.lr_num_cycles,
-        power=args.lr_power,
-        last_epoch=init_steps - 1,
-    )
-
-    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t,
-                                  args.cfg)
-    uncond_prompt_embed = train_dataset.uncond_prompt_embed
-    uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = (LengthGroupedSampler(
-        args.train_batch_size,
-        rank=rank,
-        world_size=world_size,
-        lengths=train_dataset.lengths,
-        group_frame=args.group_frame,
-        group_resolution=args.group_resolution,
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
-        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
-
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        collate_fn=latent_collate_function,
-        pin_memory=True,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-        drop_last=True,
-    )
-
-    num_update_steps_per_epoch = math.ceil(
-        len(train_dataloader) / args.gradient_accumulation_steps *
-        args.sp_size / args.train_sp_batch_size)
-    args.num_train_epochs = math.ceil(args.max_train_steps /
-                                      num_update_steps_per_epoch)
-
-    if rank <= 0:
-        project = args.tracker_project_name or "fastvideo"
-        wandb.init(project=project, config=args)
-
-    # Train!
-    total_batch_size = (world_size * args.gradient_accumulation_steps /
-                        args.sp_size * args.train_sp_batch_size)
-    logger.info("***** Running training *****")
-    logger.info("  Num examples = %s", len(train_dataset))
-    logger.info("  Dataloader size = %s", len(train_dataloader))
-    logger.info("  Num Epochs = %s", args.num_train_epochs)
-    logger.info("  Resume training from step %s", init_steps)
-    logger.info("  Instantaneous batch size per device = %s",
-                args.train_batch_size)
-    logger.info(
-        "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
-        total_batch_size)
-    logger.info("  Gradient Accumulation steps = %s",
-                args.gradient_accumulation_steps)
-    logger.info("  Total optimization steps = %s", args.max_train_steps)
-    logger.info(
-        "  Total training parameters per FSDP shard = %s B",
-        sum(p.numel()
-            for p in transformer.parameters() if p.requires_grad) / 1e9)
-    # print dtype
-    logger.info("  Master weight dtype: %s",
-                transformer.parameters().__next__().dtype)
-
-    # Potentially load in the weights and states from a previous save
-    if args.resume_from_checkpoint:
-        assert NotImplementedError(
-            "resume_from_checkpoint is not supported now.")
-        # TODO
-
-    progress_bar = tqdm(
-        range(0, args.max_train_steps),
-        initial=init_steps,
-        desc="Steps",
-        # Only show the progress bar once on each machine.
-        disable=local_rank > 0,
-    )
-
-    loader = sp_parallel_dataloader_wrapper(
-        train_dataloader,
-        device,
-        args.train_batch_size,
-        args.sp_size,
-        args.train_sp_batch_size,
-    )
-
-    step_times = deque(maxlen=100)
-
-    # todo future
-    for i in range(init_steps):
-        next(loader)
-
-    # log_validation(args, transformer, device,
-    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
-    def get_num_phases(multi_phased_distill_schedule, step):
-        # step-phase,step-phase
-        multi_phases = multi_phased_distill_schedule.split(",")
-        phase = multi_phases[-1].split("-")[-1]
-        for step_phases in multi_phases:
-            phase_step, phase = step_phases.split("-")
-            if step <= int(phase_step):
-                return int(phase)
-        return phase
-
-    for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.time()
-        assert args.multi_phased_distill_schedule is not None
-        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
-
-        loss, grad_norm, pred_norm = distill_one_step(
-            transformer,
-            args.model_type,
-            teacher_transformer,
-            ema_transformer,
-            optimizer,
-            lr_scheduler,
-            loader,
-            noise_scheduler,
-            solver,
-            noise_random_generator,
-            args.gradient_accumulation_steps,
-            args.sp_size,
-            args.max_grad_norm,
-            uncond_prompt_embed,
-            uncond_prompt_mask,
-            args.num_euler_timesteps,
-            num_phases,
-            args.not_apply_cfg_solver,
-            args.distill_cfg,
-            args.ema_decay,
-            args.pred_decay_weight,
-            args.pred_decay_type,
-            args.hunyuan_teacher_disable_cfg,
-        )
-
-        step_time = time.time() - start_time
-        step_times.append(step_time)
-        avg_step_time = sum(step_times) / len(step_times)
-
-        progress_bar.set_postfix({
-            "loss": f"{loss:.4f}",
-            "step_time": f"{step_time:.2f}s",
-            "grad_norm": grad_norm,
-            "phases": num_phases,
-        })
-        progress_bar.update(1)
-        if rank <= 0:
-            wandb.log(
-                {
-                    "train_loss":
-                    loss,
-                    "learning_rate":
-                    lr_scheduler.get_last_lr()[0],
-                    "step_time":
-                    step_time,
-                    "avg_step_time":
-                    avg_step_time,
-                    "grad_norm":
-                    grad_norm,
-                    "pred_fro_norm":
-                    pred_norm["fro"],  # codespell:ignore
-                    "pred_largest_singular_value":
-                    pred_norm["largest singular value"],
-                    "pred_absolute_mean":
-                    pred_norm["absolute mean"],
-                    "pred_absolute_max":
-                    pred_norm["absolute max"],
-                },
-                step=step,
-            )
-        if step % args.checkpointing_steps == 0:
-            if args.use_lora:
-                # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank,
-                                     args.output_dir, step)
-            else:
-                # Your existing checkpoint saving code
-                if args.use_ema:
-                    save_checkpoint(ema_transformer, rank, args.output_dir,
-                                    step)
-                else:
-                    save_checkpoint(transformer, rank, args.output_dir, step)
-            dist.barrier()
-        if args.log_validation and step % args.validation_steps == 0:
-            log_validation(
-                args,
-                transformer,
-                device,
-                torch.bfloat16,
-                step,
-                scheduler_type=args.scheduler_type,
-                shift=args.shift,
-                num_euler_timesteps=args.num_euler_timesteps,
-                linear_quadratic_threshold=args.linear_quadratic_threshold,
-                linear_range=args.linear_range,
-                ema=False,
-            )
-            if args.use_ema:
-                log_validation(
-                    args,
-                    ema_transformer,
-                    device,
-                    torch.bfloat16,
-                    step,
-                    scheduler_type=args.scheduler_type,
-                    shift=args.shift,
-                    num_euler_timesteps=args.num_euler_timesteps,
-                    linear_quadratic_threshold=args.linear_quadratic_threshold,
-                    linear_range=args.linear_range,
-                    ema=True,
-                )
-
-    if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir,
-                             args.max_train_steps)
-    else:
-        save_checkpoint(transformer, rank, args.output_dir,
-                        args.max_train_steps)
-
-    if get_sequence_parallel_state():
-        destroy_sequence_parallel_group()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument("--model_type",
-                        type=str,
-                        default="mochi",
-                        help="The type of model to train.")
-
-    # dataset & dataloader
-    parser.add_argument("--data_json_path", type=str, required=True)
-    parser.add_argument("--num_height", type=int, default=480)
-    parser.add_argument("--num_width", type=int, default=848)
-    parser.add_argument("--num_frames", type=int, default=163)
-    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("--num_latent_t",
-                        type=int,
-                        default=28,
-                        help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true")  # TODO
-    parser.add_argument("--group_resolution", action="store_true")  # TODO
-
-    # text encoder & vae & diffusion model
-    parser.add_argument("--pretrained_model_name_or_path", type=str)
-    parser.add_argument("--dit_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-
-    # diffusion setting
-    parser.add_argument("--ema_decay", type=float, default=0.95)
-    parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument("--cfg", type=float, default=0.1)
-
-    # validation & logs
-    parser.add_argument("--validation_prompt_dir", type=str)
-    parser.add_argument("--validation_sampling_steps", type=str, default="64")
-    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
-
-    parser.add_argument("--validation_steps", type=float, default=64)
-    parser.add_argument("--log_validation", action="store_true")
-    parser.add_argument("--tracker_project_name", type=str, default=None)
-    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("--shift", type=float, default=1.0)
-    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(
-        "--resume_from_lora_checkpoint",
-        type=str,
-        default=None,
-        help=
-        ("Whether training should be resumed from a previous lora 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***."),
-    )
-
-    # optimizer & scheduler & Training
-    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(
-        "--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=10,
-        help="Number of steps for the warmup in the lr scheduler.",
-    )
-    parser.add_argument("--max_grad_norm",
-                        default=1.0,
-                        type=float,
-                        help="Max gradient norm.")
-    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("--selective_checkpointing", type=float, default=1.0)
-    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(
-        "--use_cpu_offload",
-        action="store_true",
-        help=
-        "Whether to use CPU offload for param & gradient & optimizer states.",
-    )
-
-    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(
-        "--use_lora",
-        action="store_true",
-        default=False,
-        help="Whether to use LoRA for finetuning.",
-    )
-    parser.add_argument("--lora_alpha",
-                        type=int,
-                        default=256,
-                        help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank",
-                        type=int,
-                        default=128,
-                        help="LoRA rank parameter. ")
-    parser.add_argument("--fsdp_sharding_startegy", default="full")
-
-    # lr_scheduler
-    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("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument(
-        "--lr_num_cycles",
-        type=int,
-        default=1,
-        help="Number of cycles in the learning rate scheduler.",
-    )
-    parser.add_argument(
-        "--lr_power",
-        type=float,
-        default=1.0,
-        help="Power factor of the polynomial scheduler.",
-    )
-    parser.add_argument(
-        "--not_apply_cfg_solver",
-        action="store_true",
-        help="Whether to apply the cfg_solver.",
-    )
-    parser.add_argument("--distill_cfg",
-                        type=float,
-                        default=3.0,
-                        help="Distillation coefficient.")
-    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument("--scheduler_type",
-                        type=str,
-                        default="pcm",
-                        help="The scheduler type to use.")
-    parser.add_argument(
-        "--linear_quadratic_threshold",
-        type=float,
-        default=0.025,
-        help="Threshold for linear quadratic scheduler.",
-    )
-    parser.add_argument(
-        "--linear_range",
-        type=float,
-        default=0.5,
-        help="Range for linear quadratic scheduler.",
-    )
-    parser.add_argument("--weight_decay",
-                        type=float,
-                        default=0.001,
-                        help="Weight decay to apply.")
-    parser.add_argument("--use_ema",
-                        action="store_true",
-                        help="Whether to use EMA.")
-    parser.add_argument("--multi_phased_distill_schedule",
-                        type=str,
-                        default=None)
-    parser.add_argument("--pred_decay_weight", type=float, default=0.0)
-    parser.add_argument("--pred_decay_type", default="l1")
-    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
-    parser.add_argument(
-        "--master_weight_type",
-        type=str,
-        default="fp32",
-        help="Weight type to use - fp32 or bf16.",
-    )
-    args = parser.parse_args()
-    main(args)

fastvideo/models/hunyuan/idle_config.py

@@ -237,7 +237,7 @@ def add_inference_args(parser: argparse.ArgumentParser):
         type=str,
         default="540p",
         choices=["540p", "720p"],
-        help="Root path of all the models, including t2v models and extra models.",
+        help="The resolution of the model.",
     )
     group.add_argument(
         "--load-key",
@@ -361,7 +361,7 @@ def add_parallel_args(parser: argparse.ArgumentParser):
         "--ring-degree",
         type=int,
         default=1,
-        help="Ulysses degree.",
+        help="Ring degree.",
     )
 
     return parser

fastvideo/models/hunyuan/inference.py

@@ -17,7 +17,7 @@ from fastvideo.models.hunyuan.vae import load_vae
 from fastvideo.utils.parallel_states import nccl_info
 
 
-class Inference(object):
+class Inference:
 
     def __init__(
         self,

fastvideo/models/hunyuan/prompt_rewrite.py

@@ -41,7 +41,7 @@ def get_rewrite_prompt(ori_prompt, mode="Normal"):
     elif mode == "Master":
         prompt = master_mode_prompt.format(input=ori_prompt)
     else:
-        raise Exception("Only supports Normal and Normal", mode)
+        raise Exception("Only supports Normal and Master mode, but got {}".format(mode))
     return prompt
 
 

fastvideo/models/stepvideo/text_encoder/stepllm.py

@@ -267,25 +267,25 @@ class Step1Model(PreTrainedModel):
 class STEP1TextEncoder(torch.nn.Module):
 
     def __init__(self, model_dir, max_length=320):
-        super(STEP1TextEncoder, self).__init__()
+        super()
         self.max_length = max_length
         self.text_tokenizer = Wrapped_StepChatTokenizer(os.path.join(model_dir, 'step1_chat_tokenizer.model'))
         text_encoder = Step1Model.from_pretrained(model_dir)
         self.text_encoder = text_encoder.eval().to(torch.bfloat16)
 
     @torch.no_grad
+    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
     def forward(self, prompts, with_mask=True, max_length=None):
         self.device = next(self.text_encoder.parameters()).device
-        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
-            if type(prompts) is str:
-                prompts = [prompts]
+        if type(prompts) is str:
+            prompts = [prompts]
 
-            txt_tokens = self.text_tokenizer(prompts,
-                                             max_length=max_length or self.max_length,
-                                             padding="max_length",
-                                             truncation=True,
-                                             return_tensors="pt")
-            y = self.text_encoder(txt_tokens.input_ids.to(self.device),
+        txt_tokens = self.text_tokenizer(prompts,
+                                            max_length=max_length or self.max_length,
+                                            padding="max_length",
+                                            truncation=True,
+                                            return_tensors="pt")
+        y = self.text_encoder(txt_tokens.input_ids.to(self.device),
                                   attention_mask=txt_tokens.attention_mask.to(self.device) if with_mask else None)
-            y_mask = txt_tokens.attention_mask
+        y_mask = txt_tokens.attention_mask
         return y.transpose(0, 1), y_mask

fastvideo/utils/checkpoint.py

@@ -11,7 +11,6 @@ from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_
 from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
-import dataclasses
 
 from fastvideo.utils.logging_ import main_print
 
@@ -45,50 +44,13 @@ def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discrimi
         optimizer_path = os.path.join(save_dir, "optimizer.pt")
         torch.save(optim_state, optimizer_path)
     else:
-        weight_path = os.path.join(save_dstate_dictir, "discriminator_pytorch_model.safetensors")
+        weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
         optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt")
         torch.save(optim_state, optimizer_path)
     main_print(f"--> checkpoint saved at step {step}")
 
 
-def save_checkpoint_v1(transformer, rank, output_dir, step):
-    # from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-# from torch.distributed.fsdp import StateDictType, FullStateDictConfig
-
-    # Configure FSDP to save full state dict
-    FSDP.set_state_dict_type(
-        transformer,
-        state_dict_type=StateDictType.FULL_STATE_DICT,
-        state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-    )
-
-    # Now get the state dict
-    cpu_state = transformer.state_dict()
-
-    # Save it (only on rank 0 since we used rank0_only=True)
-    # if torch.distributed.get_rank() == 0:
-    #     torch.save(state_dict, "model_checkpoint.pt")
-    if rank <= 0:
-        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
-        os.makedirs(save_dir, exist_ok=True)
-        # save using safetensors
-        # weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
-        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt")
-        print(weight_path)
-        # save_file(cpu_state, weight_path)
-        torch.save(cpu_state, weight_path)
-        config_dict = transformer.hf_config
-        if "dtype" in config_dict:
-            del config_dict["dtype"]  # TODO
-        config_path = os.path.join(save_dir, "config.json")
-        # save dict as json
-        with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
-    main_print(f"--> checkpoint saved at step {step}")
-
-
-
 def save_checkpoint(transformer, rank, output_dir, step):
     main_print(f"--> saving checkpoint at step {step}")
     with FSDP.state_dict_type(

fastvideo/utils/collect_env.py

@@ -70,6 +70,8 @@ DEFAULT_CONDA_PATTERNS = {
     "optree",
     "nccl",
     "transformers",
+    "accelerate",
+    "peft",
     "zmq",
     "nvidia",
     "pynvml",
@@ -85,6 +87,8 @@ DEFAULT_PIP_PATTERNS = {
     "onnx",
     "nccl",
     "transformers",
+    "accelerate",
+    "peft",
     "zmq",
     "nvidia",
     "pynvml",

fastvideo/utils/env_utils.py

@@ -1,38 +0,0 @@
-import platform
-
-import accelerate
-import peft
-import torch
-import transformers
-from transformers.utils import is_torch_cuda_available, is_torch_npu_available
-
-VERSION = "1.2.0"
-
-if __name__ == "__main__":
-    info = {
-        "FastVideo version": VERSION,
-        "Platform": platform.platform(),
-        "Python version": platform.python_version(),
-        "PyTorch version": torch.__version__,
-        "Transformers version": transformers.__version__,
-        "Accelerate version": accelerate.__version__,
-        "PEFT version": peft.__version__,
-    }
-
-    if is_torch_cuda_available():
-        info["PyTorch version"] += " (GPU)"
-        info["GPU type"] = torch.cuda.get_device_name()
-
-    if is_torch_npu_available():
-        info["PyTorch version"] += " (NPU)"
-        info["NPU type"] = torch.npu.get_device_name()
-        info["CANN version"] = torch.version.cann  # codespell:ignore
-
-    try:
-        import bitsandbytes
-
-        info["Bitsandbytes version"] = bitsandbytes.__version__
-    except Exception:
-        pass
-
-    print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n")

fastvideo/v1/dataset/__init__.py

@@ -7,7 +7,7 @@ from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                             TemporalRandomCrop)
 
 
-def getdataset(args, start_idx=0):
+def getdataset(args, start_idx=0) -> T2V_dataset:
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
     norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
     resize_topcrop = [
@@ -37,66 +37,3 @@ def getdataset(args, start_idx=0):
                            start_idx=start_idx)
 
     raise NotImplementedError(args.dataset)
-
-
-if __name__ == "__main__":
-    import random
-
-    from accelerate import Accelerator
-    from tqdm import tqdm
-
-    from fastvideo.v1.dataset.t2v_datasets import dataset_prog
-
-    args = type(
-        "args",
-        (),
-        {
-            "ae": "CausalVAEModel_4x8x8",
-            "dataset": "t2v",
-            "attention_mode": "xformers",
-            "use_rope": True,
-            "text_max_length": 300,
-            "max_height": 320,
-            "max_width": 240,
-            "num_frames": 1,
-            "use_image_num": 0,
-            "interpolation_scale_t": 1,
-            "interpolation_scale_h": 1,
-            "interpolation_scale_w": 1,
-            "cache_dir": "../cache_dir",
-            "image_data":
-            "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
-            "video_data": "1",
-            "train_fps": 24,
-            "drop_short_ratio": 1.0,
-            "use_img_from_vid": False,
-            "speed_factor": 1.0,
-            "cfg": 0.1,
-            "text_encoder_name": "google/mt5-xxl",
-            "dataloader_num_workers": 10,
-        },
-    )
-    accelerator = Accelerator()
-    dataset = getdataset(args)
-    num = len(dataset_prog.img_cap_list)
-    zero = 0
-    for idx in tqdm(range(num)):
-        image_data = dataset_prog.img_cap_list[idx]
-        caps = [
-            i["cap"] if isinstance(i["cap"], list) else [i["cap"]]
-            for i in image_data
-        ]
-        try:
-            caps = [[random.choice(i)] for i in caps]
-        except Exception as e:
-            print(e)
-            # import ipdb;ipdb.set_trace()
-            print(image_data)
-            zero += 1
-            continue
-        assert caps[0] is not None and len(caps[0]) > 0
-    print(num, zero)
-    import ipdb
-
-    ipdb.set_trace()
-    print("end")

fastvideo/v1/dataset/latent_datasets.py

@@ -13,7 +13,7 @@ class LatentDataset(Dataset):
         json_path,
         num_latent_t,
         cfg_rate,
-    ):
+    ) -> None:
         # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
         self.json_path = json_path
         self.cfg_rate = cfg_rate
@@ -29,13 +29,12 @@ class LatentDataset(Dataset):
         # json.load(f) already keeps the order
         # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
         self.num_latent_t = num_latent_t
-        # just zero embeddings [256, 4096]
+
         self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
-        # 256 zeros
+
         self.uncond_prompt_mask = torch.zeros(256).bool()
         self.lengths = [
-            data_item["length"] if "length" in data_item else 1
-            for data_item in self.data_anno
+            data_item.get("length", 1) for data_item in self.data_anno
         ]
 
     def __getitem__(self, idx):
@@ -83,7 +82,7 @@ def latent_collate_function(batch):
     max_w = max([latent.shape[3] for latent in latents])
 
     # padding
-    latents = [
+    latent_list: list[torch.Tensor] = [
         torch.nn.functional.pad(
             latent,
             (
@@ -97,22 +96,23 @@ def latent_collate_function(batch):
         ) for latent in latents
     ]
     # attn mask
-    latent_attn_mask = torch.ones(len(latents), max_t, max_h, max_w)
+    latent_attn_mask = torch.ones(len(latent_list), max_t, max_h, max_w)
     # set to 0 if padding
-    for i, latent in enumerate(latents):
+    for i, latent in enumerate(latent_list):
         latent_attn_mask[i, latent.shape[1]:, :, :] = 0
         latent_attn_mask[i, :, latent.shape[2]:, :] = 0
         latent_attn_mask[i, :, :, latent.shape[3]:] = 0
 
     prompt_embeds = torch.stack(prompt_embeds, dim=0)
     prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0)
-    latents = torch.stack(latents, dim=0)
+    latents = torch.stack(latent_list, dim=0)
     return latents, prompt_embeds, latent_attn_mask, prompt_attention_masks
 
 
 if __name__ == "__main__":
     dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt",
-                            num_latent_t=28)
+                            num_latent_t=28,
+                            cfg_rate=0.0)
     dataloader = torch.utils.data.DataLoader(dataset,
                                              batch_size=2,
                                              shuffle=False,

fastvideo/v1/dataset/parquet_datasets.py

@@ -1,21 +1,28 @@
 import argparse
+import json
 import os
 import random
 import time
+from collections import defaultdict
+from typing import Any, Dict, List
 
 import numpy as np
 import pyarrow.parquet as pq
 import torch
 import tqdm
+from einops import rearrange
 from torch import distributed as dist
-from torch.utils.data import IterableDataset, get_worker_info
+from torch.utils.data import Dataset
 from torchdata.stateful_dataloader import StatefulDataLoader
 
-# Path to your dataset
-dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/train/"
+from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
+                                      get_sp_group)
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
 
 
-class ParquetVideoTextDataset(IterableDataset):
+class ParquetVideoTextDataset(Dataset):
     """Efficient loader for video-text data from a directory of Parquet files."""
 
     def __init__(self,
@@ -24,237 +31,184 @@ class ParquetVideoTextDataset(IterableDataset):
                  rank: int = 0,
                  world_size: int = 1,
                  cfg_rate: float = 0.0,
-                 num_latent_t: int = 2):
+                 num_latent_t: int = 2,
+                 seed: int = 0):
         super().__init__()
         self.path = str(path)
         self.batch_size = batch_size
         self.rank = rank
-        self.world_size = world_size
+        self.local_rank = get_sequence_model_parallel_rank()
+        self.sp_world_size = world_size
+        self.world_size = int(os.getenv("WORLD_SIZE", 1))
         self.cfg_rate = cfg_rate
         self.num_latent_t = num_latent_t
+        self.local_indices = None
+        self.plan_output_dir = os.path.join(self.path, "data_plan.json")
 
-        # Find all parquet files recursively
-        print(f"Scanning for parquet files in {self.path}")
-        self.parquet_files = []
-        for root, _, files in os.walk(self.path):
-            for file in files:
-                if file.endswith('.parquet'):
-                    self.parquet_files.append(os.path.join(root, file))
-        # Sort files for consistent ordering
-        self.parquet_files.sort()
+        ranks = get_sp_group().ranks
+        group_ranks: List[List] = [[] for _ in range(self.world_size)]
+        torch.distributed.all_gather_object(group_ranks, ranks)
 
-        # Distribute files among workers
-        # drop last unenven files
-        print(f"Total files: {len(self.parquet_files)}")
-        total_files = len(self.parquet_files)
-        base_count = total_files // world_size
-        extra_files = total_files % world_size
+        if rank == 0:
+            # If a plan already exists, then skip creating a new plan
+            # This will be useful when resume training
+            if os.path.exists(self.plan_output_dir):
+                print(f"Using existing plan from {self.plan_output_dir}")
+                return
 
-        if rank < extra_files:
-            start_idx = rank * (base_count + 1)
-            end_idx = start_idx + base_count + 1
-        else:
-            start_idx = rank * base_count + extra_files
-            end_idx = start_idx + base_count
+            # Find all parquet files recursively, and record num_rows for each file
+            print(f"Scanning for parquet files in {self.path}")
+            metadatas = []
+            for root, _, files in os.walk(self.path):
+                for file in sorted(files):
+                    if file.endswith('.parquet'):
+                        file_path = os.path.join(root, file)
+                        num_rows = pq.ParquetFile(file_path).metadata.num_rows
+                        for row_idx in range(num_rows):
+                            metadatas.append((file_path, row_idx))
 
-        self.parquet_files = self.parquet_files[start_idx:end_idx]
+            # Generate the plan that distribute rows among workers
+            random.seed(seed)
+            random.shuffle(metadatas)
 
-        print(f"Files assigned to rank {rank}: {len(self.parquet_files)}")
-        if len(self.parquet_files) > 0:
-            print(f"First file: {self.parquet_files[0]}")
-            print(f"Last file: {self.parquet_files[-1]}")
+            # Get all sp groups
+            # e.g. if num_gpus = 4, sp_size = 2
+            # group_ranks = [(0, 1), (2, 3)]
+            # We will assign the same batches of data to ranks in the same sp group, and we'll assign different batches to ranks in different sp groups
+            # e.g. plan = {0: [row 1, row 4], 1: [row 1, row 4], 2: [row 2, row 3], 3: [row 2, row 3]}
+            group_ranks_list: List[Any] = list(
+                set(tuple(r) for r in group_ranks))
+            num_sp_groups = len(group_ranks_list)
+            plan = defaultdict(list)
+            for idx, metadata in enumerate(metadatas):
+                sp_group_idx = idx % num_sp_groups
+                for global_rank in group_ranks_list[sp_group_idx]:
+                    plan[global_rank].append(metadata)
 
-        # Initialize current file index
-        self.current_file_idx = 0
-        self.current_reader = None
-        self.current_batches = None
-        self.total_samples = 0
+            with open(self.plan_output_dir, "w") as f:
+                json.dump(plan, f)
 
-    def _open_next_file(self):
-        """Open the next parquet file for reading."""
-        num_workers = get_worker_info().num_workers
-        worker_id = get_worker_info().id
-        total_files = len(self.parquet_files)
-        base_count = total_files // num_workers
-        extra_files = total_files % num_workers
-
-        if worker_id < extra_files:
-            start_idx = worker_id * (base_count + 1)
-            end_idx = start_idx + base_count + 1
-        else:
-            start_idx = worker_id * base_count + extra_files
-            end_idx = start_idx + base_count
-
-        worker_parquet_files = self.parquet_files[start_idx:end_idx]
-        if self.current_file_idx >= len(worker_parquet_files):
-            print(
-                f"Rank {self.rank}, Worker {worker_id}: No more files to open (current_idx={self.current_file_idx}, total_files={len(worker_parquet_files)})"
-            )
-            return False
-
-        if self.current_reader is not None:
-            self.current_reader.close()
-
-        file_path = worker_parquet_files[self.current_file_idx]
-        print(
-            f"Rank {self.rank}, Worker {worker_id}: Opening file {self.current_file_idx + 1}/{len(worker_parquet_files)}: {file_path}"
-        )
-
-        try:
-            self.current_reader = pq.ParquetFile(file_path)
-            self.current_batches = self.current_reader.iter_batches(
-                batch_size=self.batch_size)
-            self.current_file_idx += 1
-            return True
-        except Exception as e:
-            print(f"Error opening file {file_path}: {str(e)}")
-            return False
-
-    def __iter__(self):
-        """Iterate over the dataset in a streaming fashion."""
-        print(f"Rank {self.rank}: Starting iteration")
-
-        # First try to open a file
-        if not self._open_next_file():
-            print(f"Rank {self.rank}: Failed to open first file")
-            return
-
-        while True:
+    def __len__(self):
+        if self.local_indices is None:
             try:
-                # Get next batch from current file
-                batch = next(self.current_batches)
-                batch_dict = batch.to_pydict()
-                processed = self._process_batch(batch_dict)
+                with open(self.plan_output_dir) as f:
+                    plan = json.load(f)
+                self.local_indices = plan[str(self.rank)]
+            except Exception as err:
+                raise Exception(
+                    "The data plan hasn't been created yet") from err
+        assert self.local_indices is not None
+        return len(self.local_indices)
 
-                # Update sample count
-                batch_size = len(processed["latents"])
-                self.total_samples += batch_size
+    def __getitem__(self, idx):
+        if self.local_indices is None:
+            try:
+                with open(self.plan_output_dir) as f:
+                    plan = json.load(f)
+                self.local_indices = plan[self.rank]
+            except Exception as err:
+                raise Exception(
+                    "The data plan hasn't been created yet") from err
+        assert self.local_indices is not None
+        file_path, row_idx = self.local_indices[idx]
+        parquet_file = pq.ParquetFile(file_path)
 
-                # Print progress
-                if self.total_samples % 1000 == 0:
-                    print(
-                        f"Rank {self.rank}: Processed {self.total_samples} samples"
-                    )
+        # Calculate the row group to read into memory and the local idx
+        # This way we can avoid reading in the entire parquet file
+        cumulative = 0
+        for i in range(parquet_file.num_row_groups):
+            num_rows = parquet_file.metadata.row_group(i).num_rows
+            if cumulative + num_rows > idx:
+                row_group_index = i
+                local_index = idx - cumulative
+                break
+            cumulative += num_rows
 
-                # Yield each item in the batch
-                for lat, emb, mask, info in zip(processed["latents"],
-                                                processed["embeddings"],
-                                                processed["masks"],
-                                                processed["info"]):
-                    if lat.numel() == 0:  # Split is validation
-                        yield lat, emb, mask, info
-                    else:
-                        yield lat[:, -self.num_latent_t:], emb, mask, info
+        row_group = parquet_file.read_row_group(row_group_index).to_pydict()
+        row_dict = {k: v[local_index] for k, v in row_group.items()}
+        del row_group
 
-            except StopIteration:
-                # Current file is exhausted, try next file
-                print(
-                    f"Rank {self.rank}: Current file exhausted, trying next file"
-                )
-                self.current_batches = None
-                if not self._open_next_file():
-                    print(
-                        f"Rank {self.rank}: No more files to process. Total samples: {self.total_samples}"
-                    )
-                    break
-            except Exception as e:
-                print(f"Error processing batch: {str(e)}")
-                self.current_batches = None
-                if not self._open_next_file():
-                    print(
-                        f"Rank {self.rank}: Failed to open next file after error"
-                    )
-                    break
+        processed = self._process_row(row_dict)
+        lat, emb, mask, info = processed["latents"], processed[
+            "embeddings"], processed["masks"], processed["info"]
+        if lat.numel() == 0:  # Validation parquet
+            return lat, emb, mask, info
+        else:
+            lat = lat[:, -self.num_latent_t:]
+            if self.sp_world_size > 1:
+                lat = rearrange(lat,
+                                "t (n s) h w -> t n s h w",
+                                n=self.sp_world_size).contiguous()
+                lat = lat[:, self.local_rank, :, :, :]
+            return lat, emb, mask, info
 
-        # Clean up
-        if self.current_reader is not None:
-            self.current_reader.close()
-
-    def _process_batch(self, batch):
+    def _process_row(self, row) -> Dict[str, Any]:
         """Process a PyArrow batch into tensors."""
-        out = {"lat": [], "emb": [], "msk": [], "info": []}
 
-        for i in range(len(batch["vae_latent_bytes"])):
-            vae_latent_bytes = batch["vae_latent_bytes"][i]
-            vae_latent_shape = batch["vae_latent_shape"][i]
-            text_embedding_bytes = batch["text_embedding_bytes"][i]
-            text_embedding_shape = batch["text_embedding_shape"][i]
-            text_attention_mask_bytes = batch["text_attention_mask_bytes"][i]
-            text_attention_mask_shape = batch["text_attention_mask_shape"][i]
+        vae_latent_bytes = row["vae_latent_bytes"]
+        vae_latent_shape = row["vae_latent_shape"]
+        text_embedding_bytes = row["text_embedding_bytes"]
+        text_embedding_shape = row["text_embedding_shape"]
+        text_attention_mask_bytes = row["text_attention_mask_bytes"]
+        text_attention_mask_shape = row["text_attention_mask_shape"]
 
-            # Process latent
-            if not vae_latent_shape:  # No VAE latent is stored. Split is validation
-                lat = np.array([])
-            else:
-                lat = np.frombuffer(vae_latent_bytes,
-                                    dtype=np.float32).reshape(vae_latent_shape)
-                # Make array writable
-                lat = np.copy(lat)
+        # Process latent
+        if not vae_latent_shape:  # No VAE latent is stored. Split is validation
+            lat = np.array([])
+        else:
+            lat = np.frombuffer(vae_latent_bytes,
+                                dtype=np.float32).reshape(vae_latent_shape)
+            # Make array writable
+            lat = np.copy(lat)
 
-            if random.random() < self.cfg_rate:
-                emb = np.zeros((512, 4096), dtype=np.float32)
-            else:
-                emb = np.frombuffer(
-                    text_embedding_bytes,
-                    dtype=np.float32).reshape(text_embedding_shape)
-                # Make array writable
-                emb = np.copy(emb)
-            if emb.shape[0] < 512:
-                padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
-                padded_emb[:emb.shape[0], :] = emb
-                emb = padded_emb
-            elif emb.shape[0] > 512:
-                emb = emb[:512, :]
+        if random.random() < self.cfg_rate:
+            emb = np.zeros((512, 4096), dtype=np.float32)
+        else:
+            emb = np.frombuffer(text_embedding_bytes,
+                                dtype=np.float32).reshape(text_embedding_shape)
+            # Make array writable
+            emb = np.copy(emb)
+        if emb.shape[0] < 512:
+            padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
+            padded_emb[:emb.shape[0], :] = emb
+            emb = padded_emb
+        elif emb.shape[0] > 512:
+            emb = emb[:512, :]
 
-            # Process mask
-            if len(text_attention_mask_bytes) > 0 and len(
-                    text_attention_mask_shape) > 0:
-                msk = np.frombuffer(text_attention_mask_bytes,
-                                    dtype=np.uint8).astype(np.bool_)
-                msk = msk.reshape(1, -1)
-                # Make array writable
-                msk = np.copy(msk)
-                if msk.shape[1] < 512:
-                    padded_msk = np.zeros((1, 512), dtype=np.bool_)
-                    padded_msk[:, :msk.shape[1]] = msk
-                    msk = padded_msk
-                elif msk.shape[1] > 512:
-                    msk = msk[:, :512]
-            else:
-                msk = np.ones((1, 512), dtype=np.bool_)
-            # to string
-            file_name = str(batch["file_name"][i])
-            # Collect metadata
-            info = {
-                "width": batch["width"][i],
-                "height": batch["height"][i],
-                "num_frames": batch["num_frames"][i],
-                "duration_sec": batch["duration_sec"][i],
-                "fps": batch["fps"][i],
-                "file_name": batch["file_name"][i],
-                "caption": batch["caption"][i],
-            }
+        # Process mask
+        if len(text_attention_mask_bytes) > 0 and len(
+                text_attention_mask_shape) > 0:
+            msk = np.frombuffer(text_attention_mask_bytes,
+                                dtype=np.uint8).astype(np.bool_)
+            msk = msk.reshape(1, -1)
+            # Make array writable
+            msk = np.copy(msk)
+            if msk.shape[1] < 512:
+                padded_msk = np.zeros((1, 512), dtype=np.bool_)
+                padded_msk[:, :msk.shape[1]] = msk
+                msk = padded_msk
+            elif msk.shape[1] > 512:
+                msk = msk[:, :512]
+        else:
+            msk = np.ones((1, 512), dtype=np.bool_)
 
-            out["lat"].append(torch.from_numpy(lat))
-            out["emb"].append(torch.from_numpy(emb))
-            out["msk"].append(torch.from_numpy(msk))
-            out["info"].append(info)
-
-        return {
-            "latents": torch.stack(out["lat"]) if out["lat"] else None,
-            "embeddings": torch.stack(out["emb"]) if out["emb"] else None,
-            "masks": torch.stack(out["msk"]) if out["msk"] else None,
-            "info": out["info"]
+        # Collect metadata
+        info = {
+            "width": row["width"],
+            "height": row["height"],
+            "num_frames": row["num_frames"],
+            "duration_sec": row["duration_sec"],
+            "fps": row["fps"],
+            "file_name": row["file_name"],
+            "caption": row["caption"],
         }
 
-
-def bind_cpu_cores(local_rank, cpu_per_process=16):
-    """根据local_rank绑定固定cpu核。"""
-    start = local_rank * cpu_per_process
-    end = start + cpu_per_process
-    cores = list(range(start, end))
-    print(f"[Rank {local_rank}] Binding to CPU cores: {cores}")
-    os.sched_setaffinity(0, cores)
+        return {
+            "latents": torch.from_numpy(lat),
+            "embeddings": torch.from_numpy(emb),
+            "masks": torch.from_numpy(msk),
+            "info": info
+        }
 
 
 if __name__ == "__main__":
@@ -262,7 +216,7 @@ if __name__ == "__main__":
         description='Benchmark Parquet dataset loading speed')
     parser.add_argument('--path',
                         type=str,
-                        default=dataset_path,
+                        default="your/dataset/path",
                         help='Path to Parquet dataset')
     parser.add_argument('--batch_size',
                         type=int,
@@ -297,9 +251,6 @@ if __name__ == "__main__":
             f"Initialized process: rank={rank}, local_rank={local_rank}, world_size={world_size}, device={device}"
         )
 
-    # Bind CPU cores after distributed initialization
-    # bind_cpu_cores(local_rank, cpu_per_process=16)
-
     # Create dataset
     dataset = ParquetVideoTextDataset(
         args.path,

fastvideo/v1/dataset/t2v_datasets.py

@@ -17,7 +17,7 @@ from fastvideo.utils.logging_ import main_print
 
 
 class SingletonMeta(type):
-    _instances = {}
+    _instances: dict[type, 'SingletonMeta'] = {}
 
     def __call__(cls, *args, **kwargs):
         if cls not in cls._instances:
@@ -28,15 +28,15 @@ class SingletonMeta(type):
 
 class DataSetProg(metaclass=SingletonMeta):
 
-    def __init__(self):
-        self.cap_list = []
-        self.elements = []
+    def __init__(self) -> None:
+        self.cap_list: list[dict] = []
+        self.elements: list[int] = []
         self.num_workers = 1
         self.n_elements = 0
-        self.worker_elements = dict()
-        self.n_used_elements = dict()
+        self.worker_elements: dict[int, list[int]] = {}
+        self.n_used_elements: dict[int, int] = {}
 
-    def set_cap_list(self, num_workers, cap_list, n_elements):
+    def set_cap_list(self, num_workers, cap_list, n_elements) -> None:
         self.num_workers = num_workers
         self.cap_list = cap_list
         self.n_elements = n_elements
@@ -52,11 +52,8 @@ class DataSetProg(metaclass=SingletonMeta):
             end = min(start + per_worker, len(self.elements))
             self.worker_elements[i] = self.elements[start:end]
 
-    def get_item(self, work_info):
-        if work_info is None:
-            worker_id = 0
-        else:
-            worker_id = work_info.id
+    def get_item(self, work_info) -> int:
+        worker_id = 0 if work_info is None else work_info.id
 
         idx = self.worker_elements[worker_id][
             self.n_used_elements[worker_id] %
@@ -68,13 +65,11 @@ class DataSetProg(metaclass=SingletonMeta):
 dataset_prog = DataSetProg()
 
 
-def filter_resolution(h,
-                      w,
-                      max_h_div_w_ratio=17 / 16,
-                      min_h_div_w_ratio=8 / 16):
-    if h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio:
-        return True
-    return False
+def filter_resolution(h: int,
+                      w: int,
+                      max_h_div_w_ratio: float = 17 / 16,
+                      min_h_div_w_ratio: float = 8 / 16) -> bool:
+    return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
 
 
 class T2V_dataset(Dataset):
@@ -85,7 +80,7 @@ class T2V_dataset(Dataset):
                  temporal_sample,
                  tokenizer,
                  transform_topcrop,
-                 start_idx=0):
+                 start_idx=0) -> None:
         self.start_idx = start_idx
         self.data = args.data_merge_path
         self.num_frames = args.num_frames
@@ -132,14 +127,14 @@ class T2V_dataset(Dataset):
         data = self.get_data(idx)
         return data
 
-    def get_data(self, idx):
+    def get_data(self, idx) -> dict:
         path = dataset_prog.cap_list[idx]["path"]
         if path.endswith(".mp4"):
             return self.get_video(idx)
         else:
             return self.get_image(idx)
 
-    def get_video(self, idx):
+    def get_video(self, idx) -> dict:
         video_path = dataset_prog.cap_list[idx]["path"]
         assert os.path.exists(video_path), f"file {video_path} do not exist!"
         frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
@@ -183,7 +178,7 @@ class T2V_dataset(Dataset):
                     fps=dataset_prog.cap_list[idx]["fps"],
                     duration=dataset_prog.cap_list[idx]["duration"])
 
-    def get_image(self, idx):
+    def get_image(self, idx) -> dict:
         image_data = dataset_prog.cap_list[
             idx]  # [{'path': path, 'cap': cap}, ...]
 
@@ -201,14 +196,14 @@ class T2V_dataset(Dataset):
 
         image = image.float() / 127.5 - 1.0
 
-        caps = (image_data["cap"]
-                if isinstance(image_data["cap"], list) else [image_data["cap"]])
+        caps: list[str] = (image_data["cap"] if isinstance(
+            image_data["cap"], list) else [image_data["cap"]])
         caps = [random.choice(caps)]
         text = caps
         input_ids, cond_mask = [], []
-        text = text[0] if random.random() > self.cfg else ""
+        single_text = text[0] if random.random() > self.cfg else ""
         text_tokens_and_mask = self.tokenizer(
-            text,
+            single_text,
             max_length=self.text_max_length,
             padding="max_length",
             truncation=True,
@@ -226,7 +221,7 @@ class T2V_dataset(Dataset):
             path=image_data["path"],
         )
 
-    def define_frame_index(self, cap_list):
+    def define_frame_index(self, cap_list) -> tuple[list[dict], list[int]]:
         new_cap_list = []
         sample_num_frames = []
         cnt_too_long = 0
@@ -326,14 +321,14 @@ class T2V_dataset(Dataset):
         )
         return new_cap_list, sample_num_frames
 
-    def decord_read(self, path, frame_indices):
+    def decord_read(self, path, frame_indices) -> torch.Tensor:
         decord_vr = self.v_decoder(path)
         video_data = decord_vr.get_batch(frame_indices).asnumpy()
         video_data = torch.from_numpy(video_data)
         video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
         return video_data
 
-    def read_jsons(self, data):
+    def read_jsons(self, data) -> list[dict]:
         cap_lists = []
         with open(data) as f:
             folder_anno = [
@@ -349,6 +344,6 @@ class T2V_dataset(Dataset):
             cap_lists += sub_list
         return cap_lists
 
-    def get_cap_list(self):
+    def get_cap_list(self) -> list:
         cap_lists = self.read_jsons(self.data)[self.start_idx:]
         return cap_lists

fastvideo/v1/dataset/transform.py

@@ -1,42 +1,19 @@
-import numbers
 import random
 
 import torch
-from PIL import Image
 
 
-def _is_tensor_video_clip(clip):
+def _is_tensor_video_clip(clip) -> bool:
     if not torch.is_tensor(clip):
-        raise TypeError("clip should be Tensor. Got %s" % type(clip))
+        raise TypeError(f"clip should be Tensor. Got {type(clip)}")
 
     if not clip.ndimension() == 4:
-        raise ValueError("clip should be 4D. Got %dD" % clip.dim())
+        raise ValueError(f"clip should be 4D. Got {clip.dim()}D")
 
     return True
 
 
-def center_crop_arr(pil_image, image_size):
-    """
-    Center cropping implementation from ADM.
-    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
-    """
-    while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size),
-                                     resample=Image.BOX)
-
-    scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(tuple(
-        round(x * scale) for x in pil_image.size),
-                                 resample=Image.BICUBIC)
-
-    arr = np.array(pil_image)
-    crop_y = (arr.shape[0] - image_size) // 2
-    crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y:crop_y + image_size,
-                               crop_x:crop_x + image_size])
-
-
-def crop(clip, i, j, h, w):
+def crop(clip, i, j, h, w) -> torch.Tensor:
     """
     Args:
         clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
@@ -46,7 +23,7 @@ def crop(clip, i, j, h, w):
     return clip[..., i:i + h, j:j + w]
 
 
-def resize(clip, target_size, interpolation_mode):
+def resize(clip, target_size, interpolation_mode) -> torch.Tensor:
     if len(target_size) != 2:
         raise ValueError(
             f"target size should be tuple (height, width), instead got {target_size}"
@@ -60,71 +37,7 @@ def resize(clip, target_size, interpolation_mode):
     )
 
 
-def resize_scale(clip, target_size, interpolation_mode):
-    if len(target_size) != 2:
-        raise ValueError(
-            f"target size should be tuple (height, width), instead got {target_size}"
-        )
-    H, W = clip.size(-2), clip.size(-1)
-    scale_ = target_size[0] / min(H, W)
-    return torch.nn.functional.interpolate(
-        clip,
-        scale_factor=scale_,
-        mode=interpolation_mode,
-        align_corners=True,
-        antialias=True,
-    )
-
-
-def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"):
-    """
-    Do spatial cropping and resizing to the video clip
-    Args:
-        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        i (int): i in (i,j) i.e coordinates of the upper left corner.
-        j (int): j in (i,j) i.e coordinates of the upper left corner.
-        h (int): Height of the cropped region.
-        w (int): Width of the cropped region.
-        size (tuple(int, int)): height and width of resized clip
-    Returns:
-        clip (torch.tensor): Resized and cropped clip. Size is (T, C, H, W)
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    clip = crop(clip, i, j, h, w)
-    clip = resize(clip, size, interpolation_mode)
-    return clip
-
-
-def center_crop(clip, crop_size):
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    h, w = clip.size(-2), clip.size(-1)
-    th, tw = crop_size
-    if h < th or w < tw:
-        raise ValueError("height and width must be no smaller than crop_size")
-
-    i = int(round((h - th) / 2.0))
-    j = int(round((w - tw) / 2.0))
-    return crop(clip, i, j, th, tw)
-
-
-def center_crop_using_short_edge(clip):
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    h, w = clip.size(-2), clip.size(-1)
-    if h < w:
-        th, tw = h, h
-        i = 0
-        j = int(round((w - tw) / 2.0))
-    else:
-        th, tw = w, w
-        i = int(round((h - th) / 2.0))
-        j = 0
-    return crop(clip, i, j, th, tw)
-
-
-def center_crop_th_tw(clip, th, tw, top_crop):
+def center_crop_th_tw(clip, th, tw, top_crop) -> torch.Tensor:
     if not _is_tensor_video_clip(clip):
         raise ValueError("clip should be a 4D torch.tensor")
 
@@ -143,27 +56,7 @@ def center_crop_th_tw(clip, th, tw, top_crop):
     return crop(clip, i, j, new_h, new_w)
 
 
-def random_shift_crop(clip):
-    """
-    Slide along the long edge, with the short edge as crop size
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    h, w = clip.size(-2), clip.size(-1)
-
-    if h <= w:
-        short_edge = h
-    else:
-        short_edge = w
-
-    th, tw = short_edge, short_edge
-
-    i = torch.randint(0, h - th + 1, size=(1, )).item()
-    j = torch.randint(0, w - tw + 1, size=(1, )).item()
-    return crop(clip, i, j, th, tw)
-
-
-def normalize_video(clip):
+def normalize_video(clip) -> torch.Tensor:
     """
     Convert tensor data type from uint8 to float, divide value by 255.0 and
     permute the dimensions of clip tensor
@@ -174,153 +67,12 @@ def normalize_video(clip):
     """
     _is_tensor_video_clip(clip)
     if not clip.dtype == torch.uint8:
-        raise TypeError("clip tensor should have data type uint8. Got %s" %
-                        str(clip.dtype))
+        raise TypeError(
+            f"clip tensor should have data type uint8. Got {clip.dtype}")
     # return clip.float().permute(3, 0, 1, 2) / 255.0
     return clip.float() / 255.0
 
 
-def normalize(clip, mean, std, inplace=False):
-    """
-    Args:
-        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
-        mean (tuple): pixel RGB mean. Size is (3)
-        std (tuple): pixel standard deviation. Size is (3)
-    Returns:
-        normalized clip (torch.tensor): Size is (T, C, H, W)
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    if not inplace:
-        clip = clip.clone()
-    mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device)
-    # print(mean)
-    std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device)
-    clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None])
-    return clip
-
-
-def hflip(clip):
-    """
-    Args:
-        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
-    Returns:
-        flipped clip (torch.tensor): Size is (T, C, H, W)
-    """
-    if not _is_tensor_video_clip(clip):
-        raise ValueError("clip should be a 4D torch.tensor")
-    return clip.flip(-1)
-
-
-class RandomCropVideo:
-
-    def __init__(self, size):
-        if isinstance(size, numbers.Number):
-            self.size = (int(size), int(size))
-        else:
-            self.size = size
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: randomly cropped video clip.
-                size is (T, C, OH, OW)
-        """
-        i, j, h, w = self.get_params(clip)
-        return crop(clip, i, j, h, w)
-
-    def get_params(self, clip):
-        h, w = clip.shape[-2:]
-        th, tw = self.size
-
-        if h < th or w < tw:
-            raise ValueError(
-                f"Required crop size {(th, tw)} is larger than input image size {(h, w)}"
-            )
-
-        if w == tw and h == th:
-            return 0, 0, h, w
-
-        i = torch.randint(0, h - th + 1, size=(1, )).item()
-        j = torch.randint(0, w - tw + 1, size=(1, )).item()
-
-        return i, j, th, tw
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size})"
-
-
-class SpatialStrideCropVideo:
-
-    def __init__(self, stride):
-        self.stride = stride
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: cropped video clip by stride.
-                size is (T, C, OH, OW)
-        """
-        i, j, h, w = self.get_params(clip)
-        return crop(clip, i, j, h, w)
-
-    def get_params(self, clip):
-        h, w = clip.shape[-2:]
-
-        th, tw = h // self.stride * self.stride, w // self.stride * self.stride
-
-        return 0, 0, th, tw  # from top-left
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size})"
-
-
-class LongSideResizeVideo:
-    """
-    First use the long side,
-    then resize to the specified size
-    """
-
-    def __init__(
-        self,
-        size,
-        skip_low_resolution=False,
-        interpolation_mode="bilinear",
-    ):
-        self.size = size
-        self.skip_low_resolution = skip_low_resolution
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: scale resized video clip.
-                size is (T, C, 512, *) or (T, C, *, 512)
-        """
-        _, _, h, w = clip.shape
-        if self.skip_low_resolution and max(h, w) <= self.size:
-            return clip
-        if h > w:
-            w = int(w * self.size / h)
-            h = self.size
-        else:
-            h = int(h * self.size / w)
-            w = self.size
-        resize_clip = resize(clip,
-                             target_size=(h, w),
-                             interpolation_mode=self.interpolation_mode)
-        return resize_clip
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
 class CenterCropResizeVideo:
     """
     First use the short side for cropping length,
@@ -332,7 +84,7 @@ class CenterCropResizeVideo:
         size,
         top_crop=False,
         interpolation_mode="bilinear",
-    ):
+    ) -> None:
         if len(size) != 2:
             raise ValueError(
                 f"size should be tuple (height, width), instead got {size}")
@@ -340,7 +92,7 @@ class CenterCropResizeVideo:
         self.top_crop = top_crop
         self.interpolation_mode = interpolation_mode
 
-    def __call__(self, clip):
+    def __call__(self, clip) -> torch.Tensor:
         """
         Args:
             clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
@@ -348,12 +100,10 @@ class CenterCropResizeVideo:
             torch.tensor: scale resized / center cropped video clip.
                 size is (T, C, crop_size, crop_size)
         """
-        # clip_center_crop = center_crop_using_short_edge(clip)
         clip_center_crop = center_crop_th_tw(clip,
                                              self.size[0],
                                              self.size[1],
                                              top_crop=self.top_crop)
-        # import ipdb;ipdb.set_trace()
         clip_center_crop_resize = resize(
             clip_center_crop,
             target_size=self.size,
@@ -365,138 +115,15 @@ class CenterCropResizeVideo:
         return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
 
 
-class UCFCenterCropVideo:
-    """
-    First scale to the specified size in equal proportion to the short edge,
-    then center cropping
-    """
-
-    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
-    ):
-        if isinstance(size, tuple):
-            if len(size) != 2:
-                raise ValueError(
-                    f"size should be tuple (height, width), instead got {size}")
-            self.size = size
-        else:
-            self.size = (size, size)
-
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: scale resized / center cropped video clip.
-                size is (T, C, crop_size, crop_size)
-        """
-        clip_resize = resize_scale(clip=clip,
-                                   target_size=self.size,
-                                   interpolation_mode=self.interpolation_mode)
-        clip_center_crop = center_crop(clip_resize, self.size)
-        return clip_center_crop
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
-class KineticsRandomCropResizeVideo:
-    """
-    Slide along the long edge, with the short edge as crop size. And resie to the desired size.
-    """
-
-    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
-    ):
-        if isinstance(size, tuple):
-            if len(size) != 2:
-                raise ValueError(
-                    f"size should be tuple (height, width), instead got {size}")
-            self.size = size
-        else:
-            self.size = (size, size)
-
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        clip_random_crop = random_shift_crop(clip)
-        clip_resize = resize(clip_random_crop, self.size,
-                             self.interpolation_mode)
-        return clip_resize
-
-
-class CenterCropVideo:
-
-    def __init__(
-        self,
-        size,
-        interpolation_mode="bilinear",
-    ):
-        if isinstance(size, tuple):
-            if len(size) != 2:
-                raise ValueError(
-                    f"size should be tuple (height, width), instead got {size}")
-            self.size = size
-        else:
-            self.size = (size, size)
-
-        self.interpolation_mode = interpolation_mode
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
-        Returns:
-            torch.tensor: center cropped video clip.
-                size is (T, C, crop_size, crop_size)
-        """
-        clip_center_crop = center_crop(clip, self.size)
-        return clip_center_crop
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
-
-
-class Normalize:
-    """
-    Normalize the video clip by mean subtraction and division by standard deviation
-    Args:
-        mean (3-tuple): pixel RGB mean
-        std (3-tuple): pixel RGB standard deviation
-        inplace (boolean): whether do in-place normalization
-    """
-
-    def __init__(self, mean, std, inplace=False):
-        self.mean = mean
-        self.std = std
-        self.inplace = inplace
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): video clip must be normalized. Size is (C, T, H, W)
-        """
-        return normalize(clip, self.mean, self.std, self.inplace)
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})"
-
-
 class Normalize255:
     """
     Convert tensor data type from uint8 to float, divide value by 255.0 and
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
-    def __call__(self, clip):
+    def __call__(self, clip) -> torch.Tensor:
         """
         Args:
             clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
@@ -509,34 +136,6 @@ class Normalize255:
         return self.__class__.__name__
 
 
-class RandomHorizontalFlipVideo:
-    """
-    Flip the video clip along the horizontal direction with a given probability
-    Args:
-        p (float): probability of the clip being flipped. Default value is 0.5
-    """
-
-    def __init__(self, p=0.5):
-        self.p = p
-
-    def __call__(self, clip):
-        """
-        Args:
-            clip (torch.tensor): Size is (T, C, H, W)
-        Return:
-            clip (torch.tensor): Size is (T, C, H, W)
-        """
-        if random.random() < self.p:
-            clip = hflip(clip)
-        return clip
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(p={self.p})"
-
-
-#  ------------------------------------------------------------
-#  ---------------------  Sampling  ---------------------------
-#  ------------------------------------------------------------
 class TemporalRandomCrop:
     """Temporally crop the given frame indices at a random location.
 
@@ -544,99 +143,11 @@ class TemporalRandomCrop:
         size (int): Desired length of frames will be seen in the model.
     """
 
-    def __init__(self, size):
+    def __init__(self, size) -> None:
         self.size = size
 
-    def __call__(self, total_frames):
+    def __call__(self, total_frames) -> tuple[int, int]:
         rand_end = max(0, total_frames - self.size - 1)
         begin_index = random.randint(0, rand_end)
         end_index = min(begin_index + self.size, total_frames)
         return begin_index, end_index
-
-
-class DynamicSampleDuration:
-    """Temporally crop the given frame indices at a random location.
-
-    Args:
-        size (int): Desired length of frames will be seen in the model.
-    """
-
-    def __init__(self, t_stride, extra_1):
-        self.t_stride = t_stride
-        self.extra_1 = extra_1
-
-    def __call__(self, t, h, w):
-        if self.extra_1:
-            t = t - 1
-        truncate_t_list = list(
-            range(t + 1))[t // 2:][::self.t_stride]  # need half at least
-        truncate_t = random.choice(truncate_t_list)
-        if self.extra_1:
-            truncate_t = truncate_t + 1
-        return 0, truncate_t
-
-
-if __name__ == "__main__":
-    import os
-
-    import numpy as np
-    import torchvision.io as io
-    from torchvision import transforms
-    from torchvision.utils import save_image
-
-    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi",
-                                           pts_unit="sec",
-                                           output_format="TCHW")
-
-    trans = transforms.Compose([
-        Normalize255(),
-        RandomHorizontalFlipVideo(),
-        UCFCenterCropVideo(512),
-        # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5],
-                             std=[0.5, 0.5, 0.5],
-                             inplace=True),
-    ])
-
-    target_video_len = 32
-    frame_interval = 1
-    total_frames = len(vframes)
-    print(total_frames)
-
-    temporal_sample = TemporalRandomCrop(target_video_len * frame_interval)
-
-    # Sampling video frames
-    start_frame_ind, end_frame_ind = temporal_sample(total_frames)
-    # print(start_frame_ind)
-    # print(end_frame_ind)
-    assert end_frame_ind - start_frame_ind >= target_video_len
-    frame_indice = np.linspace(start_frame_ind,
-                               end_frame_ind - 1,
-                               target_video_len,
-                               dtype=int)
-    print(frame_indice)
-
-    select_vframes = vframes[frame_indice]
-    print(select_vframes.shape)
-    print(select_vframes.dtype)
-
-    select_vframes_trans = trans(select_vframes)
-    print(select_vframes_trans.shape)
-    print(select_vframes_trans.dtype)
-
-    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) *
-                                255).to(dtype=torch.uint8)
-    print(select_vframes_trans_int.dtype)
-    print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
-
-    io.write_video("./test.avi",
-                   select_vframes_trans_int.permute(0, 2, 3, 1),
-                   fps=8)
-
-    for i in range(target_video_len):
-        save_image(
-            select_vframes_trans[i],
-            os.path.join("./test000", "%04d.png" % i),
-            normalize=True,
-            value_range=(-1, 1),
-        )

fastvideo/v1/fastvideo_args.py

@@ -438,6 +438,11 @@ def get_current_fastvideo_args() -> FastVideoArgs:
 
 @dataclasses.dataclass
 class TrainingArgs(FastVideoArgs):
+    """
+    Training arguments. Inherits from FastVideoArgs and adds training-specific
+    arguments. If there are any conflicts, the training arguments will take
+    precedence.
+    """
     data_path: str = ""
     dataloader_num_workers: int = 0
     num_height: int = 0
@@ -474,7 +479,6 @@ class TrainingArgs(FastVideoArgs):
     checkpoints_total_limit: int = 0
     checkpointing_steps: int = 0
     resume_from_checkpoint: str = ""
-    resume_from_lora_checkpoint: str = ""
     logging_dir: str = ""
 
     # optimizer & scheduler
@@ -490,13 +494,7 @@ class TrainingArgs(FastVideoArgs):
     selective_checkpointing: float = 0.0
     allow_tf32: bool = False
     mixed_precision: str = ""
-    use_cpu_offload: bool = False
-    # fp16_full_eval: bool
-    # fp16_backend: str
     train_sp_batch_size: int = 0
-    use_lora: bool = False
-    lora_alpha: int = 0
-    lora_rank: int = 0
     fsdp_sharding_startegy: str = ""
 
     weighting_scheme: str = ""
@@ -504,8 +502,6 @@ class TrainingArgs(FastVideoArgs):
     logit_std: float = 1.0
     mode_scale: float = 0.0
 
-    # lr_scheduler
-    lr_scheduler: str = ""
     num_euler_timesteps: int = 0
     lr_num_cycles: int = 0
     lr_power: float = 0.0
@@ -649,9 +645,6 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--resume-from-checkpoint",
                             type=str,
                             help="Path to checkpoint to resume from")
-        parser.add_argument("--resume-from-lora-checkpoint",
-                            type=str,
-                            help="Path to LoRA checkpoint to resume from")
         parser.add_argument("--logging-dir",
                             type=str,
                             help="Directory for logging")
@@ -700,14 +693,6 @@ class TrainingArgs(FastVideoArgs):
                             type=int,
                             help="Training spatial parallelism batch size")
 
-        # LoRA configuration
-        parser.add_argument("--use-lora",
-                            action=StoreBoolean,
-                            help="Whether to use LoRA")
-        parser.add_argument("--lora-alpha",
-                            type=int,
-                            help="LoRA alpha parameter")
-        parser.add_argument("--lora-rank", type=int, help="LoRA rank")
         parser.add_argument("--fsdp-sharding-strategy",
                             type=str,
                             help="FSDP sharding strategy")

fastvideo/v1/models/loader/fsdp_load.py

@@ -138,7 +138,6 @@ def load_fsdp_model(
                 f"Unexpected param or buffer {n} on meta device.")
     for p in model.parameters():
         p.requires_grad = False
-    # set_state_dict(model, StateDictType.LOCAL_STATE_DICT)
     return model
 
 

fastvideo/v1/models/vaes/common.py

@@ -39,6 +39,9 @@ class ParallelTiledVAE(ABC):
         self.use_temporal_tiling = config.use_temporal_tiling
         self.use_parallel_tiling = config.use_parallel_tiling
 
+    def to(self, device) -> 'ParallelTiledVAE':
+        return self
+
     @property
     def temporal_compression_ratio(self) -> int:
         return cast(int, self.config.temporal_compression_ratio)

fastvideo/v1/pipelines/distill_pipeline.py

@@ -1,303 +0,0 @@
-import os
-import sys
-import time
-from collections import deque
-
-import torch
-from tqdm.auto import tqdm
-
-# import torch.distributed as dist
-import wandb
-from fastvideo.utils.checkpoint import save_checkpoint
-from fastvideo.utils.communications import sp_parallel_dataloader_wrapper
-from fastvideo.utils.validation import log_validation
-from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines import ComposedPipelineBase
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
-                                           DenoisingStage, InputValidationStage,
-                                           LatentPreparationStage,
-                                           TextEncodingStage,
-                                           TimestepPreparationStage)
-
-logger = init_logger(__name__)
-
-
-class WanTrainingPipeline(ComposedPipelineBase):  # == distill_one_step
-    _required_config_modules = ["scheduler", "transformer"]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        """Set up pipeline stages with proper dependency injection."""
-
-        self.add_stage(stage_name="input_validation_stage",
-                       stage=InputValidationStage())
-
-        self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
-                       ))
-
-        self.add_stage(stage_name="conditioning_stage",
-                       stage=ConditioningStage())
-
-        self.add_stage(stage_name="timestep_preparation_stage",
-                       stage=TimestepPreparationStage(
-                           scheduler=self.get_module("scheduler")))
-
-        self.add_stage(stage_name="latent_preparation_stage",
-                       stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler"),
-                           transformer=self.get_module("transformer")))
-
-        self.add_stage(stage_name="denoising_stage",
-                       stage=DenoisingStage(
-                           transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler")))
-
-        self.add_stage(stage_name="decoding_stage",
-                       stage=DecodingStage(vae=self.get_module("vae")))
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ):
-        device = fastvideo_args.device
-        local_rank = int(os.environ.get("LOCAL_RANK", -1))
-        rank = int(os.environ.get("RANK", -1))
-        assert rank != -1
-        assert local_rank != -1
-        sp_group = get_sp_group()
-        world_size = sp_group.world_size
-        rank = sp_group.rank
-        args = fastvideo_args
-        transformer = self.get_module("transformer")
-        teacher_transformer = self.get_module("teacher_transformer")
-        ema_transformer = None
-        assert not fastvideo_args.use_ema, "ema is not supported now"
-        assert teacher_transformer is not None
-        assert transformer is not None
-        train_dataset = self.train_dataset
-        train_dataloader = self.train_dataloader
-        init_steps = self.init_steps
-        lr_scheduler = self.lr_scheduler
-        optimizer = self.optimizer
-        noise_scheduler = self.noise_scheduler
-        solver = self.solver
-        noise_random_generator = None
-        uncond_prompt_embed = self.uncond_prompt_embed
-        uncond_prompt_mask = self.uncond_prompt_mask
-
-        # Train!
-        total_batch_size = (world_size * args.gradient_accumulation_steps /
-                            args.sp_size * args.train_sp_batch_size)
-        logger.info("***** Running training *****")
-        logger.info("  Num examples = %s", len(train_dataset))
-        logger.info("  Dataloader size = %s", len(train_dataloader))
-        logger.info("  Num Epochs = %s", args.num_train_epochs)
-        logger.info("  Resume training from step %s", init_steps)
-        logger.info("  Instantaneous batch size per device = %s",
-                    args.train_batch_size)
-        logger.info(
-            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
-            total_batch_size)
-        logger.info("  Gradient Accumulation steps = %s",
-                    args.gradient_accumulation_steps)
-        logger.info("  Total optimization steps = %s", args.max_train_steps)
-        logger.info(
-            "  Total training parameters per FSDP shard = %s B",
-            sum(p.numel()
-                for p in transformer.parameters() if p.requires_grad) / 1e9)
-        # print dtype
-        logger.info("  Master weight dtype: %s",
-                    transformer.parameters().__next__().dtype)
-
-        # Potentially load in the weights and states from a previous save
-        if args.resume_from_checkpoint:
-            assert NotImplementedError(
-                "resume_from_checkpoint is not supported now.")
-            # TODO
-
-        progress_bar = tqdm(
-            range(0, args.max_train_steps),
-            initial=init_steps,
-            desc="Steps",
-            # Only show the progress bar once on each machine.
-            disable=local_rank > 0,
-        )
-
-        loader = sp_parallel_dataloader_wrapper(
-            train_dataloader,
-            device,
-            args.train_batch_size,
-            args.sp_size,
-            args.train_sp_batch_size,
-        )
-        # loader = self.get_module("train_dataloader")
-
-        step_times = deque(maxlen=100)
-
-        # todo future
-        for i in range(init_steps):
-            next(loader)
-
-        # log_validation(args, transformer, device,
-        #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
-        def get_num_phases(multi_phased_distill_schedule, step):
-            # step-phase,step-phase
-            multi_phases = multi_phased_distill_schedule.split(",")
-            phase = multi_phases[-1].split("-")[-1]
-            for step_phases in multi_phases:
-                phase_step, phase = step_phases.split("-")
-                if step <= int(phase_step):
-                    return int(phase)
-            return phase
-
-        for step in range(init_steps + 1, args.max_train_steps + 1):
-            start_time = time.time()
-            assert args.multi_phased_distill_schedule is not None
-            num_phases = get_num_phases(args.multi_phased_distill_schedule,
-                                        step)
-
-            loss, grad_norm, pred_norm = self.distill_one_step(
-                transformer,
-                args.model_type,
-                teacher_transformer,
-                ema_transformer,
-                optimizer,
-                lr_scheduler,
-                loader,
-                noise_scheduler,
-                solver,
-                noise_random_generator,
-                args.gradient_accumulation_steps,
-                args.sp_size,
-                args.max_grad_norm,
-                uncond_prompt_embed,
-                uncond_prompt_mask,
-                args.num_euler_timesteps,
-                num_phases,
-                args.not_apply_cfg_solver,
-                args.distill_cfg,
-                args.ema_decay,
-                args.pred_decay_weight,
-                args.pred_decay_type,
-                args.hunyuan_teacher_disable_cfg,
-            )
-
-            step_time = time.time() - start_time
-            step_times.append(step_time)
-            avg_step_time = sum(step_times) / len(step_times)
-
-            progress_bar.set_postfix({
-                "loss": f"{loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
-                "phases": num_phases,
-            })
-            progress_bar.update(1)
-            if rank <= 0:
-                wandb.log(
-                    {
-                        "train_loss":
-                        loss,
-                        "learning_rate":
-                        lr_scheduler.get_last_lr()[0],
-                        "step_time":
-                        step_time,
-                        "avg_step_time":
-                        avg_step_time,
-                        "grad_norm":
-                        grad_norm,
-                        "pred_fro_norm":
-                        pred_norm["fro"],  # codespell:ignore
-                        "pred_largest_singular_value":
-                        pred_norm["largest singular value"],
-                        "pred_absolute_mean":
-                        pred_norm["absolute mean"],
-                        "pred_absolute_max":
-                        pred_norm["absolute max"],
-                    },
-                    step=step,
-                )
-            if step % args.checkpointing_steps == 0:
-                if args.use_lora:
-                    # Save LoRA weights
-                    raise NotImplementedError("lora is not supported now")
-                    # save_lora_checkpoint(transformer, optimizer, rank,
-                    #                      args.output_dir, step)
-                else:
-                    # Your existing checkpoint saving code
-                    if args.use_ema:
-                        raise NotImplementedError("ema is not supported now")
-                        save_checkpoint(ema_transformer, rank, args.output_dir,
-                                        step)
-                    else:
-                        save_checkpoint(transformer, rank, args.output_dir,
-                                        step)
-
-                sp_group.barrier()
-            if args.log_validation and step % args.validation_steps == 0:
-                log_validation(
-                    args,
-                    transformer,
-                    device,
-                    torch.bfloat16,
-                    step,
-                    scheduler_type=args.scheduler_type,
-                    shift=args.shift,
-                    num_euler_timesteps=args.num_euler_timesteps,
-                    linear_quadratic_threshold=args.linear_quadratic_threshold,
-                    linear_range=args.linear_range,
-                    ema=False,
-                )
-                if args.use_ema:
-                    log_validation(
-                        args,
-                        ema_transformer,
-                        device,
-                        torch.bfloat16,
-                        step,
-                        scheduler_type=args.scheduler_type,
-                        shift=args.shift,
-                        num_euler_timesteps=args.num_euler_timesteps,
-                        linear_quadratic_threshold=args.
-                        linear_quadratic_threshold,
-                        linear_range=args.linear_range,
-                        ema=True,
-                    )
-
-        if args.use_lora:
-            raise NotImplementedError("lora is not supported now")
-            save_lora_checkpoint(transformer, optimizer, rank, args.output_dir,
-                                 args.max_train_steps)
-        else:
-            save_checkpoint(transformer, rank, args.output_dir,
-                            args.max_train_steps)
-
-        if get_sp_group():
-            cleanup_dist_env_and_memory()
-
-
-def main(args):
-    logger.info("Starting training pipeline...")
-    pipeline = WanTrainingPipeline.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", args=args)
-    args = pipeline.fastvideo_args
-    pipeline.forward(None, args)
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.v1.fastvideo_args import TrainingArgs
-    from fastvideo.v1.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    print(args)
-    main(args)

fastvideo/v1/pipelines/preprocess_pipeline.py

@@ -140,7 +140,6 @@ class PreprocessPipeline(ComposedPipelineBase):
                     0], result_batch.prompt_attention_mask[0]
                 assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
 
-                # Remove padding from prompt_embeds using attention mask for all batches
                 # Get sequence lengths from attention masks (number of 1s)
                 seq_lens = prompt_attention_mask.sum(dim=1)
                 # Create a list to store non-padded embeddings and masks
@@ -354,9 +353,6 @@ class PreprocessPipeline(ComposedPipelineBase):
                                               "validation_parquet_dataset")
         os.makedirs(validation_parquet_dir, exist_ok=True)
 
-        # Initialize Parquet dataset
-        validation_parquet_path = os.path.join(validation_parquet_dir,
-                                               "data.parquet")
 
         with open(args.validation_prompt_txt, encoding="utf-8") as file:
             lines = file.readlines()

fastvideo/v1/pipelines/stages/decoding.py

@@ -7,6 +7,7 @@ import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.utils import PRECISION_TO_TYPE
@@ -23,7 +24,7 @@ class DecodingStage(PipelineStage):
     """
 
     def __init__(self, vae) -> None:
-        self.vae = vae
+        self.vae: ParallelTiledVAE = vae
 
     def forward(
         self,

fastvideo/v1/pipelines/stages/encoding.py

@@ -46,6 +46,8 @@ class EncodingStage(PipelineStage):
         Returns:
             The batch with encoded outputs.
         """
+        self.vae = self.vae.to(fastvideo_args.device)
+
         image_path = batch.image_path
         # TODO(will): remove this once we add input/output validation for stages
         if image_path is None:

fastvideo/v1/pipelines/training_pipeline.py

@@ -19,7 +19,6 @@ from tqdm.auto import tqdm
 # import torch.distributed as dist
 import wandb
 from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
-from fastvideo.utils.checkpoint import save_checkpoint_v1
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
 from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
@@ -30,7 +29,7 @@ from fastvideo.v1.pipelines import ComposedPipelineBase
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.training_utils import (
     _clip_grad_norm_while_handling_failing_dtensor_cases,
-    compute_density_for_timestep_sampling, get_sigmas)
+    compute_density_for_timestep_sampling, get_sigmas, save_checkpoint)
 from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
 
 logger = init_logger(__name__)
@@ -201,6 +200,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
             logger.info('embed dtype', prompt_embeds.dtype)
 
+            num_frames = (fastvideo_args.num_latent_t - 1) * 4 + 1
+            logger.info(f"validation num_frames: {num_frames}")
             # Prepare batch for validation
             # print('shape of embeddings', prompt_embeds.shape)
             batch = ForwardBatch(
@@ -214,7 +215,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                 # make sure we use the same height, width, and num_frames as the training pipeline
                 height=args.num_height,
                 width=args.num_width,
-                num_frames=args.num_frames,
+                num_frames=num_frames,
                 # num_inference_steps=fastvideo_args.validation_sampling_steps,
                 num_inference_steps=50,
                 # guidance_scale=fastvideo_args.validation_guidance_scale,
@@ -249,6 +250,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             video_captions = []
             for i, video in enumerate(videos):
                 caption = captions[i]
+                os.makedirs(fastvideo_args.output_dir, exist_ok=True)
                 filename = os.path.join(
                     fastvideo_args.output_dir,
                     f"validation_step_{global_step}_video_{i}.mp4")
@@ -351,24 +353,46 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             # Check gradients for selected parameters
             absolute_errors = []
             param_count = 0
+            rank = int(os.environ.get("RANK", 0))
 
+            sp_group = get_sp_group()
             for name, param in transformer.named_parameters():
+                sp_group.barrier()
+                # skip scale_shift_table because it is not sharded
+                if 'scale_shift_table' in name:
+                    continue
+                if isinstance(param.grad, torch.distributed.tensor.DTensor):
+                    l = param.grad.full_tensor()
+                    distributed = True
+                else:
+                    l = param.grad
+                    distributed = False
+                    continue
+                # logger.info(f"rank: {rank}, name: {name}, param: {param.shape}, grad: {param.grad.shape}, distributed: {distributed}", local_main_process_only=False)
+                # logger.info(f"rank: {rank}, name: {name}, type of param: {type(param)}, type of grad: {type(param.grad)}", local_main_process_only=False)
+                # logger.info(f"rank: {rank}, name: {name}, param: {param}, grad: {param.grad}", local_main_process_only=False)
                 if not (param.requires_grad and param.grad is not None
                         and param_count < max_params_to_check
-                        and param.grad.abs().max() > 5e-4):
+                        and l.abs().max() > 5e-4):
                     continue
+                if not distributed:
+                    if rank != 0:
+                        continue
 
                 # Get local parameter and gradient tensors
                 local_param = param._local_tensor if hasattr(
                     param, '_local_tensor') else param
                 local_grad = param.grad._local_tensor if hasattr(
                     param.grad, '_local_tensor') else param.grad
+                # logger.info(f"rank: {rank}, local_param: {local_param.shape}, local_grad: {local_grad.shape}", local_main_process_only=False)
 
                 # Find first significant gradient element
                 flat_param = local_param.data.view(-1)
                 flat_grad = local_grad.view(-1)
+                # logger.info(f"rank: {rank}, flat_param: {flat_param.shape}, flat_grad: {flat_grad.shape}", local_main_process_only=False)
                 check_idx = next((i for i in range(min(10, flat_param.numel()))
                                   if abs(flat_grad[i]) > 1e-4), 0)
+                # logger.info(f"rank: {rank}, check_idx: {check_idx}", local_main_process_only=False)
 
                 # Store original values
                 orig_value = flat_param[check_idx].item()
@@ -377,7 +401,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                 # Compute numerical gradient
                 for delta in [eps, -eps]:
                     with torch.no_grad():
-                        flat_param[check_idx] = orig_value + delta
+                        # only have a single rank modify the parameter
+                        if rank == 0:
+                            flat_param[check_idx] = orig_value + delta
                         loss = compute_loss()
                         if delta > 0: loss_plus = loss.item()
                         else: loss_minus = loss.item()
@@ -392,23 +418,24 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                                             abs(numerical_grad), 1e-3)
                 absolute_errors.append(abs_error)
 
-                logger.info(
-                    f"{name}[{check_idx}]: analytical={analytical_grad:.6f}, "
-                    f"numerical={numerical_grad:.6f}, abs_error={abs_error:.2e}, rel_error={rel_error:.2%}"
-                )
+                if self.rank <= 0:
+                    logger.info(
+                        f"{name}[{check_idx}]: analytical={analytical_grad:.6f}, "
+                        f"numerical={numerical_grad:.6f}, abs_error={abs_error:.2e}, rel_error={rel_error:.2%}"
+                    )
 
                 # param_count += 1
 
             # Compute and log statistics
-            if absolute_errors:
-                min_err, max_err, mean_err = min(absolute_errors), max(
-                    absolute_errors
-                ), sum(absolute_errors) / len(absolute_errors)
-                logger.info(
-                    f"Gradient check stats: min={min_err:.2e}, max={max_err:.2e}, mean={mean_err:.2e}"
-                )
+            if self.rank <= 0:
+                if absolute_errors:
+                    min_err, max_err, mean_err = min(absolute_errors), max(
+                        absolute_errors
+                    ), sum(absolute_errors) / len(absolute_errors)
+                    logger.info(
+                        f"Gradient check stats: min={min_err:.2e}, max={max_err:.2e}, mean={mean_err:.2e}"
+                    )
 
-                if self.rank <= 0:
                     wandb.log({
                         "grad_check/min_abs_error":
                         min_err,
@@ -526,6 +553,7 @@ class WanTrainingPipeline(TrainingPipeline):
         pass
 
     def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.latents = None
         logger.info("Initializing validation pipeline...")
         args_copy = deepcopy(fastvideo_args)
 
@@ -562,12 +590,19 @@ class WanTrainingPipeline(TrainingPipeline):
         total_loss = 0.0
         optimizer.zero_grad()
         for _ in range(gradient_accumulation_steps):
-            (
-                latents,
-                encoder_hidden_states,
-                encoder_attention_mask,
-                infos,
-            ) = next(loader_iter)
+            logger.info(f"Rank {self.rank}: Training step {_}", local_main_process_only=False)
+            if self.latents is None:
+                (
+                    self.latents,
+                    self.encoder_hidden_states,
+                    self.encoder_attention_mask,
+                    self.infos,
+                ) = next(loader_iter)
+            latents = self.latents
+            encoder_hidden_states = self.encoder_hidden_states
+            encoder_attention_mask = self.encoder_attention_mask
+            infos = self.infos
+            logger.info(f"Rank {self.rank}: Training step {_} loaded data", local_main_process_only=False)
             latents = latents.to(self.fastvideo_args.device,
                                  dtype=torch.bfloat16)
             encoder_hidden_states = encoder_hidden_states.to(
@@ -598,8 +633,6 @@ class WanTrainingPipeline(TrainingPipeline):
                 dtype=latents.dtype,
             )
             noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
-            print('device before forward ',
-                  next(transformer.named_parameters())[1].device)
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 input_kwargs = {
                     "hidden_states": noisy_model_input,
@@ -626,14 +659,8 @@ class WanTrainingPipeline(TrainingPipeline):
 
                 loss = (torch.mean((model_pred.float() - target.float())**2) /
                         gradient_accumulation_steps)
-                print('device before backwardin context',
-                      next(transformer.named_parameters())[1].device)
 
-            print('device before backward out context',
-                  next(transformer.named_parameters())[1].device)
             loss.backward()
-            print('device after backward out context',
-                  next(transformer.named_parameters())[1].device)
 
             avg_loss = loss.detach().clone()
             sp_group = get_sp_group()
@@ -776,26 +803,16 @@ class WanTrainingPipeline(TrainingPipeline):
                     step=step,
                 )
             if step % args.checkpointing_steps == 0:
-                if args.use_lora:
-                    raise NotImplementedError("LoRA is not supported now")
-                    # Save LoRA weights
-                    save_lora_checkpoint(transformer, optimizer, rank,
-                                         args.output_dir, step, pipe)
-                else:
-                    # Your existing checkpoint saving code
-                    save_checkpoint_v1(self.transformer, self.rank,
-                                       args.output_dir, step)
-                    self.transformer.train()
+                # Your existing checkpoint saving code
+                save_checkpoint(self.transformer, self.rank, args.output_dir,
+                                step)
+                self.transformer.train()
                 self.sp_group.barrier()
             if args.log_validation and step % args.validation_steps == 0:
                 self.log_validation(self.transformer, args, step)
 
-        if args.use_lora:
-            raise NotImplementedError("LoRA is not supported now")
-            # save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps, pipe)
-        else:
-            save_checkpoint_v1(self.transformer, self.rank, args.output_dir,
-                               args.max_train_steps)
+        save_checkpoint(self.transformer, self.rank, args.output_dir,
+                        args.max_train_steps)
 
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/v1/pipelines/training_utils.py

@@ -1,14 +1,21 @@
+import json
 import math
+import os
 from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
 import torch.distributed.tensor
+from torch.distributed.fsdp import FullStateDictConfig
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import StateDictType
 
 from fastvideo.v1.logger import init_logger
 
 _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
 
+logger = init_logger(__name__)
+
 
 def compute_density_for_timestep_sampling(
     weighting_scheme: str,
@@ -60,7 +67,32 @@ def get_sigmas(noise_scheduler,
     return sigma
 
 
-logger = init_logger(__name__)
+def save_checkpoint(transformer, rank, output_dir, step):
+    # Configure FSDP to save full state dict
+    FSDP.set_state_dict_type(
+        transformer,
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        state_dict_config=FullStateDictConfig(offload_to_cpu=True,
+                                              rank0_only=True),
+    )
+
+    # Now get the state dict
+    cpu_state = transformer.state_dict()
+
+    # Save it (only on rank 0 since we used rank0_only=True)
+    if rank <= 0:
+        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+        os.makedirs(save_dir, exist_ok=True)
+        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt")
+        torch.save(cpu_state, weight_path)
+        config_dict = transformer.hf_config
+        if "dtype" in config_dict:
+            del config_dict["dtype"]  # TODO
+        config_path = os.path.join(save_dir, "config.json")
+        # save dict as json
+        with open(config_path, "w") as f:
+            json.dump(config_dict, f, indent=4)
+        logger.info("--> checkpoint saved at step %s to %s", step, weight_path)
 
 
 def _clip_grad_norm_while_handling_failing_dtensor_cases(

fastvideo/v1_distill.py

@@ -1,953 +0,0 @@
-# !/bin/python3
-# isort: skip_file
-import argparse
-import math
-import os
-import time
-from collections import deque
-from copy import deepcopy
-
-import torch
-import torch.distributed as dist
-import wandb
-from accelerate.utils import set_seed
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.optimization import get_scheduler
-from diffusers.utils import check_min_version
-from peft import LoraConfig
-# from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-# from torch.distributed.fsdp import ShardingStrategy
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from tqdm.auto import tqdm
-
-from fastvideo.dataset.latent_datasets import (LatentDataset,
-                                               latent_collate_function)
-from fastvideo.distill.solver import EulerSolver, extract_into_tensor
-from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-from fastvideo.utils.checkpoint import (resume_lora_optimizer, save_checkpoint,
-                                        save_lora_checkpoint)
-from fastvideo.utils.communications import (broadcast,
-                                            sp_parallel_dataloader_wrapper)
-from fastvideo.utils.dataset_utils import LengthGroupedSampler
-from fastvideo.utils.fsdp_util import (apply_fsdp_checkpointing)
-# from fastvideo.utils.load import load_transformer
-# from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state,
-#                                              initialize_sequence_parallel_state)
-from fastvideo.utils.validation import log_validation
-
-from fastvideo.v1.models.loader.component_loader import TransformerLoader
-from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
-from fastvideo.v1.configs.models.dits import WanVideoConfig
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.utils import maybe_download_model
-from fastvideo.v1.forward_context import set_forward_context
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.31.0")
-
-
-def main_print(content):
-    if int(os.environ["LOCAL_RANK"]) <= 0:
-        print(content)
-
-
-# def reshard_fsdp(model):
-#     for m in FSDP.fsdp_modules(model):
-#         if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
-#             torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
-
-
-def get_norm(model_pred, norms, gradient_accumulation_steps):
-    fro_norm = (
-        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
-        gradient_accumulation_steps)
-    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) /
-                              gradient_accumulation_steps)
-    absolute_mean = torch.mean(
-        torch.abs(model_pred)) / gradient_accumulation_steps
-    absolute_max = torch.max(
-        torch.abs(model_pred)) / gradient_accumulation_steps
-    dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
-    dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
-    dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
-    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
-    norms["largest singular value"] += torch.mean(largest_singular_value).item()
-    norms["absolute mean"] += absolute_mean.item()
-    norms["absolute max"] += absolute_max.item()
-
-
-def distill_one_step(
-    transformer,
-    model_type,
-    teacher_transformer,
-    ema_transformer,
-    optimizer,
-    lr_scheduler,
-    loader,
-    noise_scheduler,
-    solver,
-    noise_random_generator,
-    gradient_accumulation_steps,
-    sp_size,
-    max_grad_norm,
-    uncond_prompt_embed,
-    uncond_prompt_mask,
-    num_euler_timesteps,
-    multiphase,
-    not_apply_cfg_solver,
-    distill_cfg,
-    ema_decay,
-    pred_decay_weight,
-    pred_decay_type,
-    hunyuan_teacher_disable_cfg,
-):
-    total_loss = 0.0
-    optimizer.zero_grad()
-    model_pred_norm = {
-        "fro": 0.0,  # codespell:ignore
-        "largest singular value": 0.0,
-        "absolute mean": 0.0,
-        "absolute max": 0.0,
-    }
-    for _ in range(gradient_accumulation_steps):
-        (
-            latents,
-            encoder_hidden_states,
-            latents_attention_mask,
-            encoder_attention_mask,
-        ) = next(loader)
-        model_input = normalize_dit_input(model_type, latents)
-        noise = torch.randn_like(model_input)
-        bsz = model_input.shape[0]
-        index = torch.randint(0,
-                              num_euler_timesteps, (bsz, ),
-                              device=model_input.device).long()
-        if sp_size > 1:
-            broadcast(index)
-        # Add noise according to flow matching.
-        # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
-        sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
-        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index,
-                                          model_input.shape)
-
-        timesteps = (sigmas *
-                     noise_scheduler.config.num_train_timesteps).view(-1)
-        # if squeeze to [], unsqueeze to [1]
-
-        timesteps_prev = (sigmas_prev *
-                          noise_scheduler.config.num_train_timesteps).view(-1)
-        noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
-        # Predict the noise residual
-        with torch.autocast("cuda", dtype=torch.bfloat16):
-            print(f"--> noisy_model_input.dtype: {noisy_model_input.dtype}")
-            print(f"--> noisy_model_input.shape: {noisy_model_input.shape}")
-            print(
-                f"--> encoder_hidden_states.dtype: {encoder_hidden_states.dtype}"
-            )
-            print(
-                f"--> encoder_hidden_states.shape: {encoder_hidden_states.shape}"
-            )
-            print(f"--> timesteps.dtype: {timesteps.dtype}")
-            print(f"--> timesteps.shape: {timesteps.shape}")
-            print(
-                f"--> encoder_attention_mask.dtype: {encoder_attention_mask.dtype}"
-            )
-            print(
-                f"--> encoder_attention_mask.shape: {encoder_attention_mask.shape}"
-            )
-            noisy_model_input = noisy_model_input.to(dtype=torch.bfloat16)
-            teacher_kwargs = {
-                "hidden_states": noisy_model_input,
-                "encoder_hidden_states": encoder_hidden_states,
-                "timestep": timesteps,
-                "encoder_attention_mask": encoder_attention_mask,  # B, L
-                "return_dict": False,
-            }
-            if hunyuan_teacher_disable_cfg:
-                teacher_kwargs["guidance"] = torch.tensor(
-                    [1000.0],
-                    device=noisy_model_input.device,
-                    dtype=torch.bfloat16)
-            # batch = ForwardBatch(
-            #     enable_teacache=False,
-            # )
-            with set_forward_context(current_timestep=0,
-                                     attn_metadata=None,
-                                     forward_batch=None,
-                                     fastvideo_args=None):
-                model_pred = transformer(**teacher_kwargs)[0]
-
-        print(f"--> model_pred shape: {model_pred.shape}")
-        huber_c = 0.001
-        target = torch.randn_like(model_pred)
-        loss = (torch.mean(
-            torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) -
-            huber_c) / gradient_accumulation_steps)
-        loss.backward()
-        print(f"--> loss: {loss.item()}")
-        assert False, "stop here"
-
-        # if accelerator.is_main_process:
-        model_pred, end_index = solver.euler_style_multiphase_pred(
-            noisy_model_input, model_pred, index, multiphase)
-        with torch.no_grad():
-            w = distill_cfg
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                cond_teacher_output = teacher_transformer(
-                    noisy_model_input,
-                    encoder_hidden_states,
-                    timesteps,
-                    encoder_attention_mask,  # B, L
-                    return_dict=False,
-                )[0].float()
-            if not_apply_cfg_solver:
-                uncond_teacher_output = cond_teacher_output
-            else:
-                # Get teacher model prediction on noisy_latents and unconditional embedding
-                with torch.autocast("cuda", dtype=torch.bfloat16):
-                    uncond_teacher_output = teacher_transformer(
-                        noisy_model_input,
-                        uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
-                        timesteps,
-                        uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
-                        return_dict=False,
-                    )[0].float()
-            teacher_output = uncond_teacher_output + w * (cond_teacher_output -
-                                                          uncond_teacher_output)
-            x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
-
-        # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
-        with torch.no_grad():
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                if ema_transformer is not None:
-                    target_pred = ema_transformer(
-                        x_prev.float(),
-                        encoder_hidden_states,
-                        timesteps_prev,
-                        encoder_attention_mask,  # B, L
-                        return_dict=False,
-                    )[0]
-                else:
-                    target_pred = transformer(
-                        x_prev.float(),
-                        encoder_hidden_states,
-                        timesteps_prev,
-                        encoder_attention_mask,  # B, L
-                        return_dict=False,
-                    )[0]
-
-            target, end_index = solver.euler_style_multiphase_pred(
-                x_prev, target_pred, index, multiphase, True)
-
-        huber_c = 0.001
-        # loss = loss.mean()
-        loss = (torch.mean(
-            torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) -
-            huber_c) / gradient_accumulation_steps)
-        if pred_decay_weight > 0:
-            if pred_decay_type == "l1":
-                pred_decay_loss = (
-                    torch.mean(torch.sqrt(model_pred.float()**2)) *
-                    pred_decay_weight / gradient_accumulation_steps)
-                loss += pred_decay_loss
-            elif pred_decay_type == "l2":
-                # essnetially k2?
-                pred_decay_loss = (torch.mean(model_pred.float()**2) *
-                                   pred_decay_weight /
-                                   gradient_accumulation_steps)
-                loss += pred_decay_loss
-            else:
-                assert NotImplementedError("pred_decay_type is not implemented")
-
-        # calculate model_pred norm and mean
-        get_norm(model_pred.detach().float(), model_pred_norm,
-                 gradient_accumulation_steps)
-        loss.backward()
-
-        avg_loss = loss.detach().clone()
-        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
-        total_loss += avg_loss.item()
-
-    # update ema
-    if ema_transformer is not None:
-        reshard_fsdp(ema_transformer)
-        for p_averaged, p_model in zip(ema_transformer.parameters(),
-                                       transformer.parameters()):
-            with torch.no_grad():
-                p_averaged.copy_(
-                    torch.lerp(p_averaged.detach(), p_model.detach(),
-                               1 - ema_decay))
-
-    grad_norm = transformer.clip_grad_norm_(max_grad_norm)
-    optimizer.step()
-    lr_scheduler.step()
-
-    return total_loss, grad_norm.item(), model_pred_norm
-
-
-def main(args):
-    torch.backends.cuda.matmul.allow_tf32 = True
-
-    local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    rank = int(os.environ.get("RANK", 0))
-    world_size = int(os.environ.get("WORLD_SIZE", 1))
-    init_distributed_environment(world_size=world_size,
-                                 rank=rank,
-                                 local_rank=local_rank)
-    initialize_model_parallel(tensor_model_parallel_size=args.sp_size,
-                              sequence_model_parallel_size=args.sp_size)
-
-    # If passed along, set the training seed now. On GPU...
-    if args.seed is not None:
-        # TODO: t within the same seq parallel group should be the same. Noise should be different.
-        set_seed(args.seed + rank)
-    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
-    noise_random_generator = None
-
-    # Handle the repository creation
-    if rank <= 0 and args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    # For mixed precision training we cast all non-trainable weights to half-precision
-    # as these weights are only used for inference, keeping weights in full precision is not required.
-
-    # Create model:
-
-    print(
-        f"--> local_rank: {local_rank}, rank: {rank}, world_size: {world_size}")
-    main_print(f"--> using model pipeline {args.pretrained_model_name_or_path}")
-
-    model_path = maybe_download_model(args.pretrained_model_name_or_path)
-    main_print(f"--> loading model from {model_path}")
-
-    transformer_path = os.path.join(model_path, "transformer")
-    main_print(f"--> loading transformer from {transformer_path}")
-
-    precision = torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16
-
-    precision_str = "fp32" if precision == torch.float32 else "bf16"
-    print(f"--> precision: {precision_str}")
-    print(f"--> precision: {precision}")
-
-    # transformer_path = os.path.join(args.pretrained_model_name_or_path, "transformer")
-
-    fastvideo_args = FastVideoArgs(model_path=transformer_path,
-                                   use_cpu_offload=False,
-                                   precision=precision_str)
-    # fastvideo_args.dit_config = HunyuanVideoConfig()
-    fastvideo_args.dit_config = WanVideoConfig()
-
-    fastvideo_args.device = torch.device(f"cuda:{local_rank}")
-    fastvideo_args.device_str = f"cuda:{local_rank}"
-    device = fastvideo_args.device
-    torch.cuda.set_device(device)
-
-    loader = TransformerLoader()
-    print(f"--> loading transformer to device {device} on rank {rank}")
-    transformer = loader.load(transformer_path, "",
-                              fastvideo_args).to(device, dtype=precision)
-
-    # teacher_transformer = deepcopy(transformer)
-    if args.use_ema:
-        raise NotImplementedError("EMA is not supported for v1 distillation.")
-        ema_transformer = deepcopy(transformer)
-    else:
-        ema_transformer = None
-
-    if args.use_lora:
-        raise NotImplementedError("LoRA is not supported for v1 distillation.")
-        assert args.model_type == "mochi", "LoRA is only supported for Mochi model."
-        transformer.requires_grad_(False)
-        transformer_lora_config = LoraConfig(
-            r=args.lora_rank,
-            lora_alpha=args.lora_alpha,
-            init_lora_weights=True,
-            target_modules=["to_k", "to_q", "to_v", "to_out.0"],
-        )
-        transformer.add_adapter(transformer_lora_config)
-
-    main_print(
-        f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M"
-    )
-    main_print(
-        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
-    )
-
-    if args.use_lora:
-        raise NotImplementedError("LoRA is not supported for v1 distillation.")
-        transformer.config.lora_rank = args.lora_rank
-        transformer.config.lora_alpha = args.lora_alpha
-        transformer.config.lora_target_modules = [
-            "to_k", "to_q", "to_v", "to_out.0"
-        ]
-        transformer._no_split_modules = no_split_modules
-        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](
-            transformer)
-
-    main_print("--> model loaded")
-
-    if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, no_split_modules,
-                                 args.selective_checkpointing)
-        apply_fsdp_checkpointing(teacher_transformer, no_split_modules,
-                                 args.selective_checkpointing)
-        if args.use_ema:
-            apply_fsdp_checkpointing(ema_transformer, no_split_modules,
-                                     args.selective_checkpointing)
-    # Set model as trainable.
-    transformer.train()
-    transformer.requires_grad_(True)
-    # teacher_transformer.requires_grad_(False)
-    if args.use_ema:
-        ema_transformer.requires_grad_(False)
-    noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift)
-    if args.scheduler_type == "pcm_linear_quadratic":
-        linear_steps = int(noise_scheduler.config.num_train_timesteps *
-                           args.linear_range)
-        sigmas = linear_quadratic_schedule(
-            noise_scheduler.config.num_train_timesteps,
-            args.linear_quadratic_threshold,
-            linear_steps,
-        )
-        sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
-    else:
-        sigmas = noise_scheduler.sigmas
-    solver = EulerSolver(
-        sigmas.numpy()[::-1],
-        noise_scheduler.config.num_train_timesteps,
-        euler_timesteps=args.num_euler_timesteps,
-    )
-    solver.to(device)
-    params_to_optimize = transformer.parameters()
-    # l = list(params_to_optimize)
-    # for p in params_to_optimize:
-    #     main_print(type(p))
-    #     main_print(f"--> p: {p.shape}")
-    #     main_print(f"--> p: {p.dtype}")
-    #     main_print(f"--> p: {p.device}")
-    #     main_print(f"--> p: {p.requires_grad}")
-    #     main_print('=------------------------')
-    #     break
-    # print(f"--> params_to_optimize: {list(params_to_optimize)}")
-
-    params_to_optimize = list(
-        filter(lambda p: p.requires_grad, params_to_optimize))
-    # print(f"--> params_to_optimize2: {params_to_optimize}")
-
-    optimizer = torch.optim.AdamW(
-        params_to_optimize,
-        lr=args.learning_rate,
-        betas=(0.9, 0.999),
-        weight_decay=args.weight_decay,
-        eps=1e-8,
-    )
-    # optimizer = None
-
-    init_steps = 0
-    if args.resume_from_lora_checkpoint:
-        transformer, optimizer, init_steps = resume_lora_optimizer(
-            transformer, args.resume_from_lora_checkpoint, optimizer)
-    main_print(f"optimizer: {optimizer}")
-
-    # todo add lr scheduler
-    lr_scheduler = get_scheduler(
-        args.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * world_size,
-        num_training_steps=args.max_train_steps * world_size,
-        num_cycles=args.lr_num_cycles,
-        power=args.lr_power,
-        last_epoch=init_steps - 1,
-    )
-
-    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t,
-                                  args.cfg)
-    uncond_prompt_embed = train_dataset.uncond_prompt_embed
-    uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = (LengthGroupedSampler(
-        args.train_batch_size,
-        rank=rank,
-        world_size=world_size,
-        lengths=train_dataset.lengths,
-        group_frame=args.group_frame,
-        group_resolution=args.group_resolution,
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
-        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
-
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        collate_fn=latent_collate_function,
-        pin_memory=True,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-        drop_last=True,
-    )
-
-    num_update_steps_per_epoch = math.ceil(
-        len(train_dataloader) / args.gradient_accumulation_steps *
-        args.sp_size / args.train_sp_batch_size)
-    args.num_train_epochs = math.ceil(args.max_train_steps /
-                                      num_update_steps_per_epoch)
-
-    # if rank <= 0:
-    # project = args.tracker_project_name or "fastvideo"
-    # wandb.init(project=project, config=args)
-
-    # Train!
-    total_batch_size = (world_size * args.gradient_accumulation_steps /
-                        args.sp_size * args.train_sp_batch_size)
-    main_print("***** Running training *****")
-    main_print(f"  Num examples = {len(train_dataset)}")
-    main_print(f"  Dataloader size = {len(train_dataloader)}")
-    main_print(f"  Num Epochs = {args.num_train_epochs}")
-    main_print(f"  Resume training from step {init_steps}")
-    main_print(
-        f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(
-        f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
-    )
-    main_print(
-        f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
-    main_print(f"  Total optimization steps = {args.max_train_steps}")
-    main_print(
-        f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B"
-    )
-    # print dtype
-    main_print(
-        f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
-
-    # Potentially load in the weights and states from a previous save
-    if args.resume_from_checkpoint:
-        assert NotImplementedError(
-            "resume_from_checkpoint is not supported now.")
-        # TODO
-
-    progress_bar = tqdm(
-        range(0, args.max_train_steps),
-        initial=init_steps,
-        desc="Steps",
-        # Only show the progress bar once on each machine.
-        disable=local_rank > 0,
-    )
-
-    loader = sp_parallel_dataloader_wrapper(
-        train_dataloader,
-        device,
-        args.train_batch_size,
-        args.sp_size,
-        args.train_sp_batch_size,
-    )
-
-    step_times = deque(maxlen=100)
-
-    # todo future
-    for i in range(init_steps):
-        next(loader)
-
-    # log_validation(args, transformer, device,
-    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
-    def get_num_phases(multi_phased_distill_schedule, step):
-        # step-phase,step-phase
-        multi_phases = multi_phased_distill_schedule.split(",")
-        phase = multi_phases[-1].split("-")[-1]
-        for step_phases in multi_phases:
-            phase_step, phase = step_phases.split("-")
-            if step <= int(phase_step):
-                return int(phase)
-        return phase
-
-    for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.time()
-        assert args.multi_phased_distill_schedule is not None
-        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
-
-        loss, grad_norm, pred_norm = distill_one_step(
-            transformer,
-            args.model_type,
-            None,  # teacher_transformer
-            ema_transformer,
-            optimizer,
-            lr_scheduler,
-            loader,
-            noise_scheduler,
-            solver,
-            noise_random_generator,
-            args.gradient_accumulation_steps,
-            args.sp_size,
-            args.max_grad_norm,
-            uncond_prompt_embed,
-            uncond_prompt_mask,
-            args.num_euler_timesteps,
-            num_phases,
-            args.not_apply_cfg_solver,
-            args.distill_cfg,
-            args.ema_decay,
-            args.pred_decay_weight,
-            args.pred_decay_type,
-            args.hunyuan_teacher_disable_cfg,
-        )
-
-        step_time = time.time() - start_time
-        step_times.append(step_time)
-        avg_step_time = sum(step_times) / len(step_times)
-
-        progress_bar.set_postfix({
-            "loss": f"{loss:.4f}",
-            "step_time": f"{step_time:.2f}s",
-            "grad_norm": grad_norm,
-            "phases": num_phases,
-        })
-        progress_bar.update(1)
-        if rank <= 0:
-            wandb.log(
-                {
-                    "train_loss":
-                    loss,
-                    "learning_rate":
-                    lr_scheduler.get_last_lr()[0],
-                    "step_time":
-                    step_time,
-                    "avg_step_time":
-                    avg_step_time,
-                    "grad_norm":
-                    grad_norm,
-                    "pred_fro_norm":
-                    pred_norm["fro"],  # codespell:ignore
-                    "pred_largest_singular_value":
-                    pred_norm["largest singular value"],
-                    "pred_absolute_mean":
-                    pred_norm["absolute mean"],
-                    "pred_absolute_max":
-                    pred_norm["absolute max"],
-                },
-                step=step,
-            )
-        if step % args.checkpointing_steps == 0:
-            if args.use_lora:
-                # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank,
-                                     args.output_dir, step)
-            else:
-                # Your existing checkpoint saving code
-                if args.use_ema:
-                    save_checkpoint(ema_transformer, rank, args.output_dir,
-                                    step)
-                else:
-                    save_checkpoint(transformer, rank, args.output_dir, step)
-            dist.barrier()
-        if args.log_validation and step % args.validation_steps == 0:
-            log_validation(
-                args,
-                transformer,
-                device,
-                torch.bfloat16,
-                step,
-                scheduler_type=args.scheduler_type,
-                shift=args.shift,
-                num_euler_timesteps=args.num_euler_timesteps,
-                linear_quadratic_threshold=args.linear_quadratic_threshold,
-                linear_range=args.linear_range,
-                ema=False,
-            )
-            if args.use_ema:
-                log_validation(
-                    args,
-                    ema_transformer,
-                    device,
-                    torch.bfloat16,
-                    step,
-                    scheduler_type=args.scheduler_type,
-                    shift=args.shift,
-                    num_euler_timesteps=args.num_euler_timesteps,
-                    linear_quadratic_threshold=args.linear_quadratic_threshold,
-                    linear_range=args.linear_range,
-                    ema=True,
-                )
-
-    if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir,
-                             args.max_train_steps)
-    else:
-        save_checkpoint(transformer, rank, args.output_dir,
-                        args.max_train_steps)
-
-    if get_sequence_parallel_state():
-        destroy_sequence_parallel_group()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument("--model_type",
-                        type=str,
-                        default="mochi",
-                        help="The type of model to train.")
-
-    # dataset & dataloader
-    parser.add_argument("--data_json_path", type=str, required=True)
-    parser.add_argument("--num_height", type=int, default=480)
-    parser.add_argument("--num_width", type=int, default=848)
-    parser.add_argument("--num_frames", type=int, default=163)
-    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("--num_latent_t",
-                        type=int,
-                        default=28,
-                        help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true")  # TODO
-    parser.add_argument("--group_resolution", action="store_true")  # TODO
-
-    # text encoder & vae & diffusion model
-    parser.add_argument("--pretrained_model_name_or_path", type=str)
-    # parser.add_argument("--dit_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-
-    # diffusion setting
-    parser.add_argument("--ema_decay", type=float, default=0.95)
-    parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument("--cfg", type=float, default=0.1)
-
-    # validation & logs
-    parser.add_argument("--validation_prompt_dir", type=str)
-    parser.add_argument("--validation_sampling_steps", type=str, default="64")
-    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
-
-    parser.add_argument("--validation_steps", type=float, default=64)
-    parser.add_argument("--log_validation", action="store_true")
-    parser.add_argument("--tracker_project_name", type=str, default=None)
-    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("--shift", type=float, default=1.0)
-    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(
-        "--resume_from_lora_checkpoint",
-        type=str,
-        default=None,
-        help=
-        ("Whether training should be resumed from a previous lora 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***."),
-    )
-
-    # optimizer & scheduler & Training
-    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(
-        "--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=10,
-        help="Number of steps for the warmup in the lr scheduler.",
-    )
-    parser.add_argument("--max_grad_norm",
-                        default=1.0,
-                        type=float,
-                        help="Max gradient norm.")
-    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("--selective_checkpointing", type=float, default=1.0)
-    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(
-        "--use_cpu_offload",
-        action="store_true",
-        help=
-        "Whether to use CPU offload for param & gradient & optimizer states.",
-    )
-
-    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(
-        "--use_lora",
-        action="store_true",
-        default=False,
-        help="Whether to use LoRA for finetuning.",
-    )
-    parser.add_argument("--lora_alpha",
-                        type=int,
-                        default=256,
-                        help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank",
-                        type=int,
-                        default=128,
-                        help="LoRA rank parameter. ")
-    parser.add_argument("--fsdp_sharding_startegy", default="full")
-
-    # lr_scheduler
-    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("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument(
-        "--lr_num_cycles",
-        type=int,
-        default=1,
-        help="Number of cycles in the learning rate scheduler.",
-    )
-    parser.add_argument(
-        "--lr_power",
-        type=float,
-        default=1.0,
-        help="Power factor of the polynomial scheduler.",
-    )
-    parser.add_argument(
-        "--not_apply_cfg_solver",
-        action="store_true",
-        help="Whether to apply the cfg_solver.",
-    )
-    parser.add_argument("--distill_cfg",
-                        type=float,
-                        default=3.0,
-                        help="Distillation coefficient.")
-    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument("--scheduler_type",
-                        type=str,
-                        default="pcm",
-                        help="The scheduler type to use.")
-    parser.add_argument(
-        "--linear_quadratic_threshold",
-        type=float,
-        default=0.025,
-        help="Threshold for linear quadratic scheduler.",
-    )
-    parser.add_argument(
-        "--linear_range",
-        type=float,
-        default=0.5,
-        help="Range for linear quadratic scheduler.",
-    )
-    parser.add_argument("--weight_decay",
-                        type=float,
-                        default=0.001,
-                        help="Weight decay to apply.")
-    parser.add_argument("--use_ema",
-                        action="store_true",
-                        help="Whether to use EMA.")
-    parser.add_argument("--multi_phased_distill_schedule",
-                        type=str,
-                        default=None)
-    parser.add_argument("--pred_decay_weight", type=float, default=0.0)
-    parser.add_argument("--pred_decay_type", default="l1")
-    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
-    parser.add_argument(
-        "--master_weight_type",
-        type=str,
-        default="fp32",
-        help="Weight type to use - fp32 or bf16.",
-    )
-    args = parser.parse_args()
-    main(args)

fastvideo/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.1.1"
+version = "0.1.0"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.8"

scripts/preprocess/preprocess_wan_data.sh

@@ -12,6 +12,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --model_path $MODEL_PATH \
     --data_merge_path $DATA_MERGE_PATH \
     --preprocess_video_batch_size=4 \
+    --preprocess_text_batch_size=4 \
     --max_height=480 \
     --max_width=832 \
     --num_frames=81 \
@@ -20,6 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
     --model_type $MODEL_TYPE \
     --text_encoder_name $TEXT_ENCODER_PATH \
     --train_fps 16 \
-    --validation_prompt_txt $VALIDATION_PATH \
-    --samples_per_file 108 \
-    --flush_frequency 108
+    --validation_prompt_txt $VALIDATION_PATH

scripts/preprocess/preprocess_wan_data_v1.sh

@@ -1,30 +0,0 @@
-# export WANDB_MODE="offline"
-GPU_NUM=1 # 2,4,8
-# MODEL_PATH="/home/ray/.cache/huggingface/hub/models--Wan-AI--Wan2.1-T2V-1.3B-Diffusers/snapshots/0fad780a534b6463e45facd96134c9f345acfa5b"
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_MERGE_PATH="data/cats_480/merge.txt"
-OUTPUT_DIR="data/cats_480_latents/"
-VALIDATION_PATH="assets/prompt.txt"
-
-# torchrun --nproc_per_node=$GPU_NUM \
-#     fastvideo/data_preprocess/preprocess_vae_latents_v1.py \
-#     --model_path $MODEL_PATH \
-#     --data_merge_path $DATA_MERGE_PATH \
-#     --train_batch_size=1 \
-#     --max_height=480 \
-#     --max_width=832 \
-#     --num_frames=81 \
-#     --dataloader_num_workers 1 \
-#     --output_dir=$OUTPUT_DIR \
-#     --train_fps 16
-
-# torchrun --nproc_per_node=$GPU_NUM \
-#     fastvideo/data_preprocess/preprocess_text_embeddings_v1.py \
-#     --model_path $MODEL_PATH \
-#     --output_dir=$OUTPUT_DIR 
-
-torchrun --nproc_per_node=1 \
-    fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py \
-    --model_path $MODEL_PATH \
-    --output_dir=$OUTPUT_DIR \
-    --validation_prompt_txt $VALIDATION_PATH

train_v1.sh

@@ -2,9 +2,9 @@ export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
-DATA_DIR=data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset
-VALIDATION_DIR=data/HD-Mixkit-Finetune-Wan/validation_parquet_dataset
-NUM_GPUS=1
+DATA_DIR=data/cats_480_single_latents_parq/combined_parquet_dataset
+VALIDATION_DIR=data/cats_480_single_latents_parq/validation_parquet_dataset
+NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
@@ -21,17 +21,17 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --data_path "$DATA_DIR"\
     --validation_prompt_dir "$VALIDATION_DIR"\
     --train_batch_size=1\
-    --num_latent_t 4 \
+    --num_latent_t 20 \
     --sp_size $NUM_GPUS \
     --tp_size $NUM_GPUS \
     --train_sp_batch_size 1\
-    --dataloader_num_workers 1\
+    --dataloader_num_workers 6\
     --gradient_accumulation_steps=1\
-    --max_train_steps=120 \
+    --max_train_steps=3000 \
     --learning_rate=1e-6\
     --mixed_precision="bf16"\
-    --checkpointing_steps=60 \
-    --validation_steps 20\
+    --checkpointing_steps=5000 \
+    --validation_steps 200\
     --validation_sampling_steps "2,4,8" \
     --log_validation \
     --checkpoints_total_limit 3\