From 1d5bf447b358c547d2c4cb9e7e95b81114520a75 Mon Sep 17 00:00:00 2001
From: zR <2448370773@qq.com>
Date: Sat, 5 Oct 2024 22:28:45 +0800
Subject: [PATCH 1/4] Multi-GPU parallel encoding support for training videos.

---
 .gitignore                  |   3 +
 README.md                   |   7 +-
 README_zh.md                |  95 +++++++++++++++++++++++
 prepare_dataset.sh          |  46 +++++------
 train_text_to_video_lora.sh |  14 ++--
 training/prepare_dataset.py | 151 +++++++++++++++++++++++++-----------
 6 files changed, 238 insertions(+), 78 deletions(-)
 create mode 100644 README_zh.md

diff --git a/.gitignore b/.gitignore
index e652a10..bcc6640 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,9 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+# JetBrains
+.idea
+
 # C extensions
 *.so
 
diff --git a/README.md b/README.md
index c1e585f..093eb41 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,9 @@
-# Finetuning CogVideoX
+# CogVideoX Factory
+
+## Introduction
+
+This is a repos for CogVideoX Fine-tuning.
+
 
 
 ## Dataset Preparation
diff --git a/README_zh.md b/README_zh.md
new file mode 100644
index 0000000..3758d84
--- /dev/null
+++ b/README_zh.md
@@ -0,0 +1,95 @@
+# CogVideoX Factory
+
+## 简介
+
+这是用于 CogVideoX 微调的仓库。
+
+## 数据集准备
+
+创建两个文件，一个文件包含以换行符分隔的提示词，另一个文件包含以换行符分隔的视频数据路径（视频文件的路径必须相对于您在指定 `--data_root` 时传递的路径）。让我们通过一个例子来更好地理解这一点！
+
+假设您将 `--data_root` 指定为 `/dataset`，并且该目录包含文件：`prompts.txt` 和 `videos.txt`。
+
+`prompts.txt` 文件应包含以换行符分隔的提示词：
+
+```
+一段黑白动画序列，主角是一只名为 Rabbity Ribfried 的兔子和一只拟人化的山羊，在一个充满音乐和趣味的环境中，展示他们不断发展的互动。
+一段黑白动画序列，场景在船甲板上，主角是一只名为 Bully Bulldoger 的斗牛犬角色，展示了夸张的面部表情和肢体语言。角色从自信到专注，再到紧张和痛苦，展示了一系列情绪，随着它克服挑战。船的内部在背景中保持静止，只有简单的细节，如钟声和开着的门。角色的动态动作和变化的表情推动了故事的发展，没有镜头移动，确保观众专注于其不断变化的反应和肢体动作。
+...
+```
+
+`videos.txt` 文件应包含以换行符分隔的视频文件路径。请注意，路径应相对于 `--data_root` 目录。
+
+```bash
+videos/00000.mp4
+videos/00001.mp4
+...
+```
+
+总体而言，如果您在数据集根目录运行 `tree` 命令，您的数据集应如下所示：
+
+```bash
+/dataset
+├── prompts.txt
+├── videos.txt
+├── videos
+    ├── videos/00000.mp4
+    ├── videos/00001.mp4
+    ├── ...
+```
+
+使用此格式时，`--caption_column` 必须是 `prompts.txt`，`--video_column` 必须是 `videos.txt`。如果您的数据存储在 CSV 文件中，您也可以指定 `--dataset_file` 为 CSV 的路径，`--caption_column` 和 `--video_column` 为 CSV 文件中的实际列名。
+
+例如，让我们使用这个 [Disney 数据集](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) 进行微调。要下载，可以使用 🤗 Hugging Face CLI。
+
+```bash
+huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir video-dataset-disney
+```
+
+## 训练
+
+TODO
+
+请查看 `training/*.sh`
+
+注意：未在 MPS 上测试
+
+## 内存需求
+
+训练支持并验证的内存优化包括：
+
+- 来自 [TorchAO](https://github.com/pytorch/ao) 的 `CPUOffloadOptimizer`。
+- 来自 [bitsandbytes](https://huggingface.co/docs/bitsandbytes/optimizers) 的低位优化器。
+
+### LoRA 微调
+
+<details>
+<summary> AdamW </summary>
+
+With `train_batch_size = 1`:
+
+|       model        | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing |
+|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:|
+| THUDM/CogVideoX-2b |    16     |          False         |         12.945         |          43.764          |         46.918          |       24.234         |
+| THUDM/CogVideoX-2b |    16     |          True          |         12.945         |          12.945          |         21.121          |       24.234         |
+| THUDM/CogVideoX-2b |    64     |          False         |         13.035         |          44.314          |         47.469          |       24.469         |
+| THUDM/CogVideoX-2b |    64     |          True          |         13.036         |          13.035          |         21.564          |       24.500         |
+| THUDM/CogVideoX-2b |    256    |          False         |         13.095         |          45.826          |         48.990          |       25.543         |
+| THUDM/CogVideoX-2b |    256    |          True          |         13.094         |          13.095          |         22.344          |       25.537         |
+| THUDM/CogVideoX-5b |    16     |          True          |         19.742         |          19.742          |         28.746          |       38.123         |
+| THUDM/CogVideoX-5b |    64     |          True          |         20.006         |          20.818          |         30.338          |       38.738         |
+| THUDM/CogVideoX-5b |    256    |          True          |         20.771         |          22.119          |         31.939          |       41.537         |
+
+With `train_batch_size = 4`:
+
+|       model        | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing |
+|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:|
+| THUDM/CogVideoX-2b |    16     |          True          |         12.945         |          21.803          |         21.814          |       24.322         |
+| THUDM/CogVideoX-2b |    64     |          True          |         13.035         |          22.254          |         22.254          |       24.572         |
+| THUDM/CogVideoX-2b |    256    |          True          |         13.094         |          22.020          |         22.033          |       25.574         |
+| THUDM/CogVideoX-5b |    16     |          True          |         19.742         |          46.492          |         46.492          |       38.197         |
+| THUDM/CogVideoX-5b |    64     |          True          |         20.006         |          47.805          |         47.805          |       39.365         |
+| THUDM/CogVideoX-5b |    256    |          True          |         20.771         |          47.268          |         47.332          |       41.008         |
+
+> [!NOTE]
+> 
\ No newline at end of file
diff --git a/prepare_dataset.sh b/prepare_dataset.sh
index b0414e9..84086a2 100755
--- a/prepare_dataset.sh
+++ b/prepare_dataset.sh
@@ -1,12 +1,11 @@
 #!/bin/bash
 
-MODEL_ID="THUDM/CogVideoX-2b"
+MODEL_ID="/share/official_pretrains/hf_home/CogVideoX-5b"
 
-# For more details on the expected data format, please refer to the README.
-DATA_ROOT="/raid/aryan/video-dataset-tom-and-jerry"  # This needs to be the path to the base directory where your videos are located.
+DATA_ROOT="/share/home/zyx/disney_cogvideox"
 CAPTION_COLUMN="prompts.txt"
 VIDEO_COLUMN="videos.txt"
-OUTPUT_DIR="/raid/aryan/video-dataset-tom-and-jerry-encoded"
+OUTPUT_DIR="/share/home/zyx/disney_cogvideox-encoded-multi"
 HEIGHT=480
 WIDTH=720
 MAX_NUM_FRAMES=49
@@ -14,29 +13,24 @@ MAX_SEQUENCE_LENGTH=226
 TARGET_FPS=8
 BATCH_SIZE=1
 DTYPE=fp32
+NUM_GPUS=8
 
-# To create a folder-style dataset structure without pre-encoding videos and captions'
-CMD_WITHOUT_PRE_ENCODING="\
-  python3 training/prepare_dataset.py \
-    --model_id $MODEL_ID \
-    --data_root $DATA_ROOT \
-    --caption_column $CAPTION_COLUMN \
-    --video_column $VIDEO_COLUMN \
-    --output_dir $OUTPUT_DIR \
-    --height $HEIGHT \
-    --width $WIDTH \
-    --max_num_frames $MAX_NUM_FRAMES \
-    --max_sequence_length $MAX_SEQUENCE_LENGTH \
-    --target_fps $TARGET_FPS \
-    --batch_size $BATCH_SIZE \
-    --dtype $DTYPE
-"
-
-CMD_WITH_PRE_ENCODING="$CMD_WITHOUT_PRE_ENCODING --save_tensors"
-
-# Select which you'd like to run
-CMD=$CMD_WITH_PRE_ENCODING
+CMD="torchrun --nproc_per_node=$NUM_GPUS \
+  training/prepare_dataset.py \
+  --model_id $MODEL_ID \
+  --data_root $DATA_ROOT \
+  --caption_column $CAPTION_COLUMN \
+  --video_column $VIDEO_COLUMN \
+  --output_dir $OUTPUT_DIR \
+  --height $HEIGHT \
+  --width $WIDTH \
+  --max_num_frames $MAX_NUM_FRAMES \
+  --max_sequence_length $MAX_SEQUENCE_LENGTH \
+  --target_fps $TARGET_FPS \
+  --batch_size $BATCH_SIZE \
+  --dtype $DTYPE \
+  --save_tensors"
 
 echo "===== Running \`$CMD\` ====="
 eval $CMD
-echo -ne "===== Finished running script =====\n"
+echo -ne "===== Finished running script =====\n"
\ No newline at end of file
diff --git a/train_text_to_video_lora.sh b/train_text_to_video_lora.sh
index 9af8187..ab9013d 100755
--- a/train_text_to_video_lora.sh
+++ b/train_text_to_video_lora.sh
@@ -4,7 +4,7 @@ export WANDB_MODE="offline"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 
-GPU_IDS="0"
+GPU_IDS="0,1,2,3,4,5,6,7"
 
 # Training Configurations
 # Experiment with as many hyperparameters as you want!
@@ -19,22 +19,26 @@ ACCELERATE_CONFIG_FILE="accelerate_configs/uncompiled_1.yaml"
 # Absolute path to where the data is located. Make sure to have read the README for how to prepare data.
 # This example assumes you downloaded an already prepared dataset from HF CLI as follows:
 #   huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir /path/to/my/datasets/disney-dataset
-DATA_ROOT="/path/to/my/datasets/disney-dataset"
-CAPTION_COLUMN="prompt.txt"
+
+DATA_ROOT="/share/home/zyx/disney_cogvideox-encoded-multi"
+CAPTION_COLUMN="prompts.txt"
 VIDEO_COLUMN="videos.txt"
+MODEL_PATH="/share/official_pretrains/hf_home/CogVideoX-5b"
+
 
 # Launch experiments with different hyperparameters
 for learning_rate in "${LEARNING_RATES[@]}"; do
   for lr_schedule in "${LR_SCHEDULES[@]}"; do
     for optimizer in "${OPTIMIZERS[@]}"; do
       for steps in "${MAX_TRAIN_STEPS[@]}"; do
-        output_dir="/path/to/my/models/cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
+        output_dir="cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
 
         cmd="accelerate launch --config_file $ACCELERATE_CONFIG_FILE --gpu_ids $GPU_IDS training/cogvideox_text_to_video_lora.py \
-          --pretrained_model_name_or_path THUDM/CogVideoX-5b \
+          --pretrained_model_name_or_path $MODEL_PATH \
           --data_root $DATA_ROOT \
           --caption_column $CAPTION_COLUMN \
           --video_column $VIDEO_COLUMN \
+          --load_tensors \
           --id_token BW_STYLE \
           --height_buckets 480 \
           --width_buckets 720 \
diff --git a/training/prepare_dataset.py b/training/prepare_dataset.py
index 9d70884..4e54fab 100644
--- a/training/prepare_dataset.py
+++ b/training/prepare_dataset.py
@@ -1,23 +1,21 @@
 #!/usr/bin/env python3
 
-# For folder structure dataset: python3 prepare_dataset.py --model_id THUDM/CogVideoX-2b --data_root /raid/aryan/video-dataset-disney/ --caption_column prompts.txt --video_column videos.txt --output_dir dump --height 480 --width 720 --max_num_frames 49 --max_sequence_length 226 --target_fps 8 --batch_size 1 --dtype fp32
-# For latent/embed structure dataset: python3 prepare_dataset.py --model_id THUDM/CogVideoX-2b --data_root /raid/aryan/video-dataset-disney/ --caption_column prompts.txt --video_column videos.txt --output_dir dump --height 480 --width 720 --max_num_frames 49 --max_sequence_length 226 --target_fps 8 --batch_size 1 --dtype fp32 --save_tensors
-
 import argparse
 import gc
+import os
 import pathlib
 import traceback
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
 import torch
+import torch.distributed as dist
 from diffusers import AutoencoderKLCogVideoX
 from diffusers.utils import export_to_video, get_logger
 from torchvision import transforms
 from transformers import T5EncoderModel, T5Tokenizer
+from tqdm import tqdm
 
-
-# Must import after importing torch, otherwise there's a nasty segfault when loading text_encoder/vae
 import decord  # isort:skip
 
 decord.bridge.set_bridge("torch")
@@ -104,7 +102,7 @@ def get_args() -> Dict[str, Any]:
 
 def load_dataset_from_local_path(
     data_root: pathlib.Path, caption_column: str, video_column: str
-) -> Tuple[List[str], List[str]]:
+) -> Tuple[List[str], List[pathlib.Path]]:
     if not data_root.exists():
         raise ValueError("Root folder for videos does not exist")
 
@@ -127,7 +125,7 @@ def load_dataset_from_local_path(
 
     if any(not path.is_file() for path in video_paths):
         raise ValueError(
-            f"Expected `{video_column=}` to be a path to a file in `{data_root=}` containing line-separated paths to video data but found atleast one path that is not a valid file."
+            f"Expected `{video_column}` to be a path to a file in `{data_root}` containing line-separated paths to video data but found at least one path that is not a valid file."
         )
 
     return prompts, video_paths
@@ -135,7 +133,7 @@ def load_dataset_from_local_path(
 
 def load_dataset_from_csv(
     data_root: pathlib.Path, dataset_file: pathlib.Path, caption_column: str, video_column: str
-) -> Tuple[List[str], List[str]]:
+) -> Tuple[List[str], List[pathlib.Path]]:
     df = pd.read_csv(dataset_file)
     prompts = df[caption_column].tolist()
     video_paths = df[video_column].tolist()
@@ -143,7 +141,7 @@ def load_dataset_from_csv(
 
     if any(not path.is_file() for path in video_paths):
         raise ValueError(
-            f"Expected `{video_column=}` to be a path to a file in `{data_root=}` containing line-separated paths to video data but found atleast one path that is not a valid file."
+            f"Expected `{video_column}` to be a path to a file in `{data_root}` containing line-separated paths to video data but found at least one path that is not a valid file."
         )
 
     return prompts, video_paths
@@ -151,7 +149,7 @@ def load_dataset_from_csv(
 
 def load_and_preprocess_video(
     path: pathlib.Path, height: int, width: int, max_num_frames: int, video_transforms, num_threads: int = 0
-) -> torch.Tensor:
+) -> Optional[torch.Tensor]:
     frames = None
 
     try:
@@ -160,11 +158,11 @@ def load_and_preprocess_video(
 
         if video_num_frames < max_num_frames:
             logger.warning(
-                f"Video at `{path.as_posix()}` should have atleast `{max_num_frames=}`, but got only `{video_num_frames=}`. Skipping it."
+                f"Video at `{path.as_posix()}` should have at least `{max_num_frames=}`, but got only `{video_num_frames=}`. Skipping it."
             )
-            return
+            return None
 
-        indices = list(range(0, video_num_frames, video_num_frames // max_num_frames))
+        indices = list(range(0, video_num_frames, max(video_num_frames // max_num_frames, 1)))
         frames: torch.Tensor = video_reader.get_batch(indices)
         frames = frames[:max_num_frames].float()
         frames = frames.permute(0, 3, 1, 2).contiguous()
@@ -241,7 +239,7 @@ def encode_prompt(
 def compute_prompt_embeddings(
     tokenizer: T5Tokenizer,
     text_encoder: T5EncoderModel,
-    prompt: str,
+    prompts: List[str],
     max_sequence_length: int,
     device: torch.device,
     dtype: torch.dtype,
@@ -251,7 +249,7 @@ def compute_prompt_embeddings(
         prompt_embeds = encode_prompt(
             tokenizer,
             text_encoder,
-            prompt,
+            prompts,
             num_videos_per_prompt=1,
             max_sequence_length=max_sequence_length,
             device=device,
@@ -262,7 +260,7 @@ def compute_prompt_embeddings(
             prompt_embeds = encode_prompt(
                 tokenizer,
                 text_encoder,
-                prompt,
+                prompts,
                 num_videos_per_prompt=1,
                 max_sequence_length=max_sequence_length,
                 device=device,
@@ -272,7 +270,7 @@ def compute_prompt_embeddings(
 
 
 def save_videos(
-    videos: torch.Tensor, video_paths: List[str], prompts: List[str], output_dir: pathlib.Path, target_fps: int = 8
+    videos: torch.Tensor, video_paths: List[pathlib.Path], prompts: List[str], output_dir: pathlib.Path, target_fps: int = 8
 ) -> None:
     assert videos.size(0) == len(video_paths)
 
@@ -289,13 +287,13 @@ def save_videos(
     videos_pil = [[to_pil_image(frame) for frame in video] for video in videos]
 
     for video, video_path in zip(videos_pil, video_paths):
-        filename = video_dir.joinpath(pathlib.Path(video_path).name)
+        filename = video_dir.joinpath(video_path.name)
         logger.debug(f"Saving video to `{filename}`")
         export_to_video(video, filename.as_posix(), fps=target_fps)
 
     with open(output_dir.joinpath("videos.txt").as_posix(), "w", encoding="utf-8") as file:
         for video_path in video_paths:
-            file.write(f"videos/{pathlib.Path(video_path).name}\n")
+            file.write(f"videos/{video_path.name}\n")
 
     with open(output_dir.joinpath("prompts.txt").as_posix(), "w", encoding="utf-8") as file:
         for prompt in prompts:
@@ -305,7 +303,7 @@ def save_videos(
 def save_latents_and_embeddings(
     latents: torch.Tensor,
     prompt_embeds: torch.Tensor,
-    video_paths: List[str],
+    video_paths: List[pathlib.Path],
     prompts: List[str],
     output_dir: pathlib.Path,
 ) -> None:
@@ -321,27 +319,20 @@ def save_latents_and_embeddings(
     embeds_dir.mkdir(parents=True, exist_ok=True)
 
     for latent, embed, video_path in zip(latents, prompt_embeds, video_paths):
-        # Need to perform the clone, otherwise the entire `latents` or `prompt_embeds` tensor is
-        # saved for every single video/prompt embedding. This is due to us viewing a slice of a
-        # large tensor when iteratively saving stuff here.
         latent = latent.clone()
         embed = embed.clone()
 
-        video_path = pathlib.Path(video_path)
-        filename_without_ext = video_path.name.split(".")[0]
+        filename_without_ext = video_path.stem
 
-        latent_filename = latents_dir.joinpath(filename_without_ext)
-        embed_filename = embeds_dir.joinpath(filename_without_ext)
-
-        latent_filename = f"{latent_filename}.pt"
-        embed_filename = f"{embed_filename}.pt"
+        latent_filename = latents_dir.joinpath(f"{filename_without_ext}.pt")
+        embed_filename = embeds_dir.joinpath(f"{filename_without_ext}.pt")
 
         torch.save(latent, latent_filename)
         torch.save(embed, embed_filename)
 
     with open(output_dir.joinpath("videos.txt").as_posix(), "w", encoding="utf-8") as file:
         for video_path in video_paths:
-            file.write(f"videos/{pathlib.Path(video_path).name}\n")
+            file.write(f"videos/{video_path.name}\n")
 
     with open(output_dir.joinpath("prompts.txt").as_posix(), "w", encoding="utf-8") as file:
         for prompt in prompts:
@@ -349,7 +340,23 @@ def save_latents_and_embeddings(
 
 
 @torch.no_grad()
-def main(args: Dict[str, Any]) -> None:
+def main():
+    args = get_args()
+
+    # Initialize distributed processing
+    if 'LOCAL_RANK' in os.environ:
+        local_rank = int(os.environ['LOCAL_RANK'])
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(backend="nccl")
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+    else:
+        # Single GPU
+        local_rank = 0
+        world_size = 1
+        rank = 0
+        torch.cuda.set_device(local_rank)
+
     data_root = pathlib.Path(args.data_root)
     dataset_file = None
     if args.dataset_file:
@@ -367,10 +374,18 @@ def main(args: Dict[str, Any]) -> None:
         ]
     )
 
+    # Preprocess videos with progress bar
     prompts_usable = []
     video_paths_usable = []
     videos = []
-    for prompt, path in zip(prompts, video_paths):
+
+    # Only show progress bar on the main process
+    if rank == 0:
+        iterator = tqdm(zip(prompts, video_paths), total=len(prompts), desc="Load and Preprocess videos")
+    else:
+        iterator = zip(prompts, video_paths)
+
+    for prompt, path in iterator:
         video = load_and_preprocess_video(
             path, args.height, args.width, args.max_num_frames, video_transforms, args.num_decode_threads
         )
@@ -378,18 +393,47 @@ def main(args: Dict[str, Any]) -> None:
             prompts_usable.append(prompt)
             video_paths_usable.append(path)
             videos.append(video)
+
+    if len(videos) == 0:
+        logger.error("No usable videos found after preprocessing.")
+        return
+
     videos = torch.stack(videos)
 
+    # Split data among GPUs
+    if world_size > 1:
+        total_samples = len(prompts_usable)
+        samples_per_gpu = total_samples // world_size
+        start_index = rank * samples_per_gpu
+        end_index = start_index + samples_per_gpu
+        if rank == world_size - 1:
+            end_index = total_samples  # Make sure the last GPU gets the remaining data
+
+        # Slice the data
+        prompts_usable = prompts_usable[start_index:end_index]
+        video_paths_usable = video_paths_usable[start_index:end_index]
+        videos = videos[start_index:end_index]
+    else:
+        pass
+
+    device = torch.device(f"cuda:{local_rank}")
+
     if not args.save_tensors:
         save_videos(videos, video_paths_usable, prompts_usable, pathlib.Path(args.output_dir), args.target_fps)
     else:
         dtype = DTYPE_MAPPING[args.dtype]
         tokenizer = T5Tokenizer.from_pretrained(args.model_id, subfolder="tokenizer")
         text_encoder = T5EncoderModel.from_pretrained(args.model_id, subfolder="text_encoder", torch_dtype=dtype)
-        text_encoder = text_encoder.to("cuda")
+        text_encoder = text_encoder.to(device)
 
         prompt_embeds_list = []
-        for start_index in range(0, len(prompts_usable), args.batch_size):
+
+        if rank == 0:
+            iterator = tqdm(range(0, len(prompts_usable), args.batch_size), desc="Encoding prompts")
+        else:
+            iterator = range(0, len(prompts_usable), args.batch_size)
+
+        for start_index in iterator:
             end_index = min(len(prompts_usable), start_index + args.batch_size)
             batch_prompts = prompts_usable[start_index:end_index]
 
@@ -398,20 +442,20 @@ def main(args: Dict[str, Any]) -> None:
                 text_encoder,
                 batch_prompts,
                 max_sequence_length=args.max_sequence_length,
-                device="cuda",
+                device=device,
                 dtype=dtype,
             )
-            prompt_embeds_list.append(prompt_embeds)
+            prompt_embeds_list.append(prompt_embeds.to("cpu"))
 
-        prompt_embeds = torch.cat(prompt_embeds_list).to("cpu")
+        prompt_embeds = torch.cat(prompt_embeds_list)
 
         del tokenizer, text_encoder
         gc.collect()
         torch.cuda.empty_cache()
-        torch.cuda.synchronize("cuda")
+        torch.cuda.synchronize(device)
 
         vae = AutoencoderKLCogVideoX.from_pretrained(args.model_id, subfolder="vae", torch_dtype=dtype)
-        vae = vae.to("cuda")
+        vae = vae.to(device)
 
         if args.use_slicing:
             vae.enable_slicing()
@@ -419,11 +463,17 @@ def main(args: Dict[str, Any]) -> None:
             vae.enable_tiling()
 
         encoded_videos = []
-        for start_index in range(0, len(video_paths_usable), args.batch_size):
+
+        if rank == 0:
+            iterator = tqdm(range(0, len(video_paths_usable), args.batch_size), desc="Encoding videos")
+        else:
+            iterator = range(0, len(video_paths_usable), args.batch_size)
+
+        for start_index in iterator:
             end_index = min(len(video_paths_usable), start_index + args.batch_size)
             batch_videos = videos[start_index:end_index]
 
-            batch_videos = batch_videos.to("cuda")
+            batch_videos = batch_videos.to(device)
             batch_videos = batch_videos.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
 
             if args.use_slicing:
@@ -432,19 +482,28 @@ def main(args: Dict[str, Any]) -> None:
             else:
                 encoded_video = vae._encode(batch_videos)
 
-            encoded_videos.append(encoded_video)
+            encoded_videos.append(encoded_video.to("cpu"))
 
-        encoded_videos = torch.cat(encoded_videos).to("cpu")
+        encoded_videos = torch.cat(encoded_videos)
 
         del vae
         gc.collect()
         torch.cuda.empty_cache()
-        torch.cuda.synchronize("cuda")
+        torch.cuda.synchronize(device)
+
+        # Ensure that only one process creates the output directories
+        if world_size > 1:
+            dist.barrier()
 
         save_latents_and_embeddings(
             encoded_videos, prompt_embeds, video_paths_usable, prompts_usable, pathlib.Path(args.output_dir)
         )
 
+    # Finalize distributed processing
+    if world_size > 1:
+        dist.barrier()
+        dist.destroy_process_group()
+
 
 if __name__ == "__main__":
     args = get_args()
@@ -455,4 +514,4 @@ if __name__ == "__main__":
         args.max_num_frames % 4 == 0 or args.max_num_frames % 4 == 1
     ), "`--max_num_frames` must be of form 4 * k or 4 * k + 1 to be compatible with VAE."
 
-    main(args)
+    main()
\ No newline at end of file

From 0bd88a4f4dab9e43e379ffb4369b183c806b54a7 Mon Sep 17 00:00:00 2001
From: zR <2448370773@qq.com>
Date: Thu, 10 Oct 2024 21:38:21 +0800
Subject: [PATCH 2/4] zh readme

---
 README_zh.md     | 119 +++++++++++++++++++++++++++++------------------
 requirements.txt |  29 ++++++------
 2 files changed, 89 insertions(+), 59 deletions(-)

diff --git a/README_zh.md b/README_zh.md
index 3758d84..250c7e3 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -1,24 +1,55 @@
-# CogVideoX Factory
+# CogVideoX Factor 🧪
 
-## 简介
+在 24GB GPU 内存下微调 Cog 系列视频模型以生成自定义视频 ⚡️📼
 
-这是用于 CogVideoX 微调的仓库。
+TODO：添加有趣的视频结果表
+
+## 快速开始
+
+确保已安装所需的依赖：`pip install -r requirements.txt`。
+
+然后下载数据集：
+
+```bash
+# 安装 `huggingface_hub`
+huggingface-cli download   --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset   --local-dir video-dataset-disney
+```
+
+然后启动文本到视频的 LoRA 微调：
+
+```bash
+TODO
+```
+
+我们现在可以使用训练好的模型进行推理：
+
+```python
+TODO
+```
+
+我们还可以使用 LoRA 微调 5B 版本：
+
+```python
+TODO
+```
+
+在下方的部分中，我们提供了有关更多选项的详细信息，这些选项旨在使视频模型的微调尽可能易于使用。
 
 ## 数据集准备
 
-创建两个文件，一个文件包含以换行符分隔的提示词，另一个文件包含以换行符分隔的视频数据路径（视频文件的路径必须相对于您在指定 `--data_root` 时传递的路径）。让我们通过一个例子来更好地理解这一点！
+创建两个文件，一个文件包含逐行分隔的提示，另一个文件包含逐行分隔的视频数据路径（视频文件的路径必须相对于您在指定 `--data_root` 时传递的路径）。让我们通过一个示例来更好地理解这一点！
 
-假设您将 `--data_root` 指定为 `/dataset`，并且该目录包含文件：`prompts.txt` 和 `videos.txt`。
+假设您指定的 `--data_root` 为 `/dataset`，并且该目录包含以下文件：`prompts.txt` 和 `videos.txt`。
 
-`prompts.txt` 文件应包含以换行符分隔的提示词：
+`prompts.txt` 文件应包含逐行分隔的提示：
 
 ```
-一段黑白动画序列，主角是一只名为 Rabbity Ribfried 的兔子和一只拟人化的山羊，在一个充满音乐和趣味的环境中，展示他们不断发展的互动。
-一段黑白动画序列，场景在船甲板上，主角是一只名为 Bully Bulldoger 的斗牛犬角色，展示了夸张的面部表情和肢体语言。角色从自信到专注，再到紧张和痛苦，展示了一系列情绪，随着它克服挑战。船的内部在背景中保持静止，只有简单的细节，如钟声和开着的门。角色的动态动作和变化的表情推动了故事的发展，没有镜头移动，确保观众专注于其不断变化的反应和肢体动作。
+一段黑白动画序列，主角是一只名为 Rabbity Ribfried 的兔子和一只拟人化的山羊，展示了它们在音乐与游戏环境中的互动演变。
+一段黑白动画序列，发生在船甲板上，主角是一只名为 Bully Bulldoger 的斗牛犬，展现了夸张的面部表情和肢体语言。角色从自信、专注逐渐转变为紧张与痛苦，展示了随着挑战出现的情感变化。船的内部在背景中保持静止，只有一些简单的细节，如钟声和敞开的门。角色的动态动作和不断变化的表情推动了叙事，没有摄像机运动来分散注意力。
 ...
 ```
 
-`videos.txt` 文件应包含以换行符分隔的视频文件路径。请注意，路径应相对于 `--data_root` 目录。
+`videos.txt` 文件应包含逐行分隔的视频文件路径。请注意，路径应相对于 `--data_root` 目录。
 
 ```bash
 videos/00000.mp4
@@ -26,7 +57,7 @@ videos/00001.mp4
 ...
 ```
 
-总体而言，如果您在数据集根目录运行 `tree` 命令，您的数据集应如下所示：
+整体而言，如果在数据集根目录运行 `tree` 命令，您的数据集应如下所示：
 
 ```bash
 /dataset
@@ -38,58 +69,54 @@ videos/00001.mp4
     ├── ...
 ```
 
-使用此格式时，`--caption_column` 必须是 `prompts.txt`，`--video_column` 必须是 `videos.txt`。如果您的数据存储在 CSV 文件中，您也可以指定 `--dataset_file` 为 CSV 的路径，`--caption_column` 和 `--video_column` 为 CSV 文件中的实际列名。
+使用此格式时，`--caption_column` 必须是 `prompts.txt`，`--video_column` 必须是 `videos.txt`。如果您将数据存储在 CSV 文件中，还可以指定 `--dataset_file` 为 CSV 的路径，`--caption_column` 和 `--video_column` 为 CSV 文件中的实际列名。
 
-例如，让我们使用这个 [Disney 数据集](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) 进行微调。要下载，可以使用 🤗 Hugging Face CLI。
+例如，让我们使用[这个](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) Disney 数据集进行微调。要下载，您可以使用 🤗 Hugging Face CLI。
 
 ```bash
 huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir video-dataset-disney
 ```
 
+TODO：添加一个关于创建和使用预计算嵌入的部分。
+
 ## 训练
 
-TODO
+我们提供了与 [Cog 系列模型](https://huggingface.co/collections/THUDM/cogvideo-66c08e62f1685a3ade464cce) 兼容的文本到视频和图像到视频生成的训练脚本。
 
-请查看 `training/*.sh`
+查看 `*.sh` 文件
 
 注意：未在 MPS 上测试
 
 ## 内存需求
 
-训练支持并验证的内存优化包括：
+<table align="center">
+<tr>
+  <td align="center"><a href="https://www.youtube.com/watch?v=UvRl4ansfCg"> 使用 PyTorch 消除 OOM</a></td>
+</tr>
+<tr>
+  <td align="center"><img src="assets/slaying-ooms.png" style="width: 480px; height: 480px;"></td>
+</tr>
+</table>
 
-- 来自 [TorchAO](https://github.com/pytorch/ao) 的 `CPUOffloadOptimizer`。
-- 来自 [bitsandbytes](https://huggingface.co/docs/bitsandbytes/optimizers) 的低位优化器。
+支持和验证的内存优化训练选项包括：
+
+- [`torchao`](https://github.com/pytorch/ao) 中的 `CPUOffloadOptimizer`。您可以阅读它的能力和限制 [此处](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload)。简而言之，它允许您使用 CPU 存储可训练的参数和梯度。这导致优化器步骤在 CPU 上进行，需要一个快速的 CPU 优化器，例如 `torch.optim.AdamW(fused=True)` 或在优化器步骤上应用 `torch.compile`。此外，建议不要将模型编译用于训练。梯度裁剪和积累尚不支持。
+- [`bitsandbytes`](https://huggingface.co/docs/bitsandbytes/optimizers) 中的低位优化器。TODO：测试并使 [`torchao`](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim) 工作
+- DeepSpeed Zero2：由于我们依赖 `accelerate`，请按照[本指南](https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed) 配置 `accelerate` 以启用 DeepSpeed Zero2 优化。
+
+> [!IMPORTANT]
+> 内存需求是在运行 `training/prepare_dataset.py` 后报告的，它将视频和字幕转换为潜变量和嵌入。在训练过程中，我们直接加载潜变量和嵌入，而不需要 VAE 或 T5 文本编码器。但是，如果您执行验证/测试，则必须加载这些内容，并增加所需的内存量。不执行验证/测试可以节省大量内存，对于使用较小 VRAM 的 GPU，这可以用于专注于训练。
+>
+> 如果您选择运行验证/测试，可以通过指定 `--enable_model_cpu_offload` 在较低 VRAM 的 GPU 上节省一些内存。
 
 ### LoRA 微调
 
-<details>
-<summary> AdamW </summary>
-
-With `train_batch_size = 1`:
-
-|       model        | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing |
-|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:|
-| THUDM/CogVideoX-2b |    16     |          False         |         12.945         |          43.764          |         46.918          |       24.234         |
-| THUDM/CogVideoX-2b |    16     |          True          |         12.945         |          12.945          |         21.121          |       24.234         |
-| THUDM/CogVideoX-2b |    64     |          False         |         13.035         |          44.314          |         47.469          |       24.469         |
-| THUDM/CogVideoX-2b |    64     |          True          |         13.036         |          13.035          |         21.564          |       24.500         |
-| THUDM/CogVideoX-2b |    256    |          False         |         13.095         |          45.826          |         48.990          |       25.543         |
-| THUDM/CogVideoX-2b |    256    |          True          |         13.094         |          13.095          |         22.344          |       25.537         |
-| THUDM/CogVideoX-5b |    16     |          True          |         19.742         |          19.742          |         28.746          |       38.123         |
-| THUDM/CogVideoX-5b |    64     |          True          |         20.006         |          20.818          |         30.338          |       38.738         |
-| THUDM/CogVideoX-5b |    256    |          True          |         20.771         |          22.119          |         31.939          |       41.537         |
-
-With `train_batch_size = 4`:
-
-|       model        | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing |
-|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:|
-| THUDM/CogVideoX-2b |    16     |          True          |         12.945         |          21.803          |         21.814          |       24.322         |
-| THUDM/CogVideoX-2b |    64     |          True          |         13.035         |          22.254          |         22.254          |       24.572         |
-| THUDM/CogVideoX-2b |    256    |          True          |         13.094         |          22.020          |         22.033          |       25.574         |
-| THUDM/CogVideoX-5b |    16     |          True          |         19.742         |          46.492          |         46.492          |       38.197         |
-| THUDM/CogVideoX-5b |    64     |          True          |         20.006         |          47.805          |         47.805          |       39.365         |
-| THUDM/CogVideoX-5b |    256    |          True          |         20.771         |          47.268          |         47.332          |       41.008         |
-
 > [!NOTE]
-> 
\ No newline at end of file
+> 图像到视频 LoRA 微调的内存需求与 `THUDM/CogVideoX-5b` 上的文本到视频类似，因此未明确报告。
+>
+> 此外，要为 I2V 微调准备测试图像，您可以通过修改脚本动态生成它们，或使用以下命令从您的训练数据中提取一些帧：
+> `ffmpeg -i input.mp4 -frames:v 1 frame.png`，
+> 或提供一个有效且可访问的图像 URL。
+
+...
+
diff --git a/requirements.txt b/requirements.txt
index e5ddcf3..b75674d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,13 +1,16 @@
-accelerate
-bitsandbytes
-diffusers
-transformers
-huggingface_hub
-hf_transfer
-peft
-decord
-wandb
-pandas
-torch
-torchvision
-torchao
+accelerate>=1.0.0
+bitsandbytes>=0.44.1
+diffusers>=0.30.4
+transformers>=0.45.2
+huggingface_hub>=0.25.2
+hf_transfer>=0.1.8
+peft>=0.13.1
+decord>=0.6.0
+wandb>=0.18.3
+pandas>=2.4.1
+torch>=2.4.0
+torchvision>=0.19.0
+torchao>=0.5.0
+sentencepiece>=0.2.0
+imageio-ffmpeg>=0.5.1
+numpy>=2.1.2
\ No newline at end of file

From 353e4fc88553600334df0fd4e4b74ea3566a6623 Mon Sep 17 00:00:00 2001
From: zR <2448370773@qq.com>
Date: Thu, 10 Oct 2024 21:49:31 +0800
Subject: [PATCH 3/4] Update README_zh.md

---
 README_zh.md | 116 +++++++++++++++++++++++++++++++--------------------
 1 file changed, 70 insertions(+), 46 deletions(-)

diff --git a/README_zh.md b/README_zh.md
index 3758d84..c0a3f1c 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -1,24 +1,55 @@
-# CogVideoX Factory
+# CogVideoX Factory 🧪
 
-## 简介
+在 24GB GPU 内存下微调 Cog 系列视频模型以生成自定义视频 ⚡️📼
 
-这是用于 CogVideoX 微调的仓库。
+TODO：添加有趣的视频结果表
+
+## 快速开始
+
+确保已安装所需的依赖：`pip install -r requirements.txt`。
+
+然后下载数据集：
+
+```bash
+# 安装 `huggingface_hub`
+huggingface-cli download   --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset   --local-dir video-dataset-disney
+```
+
+然后启动文本到视频的 LoRA 微调：
+
+```bash
+TODO
+```
+
+我们现在可以使用训练好的模型进行推理：
+
+```python
+TODO
+```
+
+我们还可以使用 LoRA 微调 5B 版本：
+
+```python
+TODO
+```
+
+在下方的部分中，我们提供了有关更多选项的详细信息，这些选项旨在使视频模型的微调尽可能易于使用。
 
 ## 数据集准备
 
-创建两个文件，一个文件包含以换行符分隔的提示词，另一个文件包含以换行符分隔的视频数据路径（视频文件的路径必须相对于您在指定 `--data_root` 时传递的路径）。让我们通过一个例子来更好地理解这一点！
+创建两个文件，一个文件包含逐行分隔的提示，另一个文件包含逐行分隔的视频数据路径（视频文件的路径必须相对于您在指定 `--data_root` 时传递的路径）。让我们通过一个示例来更好地理解这一点！
 
-假设您将 `--data_root` 指定为 `/dataset`，并且该目录包含文件：`prompts.txt` 和 `videos.txt`。
+假设您指定的 `--data_root` 为 `/dataset`，并且该目录包含以下文件：`prompts.txt` 和 `videos.txt`。
 
-`prompts.txt` 文件应包含以换行符分隔的提示词：
+`prompts.txt` 文件应包含逐行分隔的提示：
 
 ```
-一段黑白动画序列，主角是一只名为 Rabbity Ribfried 的兔子和一只拟人化的山羊，在一个充满音乐和趣味的环境中，展示他们不断发展的互动。
-一段黑白动画序列，场景在船甲板上，主角是一只名为 Bully Bulldoger 的斗牛犬角色，展示了夸张的面部表情和肢体语言。角色从自信到专注，再到紧张和痛苦，展示了一系列情绪，随着它克服挑战。船的内部在背景中保持静止，只有简单的细节，如钟声和开着的门。角色的动态动作和变化的表情推动了故事的发展，没有镜头移动，确保观众专注于其不断变化的反应和肢体动作。
+一段黑白动画序列，主角是一只名为 Rabbity Ribfried 的兔子和一只拟人化的山羊，展示了它们在音乐与游戏环境中的互动演变。
+一段黑白动画序列，发生在船甲板上，主角是一只名为 Bully Bulldoger 的斗牛犬，展现了夸张的面部表情和肢体语言。角色从自信、专注逐渐转变为紧张与痛苦，展示了随着挑战出现的情感变化。船的内部在背景中保持静止，只有一些简单的细节，如钟声和敞开的门。角色的动态动作和不断变化的表情推动了叙事，没有摄像机运动来分散注意力。
 ...
 ```
 
-`videos.txt` 文件应包含以换行符分隔的视频文件路径。请注意，路径应相对于 `--data_root` 目录。
+`videos.txt` 文件应包含逐行分隔的视频文件路径。请注意，路径应相对于 `--data_root` 目录。
 
 ```bash
 videos/00000.mp4
@@ -26,7 +57,7 @@ videos/00001.mp4
 ...
 ```
 
-总体而言，如果您在数据集根目录运行 `tree` 命令，您的数据集应如下所示：
+整体而言，如果在数据集根目录运行 `tree` 命令，您的数据集应如下所示：
 
 ```bash
 /dataset
@@ -38,58 +69,51 @@ videos/00001.mp4
     ├── ...
 ```
 
-使用此格式时，`--caption_column` 必须是 `prompts.txt`，`--video_column` 必须是 `videos.txt`。如果您的数据存储在 CSV 文件中，您也可以指定 `--dataset_file` 为 CSV 的路径，`--caption_column` 和 `--video_column` 为 CSV 文件中的实际列名。
+使用此格式时，`--caption_column` 必须是 `prompts.txt`，`--video_column` 必须是 `videos.txt`。如果您将数据存储在 CSV 文件中，还可以指定 `--dataset_file` 为 CSV 的路径，`--caption_column` 和 `--video_column` 为 CSV 文件中的实际列名。
 
-例如，让我们使用这个 [Disney 数据集](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) 进行微调。要下载，可以使用 🤗 Hugging Face CLI。
+例如，让我们使用[这个](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) Disney 数据集进行微调。要下载，您可以使用 🤗 Hugging Face CLI。
 
 ```bash
 huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir video-dataset-disney
 ```
 
+TODO：添加一个关于创建和使用预计算嵌入的部分。
+
 ## 训练
 
-TODO
+我们提供了与 [Cog 系列模型](https://huggingface.co/collections/THUDM/cogvideo-66c08e62f1685a3ade464cce) 兼容的文本到视频和图像到视频生成的训练脚本。
 
-请查看 `training/*.sh`
+查看 `*.sh` 文件
 
 注意：未在 MPS 上测试
 
 ## 内存需求
 
-训练支持并验证的内存优化包括：
+<table align="center">
+<tr>
+  <td align="center"><a href="https://www.youtube.com/watch?v=UvRl4ansfCg"> 使用 PyTorch 消除 OOM</a></td>
+</tr>
+<tr>
+  <td align="center"><img src="assets/slaying-ooms.png" style="width: 480px; height: 480px;"></td>
+</tr>
+</table>
 
-- 来自 [TorchAO](https://github.com/pytorch/ao) 的 `CPUOffloadOptimizer`。
-- 来自 [bitsandbytes](https://huggingface.co/docs/bitsandbytes/optimizers) 的低位优化器。
+支持和验证的内存优化训练选项包括：
+
+- [`torchao`](https://github.com/pytorch/ao) 中的 `CPUOffloadOptimizer`。您可以阅读它的能力和限制 [此处](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload)。简而言之，它允许您使用 CPU 存储可训练的参数和梯度。这导致优化器步骤在 CPU 上进行，需要一个快速的 CPU 优化器，例如 `torch.optim.AdamW(fused=True)` 或在优化器步骤上应用 `torch.compile`。此外，建议不要将模型编译用于训练。梯度裁剪和积累尚不支持。
+- [`bitsandbytes`](https://huggingface.co/docs/bitsandbytes/optimizers) 中的低位优化器。TODO：测试并使 [`torchao`](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim) 工作
+- DeepSpeed Zero2：由于我们依赖 `accelerate`，请按照[本指南](https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed) 配置 `accelerate` 以启用 DeepSpeed Zero2 优化。
+
+> [!IMPORTANT]
+> 内存需求是在运行 `training/prepare_dataset.py` 后报告的，它将视频和字幕转换为潜变量和嵌入。在训练过程中，我们直接加载潜变量和嵌入，而不需要 VAE 或 T5 文本编码器。但是，如果您执行验证/测试，则必须加载这些内容，并增加所需的内存量。不执行验证/测试可以节省大量内存，对于使用较小 VRAM 的 GPU，这可以用于专注于训练。
+>
+> 如果您选择运行验证/测试，可以通过指定 `--enable_model_cpu_offload` 在较低 VRAM 的 GPU 上节省一些内存。
 
 ### LoRA 微调
 
-<details>
-<summary> AdamW </summary>
-
-With `train_batch_size = 1`:
-
-|       model        | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing |
-|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:|
-| THUDM/CogVideoX-2b |    16     |          False         |         12.945         |          43.764          |         46.918          |       24.234         |
-| THUDM/CogVideoX-2b |    16     |          True          |         12.945         |          12.945          |         21.121          |       24.234         |
-| THUDM/CogVideoX-2b |    64     |          False         |         13.035         |          44.314          |         47.469          |       24.469         |
-| THUDM/CogVideoX-2b |    64     |          True          |         13.036         |          13.035          |         21.564          |       24.500         |
-| THUDM/CogVideoX-2b |    256    |          False         |         13.095         |          45.826          |         48.990          |       25.543         |
-| THUDM/CogVideoX-2b |    256    |          True          |         13.094         |          13.095          |         22.344          |       25.537         |
-| THUDM/CogVideoX-5b |    16     |          True          |         19.742         |          19.742          |         28.746          |       38.123         |
-| THUDM/CogVideoX-5b |    64     |          True          |         20.006         |          20.818          |         30.338          |       38.738         |
-| THUDM/CogVideoX-5b |    256    |          True          |         20.771         |          22.119          |         31.939          |       41.537         |
-
-With `train_batch_size = 4`:
-
-|       model        | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing |
-|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:|
-| THUDM/CogVideoX-2b |    16     |          True          |         12.945         |          21.803          |         21.814          |       24.322         |
-| THUDM/CogVideoX-2b |    64     |          True          |         13.035         |          22.254          |         22.254          |       24.572         |
-| THUDM/CogVideoX-2b |    256    |          True          |         13.094         |          22.020          |         22.033          |       25.574         |
-| THUDM/CogVideoX-5b |    16     |          True          |         19.742         |          46.492          |         46.492          |       38.197         |
-| THUDM/CogVideoX-5b |    64     |          True          |         20.006         |          47.805          |         47.805          |       39.365         |
-| THUDM/CogVideoX-5b |    256    |          True          |         20.771         |          47.268          |         47.332          |       41.008         |
-
 > [!NOTE]
-> 
\ No newline at end of file
+> 图像到视频 LoRA 微调的内存需求与 `THUDM/CogVideoX-5b` 上的文本到视频类似，因此未明确报告。
+>
+> 此外，要为 I2V 微调准备测试图像，您可以通过修改脚本动态生成它们，或使用以下命令从您的训练数据中提取一些帧：
+> `ffmpeg -i input.mp4 -frames:v 1 frame.png`，
+> 或提供一个有效且可访问的图像 URL。

From 13083b5bbf6afdc33d375daabfe2c4c5e6f96127 Mon Sep 17 00:00:00 2001
From: zR <2448370773@qq.com>
Date: Thu, 10 Oct 2024 22:14:25 +0800
Subject: [PATCH 4/4] .github page

---
 .github/ISSUE_TEMPLATE/bug_report.yaml      | 51 +++++++++++++++++++
 .github/ISSUE_TEMPLATE/feature-request.yaml | 34 +++++++++++++
 README.md                                   |  2 +
 README_zh.md                                | 56 ++++++++++++++-------
 assets/contribute.md                        |  3 ++
 train_text_to_video_lora.sh                 | 14 ++----
 6 files changed, 134 insertions(+), 26 deletions(-)
 create mode 100644 .github/ISSUE_TEMPLATE/bug_report.yaml
 create mode 100644 .github/ISSUE_TEMPLATE/feature-request.yaml
 create mode 100644 assets/contribute.md

diff --git a/.github/ISSUE_TEMPLATE/bug_report.yaml b/.github/ISSUE_TEMPLATE/bug_report.yaml
new file mode 100644
index 0000000..89ccb8b
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug_report.yaml
@@ -0,0 +1,51 @@
+name: "\U0001F41B Bug Report"
+description: Submit a bug report to help us improve CogVideoX-Factory / 提交一个 Bug 问题报告来帮助我们改进 CogVideoX-Factory 开源框架
+body:
+  - type: textarea
+    id: system-info
+    attributes:
+      label: System Info / 系統信息
+      description: Your operating environment / 您的运行环境信息
+      placeholder: Includes Cuda version, Diffusers version, Python version, operating system, hardware information (if you suspect a hardware problem)... / 包括Cuda版本，Diffusers，Python版本，操作系统，硬件信息(如果您怀疑是硬件方面的问题)...
+    validations:
+      required: true
+
+  - type: checkboxes
+    id: information-scripts-examples
+    attributes:
+      label: Information / 问题信息
+      description: 'The problem arises when using: / 问题出现在'
+      options:
+        - label: "The official example scripts / 官方的示例脚本"
+        - label: "My own modified scripts / 我自己修改的脚本和任务"
+
+  - type: textarea
+    id: reproduction
+    validations:
+      required: true
+    attributes:
+      label: Reproduction / 复现过程
+      description: |
+        Please provide a code example that reproduces the problem you encountered, preferably with a minimal reproduction unit.
+        If you have code snippets, error messages, stack traces, please provide them here as well.
+        Please format your code correctly using code tags. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        Do not use screenshots, as they are difficult to read and (more importantly) do not allow others to copy and paste your code.
+        
+        请提供能重现您遇到的问题的代码示例,最好是最小复现单元。
+        如果您有代码片段、错误信息、堆栈跟踪，也请在此提供。
+        请使用代码标签正确格式化您的代码。请参见 https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        请勿使用截图，因为截图难以阅读，而且（更重要的是）不允许他人复制粘贴您的代码。
+      placeholder: |
+        Steps to reproduce the behavior/复现Bug的步骤:
+          
+          1.
+          2.
+          3.
+
+  - type: textarea
+    id: expected-behavior
+    validations:
+      required: true
+    attributes:
+      label: Expected behavior / 期待表现
+      description: "A clear and concise description of what you would expect to happen. /简单描述您期望发生的事情。"
\ No newline at end of file
diff --git a/.github/ISSUE_TEMPLATE/feature-request.yaml b/.github/ISSUE_TEMPLATE/feature-request.yaml
new file mode 100644
index 0000000..ac2f0fc
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature-request.yaml
@@ -0,0 +1,34 @@
+name: "\U0001F680 Feature request"
+description: Submit a request for a new CogVideoX-Factory feature / 提交一个新的 CogVideoX-Factory 开源项目的功能建议
+labels: [ "feature" ]
+body:
+  - type: textarea
+    id: feature-request
+    validations:
+      required: true
+    attributes:
+      label: Feature request  / 功能建议
+      description: |
+        A brief description of the functional proposal. Links to corresponding papers and code are desirable.
+        对功能建议的简述。最好提供对应的论文和代码链接。
+
+  - type: textarea
+    id: motivation
+    validations:
+      required: true
+    attributes:
+      label: Motivation / 动机
+      description: |
+        Your motivation for making the suggestion. If that motivation is related to another GitHub issue, link to it here.
+        您提出建议的动机。如果该动机与另一个 GitHub 问题有关，请在此处提供对应的链接。
+
+  - type: textarea
+    id: contribution
+    validations:
+      required: true
+    attributes:
+      label: Your contribution / 您的贡献
+      description: |
+        
+        Your PR link or any other link you can help with.
+        您的PR链接或者其他您能提供帮助的链接。
\ No newline at end of file
diff --git a/README.md b/README.md
index b150494..f28a905 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,7 @@
 # CogVideoX Factory 🧪
 
+[中文阅读](./README_zh.md)
+
 Fine-tune Cog family of video models for custom video generation under 24GB of GPU memory ⚡️📼
 
 <table align="center">
diff --git a/README_zh.md b/README_zh.md
index c0a3f1c..a3d918c 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -1,12 +1,19 @@
 # CogVideoX Factory 🧪
 
+[Read this in English](./README_zh.md)
+
 在 24GB GPU 内存下微调 Cog 系列视频模型以生成自定义视频 ⚡️📼
 
-TODO：添加有趣的视频结果表
+<table align="center">
+<tr>
+  <td align="center"><video src="https://github.com/user-attachments/assets/aad07161-87cb-4784-9e6b-16d06581e3e5">Your browser does not support the video tag.</video></td>
+</tr>
+</table>
+
 
 ## 快速开始
 
-确保已安装所需的依赖：`pip install -r requirements.txt`。
+克隆此仓库并确保已安装所有依赖：`pip install -r requirements.txt`。
 
 然后下载数据集：
 
@@ -15,25 +22,39 @@ TODO：添加有趣的视频结果表
 huggingface-cli download   --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset   --local-dir video-dataset-disney
 ```
 
-然后启动文本到视频的 LoRA 微调：
+然后启动文本到视频的 LoRA 微调（根据您的需求修改不同的超参数、数据集根目录和其他配置选项）：
 
 ```bash
-TODO
+# 对 CogVideoX 文本到视频模型进行 LoRA 微调
+./train_text_to_video_lora.sh
+
+# 对 CogVideoX 文本到视频模型进行全微调
+./train_text_to_video_sft.sh
+
+# 对 CogVideoX 图像到视频模型进行 LoRA 微调
+./train_image_to_video_lora.sh
 ```
 
-我们现在可以使用训练好的模型进行推理：
+假设您的 LoRA 已保存并推送到 HF Hub，并命名为 `my-awesome-name/my-awesome-lora`，我们现在可以使用微调后的模型进行推理：
 
-```python
-TODO
+```diff
+import torch
+from diffusers import CogVideoXPipeline
+from diffusers import export_to_video
+
+pipe = CogVideoXPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16
+).to("cuda")
++ pipe.load_lora_weights("my-awesome-name/my-awesome-lora", adapter_name=["cogvideox-lora"])
++ pipe.set_adapters(["cogvideox-lora"], [1.0])
+
+video = pipe("<my-awesome-prompt>").frames[0]
+export_to_video(video, "output.mp4", fps=8)
 ```
 
-我们还可以使用 LoRA 微调 5B 版本：
+**注意：** 对于图像到视频的微调，您必须从 [此](https://github.com/huggingface/diffusers/pull/9482) 分支安装 diffusers（该分支添加了 CogVideoX 图像到视频的 LoRA 加载支持），直到它被合并。
 
-```python
-TODO
-```
-
-在下方的部分中，我们提供了有关更多选项的详细信息，这些选项旨在使视频模型的微调尽可能易于使用。
+在下方的部分中，我们提供了在本仓库中探索的更多选项的详细信息。它们都试图通过尽可能减少内存需求，使视频模型的微调变得尽可能容易。
 
 ## 数据集准备
 
@@ -83,9 +104,9 @@ TODO：添加一个关于创建和使用预计算嵌入的部分。
 
 我们提供了与 [Cog 系列模型](https://huggingface.co/collections/THUDM/cogvideo-66c08e62f1685a3ade464cce) 兼容的文本到视频和图像到视频生成的训练脚本。
 
-查看 `*.sh` 文件
+查看 `*.sh` 文件。
 
-注意：未在 MPS 上测试
+注意：本代码未在 MPS 上测试，建议在 Linux 环境下使用 CUDA文件测试。
 
 ## 内存需求
 
@@ -101,7 +122,8 @@ TODO：添加一个关于创建和使用预计算嵌入的部分。
 支持和验证的内存优化训练选项包括：
 
 - [`torchao`](https://github.com/pytorch/ao) 中的 `CPUOffloadOptimizer`。您可以阅读它的能力和限制 [此处](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload)。简而言之，它允许您使用 CPU 存储可训练的参数和梯度。这导致优化器步骤在 CPU 上进行，需要一个快速的 CPU 优化器，例如 `torch.optim.AdamW(fused=True)` 或在优化器步骤上应用 `torch.compile`。此外，建议不要将模型编译用于训练。梯度裁剪和积累尚不支持。
-- [`bitsandbytes`](https://huggingface.co/docs/bitsandbytes/optimizers) 中的低位优化器。TODO：测试并使 [`torchao`](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim) 工作
+- [`bitsandbytes`](https://huggingface.co/docs/bitsandbytes/optimizers) 中的低位优化器。
+  - TODO：测试并使 [`torchao`](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim) 工作
 - DeepSpeed Zero2：由于我们依赖 `accelerate`，请按照[本指南](https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed) 配置 `accelerate` 以启用 DeepSpeed Zero2 优化。
 
 > [!IMPORTANT]
@@ -114,6 +136,6 @@ TODO：添加一个关于创建和使用预计算嵌入的部分。
 > [!NOTE]
 > 图像到视频 LoRA 微调的内存需求与 `THUDM/CogVideoX-5b` 上的文本到视频类似，因此未明确报告。
 >
-> 此外，要为 I2V 微调准备测试图像，您可以通过修改脚本动态生成它们，或使用以下命令从您的训练数据中提取一些帧：
+> I2V训练会使用视频的第一帧进行微调。 要为 I2V 微调准备测试图像，您可以通过修改脚本动态生成它们，或使用以下命令从您的训练数据中提取一些帧：
 > `ffmpeg -i input.mp4 -frames:v 1 frame.png`，
 > 或提供一个有效且可访问的图像 URL。
diff --git a/assets/contribute.md b/assets/contribute.md
new file mode 100644
index 0000000..9959241
--- /dev/null
+++ b/assets/contribute.md
@@ -0,0 +1,3 @@
+# 欢迎你们的贡献
+
+本项目属于非常初级的阶段
\ No newline at end of file
diff --git a/train_text_to_video_lora.sh b/train_text_to_video_lora.sh
index 4aac214..2892c75 100755
--- a/train_text_to_video_lora.sh
+++ b/train_text_to_video_lora.sh
@@ -4,7 +4,7 @@ export WANDB_MODE="offline"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 
-GPU_IDS="0,1,2,3,4,5,6,7"
+GPU_IDS="0"
 
 # Training Configurations
 # Experiment with as many hyperparameters as you want!
@@ -19,26 +19,22 @@ ACCELERATE_CONFIG_FILE="accelerate_configs/uncompiled_1.yaml"
 # Absolute path to where the data is located. Make sure to have read the README for how to prepare data.
 # This example assumes you downloaded an already prepared dataset from HF CLI as follows:
 #   huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir /path/to/my/datasets/disney-dataset
-
-DATA_ROOT="/share/home/zyx/disney_cogvideox-encoded-multi"
-CAPTION_COLUMN="prompts.txt"
+DATA_ROOT="/path/to/my/datasets/disney-dataset"
+CAPTION_COLUMN="prompt.txt"
 VIDEO_COLUMN="videos.txt"
-MODEL_PATH="/share/official_pretrains/hf_home/CogVideoX-5b"
-
 
 # Launch experiments with different hyperparameters
 for learning_rate in "${LEARNING_RATES[@]}"; do
   for lr_schedule in "${LR_SCHEDULES[@]}"; do
     for optimizer in "${OPTIMIZERS[@]}"; do
       for steps in "${MAX_TRAIN_STEPS[@]}"; do
-        output_dir="cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
+        output_dir="/path/to/my/models/cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
 
         cmd="accelerate launch --config_file $ACCELERATE_CONFIG_FILE --gpu_ids $GPU_IDS training/cogvideox_text_to_video_lora.py \
-          --pretrained_model_name_or_path $MODEL_PATH \
+          --pretrained_model_name_or_path THUDM/CogVideoX-5b \
           --data_root $DATA_ROOT \
           --caption_column $CAPTION_COLUMN \
           --video_column $VIDEO_COLUMN \
-          --load_tensors \
           --id_token BW_STYLE \
           --height_buckets 480 \
           --width_buckets 720 \