mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-17 00:23:35 +00:00
CogVideoX LoRA and full finetuning (#1)
* update * update * update * update * update * update * update * update * update * update * update * update * update * update * update
This commit is contained in:
Aryan
@@ -0,0 +1,289 @@
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
from accelerate.logging import get_logger
|
||||
from torch.utils.data import Dataset, Sampler
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms.functional import resize
|
||||
|
||||
|
||||
# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error
|
||||
# Very few bug reports but it happens. Look in decord Github issues for more relevant information.
|
||||
import decord # isort:skip
|
||||
|
||||
decord.bridge.set_bridge("torch")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
HEIGHT_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536]
|
||||
WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536]
|
||||
FRAME_BUCKETS = [16, 24, 32, 48, 64, 80]
|
||||
|
||||
|
||||
class VideoDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root: str,
|
||||
dataset_file: Optional[str] = None,
|
||||
caption_column: str = "text",
|
||||
video_column: str = "video",
|
||||
max_num_frames: int = 49,
|
||||
id_token: Optional[str] = None,
|
||||
height_buckets: List[int] = None,
|
||||
width_buckets: List[int] = None,
|
||||
frame_buckets: List[int] = None,
|
||||
load_tensors: bool = False,
|
||||
random_flip: Optional[float] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.data_root = Path(data_root)
|
||||
self.dataset_file = dataset_file
|
||||
self.caption_column = caption_column
|
||||
self.video_column = video_column
|
||||
self.max_num_frames = max_num_frames
|
||||
self.id_token = id_token or ""
|
||||
self.height_buckets = height_buckets or HEIGHT_BUCKETS
|
||||
self.width_buckets = width_buckets or WIDTH_BUCKETS
|
||||
self.frame_buckets = frame_buckets or FRAME_BUCKETS
|
||||
self.load_tensors = load_tensors
|
||||
self.random_flip = random_flip
|
||||
|
||||
self.resolutions = [
|
||||
(f, h, w) for h in self.height_buckets for w in self.width_buckets for f in self.frame_buckets
|
||||
]
|
||||
|
||||
# Two methods of loading data are supported.
|
||||
# - Using a CSV: caption_column and video_column must be some column in the CSV. One could
|
||||
# make use of other columns too, such as a motion score or aesthetic score, by modifying the
|
||||
# logic in CSV processing.
|
||||
# - Using two files containing line-separate captions and relative paths to videos.
|
||||
# For a more detailed explanation about preparing dataset format, checkout the README.
|
||||
if dataset_file is None:
|
||||
(
|
||||
self.prompts,
|
||||
self.video_paths,
|
||||
) = self._load_dataset_from_local_path()
|
||||
else:
|
||||
(
|
||||
self.prompts,
|
||||
self.video_paths,
|
||||
) = self._load_dataset_from_csv()
|
||||
|
||||
self.num_videos = len(self.video_paths)
|
||||
if self.num_videos != len(self.prompts):
|
||||
raise ValueError(
|
||||
f"Expected length of prompts and videos to be the same but found {len(self.prompts)=} and {len(self.video_paths)=}. Please ensure that the number of caption prompts and videos match in your dataset."
|
||||
)
|
||||
|
||||
self.video_transforms = transforms.Compose(
|
||||
[
|
||||
transforms.RandomHorizontalFlip(random_flip) if random_flip else transforms.Lambda(lambda x: x),
|
||||
transforms.Lambda(lambda x: x / 255.0),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
|
||||
]
|
||||
)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_videos
|
||||
|
||||
def __getitem__(self, index: int) -> Dict[str, Any]:
|
||||
if isinstance(index, list):
|
||||
# Here, index is actually a list of data objects that we need to return.
|
||||
# The BucketSampler should ideally return indices. But, in the sampler, we'd like
|
||||
# to have information about num_frames, height and width. Since this is not stored
|
||||
# as metadata, we need to read the video to get this information. You could read this
|
||||
# information without loading the full video in memory, but we do it anyway. In order
|
||||
# to not load the video twice (once to get the metadata, and once to return the loaded video
|
||||
# based on sampled indices), we cache it in the BucketSampler. When the sampler is
|
||||
# to yield, we yield the cache data instead of indices. So, this special check ensures
|
||||
# that data is not loaded a second time. PRs are welcome for improvements.
|
||||
return index
|
||||
|
||||
if self.load_tensors:
|
||||
latents, prompt_embeds = self._preprocess_video(self.video_paths[index])
|
||||
|
||||
# This is hardcoded for now.
|
||||
# The VAE's temporal compression ratio is 4.
|
||||
# The VAE's spatial compression ratio is 8.
|
||||
latent_num_frames = latents.size(1)
|
||||
if latent_num_frames % 2 == 0:
|
||||
num_frames = latent_num_frames * 4
|
||||
else:
|
||||
num_frames = (latent_num_frames - 1) * 4 + 1
|
||||
|
||||
height = latents.size(2) * 8
|
||||
width = latents.size(3) * 8
|
||||
|
||||
return {
|
||||
"prompt": prompt_embeds,
|
||||
"video": latents,
|
||||
"video_metadata": {
|
||||
"num_frames": num_frames,
|
||||
"height": height,
|
||||
"width": width,
|
||||
},
|
||||
}
|
||||
else:
|
||||
video, _ = self._preprocess_video(self.video_paths[index])
|
||||
|
||||
return {
|
||||
"prompt": self.id_token + self.prompts[index],
|
||||
"video": video,
|
||||
"video_metadata": {
|
||||
"num_frames": video.shape[0],
|
||||
"height": video.shape[2],
|
||||
"width": video.shape[3],
|
||||
},
|
||||
}
|
||||
|
||||
def _load_dataset_from_local_path(self) -> Tuple[List[str], List[str]]:
|
||||
if not self.data_root.exists():
|
||||
raise ValueError("Root folder for videos does not exist")
|
||||
|
||||
prompt_path = self.data_root.joinpath(self.caption_column)
|
||||
video_path = self.data_root.joinpath(self.video_column)
|
||||
|
||||
if not prompt_path.exists() or not prompt_path.is_file():
|
||||
raise ValueError(
|
||||
"Expected `--caption_column` to be path to a file in `--data_root` containing line-separated text prompts."
|
||||
)
|
||||
if not video_path.exists() or not video_path.is_file():
|
||||
raise ValueError(
|
||||
"Expected `--video_column` to be path to a file in `--data_root` containing line-separated paths to video data in the same directory."
|
||||
)
|
||||
|
||||
with open(prompt_path, "r", encoding="utf-8") as file:
|
||||
prompts = [line.strip() for line in file.readlines() if len(line.strip()) > 0]
|
||||
with open(video_path, "r", encoding="utf-8") as file:
|
||||
video_paths = [self.data_root.joinpath(line.strip()) for line in file.readlines() if len(line.strip()) > 0]
|
||||
|
||||
if not self.load_tensors and any(not path.is_file() for path in video_paths):
|
||||
raise ValueError(
|
||||
f"Expected `{self.video_column=}` to be a path to a file in `{self.data_root=}` containing line-separated paths to video data but found atleast one path that is not a valid file."
|
||||
)
|
||||
|
||||
return prompts, video_paths
|
||||
|
||||
def _load_dataset_from_csv(self) -> Tuple[List[str], List[str]]:
|
||||
df = pd.read_csv(self.dataset_file)
|
||||
prompts = df[self.caption_column].tolist()
|
||||
video_paths = df[self.video_column].tolist()
|
||||
video_paths = [self.data_root.joinpath(line.strip()) for line in video_paths]
|
||||
|
||||
if any(not path.is_file() for path in video_paths):
|
||||
raise ValueError(
|
||||
f"Expected `{self.video_column=}` to be a path to a file in `{self.data_root=}` containing line-separated paths to video data but found atleast one path that is not a valid file."
|
||||
)
|
||||
|
||||
return prompts, video_paths
|
||||
|
||||
def _preprocess_video(self, path: Path) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||
r"""
|
||||
Loads a single video, or latent and prompt embedding, based on initialization parameters.
|
||||
|
||||
If returning a video, returns a [F, C, H, W] video tensor, and None for the prompt embedding. Here,
|
||||
F, C, H and W are the frames, channels, height and width of the input video.
|
||||
|
||||
If returning latent/embedding, returns a [F, C, H, W] latent, and the prompt embedding of shape [S, D].
|
||||
F, C, H and W are the frames, channels, height and width of the latent, and S, D are the sequence length
|
||||
and embedding dimension of prompt embeddings.
|
||||
"""
|
||||
if self.load_tensors:
|
||||
return self._load_preprocessed_latents_and_embeds(path)
|
||||
else:
|
||||
video_reader = decord.VideoReader(uri=path.as_posix())
|
||||
video_num_frames = len(video_reader)
|
||||
|
||||
indices = list(range(0, video_num_frames, video_num_frames // self.max_num_frames))
|
||||
frames = video_reader.get_batch(indices)
|
||||
frames = frames[: self.max_num_frames].float()
|
||||
frames = frames.permute(0, 3, 1, 2).contiguous()
|
||||
frames = torch.stack([self.video_transforms(frame) for frame in frames], dim=0)
|
||||
|
||||
return frames, None
|
||||
|
||||
def _load_preprocessed_latents_and_embeds(self, path: Path) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
filename_without_ext = path.name.split(".")[0]
|
||||
pt_filename = f"{filename_without_ext}.pt"
|
||||
|
||||
# The current path is something like: /a/b/c/d/videos/00001.mp4
|
||||
# We need to reach: /a/b/c/d/latents/00001.pt
|
||||
latents_path = path.parent.parent.joinpath("latents")
|
||||
embeds_path = path.parent.parent.joinpath("embeddings")
|
||||
|
||||
if not latents_path.exists() or not embeds_path.exists():
|
||||
raise ValueError(
|
||||
f"When setting the load_tensors parameter to `True`, it is expected that the `{self.data_root=}` contains two folders named `latents` and `embeddings`. However, these folders were not found. Please make sure to have prepared your data correctly using `prepare_data.py`."
|
||||
)
|
||||
|
||||
latent_filepath = latents_path.joinpath(pt_filename)
|
||||
embeds_filepath = embeds_path.joinpath(pt_filename)
|
||||
|
||||
if not latent_filepath.is_file() or not embeds_filepath.is_file():
|
||||
latent_filepath = latent_filepath.as_posix()
|
||||
embeds_filepath = embeds_filepath.as_posix()
|
||||
raise ValueError(
|
||||
f"The file {latent_filepath=} or {embeds_filepath=} could not be found. Please ensure that you've correctly executed `prepare_dataset.py`."
|
||||
)
|
||||
|
||||
latents = torch.load(latent_filepath, map_location="cpu", weights_only=True)
|
||||
embeds = torch.load(embeds_filepath, map_location="cpu", weights_only=True)
|
||||
|
||||
return latents, embeds
|
||||
|
||||
|
||||
class VideoDatasetWithResizing(VideoDataset):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def _preprocess_video(self, path: Path) -> torch.Tensor:
|
||||
if self.load_tensors:
|
||||
return self._load_preprocessed_latents_and_embeds(path)
|
||||
else:
|
||||
video_reader = decord.VideoReader(uri=path.as_posix())
|
||||
video_num_frames = len(video_reader)
|
||||
nearest_frame_bucket = min(
|
||||
self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames))
|
||||
)
|
||||
|
||||
frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket))
|
||||
|
||||
frames = video_reader.get_batch(frame_indices)
|
||||
frames = frames[:nearest_frame_bucket].float()
|
||||
frames = frames.permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
nearest_res = self._find_nearest_resolution(frames.shape[2], frames.shape[3])
|
||||
frames_resized = torch.stack([resize(frame, nearest_res) for frame in frames], dim=0)
|
||||
|
||||
frames = torch.stack([self.video_transforms(frame) for frame in frames_resized], dim=0)
|
||||
return frames, None
|
||||
|
||||
def _find_nearest_resolution(self, height, width):
|
||||
nearest_res = min(self.resolutions, key=lambda x: abs(x[1] - height) + abs(x[2] - width))
|
||||
return nearest_res[1], nearest_res[2]
|
||||
|
||||
|
||||
class BucketSampler(Sampler):
|
||||
def __init__(self, data_source: VideoDataset, batch_size: int = 8, shuffle: bool = True) -> None:
|
||||
self.data_source = data_source
|
||||
self.batch_size = batch_size
|
||||
self.shuffle = shuffle
|
||||
|
||||
self.buckets = {resolution: [] for resolution in data_source.resolutions}
|
||||
|
||||
def __iter__(self):
|
||||
for index, data in enumerate(self.data_source):
|
||||
video_metadata = data["video_metadata"]
|
||||
f, h, w = video_metadata["num_frames"], video_metadata["height"], video_metadata["width"]
|
||||
|
||||
self.buckets[(f, h, w)].append(data)
|
||||
if len(self.buckets[(f, h, w)]) == self.batch_size:
|
||||
if self.shuffle:
|
||||
random.shuffle(self.buckets[(f, h, w)])
|
||||
yield self.buckets[(f, h, w)]
|
||||
del self.buckets[(f, h, w)]
|
||||
self.buckets[(f, h, w)] = []
|
||||
Reference in New Issue
Block a user