mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-18 00:33:52 +00:00
LTX Video (#123)
* rename files * ltx finetuning * update * update * improvements * make style * gradient clipping * update * fix distributed inference * update * update
This commit is contained in:
Aryan
@@ -0,0 +1,2 @@
|
||||
from .args import Args, parse_arguments
|
||||
from .trainer import Trainer
|
||||
@@ -0,0 +1,779 @@
|
||||
import argparse
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .constants import DEFAULT_IMAGE_RESOLUTION_BUCKETS, DEFAULT_VIDEO_RESOLUTION_BUCKETS
|
||||
|
||||
|
||||
class Args:
|
||||
r"""
|
||||
The arguments for the finetrainers training script.
|
||||
|
||||
Args:
|
||||
flow_resolution_shifting (`bool`, defaults to `False`):
|
||||
Resolution-dependant shifting of timestep schedules.
|
||||
[Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206)
|
||||
"""
|
||||
|
||||
# Model arguments
|
||||
model_name: str = None
|
||||
pretrained_model_name_or_path: str = None
|
||||
revision: Optional[str] = None
|
||||
variant: Optional[str] = None
|
||||
cache_dir: Optional[str] = None
|
||||
|
||||
# Dataset arguments
|
||||
data_root: str = None
|
||||
dataset_file: Optional[str] = None
|
||||
video_column: str = None
|
||||
caption_column: str = None
|
||||
id_token: Optional[str] = None
|
||||
image_resolution_buckets: List[Tuple[int, int]] = None
|
||||
video_resolution_buckets: List[Tuple[int, int, int]] = None
|
||||
video_reshape_mode: Optional[str] = None
|
||||
caption_dropout_p: float = 0.00
|
||||
caption_dropout_technique: str = "empty"
|
||||
|
||||
# Dataloader arguments
|
||||
dataloader_num_workers: int = 0
|
||||
pin_memory: bool = False
|
||||
|
||||
# Diffusion arguments
|
||||
flow_resolution_shifting: bool = False
|
||||
flow_base_image_seq_len: int = 256
|
||||
flow_max_image_seq_len: int = 4096
|
||||
flow_base_shift: float = 0.5
|
||||
flow_max_shift: float = 1.15
|
||||
flow_shift: float = 1.0
|
||||
flow_weighting_scheme: str = "none"
|
||||
flow_logit_mean: float = 0.0
|
||||
flow_logit_std: float = 1.0
|
||||
flow_mode_scale: float = 1.29
|
||||
|
||||
# Training arguments
|
||||
training_type: str = None
|
||||
seed: int = 42
|
||||
mixed_precision: str = None
|
||||
batch_size: int = 1
|
||||
train_epochs: int = 1
|
||||
train_steps: int = None
|
||||
rank: int = 128
|
||||
lora_alpha: float = 64
|
||||
target_modules: List[str] = ["to_k", "to_q", "to_v", "to_out.0"]
|
||||
gradient_accumulation_steps: int = 1
|
||||
gradient_checkpointing: bool = False
|
||||
checkpointing_steps: int = 500
|
||||
checkpointing_limit: Optional[int] = None
|
||||
resume_from_checkpoint: Optional[str] = None
|
||||
enable_slicing: bool = False
|
||||
enable_tiling: bool = False
|
||||
|
||||
# Optimizer arguments
|
||||
optimizer: str = "adamw"
|
||||
lr: float = 1e-4
|
||||
scale_lr: bool = False
|
||||
lr_scheduler: str = "cosine_with_restarts"
|
||||
lr_warmup_steps: int = 0
|
||||
lr_num_cycles: int = 1
|
||||
lr_power: float = 1.0
|
||||
beta1: float = 0.9
|
||||
beta2: float = 0.95
|
||||
beta3: float = 0.999
|
||||
weight_decay: float = 0.0001
|
||||
epsilon: float = 1e-8
|
||||
max_grad_norm: float = 1.0
|
||||
|
||||
# Validation arguments
|
||||
validation_prompts: List[str] = None
|
||||
validation_images: List[str] = None
|
||||
validation_videos: List[str] = None
|
||||
validation_heights: List[int] = None
|
||||
validation_widths: List[int] = None
|
||||
validation_num_frames: List[int] = None
|
||||
num_validation_videos_per_prompt: int = 1
|
||||
validation_every_n_epochs: Optional[int] = None
|
||||
validation_every_n_steps: Optional[int] = None
|
||||
enable_model_cpu_offload: bool = False
|
||||
|
||||
# Miscellaneous arguments
|
||||
tracker_name: str = "finetrainers"
|
||||
push_to_hub: bool = False
|
||||
hub_token: Optional[str] = None
|
||||
hub_model_id: Optional[str] = None
|
||||
output_dir: str = None
|
||||
logging_dir: Optional[str] = "logs"
|
||||
allow_tf32: bool = False
|
||||
nccl_timeout: int = 1800 # 30 minutes
|
||||
report_to: str = "wandb"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"model_arguments": {
|
||||
"model_name": self.model_name,
|
||||
"pretrained_model_name_or_path": self.pretrained_model_name_or_path,
|
||||
"revision": self.revision,
|
||||
"variant": self.variant,
|
||||
"cache_dir": self.cache_dir,
|
||||
},
|
||||
"dataset_arguments": {
|
||||
"data_root": self.data_root,
|
||||
"dataset_file": self.dataset_file,
|
||||
"video_column": self.video_column,
|
||||
"caption_column": self.caption_column,
|
||||
"id_token": self.id_token,
|
||||
"image_resolution_buckets": self.image_resolution_buckets,
|
||||
"video_resolution_buckets": self.video_resolution_buckets,
|
||||
"video_reshape_mode": self.video_reshape_mode,
|
||||
"caption_dropout_p": self.caption_dropout_p,
|
||||
},
|
||||
"dataloader_arguments": {
|
||||
"dataloader_num_workers": self.dataloader_num_workers,
|
||||
"pin_memory": self.pin_memory,
|
||||
},
|
||||
"training_arguments": {
|
||||
"training_type": self.training_type,
|
||||
"seed": self.seed,
|
||||
"mixed_precision": self.mixed_precision,
|
||||
"batch_size": self.batch_size,
|
||||
"train_epochs": self.train_epochs,
|
||||
"train_steps": self.train_steps,
|
||||
"rank": self.rank,
|
||||
"lora_alpha": self.lora_alpha,
|
||||
"target_modules": self.target_modules,
|
||||
"gradient_accumulation_steps": self.gradient_accumulation_steps,
|
||||
"gradient_checkpointing": self.gradient_checkpointing,
|
||||
"checkpointing_steps": self.checkpointing_steps,
|
||||
"checkpointing_limit": self.checkpointing_limit,
|
||||
"resume_from_checkpoint": self.resume_from_checkpoint,
|
||||
"enable_slicing": self.enable_slicing,
|
||||
"enable_tiling": self.enable_tiling,
|
||||
},
|
||||
"optimizer_arguments": {
|
||||
"optimizer": self.optimizer,
|
||||
"lr": self.lr,
|
||||
"scale_lr": self.scale_lr,
|
||||
"lr_scheduler": self.lr_scheduler,
|
||||
"lr_warmup_steps": self.lr_warmup_steps,
|
||||
"lr_num_cycles": self.lr_num_cycles,
|
||||
"lr_power": self.lr_power,
|
||||
"beta1": self.beta1,
|
||||
"beta2": self.beta2,
|
||||
"beta3": self.beta3,
|
||||
"weight_decay": self.weight_decay,
|
||||
"epsilon": self.epsilon,
|
||||
"max_grad_norm": self.max_grad_norm,
|
||||
},
|
||||
"validation_arguments": {
|
||||
"validation_prompts": self.validation_prompts,
|
||||
"validation_images": self.validation_images,
|
||||
"validation_videos": self.validation_videos,
|
||||
"num_validation_videos_per_prompt": self.num_validation_videos_per_prompt,
|
||||
"validation_every_n_epochs": self.validation_every_n_epochs,
|
||||
"validation_every_n_steps": self.validation_every_n_steps,
|
||||
"enable_model_cpu_offload": self.enable_model_cpu_offload,
|
||||
},
|
||||
"miscellaneous_arguments": {
|
||||
"tracker_name": self.tracker_name,
|
||||
"push_to_hub": self.push_to_hub,
|
||||
"hub_token": self.hub_token,
|
||||
"hub_model_id": self.hub_model_id,
|
||||
"output_dir": self.output_dir,
|
||||
"logging_dir": self.logging_dir,
|
||||
"allow_tf32": self.allow_tf32,
|
||||
"nccl_timeout": self.nccl_timeout,
|
||||
"report_to": self.report_to,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def parse_arguments() -> Args:
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
_add_model_arguments(parser)
|
||||
_add_dataset_arguments(parser)
|
||||
_add_dataloader_arguments(parser)
|
||||
_add_diffusion_arguments(parser)
|
||||
_add_training_arguments(parser)
|
||||
_add_optimizer_arguments(parser)
|
||||
_add_validation_arguments(parser)
|
||||
_add_miscellaneous_arguments(parser)
|
||||
|
||||
args = parser.parse_args()
|
||||
return _map_to_args_type(args)
|
||||
|
||||
|
||||
def validate_args(args: Args):
|
||||
_validate_validation_args(args)
|
||||
|
||||
|
||||
def _add_model_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument("--model_name", type=str, required=True, choices=["ltx_video"], help="Name of model to train.")
|
||||
parser.add_argument(
|
||||
"--pretrained_model_name_or_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--revision",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="Revision of pretrained model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--variant",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The directory where the downloaded models and datasets will be stored.",
|
||||
)
|
||||
|
||||
|
||||
def _add_dataset_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
def parse_resolution_bucket(resolution_bucket: str) -> Tuple[int, ...]:
|
||||
return tuple(map(int, resolution_bucket.split("x")))
|
||||
|
||||
def parse_image_resolution_bucket(resolution_bucket: str) -> Tuple[int, int]:
|
||||
resolution_bucket = parse_resolution_bucket(resolution_bucket)
|
||||
assert (
|
||||
len(resolution_bucket) == 2
|
||||
), f"Expected 2D resolution bucket, got {len(resolution_bucket)}D resolution bucket"
|
||||
return resolution_bucket
|
||||
|
||||
def parse_video_resolution_bucket(resolution_bucket: str) -> Tuple[int, int, int]:
|
||||
resolution_bucket = parse_resolution_bucket(resolution_bucket)
|
||||
assert (
|
||||
len(resolution_bucket) == 3
|
||||
), f"Expected 3D resolution bucket, got {len(resolution_bucket)}D resolution bucket"
|
||||
return resolution_bucket
|
||||
|
||||
parser.add_argument(
|
||||
"--data_root",
|
||||
type=str,
|
||||
default=None,
|
||||
help=("A folder containing the training data."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset_file",
|
||||
type=str,
|
||||
default=None,
|
||||
help=("Path to a CSV file if loading prompts/video paths using this format."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_column",
|
||||
type=str,
|
||||
default="video",
|
||||
help="The column of the dataset containing videos. Or, the name of the file in `--data_root` folder containing the line-separated path to video data.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--caption_column",
|
||||
type=str,
|
||||
default="text",
|
||||
help="The column of the dataset containing the instance prompt for each video. Or, the name of the file in `--data_root` folder containing the line-separated instance prompts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--id_token",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Identifier token appended to the start of each prompt if provided.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_resolution_buckets",
|
||||
type=parse_image_resolution_bucket,
|
||||
default=None,
|
||||
nargs="+",
|
||||
help="Resolution buckets for images.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_resolution_buckets",
|
||||
type=parse_video_resolution_bucket,
|
||||
default=None,
|
||||
nargs="+",
|
||||
help="Resolution buckets for videos.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_reshape_mode",
|
||||
type=str,
|
||||
default=None,
|
||||
help="All input videos are reshaped to this mode. Choose between ['center', 'random', 'none']",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--caption_dropout_p",
|
||||
type=float,
|
||||
default=0.00,
|
||||
help="Probability of dropout for the caption tokens.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--caption_dropout_technique",
|
||||
type=str,
|
||||
default="empty",
|
||||
choices=["empty", "zero"],
|
||||
help="Technique to use for caption dropout.",
|
||||
)
|
||||
|
||||
|
||||
def _add_dataloader_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--dataloader_num_workers",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pin_memory",
|
||||
action="store_true",
|
||||
help="Whether or not to use the pinned memory setting in pytorch dataloader.",
|
||||
)
|
||||
|
||||
|
||||
def _add_diffusion_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--flow_resolution_shifting",
|
||||
action="store_true",
|
||||
help="Resolution-dependant shifting of timestep schedules.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow_weighting_scheme",
|
||||
type=str,
|
||||
default="none",
|
||||
choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"],
|
||||
help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow_logit_mean",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="mean to use when using the `'logit_normal'` weighting scheme.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow_logit_std",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="std to use when using the `'logit_normal'` weighting scheme.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow_mode_scale",
|
||||
type=float,
|
||||
default=1.29,
|
||||
help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
|
||||
)
|
||||
|
||||
|
||||
def _add_training_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
# TODO: support full finetuning and other kinds
|
||||
parser.add_argument(
|
||||
"--training_type",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Type of training to perform. Choose between ['lora']",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
|
||||
parser.add_argument(
|
||||
"--mixed_precision",
|
||||
type=str,
|
||||
default="no",
|
||||
choices=["no", "fp8", "fp16", "bf16"],
|
||||
help=(
|
||||
"Whether to use mixed precision. Defaults 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(
|
||||
"--batch_size",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Batch size (per device) for the training dataloader.",
|
||||
)
|
||||
parser.add_argument("--train_epochs", type=int, default=1)
|
||||
parser.add_argument(
|
||||
"--train_steps",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Total number of training steps to perform. If provided, overrides `--num_train_epochs`.",
|
||||
)
|
||||
parser.add_argument("--rank", type=int, default=64, help="The rank for LoRA matrices.")
|
||||
parser.add_argument(
|
||||
"--lora_alpha",
|
||||
type=int,
|
||||
default=64,
|
||||
help="The lora_alpha to compute scaling factor (lora_alpha / rank) for LoRA matrices.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_modules", type=str, default="to_k to_q to_v to_out.0", nargs="+", help="The target modules for LoRA."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gradient_accumulation_steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of updates steps to accumulate before performing a backward/update pass.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gradient_checkpointing",
|
||||
action="store_true",
|
||||
help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--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(
|
||||
"--checkpointing_limit",
|
||||
type=int,
|
||||
default=None,
|
||||
help=("Max number of checkpoints to store."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resume_from_checkpoint",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Whether training should be resumed from a previous checkpoint. Use a path saved by"
|
||||
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable_slicing",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Whether or not to use VAE slicing for saving memory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable_tiling",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Whether or not to use VAE tiling for saving memory.",
|
||||
)
|
||||
|
||||
|
||||
def _add_optimizer_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--lr",
|
||||
type=float,
|
||||
default=1e-4,
|
||||
help="Initial learning rate (after the potential warmup period) to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--scale_lr",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr_scheduler",
|
||||
type=str,
|
||||
default="constant",
|
||||
help=(
|
||||
'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
|
||||
' "constant", "constant_with_warmup"]'
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr_warmup_steps",
|
||||
type=int,
|
||||
default=500,
|
||||
help="Number of steps for the warmup in the lr scheduler.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr_num_cycles",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr_power",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Power factor of the polynomial scheduler.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--optimizer",
|
||||
type=lambda s: s.lower(),
|
||||
default="adam",
|
||||
choices=["adam", "adamw"],
|
||||
help=("The optimizer type to use."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--beta1",
|
||||
type=float,
|
||||
default=0.9,
|
||||
help="The beta1 parameter for the Adam and Prodigy optimizers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--beta2",
|
||||
type=float,
|
||||
default=0.95,
|
||||
help="The beta2 parameter for the Adam and Prodigy optimizers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--beta3",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Coefficients for computing the Prodigy optimizer's stepsize using running averages. If set to None, uses the value of square root of beta2.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--weight_decay",
|
||||
type=float,
|
||||
default=1e-04,
|
||||
help="Weight decay to use for optimizer.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--epsilon",
|
||||
type=float,
|
||||
default=1e-8,
|
||||
help="Epsilon value for the Adam optimizer and Prodigy optimizers.",
|
||||
)
|
||||
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
|
||||
|
||||
|
||||
def _add_validation_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--validation_prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_seperator' string.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validation_images",
|
||||
type=str,
|
||||
default=None,
|
||||
help="One or more image path(s)/URLs that is used during validation to verify that the model is learning. Multiple validation paths should be separated by the '--validation_prompt_seperator' string. These should correspond to the order of the validation prompts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validation_videos",
|
||||
type=str,
|
||||
default=None,
|
||||
help="One or more video path(s)/URLs that is used during validation to verify that the model is learning. Multiple validation paths should be separated by the '--validation_prompt_seperator' string. These should correspond to the order of the validation prompts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validation_separator",
|
||||
type=str,
|
||||
default=":::",
|
||||
help="String that separates multiple validation prompts",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_validation_videos",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of videos that should be generated during validation per `validation_prompt`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validation_epochs",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Run validation every X training epochs. Validation consists of running the validation prompt `args.num_validation_videos` times.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validation_steps",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Run validation every X training steps. Validation consists of running the validation prompt `args.num_validation_videos` times.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable_model_cpu_offload",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Whether or not to enable model-wise CPU offloading when performing validation/testing to save memory.",
|
||||
)
|
||||
|
||||
|
||||
def _add_miscellaneous_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument("--tracker_name", type=str, default="finetrainers", help="Project tracker name")
|
||||
parser.add_argument(
|
||||
"--push_to_hub",
|
||||
action="store_true",
|
||||
help="Whether or not to push the model to the Hub.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hub_token",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The token to use to push to the Model Hub.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hub_model_id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The name of the repository to keep in sync with the local `output_dir`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="finetrainer-training",
|
||||
help="The output directory where the model predictions and checkpoints will be written.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging_dir",
|
||||
type=str,
|
||||
default="logs",
|
||||
help="Directory where logs are stored.",
|
||||
)
|
||||
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(
|
||||
"--nccl_timeout",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Maximum timeout duration before which allgather, or related, operations fail in multi-GPU/multi-node training settings.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report_to",
|
||||
type=str,
|
||||
default="none",
|
||||
help=(
|
||||
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
|
||||
' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _map_to_args_type(args: Dict[str, Any]) -> Args:
|
||||
result_args = Args()
|
||||
|
||||
# Model arguments
|
||||
result_args.model_name = args.model_name
|
||||
result_args.pretrained_model_name_or_path = args.pretrained_model_name_or_path
|
||||
result_args.revision = args.revision
|
||||
result_args.variant = args.variant
|
||||
result_args.cache_dir = args.cache_dir
|
||||
|
||||
# Dataset arguments
|
||||
if args.data_root is None and args.dataset_file is None:
|
||||
raise ValueError("At least one of `data_root` or `dataset_file` should be provided.")
|
||||
|
||||
result_args.data_root = args.data_root
|
||||
result_args.dataset_file = args.dataset_file
|
||||
result_args.video_column = args.video_column
|
||||
result_args.caption_column = args.caption_column
|
||||
result_args.id_token = args.id_token
|
||||
result_args.image_resolution_buckets = args.image_resolution_buckets or DEFAULT_IMAGE_RESOLUTION_BUCKETS
|
||||
result_args.video_resolution_buckets = args.video_resolution_buckets or DEFAULT_VIDEO_RESOLUTION_BUCKETS
|
||||
result_args.video_reshape_mode = args.video_reshape_mode
|
||||
result_args.caption_dropout_p = args.caption_dropout_p
|
||||
|
||||
# Dataloader arguments
|
||||
result_args.dataloader_num_workers = args.dataloader_num_workers
|
||||
result_args.pin_memory = args.pin_memory
|
||||
|
||||
# Diffusion arguments
|
||||
result_args.flow_resolution_shifting = args.flow_resolution_shifting
|
||||
result_args.flow_weighting_scheme = args.flow_weighting_scheme
|
||||
result_args.flow_logit_mean = args.flow_logit_mean
|
||||
result_args.flow_logit_std = args.flow_logit_std
|
||||
result_args.flow_mode_scale = args.flow_mode_scale
|
||||
|
||||
# Training arguments
|
||||
result_args.training_type = args.training_type
|
||||
result_args.seed = args.seed
|
||||
result_args.mixed_precision = args.mixed_precision
|
||||
result_args.batch_size = args.batch_size
|
||||
result_args.train_epochs = args.train_epochs
|
||||
result_args.train_steps = args.train_steps
|
||||
result_args.rank = args.rank
|
||||
result_args.lora_alpha = args.lora_alpha
|
||||
result_args.gradient_accumulation_steps = args.gradient_accumulation_steps
|
||||
result_args.gradient_checkpointing = args.gradient_checkpointing
|
||||
result_args.checkpointing_steps = args.checkpointing_steps
|
||||
result_args.checkpointing_limit = args.checkpointing_limit
|
||||
result_args.resume_from_checkpoint = args.resume_from_checkpoint
|
||||
result_args.enable_slicing = args.enable_slicing
|
||||
result_args.enable_tiling = args.enable_tiling
|
||||
|
||||
# Optimizer arguments
|
||||
result_args.optimizer = args.optimizer or "adamw"
|
||||
result_args.lr = args.lr or 1e-4
|
||||
result_args.scale_lr = args.scale_lr
|
||||
result_args.lr_scheduler = args.lr_scheduler
|
||||
result_args.lr_warmup_steps = args.lr_warmup_steps
|
||||
result_args.lr_num_cycles = args.lr_num_cycles
|
||||
result_args.lr_power = args.lr_power
|
||||
result_args.beta1 = args.beta1
|
||||
result_args.beta2 = args.beta2
|
||||
result_args.beta3 = args.beta3
|
||||
result_args.weight_decay = args.weight_decay
|
||||
result_args.epsilon = args.epsilon
|
||||
result_args.max_grad_norm = args.max_grad_norm
|
||||
|
||||
# Validation arguments
|
||||
validation_prompts = args.validation_prompts.split(args.validation_separator) if args.validation_prompts else []
|
||||
validation_images = args.validation_images.split(args.validation_separator) if args.validation_images else None
|
||||
validation_videos = args.validation_videos.split(args.validation_separator) if args.validation_videos else None
|
||||
stripped_validation_prompts = []
|
||||
validation_heights = []
|
||||
validation_widths = []
|
||||
validation_num_frames = []
|
||||
for prompt in validation_prompts:
|
||||
prompt: str
|
||||
prompt = prompt.strip()
|
||||
actual_prompt, separator, resolution = prompt.rpartition("@@@")
|
||||
stripped_validation_prompts.append(actual_prompt)
|
||||
num_frames, height, width = None, None, None
|
||||
if len(resolution) > 0:
|
||||
num_frames, height, width = map(int, resolution.split("x"))
|
||||
validation_num_frames.append(num_frames)
|
||||
validation_heights.append(height)
|
||||
validation_widths.append(width)
|
||||
|
||||
if validation_images is None:
|
||||
validation_images = [None] * len(validation_prompts)
|
||||
if validation_videos is None:
|
||||
validation_videos = [None] * len(validation_prompts)
|
||||
|
||||
result_args.validation_prompts = stripped_validation_prompts
|
||||
result_args.validation_heights = validation_heights
|
||||
result_args.validation_widths = validation_widths
|
||||
result_args.validation_num_frames = validation_num_frames
|
||||
result_args.validation_images = validation_images
|
||||
result_args.validation_videos = validation_videos
|
||||
|
||||
result_args.num_validation_videos_per_prompt = args.num_validation_videos
|
||||
result_args.validation_every_n_epochs = args.validation_epochs
|
||||
result_args.validation_every_n_steps = args.validation_steps
|
||||
result_args.enable_model_cpu_offload = args.enable_model_cpu_offload
|
||||
|
||||
# Miscellaneous arguments
|
||||
result_args.tracker_name = args.tracker_name
|
||||
result_args.push_to_hub = args.push_to_hub
|
||||
result_args.hub_token = args.hub_token
|
||||
result_args.hub_model_id = args.hub_model_id
|
||||
result_args.output_dir = args.output_dir
|
||||
result_args.logging_dir = args.logging_dir
|
||||
result_args.allow_tf32 = args.allow_tf32
|
||||
result_args.nccl_timeout = args.nccl_timeout
|
||||
result_args.report_to = args.report_to
|
||||
|
||||
return result_args
|
||||
|
||||
|
||||
def _validate_validation_args(args: Args):
|
||||
assert args.validation_prompts is not None, "Validation prompts are required for validation"
|
||||
if args.validation_images is not None:
|
||||
assert len(args.validation_images) == len(
|
||||
args.validation_prompts
|
||||
), "Validation images and prompts should be of same length"
|
||||
if args.validation_videos is not None:
|
||||
assert len(args.validation_videos) == len(
|
||||
args.validation_prompts
|
||||
), "Validation videos and prompts should be of same length"
|
||||
assert len(args.validation_prompts) == len(
|
||||
args.validation_heights
|
||||
), "Validation prompts and heights should be of same length"
|
||||
assert len(args.validation_prompts) == len(
|
||||
args.validation_widths
|
||||
), "Validation prompts and widths should be of same length"
|
||||
@@ -0,0 +1,50 @@
|
||||
import os
|
||||
|
||||
|
||||
DEFAULT_HEIGHT_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536]
|
||||
DEFAULT_WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536]
|
||||
DEFAULT_FRAME_BUCKETS = [49]
|
||||
|
||||
DEFAULT_IMAGE_RESOLUTION_BUCKETS = []
|
||||
for height in DEFAULT_HEIGHT_BUCKETS:
|
||||
for width in DEFAULT_WIDTH_BUCKETS:
|
||||
DEFAULT_IMAGE_RESOLUTION_BUCKETS.append((height, width))
|
||||
|
||||
DEFAULT_VIDEO_RESOLUTION_BUCKETS = []
|
||||
for frames in DEFAULT_FRAME_BUCKETS:
|
||||
for height in DEFAULT_HEIGHT_BUCKETS:
|
||||
for width in DEFAULT_WIDTH_BUCKETS:
|
||||
DEFAULT_VIDEO_RESOLUTION_BUCKETS.append((frames, height, width))
|
||||
|
||||
|
||||
FINETRAINERS_LOG_LEVEL = os.environ.get("FINETRAINERS_LOG_LEVEL", "INFO")
|
||||
|
||||
MODEL_DESCRIPTION = r"""
|
||||
\# {model_id} {training_type} finetune
|
||||
|
||||
<Gallery />
|
||||
|
||||
\#\# Model Description
|
||||
|
||||
This model is a {training_type} of the `{model_id}` model.
|
||||
|
||||
This model was trained using the `fine-video-trainers` library - a repository containing memory-optimized scripts for training video models with [Diffusers](https://github.com/huggingface/diffusers).
|
||||
|
||||
\#\# Download model
|
||||
|
||||
[Download LoRA]({repo_id}/tree/main) in the Files & Versions tab.
|
||||
|
||||
\#\# Usage
|
||||
|
||||
Requires [🧨 Diffusers](https://github.com/huggingface/diffusers) installed.
|
||||
|
||||
```python
|
||||
{model_example}
|
||||
```
|
||||
|
||||
For more details, including weighting, merging and fusing LoRAs, check the [documentation](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters) on loading LoRAs in diffusers.
|
||||
|
||||
\#\# License
|
||||
|
||||
Please adhere to the license of the base model.
|
||||
""".strip()
|
||||
@@ -0,0 +1,321 @@
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torchvision.transforms as TT
|
||||
from accelerate.logging import get_logger
|
||||
from torch.utils.data import Dataset, Sampler
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import InterpolationMode
|
||||
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__)
|
||||
|
||||
|
||||
class VideoDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root: str,
|
||||
caption_column: str,
|
||||
video_column: str,
|
||||
resolution_buckets: List[Tuple[int, int, int]],
|
||||
dataset_file: Optional[str] = None,
|
||||
id_token: Optional[str] = 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.id_token = f"{id_token.strip()} " if id_token else ""
|
||||
self.resolution_buckets = resolution_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()
|
||||
|
||||
if len(self.video_paths) != 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.Lambda(self.scale_transform),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
|
||||
]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def scale_transform(x):
|
||||
return x / 255.0
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.video_paths)
|
||||
|
||||
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
|
||||
|
||||
prompt = self.id_token + self.prompts[index]
|
||||
video = self._preprocess_video(self.video_paths[index])
|
||||
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"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 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.
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
class VideoDatasetWithResizing(VideoDataset):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.max_num_frames = max(self.resolution_buckets, key=lambda x: x[0])[0]
|
||||
|
||||
def _preprocess_video(self, path: Path) -> torch.Tensor:
|
||||
video_reader = decord.VideoReader(uri=path.as_posix())
|
||||
video_num_frames = len(video_reader)
|
||||
nearest_frame_bucket = min(
|
||||
[bucket for bucket in self.resolution_buckets if bucket[0] <= video_num_frames],
|
||||
key=lambda x: abs(x[0] - min(video_num_frames, self.max_num_frames)),
|
||||
default=1,
|
||||
)[0]
|
||||
|
||||
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
|
||||
|
||||
def _find_nearest_resolution(self, height, width):
|
||||
nearest_res = min(self.resolution_buckets, key=lambda x: abs(x[1] - height) + abs(x[2] - width))
|
||||
return nearest_res[1], nearest_res[2]
|
||||
|
||||
|
||||
class VideoDatasetWithResizeAndRectangleCrop(VideoDataset):
|
||||
def __init__(self, video_reshape_mode: str = "center", *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.video_reshape_mode = video_reshape_mode
|
||||
self.max_num_frames = max(self.resolution_buckets, key=lambda x: x[0])[0]
|
||||
|
||||
def _resize_for_rectangle_crop(self, arr, image_size):
|
||||
reshape_mode = self.video_reshape_mode
|
||||
if arr.shape[3] / arr.shape[2] > image_size[1] / image_size[0]:
|
||||
arr = resize(
|
||||
arr,
|
||||
size=[image_size[0], int(arr.shape[3] * image_size[0] / arr.shape[2])],
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
)
|
||||
else:
|
||||
arr = resize(
|
||||
arr,
|
||||
size=[int(arr.shape[2] * image_size[1] / arr.shape[3]), image_size[1]],
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
)
|
||||
|
||||
h, w = arr.shape[2], arr.shape[3]
|
||||
arr = arr.squeeze(0)
|
||||
|
||||
delta_h = h - image_size[0]
|
||||
delta_w = w - image_size[1]
|
||||
|
||||
if reshape_mode == "random" or reshape_mode == "none":
|
||||
top = np.random.randint(0, delta_h + 1)
|
||||
left = np.random.randint(0, delta_w + 1)
|
||||
elif reshape_mode == "center":
|
||||
top, left = delta_h // 2, delta_w // 2
|
||||
else:
|
||||
raise NotImplementedError
|
||||
arr = TT.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])
|
||||
return arr
|
||||
|
||||
def _preprocess_video(self, path: Path) -> torch.Tensor:
|
||||
video_reader = decord.VideoReader(uri=path.as_posix())
|
||||
video_num_frames = len(video_reader)
|
||||
nearest_frame_bucket = min(
|
||||
[bucket for bucket in self.resolution_buckets if bucket <= video_num_frames],
|
||||
key=lambda x: abs(x[0] - min(video_num_frames, self.max_num_frames)),
|
||||
default=1,
|
||||
)[0]
|
||||
|
||||
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 = self._resize_for_rectangle_crop(frames, nearest_res)
|
||||
frames = torch.stack([self.video_transforms(frame) for frame in frames_resized], dim=0)
|
||||
return frames
|
||||
|
||||
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):
|
||||
r"""
|
||||
PyTorch Sampler that groups 3D data by height, width and frames.
|
||||
|
||||
Args:
|
||||
data_source (`VideoDataset`):
|
||||
A PyTorch dataset object that is an instance of `VideoDataset`.
|
||||
batch_size (`int`, defaults to `8`):
|
||||
The batch size to use for training.
|
||||
shuffle (`bool`, defaults to `True`):
|
||||
Whether or not to shuffle the data in each batch before dispatching to dataloader.
|
||||
drop_last (`bool`, defaults to `False`):
|
||||
Whether or not to drop incomplete buckets of data after completely iterating over all data
|
||||
in the dataset. If set to True, only batches that have `batch_size` number of entries will
|
||||
be yielded. If set to False, it is guaranteed that all data in the dataset will be processed
|
||||
and batches that do not have `batch_size` number of entries will also be yielded.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, data_source: VideoDataset, batch_size: int = 8, shuffle: bool = True, drop_last: bool = False
|
||||
) -> None:
|
||||
self.data_source = data_source
|
||||
self.batch_size = batch_size
|
||||
self.shuffle = shuffle
|
||||
self.drop_last = drop_last
|
||||
|
||||
self.buckets = {resolution: [] for resolution in data_source.resolution_buckets}
|
||||
|
||||
self._raised_warning_for_drop_last = False
|
||||
|
||||
def __len__(self):
|
||||
if self.drop_last and not self._raised_warning_for_drop_last:
|
||||
self._raised_warning_for_drop_last = True
|
||||
logger.warning(
|
||||
"Calculating the length for bucket sampler is not possible when `drop_last` is set to True. This may cause problems when setting the number of epochs used for training."
|
||||
)
|
||||
return (len(self.data_source) + self.batch_size - 1) // self.batch_size
|
||||
|
||||
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)] = []
|
||||
|
||||
if self.drop_last:
|
||||
return
|
||||
|
||||
for fhw, bucket in list(self.buckets.items()):
|
||||
if len(bucket) == 0:
|
||||
continue
|
||||
if self.shuffle:
|
||||
random.shuffle(bucket)
|
||||
yield bucket
|
||||
del self.buckets[fhw]
|
||||
self.buckets[fhw] = []
|
||||
@@ -0,0 +1 @@
|
||||
from .ltx_video import LTX_VIDEO_T2V_CONFIG
|
||||
@@ -0,0 +1,264 @@
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from accelerate.logging import get_logger
|
||||
from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
|
||||
from diffusers.utils import logging
|
||||
from transformers import T5EncoderModel, T5Tokenizer
|
||||
from PIL import Image
|
||||
|
||||
|
||||
logger = get_logger("finetrainers") # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def load_components(
|
||||
model_id: str = "Lightricks/LTX-Video",
|
||||
text_encoder_dtype: torch.dtype = torch.bfloat16,
|
||||
transformer_dtype: torch.dtype = torch.bfloat16,
|
||||
vae_dtype: torch.dtype = torch.bfloat16,
|
||||
cache_dir: Optional[str] = None,
|
||||
) -> Dict[str, nn.Module]:
|
||||
tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer", cache_dir=cache_dir)
|
||||
text_encoder = T5EncoderModel.from_pretrained(
|
||||
model_id, subfolder="text_encoder", torch_dtype=text_encoder_dtype, cache_dir=cache_dir
|
||||
)
|
||||
transformer = LTXVideoTransformer3DModel.from_pretrained(
|
||||
model_id, subfolder="transformer", torch_dtype=transformer_dtype, cache_dir=cache_dir
|
||||
)
|
||||
vae = AutoencoderKLLTXVideo.from_pretrained(model_id, subfolder="vae", torch_dtype=vae_dtype, cache_dir=cache_dir)
|
||||
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler", cache_dir=cache_dir)
|
||||
return {
|
||||
"tokenizer": tokenizer,
|
||||
"text_encoder": text_encoder,
|
||||
"transformer": transformer,
|
||||
"vae": vae,
|
||||
"scheduler": scheduler,
|
||||
}
|
||||
|
||||
|
||||
def initialize_pipeline(
|
||||
model_id: str = "Lightricks/LTX-Video",
|
||||
text_encoder_dtype: torch.dtype = torch.bfloat16,
|
||||
transformer_dtype: torch.dtype = torch.bfloat16,
|
||||
vae_dtype: torch.dtype = torch.bfloat16,
|
||||
tokenizer: Optional[T5Tokenizer] = None,
|
||||
text_encoder: Optional[T5EncoderModel] = None,
|
||||
transformer: Optional[LTXVideoTransformer3DModel] = None,
|
||||
vae: Optional[AutoencoderKLLTXVideo] = None,
|
||||
scheduler: Optional[FlowMatchEulerDiscreteScheduler] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
enable_slicing: bool = False,
|
||||
enable_tiling: bool = False,
|
||||
enable_model_cpu_offload: bool = False,
|
||||
) -> LTXPipeline:
|
||||
component_name_pairs = [
|
||||
("tokenizer", tokenizer),
|
||||
("text_encoder", text_encoder),
|
||||
("transformer", transformer),
|
||||
("vae", vae),
|
||||
("scheduler", scheduler),
|
||||
]
|
||||
components = {}
|
||||
for name, component in component_name_pairs:
|
||||
if component is not None:
|
||||
components[name] = component
|
||||
|
||||
pipe = LTXPipeline.from_pretrained(model_id, **components, cache_dir=cache_dir)
|
||||
pipe.text_encoder = pipe.text_encoder.to(dtype=text_encoder_dtype)
|
||||
pipe.transformer = pipe.transformer.to(dtype=transformer_dtype)
|
||||
pipe.vae = pipe.vae.to(dtype=vae_dtype)
|
||||
|
||||
if enable_slicing:
|
||||
pipe.vae.enable_slicing()
|
||||
if enable_tiling:
|
||||
pipe.vae.enable_tiling()
|
||||
|
||||
if enable_model_cpu_offload:
|
||||
pipe.enable_model_cpu_offload(device=device)
|
||||
else:
|
||||
pipe.to(device=device)
|
||||
|
||||
return pipe
|
||||
|
||||
|
||||
def prepare_conditions(
|
||||
tokenizer: T5Tokenizer,
|
||||
text_encoder: T5EncoderModel,
|
||||
prompt: Union[str, List[str]],
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
max_sequence_length: int = 128,
|
||||
) -> torch.Tensor:
|
||||
device = device or text_encoder.device
|
||||
dtype = dtype or text_encoder.dtype
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
|
||||
return _encode_prompt_t5(tokenizer, text_encoder, prompt, device, dtype, max_sequence_length)
|
||||
|
||||
|
||||
def prepare_latents(
|
||||
vae: AutoencoderKLLTXVideo,
|
||||
image_or_video: torch.Tensor,
|
||||
patch_size: int = 1,
|
||||
patch_size_t: int = 1,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> torch.Tensor:
|
||||
device = device or vae.device
|
||||
dtype = dtype or vae.dtype
|
||||
|
||||
if image_or_video.ndim == 4:
|
||||
image_or_video = image_or_video.unsqueeze(2)
|
||||
assert image_or_video.ndim == 5, f"Expected 5D tensor, got {image_or_video.ndim}D tensor"
|
||||
|
||||
image_or_video = image_or_video.to(device=device, dtype=dtype)
|
||||
image_or_video = image_or_video.permute(0, 2, 1, 3, 4).contiguous() # [B, C, F, H, W] -> [B, F, C, H, W]
|
||||
latents = vae.encode(image_or_video).latent_dist.sample(generator=generator)
|
||||
_, _, num_frames, height, width = latents.shape
|
||||
latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std)
|
||||
latents = _pack_latents(latents, patch_size, patch_size_t)
|
||||
return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
|
||||
|
||||
|
||||
def collate_fn_t2v(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
|
||||
return {
|
||||
"prompts": [x["prompt"] for x in batch[0]],
|
||||
"videos": torch.stack([x["video"] for x in batch[0]]),
|
||||
}
|
||||
|
||||
|
||||
def forward_pass(
|
||||
transformer: LTXVideoTransformer3DModel,
|
||||
prompt_embeds: torch.Tensor,
|
||||
prompt_attention_mask: torch.Tensor,
|
||||
latents: torch.Tensor,
|
||||
noisy_latents: torch.Tensor,
|
||||
timesteps: torch.LongTensor,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
) -> torch.Tensor:
|
||||
# TODO(aryan): make configurable
|
||||
rope_interpolation_scale = [1 / 25, 32, 32]
|
||||
|
||||
denoised_latents = transformer(
|
||||
hidden_states=noisy_latents,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
timestep=timesteps,
|
||||
encoder_attention_mask=prompt_attention_mask,
|
||||
num_frames=num_frames,
|
||||
height=height,
|
||||
width=width,
|
||||
rope_interpolation_scale=rope_interpolation_scale,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
return {"latents": denoised_latents}
|
||||
|
||||
|
||||
def validation(
|
||||
pipeline: LTXPipeline,
|
||||
prompt: str,
|
||||
image: Optional[Image.Image] = None,
|
||||
video: Optional[List[Image.Image]] = None,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_frames: Optional[int] = None,
|
||||
frame_rate: int = 25,
|
||||
num_videos_per_prompt: int = 1,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
**kwargs,
|
||||
):
|
||||
generation_kwargs = {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"num_frames": num_frames,
|
||||
"frame_rate": frame_rate,
|
||||
"num_videos_per_prompt": num_videos_per_prompt,
|
||||
"generator": generator,
|
||||
"return_dict": True,
|
||||
"output_type": "pil",
|
||||
}
|
||||
generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
|
||||
video = pipeline(**generation_kwargs).frames[0]
|
||||
return [("video", video)]
|
||||
|
||||
|
||||
def _encode_prompt_t5(
|
||||
tokenizer: T5Tokenizer,
|
||||
text_encoder: T5EncoderModel,
|
||||
prompt: List[str],
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
max_sequence_length,
|
||||
) -> torch.Tensor:
|
||||
batch_size = len(prompt)
|
||||
|
||||
text_inputs = tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
prompt_attention_mask = text_inputs.attention_mask
|
||||
prompt_attention_mask = prompt_attention_mask.bool().to(device)
|
||||
|
||||
prompt_embeds = text_encoder(text_input_ids.to(device))[0]
|
||||
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
||||
prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
|
||||
|
||||
return {"prompt_embeds": prompt_embeds, "prompt_attention_mask": prompt_attention_mask}
|
||||
|
||||
|
||||
def _normalize_latents(
|
||||
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
|
||||
) -> torch.Tensor:
|
||||
# Normalize latents across the channel dimension [B, C, F, H, W]
|
||||
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
latents = (latents - latents_mean) * scaling_factor / latents_std
|
||||
return latents
|
||||
|
||||
|
||||
def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
|
||||
# Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
|
||||
# The patch dimensions are then permuted and collapsed into the channel dimension of shape:
|
||||
# [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
|
||||
# dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
|
||||
batch_size, num_channels, num_frames, height, width = latents.shape
|
||||
post_patch_num_frames = num_frames // patch_size_t
|
||||
post_patch_height = height // patch_size
|
||||
post_patch_width = width // patch_size
|
||||
latents = latents.reshape(
|
||||
batch_size,
|
||||
-1,
|
||||
post_patch_num_frames,
|
||||
patch_size_t,
|
||||
post_patch_height,
|
||||
patch_size,
|
||||
post_patch_width,
|
||||
patch_size,
|
||||
)
|
||||
latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
|
||||
return latents
|
||||
|
||||
|
||||
LTX_VIDEO_T2V_CONFIG = {
|
||||
"pipeline_cls": LTXPipeline,
|
||||
"load_components": load_components,
|
||||
"initialize_pipeline": initialize_pipeline,
|
||||
"prepare_conditions": prepare_conditions,
|
||||
"prepare_latents": prepare_latents,
|
||||
"collate_fn": collate_fn_t2v,
|
||||
"forward_pass": forward_pass,
|
||||
"validation": validation,
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
from typing import Any, Dict
|
||||
|
||||
from .ltx_video import LTX_VIDEO_T2V_CONFIG
|
||||
|
||||
|
||||
SUPPORTED_MODEL_CONFIGS = {
|
||||
"ltx_video": LTX_VIDEO_T2V_CONFIG,
|
||||
}
|
||||
|
||||
|
||||
def get_config_from_model_name(model_name: str) -> Dict[str, Any]:
|
||||
if model_name not in SUPPORTED_MODEL_CONFIGS:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not supported. Supported models are: {list(SUPPORTED_MODEL_CONFIGS.keys())}"
|
||||
)
|
||||
return SUPPORTED_MODEL_CONFIGS[model_name]
|
||||
@@ -0,0 +1,23 @@
|
||||
import torch
|
||||
from accelerate import Accelerator
|
||||
|
||||
|
||||
class State:
|
||||
# Training state
|
||||
seed: int = None
|
||||
model_name: str = None
|
||||
accelerator: Accelerator = None
|
||||
weight_dtype: torch.dtype = None
|
||||
train_epochs: int = None
|
||||
train_steps: int = None
|
||||
overwrote_max_train_steps: bool = False
|
||||
num_trainable_parameters: int = 0
|
||||
learning_rate: float = None
|
||||
train_batch_size: int = None
|
||||
generator: torch.Generator = None
|
||||
|
||||
# Hub state
|
||||
repo_id: str = None
|
||||
|
||||
# Artifacts state
|
||||
output_dir: str = None
|
||||
@@ -0,0 +1,679 @@
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
from datetime import timedelta
|
||||
from typing import Any, Dict
|
||||
from pathlib import Path
|
||||
|
||||
import diffusers
|
||||
import torch
|
||||
import torch.backends
|
||||
import transformers
|
||||
import wandb
|
||||
from accelerate import Accelerator, DistributedType
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import (
|
||||
DistributedDataParallelKwargs,
|
||||
InitProcessGroupKwargs,
|
||||
ProjectConfiguration,
|
||||
set_seed,
|
||||
gather_object,
|
||||
)
|
||||
from diffusers.optimization import get_scheduler
|
||||
from diffusers.training_utils import (
|
||||
cast_training_params,
|
||||
compute_density_for_timestep_sampling,
|
||||
compute_loss_weighting_for_sd3,
|
||||
)
|
||||
from diffusers.utils import export_to_video, load_image, load_video
|
||||
from huggingface_hub import create_repo, upload_folder
|
||||
from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict
|
||||
from tqdm import tqdm
|
||||
|
||||
from .args import Args, validate_args
|
||||
from .constants import FINETRAINERS_LOG_LEVEL
|
||||
from .dataset import BucketSampler, VideoDatasetWithResizing
|
||||
from .models import get_config_from_model_name
|
||||
from .state import State
|
||||
from .utils.file_utils import find_files, delete_files, string_to_filename
|
||||
from .utils.optimizer_utils import get_optimizer, gradient_norm
|
||||
from .utils.memory_utils import get_memory_statistics, free_memory, make_contiguous
|
||||
from .utils.torch_utils import unwrap_model
|
||||
|
||||
|
||||
logger = get_logger("finetrainers")
|
||||
logger.setLevel(FINETRAINERS_LOG_LEVEL)
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(self, args: Args) -> None:
|
||||
validate_args(args)
|
||||
|
||||
self.args = args
|
||||
self.state = State()
|
||||
|
||||
# Tokenizers
|
||||
self.tokenizer = None
|
||||
self.tokenizer_2 = None
|
||||
self.tokenizer_3 = None
|
||||
|
||||
# Text encoders
|
||||
self.text_encoder = None
|
||||
self.text_encoder_2 = None
|
||||
self.text_encoder_3 = None
|
||||
|
||||
# Denoisers
|
||||
self.transformer = None
|
||||
self.unet = None
|
||||
|
||||
# Autoencoders
|
||||
self.vae = None
|
||||
|
||||
self._init_distributed()
|
||||
self._init_logging()
|
||||
self._init_directories_and_repositories()
|
||||
|
||||
self.state.model_name = self.args.model_name
|
||||
self.model_config = get_config_from_model_name(self.args.model_name)
|
||||
|
||||
def prepare_models(self) -> None:
|
||||
logger.info("Initializing models")
|
||||
|
||||
# TODO(aryan): refactor in future
|
||||
load_components_kwargs = {
|
||||
"text_encoder_dtype": torch.bfloat16,
|
||||
"transformer_dtype": torch.bfloat16,
|
||||
"vae_dtype": torch.bfloat16,
|
||||
"cache_dir": self.args.cache_dir,
|
||||
}
|
||||
if self.args.pretrained_model_name_or_path is not None:
|
||||
load_components_kwargs["model_id"] = self.args.pretrained_model_name_or_path
|
||||
components = self._model_config_call(self.model_config["load_components"], load_components_kwargs)
|
||||
|
||||
self.tokenizer = components.get("tokenizer", None)
|
||||
self.text_encoder = components.get("text_encoder", None)
|
||||
self.transformer = components.get("transformer", None)
|
||||
self.vae = components.get("vae", None)
|
||||
self.scheduler = components.get("scheduler", None)
|
||||
|
||||
self.transformer_config = self.transformer.config if self.transformer is not None else None
|
||||
|
||||
def prepare_dataset(self) -> None:
|
||||
logger.info("Initializing dataset and dataloader")
|
||||
|
||||
self.dataset = VideoDatasetWithResizing(
|
||||
data_root=self.args.data_root,
|
||||
caption_column=self.args.caption_column,
|
||||
video_column=self.args.video_column,
|
||||
resolution_buckets=self.args.video_resolution_buckets,
|
||||
dataset_file=self.args.dataset_file,
|
||||
id_token=self.args.id_token,
|
||||
)
|
||||
self.dataloader = torch.utils.data.DataLoader(
|
||||
self.dataset,
|
||||
batch_size=1,
|
||||
sampler=BucketSampler(self.dataset, batch_size=self.args.batch_size, shuffle=True),
|
||||
collate_fn=self.model_config.get("collate_fn"),
|
||||
num_workers=self.args.dataloader_num_workers,
|
||||
pin_memory=self.args.pin_memory,
|
||||
)
|
||||
|
||||
def prepare_trainable_parameters(self) -> None:
|
||||
logger.info("Initializing trainable parameters")
|
||||
|
||||
# TODO(aryan): refactor later. for now only lora is supported
|
||||
self.text_encoder.requires_grad_(False)
|
||||
self.transformer.requires_grad_(False)
|
||||
self.vae.requires_grad_(False)
|
||||
|
||||
# For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
|
||||
# as these weights are only used for inference, keeping weights in full precision is not required.
|
||||
weight_dtype = torch.float32
|
||||
if self.state.accelerator.mixed_precision == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif self.state.accelerator.mixed_precision == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
|
||||
# due to pytorch#99272, MPS does not yet support bfloat16.
|
||||
raise ValueError(
|
||||
"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
|
||||
)
|
||||
|
||||
# TODO(aryan): handle torch dtype from accelerator vs model dtype
|
||||
self.state.weight_dtype = weight_dtype
|
||||
self.text_encoder.to(self.state.accelerator.device, dtype=weight_dtype)
|
||||
self.transformer.to(self.state.accelerator.device, dtype=weight_dtype)
|
||||
self.vae.to(self.state.accelerator.device, dtype=weight_dtype)
|
||||
|
||||
if self.args.gradient_checkpointing:
|
||||
self.transformer.enable_gradient_checkpointing()
|
||||
|
||||
transformer_lora_config = LoraConfig(
|
||||
r=self.args.rank,
|
||||
lora_alpha=self.args.lora_alpha,
|
||||
init_lora_weights=True,
|
||||
target_modules=self.args.target_modules,
|
||||
)
|
||||
self.transformer.add_adapter(transformer_lora_config)
|
||||
|
||||
# TODO: refactor
|
||||
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
|
||||
def save_model_hook(models, weights, output_dir):
|
||||
if self.state.accelerator.is_main_process:
|
||||
transformer_lora_layers_to_save = None
|
||||
|
||||
for model in models:
|
||||
if isinstance(
|
||||
unwrap_model(self.state.accelerator, model),
|
||||
type(unwrap_model(self.state.accelerator, self.transformer)),
|
||||
):
|
||||
model = unwrap_model(self.state.accelerator, model)
|
||||
transformer_lora_layers_to_save = get_peft_model_state_dict(model)
|
||||
else:
|
||||
raise ValueError(f"Unexpected save model: {model.__class__}")
|
||||
|
||||
# make sure to pop weight so that corresponding model is not saved again
|
||||
if weights:
|
||||
weights.pop()
|
||||
|
||||
self.model_config["pipeline_cls"].save_lora_weights(
|
||||
output_dir,
|
||||
transformer_lora_layers=transformer_lora_layers_to_save,
|
||||
)
|
||||
|
||||
def load_model_hook(models, input_dir):
|
||||
transformer_ = self.model_config["pipeline_cls"].from_pretrained(
|
||||
self.args.pretrained_model_name_or_path, subfolder="transformer"
|
||||
)
|
||||
transformer_.add_adapter(transformer_lora_config)
|
||||
|
||||
lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir)
|
||||
|
||||
transformer_state_dict = {
|
||||
f'{k.replace("transformer.", "")}': v
|
||||
for k, v in lora_state_dict.items()
|
||||
if k.startswith("transformer.")
|
||||
}
|
||||
incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
|
||||
if incompatible_keys is not None:
|
||||
# check only for unexpected keys
|
||||
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
|
||||
if unexpected_keys:
|
||||
logger.warning(
|
||||
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
|
||||
f" {unexpected_keys}. "
|
||||
)
|
||||
|
||||
# Make sure the trainable params are in float32. This is again needed since the base models
|
||||
# are in `weight_dtype`. More details:
|
||||
# https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
|
||||
if self.args.mixed_precision == "fp16":
|
||||
# only upcast trainable parameters (LoRA) into fp32
|
||||
cast_training_params([transformer_])
|
||||
|
||||
self.state.accelerator.register_save_state_pre_hook(save_model_hook)
|
||||
self.state.accelerator.register_load_state_pre_hook(load_model_hook)
|
||||
|
||||
# Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
||||
if self.args.allow_tf32 and torch.cuda.is_available():
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
def prepare_optimizer(self) -> None:
|
||||
logger.info("Initializing optimizer and lr scheduler")
|
||||
|
||||
self.state.train_epochs = self.args.train_epochs
|
||||
self.state.train_steps = self.args.train_steps
|
||||
|
||||
# Make sure the trainable params are in float32
|
||||
if self.args.mixed_precision == "fp16":
|
||||
# only upcast trainable parameters (LoRA) into fp32
|
||||
cast_training_params([self.transformer], dtype=torch.float32)
|
||||
|
||||
self.state.learning_rate = self.args.lr
|
||||
if self.args.scale_lr:
|
||||
self.state.learning_rate = (
|
||||
self.state.learning_rate
|
||||
* self.args.gradient_accumulation_steps
|
||||
* self.args.batch_size
|
||||
* self.state.accelerator.num_processes
|
||||
)
|
||||
|
||||
transformer_lora_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters()))
|
||||
transformer_parameters_with_lr = {
|
||||
"params": transformer_lora_parameters,
|
||||
"lr": self.state.learning_rate,
|
||||
}
|
||||
params_to_optimize = [transformer_parameters_with_lr]
|
||||
self.state.num_trainable_parameters = sum(p.numel() for p in transformer_lora_parameters)
|
||||
|
||||
# TODO(aryan): add deepspeed support
|
||||
optimizer = get_optimizer(
|
||||
params_to_optimize=params_to_optimize,
|
||||
optimizer_name=self.args.optimizer,
|
||||
learning_rate=self.args.lr,
|
||||
beta1=self.args.beta1,
|
||||
beta2=self.args.beta2,
|
||||
beta3=self.args.beta3,
|
||||
epsilon=self.args.epsilon,
|
||||
weight_decay=self.args.weight_decay,
|
||||
)
|
||||
|
||||
num_update_steps_per_epoch = math.ceil(len(self.dataloader) / self.args.gradient_accumulation_steps)
|
||||
if self.state.train_steps is None:
|
||||
self.state.train_steps = self.state.train_epochs * num_update_steps_per_epoch
|
||||
self.state.overwrote_max_train_steps = True
|
||||
|
||||
lr_scheduler = get_scheduler(
|
||||
name=self.args.lr_scheduler,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=self.args.lr_warmup_steps * self.state.accelerator.num_processes,
|
||||
num_training_steps=self.state.train_steps * self.state.accelerator.num_processes,
|
||||
num_cycles=self.args.lr_num_cycles,
|
||||
power=self.args.lr_power,
|
||||
)
|
||||
|
||||
self.optimizer = optimizer
|
||||
self.lr_scheduler = lr_scheduler
|
||||
|
||||
def prepare_for_training(self) -> None:
|
||||
self.transformer, self.optimizer, self.dataloader, self.lr_scheduler = self.state.accelerator.prepare(
|
||||
self.transformer, self.optimizer, self.dataloader, self.lr_scheduler
|
||||
)
|
||||
|
||||
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
||||
num_update_steps_per_epoch = math.ceil(len(self.dataloader) / self.args.gradient_accumulation_steps)
|
||||
if self.state.overwrote_max_train_steps:
|
||||
self.state.train_steps = self.state.train_epochs * num_update_steps_per_epoch
|
||||
# Afterwards we recalculate our number of training epochs
|
||||
self.state.train_epochs = math.ceil(self.state.train_steps / num_update_steps_per_epoch)
|
||||
|
||||
def prepare_trackers(self) -> None:
|
||||
logger.info("Initializing trackers")
|
||||
|
||||
tracker_name = self.args.tracker_name or "finetrainers-experiment"
|
||||
self.state.accelerator.init_trackers(tracker_name, config=self.args.to_dict())
|
||||
|
||||
def train(self) -> None:
|
||||
logger.info("Starting training")
|
||||
|
||||
memory_statistics = get_memory_statistics()
|
||||
logger.info(f"Memory before training start: {json.dumps(memory_statistics, indent=4)}")
|
||||
|
||||
self.state.train_batch_size = (
|
||||
self.args.batch_size * self.state.accelerator.num_processes * self.args.gradient_accumulation_steps
|
||||
)
|
||||
info = {
|
||||
"trainable parameters": self.state.num_trainable_parameters,
|
||||
"total samples": len(self.dataset),
|
||||
"train epochs": self.state.train_epochs,
|
||||
"train steps": self.state.train_steps,
|
||||
"batches per device": self.args.batch_size,
|
||||
"total batches observed per epoch": len(self.dataloader),
|
||||
"train batch size": self.state.train_batch_size,
|
||||
"gradient accumulation steps": self.args.gradient_accumulation_steps,
|
||||
}
|
||||
logger.info(f"Training configuration: {json.dumps(info, indent=4)}")
|
||||
|
||||
# TODO(aryan): handle resume from checkpoint
|
||||
|
||||
global_step = 0
|
||||
first_epoch = 0
|
||||
initial_global_step = 0
|
||||
progress_bar = tqdm(
|
||||
range(0, self.state.train_steps),
|
||||
initial=initial_global_step,
|
||||
desc="Training steps",
|
||||
disable=not self.state.accelerator.is_local_main_process,
|
||||
)
|
||||
|
||||
accelerator = self.state.accelerator
|
||||
weight_dtype = self.state.weight_dtype
|
||||
scheduler_sigmas = self.scheduler.sigmas.clone().to(device=accelerator.device, dtype=weight_dtype)
|
||||
generator = torch.Generator(device=accelerator.device)
|
||||
if self.args.seed is not None:
|
||||
generator = generator.manual_seed(self.args.seed)
|
||||
self.state.generator = generator
|
||||
|
||||
for epoch in range(first_epoch, self.state.train_epochs):
|
||||
logger.debug(f"Starting epoch ({epoch + 1}/{self.state.train_epochs})")
|
||||
|
||||
self.transformer.train()
|
||||
models_to_accumulate = [self.transformer]
|
||||
|
||||
for step, batch in enumerate(self.dataloader):
|
||||
logger.debug(f"Starting step {step + 1}")
|
||||
logs = {}
|
||||
|
||||
with accelerator.accumulate(models_to_accumulate):
|
||||
videos = batch["videos"]
|
||||
prompts = batch["prompts"]
|
||||
batch_size = len(prompts)
|
||||
|
||||
if self.args.caption_dropout_technique == "empty":
|
||||
if random.random() < self.args.caption_dropout_p:
|
||||
prompts = [""] * batch_size
|
||||
|
||||
latent_conditions = self.model_config["prepare_latents"](
|
||||
vae=self.vae,
|
||||
image_or_video=videos,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=generator,
|
||||
)
|
||||
latent_conditions = make_contiguous(latent_conditions)
|
||||
|
||||
other_conditions = self.model_config["prepare_conditions"](
|
||||
tokenizer=self.tokenizer,
|
||||
text_encoder=self.text_encoder,
|
||||
prompt=prompts,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
)
|
||||
other_conditions = make_contiguous(other_conditions)
|
||||
|
||||
if self.args.caption_dropout_technique == "zero":
|
||||
if random.random() < self.args.caption_dropout_p:
|
||||
other_conditions["prompt_embeds"].fill_(0)
|
||||
other_conditions["prompt_attention_mask"].fill_(False)
|
||||
|
||||
# These weighting schemes use a uniform timestep sampling and instead post-weight the loss
|
||||
weights = compute_density_for_timestep_sampling(
|
||||
weighting_scheme=self.args.flow_weighting_scheme,
|
||||
batch_size=batch_size,
|
||||
logit_mean=self.args.flow_logit_mean,
|
||||
logit_std=self.args.flow_logit_std,
|
||||
mode_scale=self.args.flow_mode_scale,
|
||||
)
|
||||
indices = (weights * self.scheduler.config.num_train_timesteps).long()
|
||||
sigmas = scheduler_sigmas[indices].flatten()
|
||||
|
||||
while sigmas.ndim < latent_conditions["latents"].ndim:
|
||||
sigmas = sigmas.unsqueeze(-1)
|
||||
|
||||
timesteps = (sigmas * 1000.0).long()
|
||||
|
||||
noise = torch.randn(
|
||||
latent_conditions["latents"].shape,
|
||||
generator=generator,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
)
|
||||
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
|
||||
|
||||
latent_conditions.update({"noisy_latents": noisy_latents})
|
||||
other_conditions.update({"timesteps": timesteps})
|
||||
|
||||
# These weighting schemes use a uniform timestep sampling and instead post-weight the loss
|
||||
weights = compute_loss_weighting_for_sd3(
|
||||
weighting_scheme=self.args.flow_weighting_scheme, sigmas=sigmas
|
||||
)
|
||||
pred = self.model_config["forward_pass"](
|
||||
transformer=self.transformer, **latent_conditions, **other_conditions
|
||||
)
|
||||
target = noise - latent_conditions["latents"]
|
||||
|
||||
loss = weights.float() * (pred["latents"].float() - target.float()).pow(2)
|
||||
# Average loss across channel dimension
|
||||
loss = loss.mean(list(range(1, loss.ndim)))
|
||||
# Average loss across batch dimension
|
||||
loss = loss.mean()
|
||||
accelerator.backward(loss)
|
||||
|
||||
if accelerator.sync_gradients and accelerator.distributed_type != DistributedType.DEEPSPEED:
|
||||
accelerator.clip_grad_norm_(self.transformer.parameters(), self.args.max_grad_norm)
|
||||
|
||||
self.optimizer.step()
|
||||
self.lr_scheduler.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Checks if the accelerator has performed an optimization step behind the scenes
|
||||
if accelerator.sync_gradients:
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
|
||||
# Checkpointing
|
||||
if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
|
||||
if global_step % self.args.checkpointing_steps == 0:
|
||||
# before saving state, check if this save would set us over the `checkpointing_limit`
|
||||
if self.args.checkpointing_limit is not None:
|
||||
checkpoints = find_files(self.args.output_dir, prefix="checkpoint")
|
||||
|
||||
# before we save the new checkpoint, we need to have at_most `checkpoints_total_limit - 1` checkpoints
|
||||
if len(checkpoints) >= self.args.checkpointing_limit:
|
||||
num_to_remove = len(checkpoints) - self.args.checkpointing_limit + 1
|
||||
checkpoints_to_remove = checkpoints[0:num_to_remove]
|
||||
delete_files(checkpoints_to_remove)
|
||||
|
||||
logger.info(f"Checkpointing at step {global_step}")
|
||||
save_path = os.path.join(self.args.output_dir, f"checkpoint-{global_step}")
|
||||
accelerator.save_state(save_path)
|
||||
logger.info(f"Saved state to {save_path}")
|
||||
|
||||
# Maybe run validation
|
||||
should_run_validation = (
|
||||
self.args.validation_every_n_steps is not None
|
||||
and global_step % self.args.validation_every_n_steps == 0
|
||||
)
|
||||
if should_run_validation:
|
||||
self.validate(global_step)
|
||||
|
||||
logs = {"loss": loss.detach().item(), "lr": self.lr_scheduler.get_last_lr()[0]}
|
||||
progress_bar.set_postfix(logs)
|
||||
accelerator.log(logs, step=global_step)
|
||||
|
||||
if global_step >= self.state.train_steps:
|
||||
break
|
||||
|
||||
memory_statistics = get_memory_statistics()
|
||||
logger.info(f"Memory after epoch {epoch + 1}: {json.dumps(memory_statistics, indent=4)}")
|
||||
|
||||
# Maybe run validation
|
||||
should_run_validation = (
|
||||
self.args.validation_every_n_epochs is not None
|
||||
and (epoch + 1) % self.args.validation_every_n_epochs == 0
|
||||
)
|
||||
if should_run_validation:
|
||||
self.validate(global_step)
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
if accelerator.is_main_process:
|
||||
self.transformer = unwrap_model(accelerator, self.transformer)
|
||||
dtype = (
|
||||
torch.float16
|
||||
if self.args.mixed_precision == "fp16"
|
||||
else torch.bfloat16
|
||||
if self.args.mixed_precision == "bf16"
|
||||
else torch.float32
|
||||
)
|
||||
self.transformer = self.transformer.to(dtype)
|
||||
transformer_lora_layers = get_peft_model_state_dict(self.transformer)
|
||||
|
||||
self.model_config["pipeline_cls"].save_lora_weights(
|
||||
save_directory=self.args.output_dir,
|
||||
transformer_lora_layers=transformer_lora_layers,
|
||||
)
|
||||
|
||||
del self.tokenizer, self.text_encoder, self.transformer, self.vae, self.scheduler
|
||||
free_memory()
|
||||
memory_statistics = get_memory_statistics()
|
||||
logger.info(f"Memory after training end: {json.dumps(memory_statistics, indent=4)}")
|
||||
|
||||
accelerator.end_training()
|
||||
|
||||
def validate(self, step: int) -> None:
|
||||
logger.info("Starting validation")
|
||||
|
||||
accelerator = self.state.accelerator
|
||||
num_validation_samples = len(self.args.validation_prompts)
|
||||
|
||||
if num_validation_samples == 0:
|
||||
logger.warning("No validation samples found. Skipping validation.")
|
||||
return
|
||||
|
||||
self.transformer.eval()
|
||||
|
||||
memory_statistics = get_memory_statistics()
|
||||
logger.info(f"Memory before validation start: {json.dumps(memory_statistics, indent=4)}")
|
||||
|
||||
pipeline = self.model_config["initialize_pipeline"](
|
||||
model_id=self.args.pretrained_model_name_or_path,
|
||||
cache_dir=self.args.cache_dir,
|
||||
tokenizer=self.tokenizer,
|
||||
text_encoder=self.text_encoder,
|
||||
transformer=unwrap_model(accelerator, self.transformer),
|
||||
vae=self.vae,
|
||||
device=accelerator.device,
|
||||
enable_slicing=self.args.enable_slicing,
|
||||
enable_tiling=self.args.enable_tiling,
|
||||
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
|
||||
)
|
||||
|
||||
all_processes_artifacts = []
|
||||
for i in range(num_validation_samples):
|
||||
# Skip current validation on all processes but one
|
||||
if i % accelerator.num_processes != accelerator.process_index:
|
||||
continue
|
||||
|
||||
prompt = self.args.validation_prompts[i]
|
||||
image = self.args.validation_images[i]
|
||||
video = self.args.validation_videos[i]
|
||||
height = self.args.validation_heights[i]
|
||||
width = self.args.validation_widths[i]
|
||||
num_frames = self.args.validation_num_frames[i]
|
||||
|
||||
if image is not None:
|
||||
image = load_image(image)
|
||||
if video is not None:
|
||||
video = load_video(video)
|
||||
|
||||
logger.debug(
|
||||
f"Validating sample {i + 1}/{num_validation_samples} on process {accelerator.process_index}. Prompt: {prompt}",
|
||||
main_process_only=False,
|
||||
)
|
||||
validation_artifacts = self.model_config["validation"](
|
||||
pipeline=pipeline,
|
||||
prompt=prompt,
|
||||
image=image,
|
||||
video=video,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_videos_per_prompt=self.args.num_validation_videos_per_prompt,
|
||||
generator=self.state.generator,
|
||||
)
|
||||
|
||||
prompt_filename = string_to_filename(prompt)[:25]
|
||||
artifacts = {
|
||||
"image": {"type": "image", "value": image},
|
||||
"video": {"type": "video", "value": video},
|
||||
}
|
||||
for i, (artifact_type, artifact_value) in enumerate(validation_artifacts):
|
||||
artifacts.update({f"artifact_{i}": {"type": artifact_type, "value": artifact_value}})
|
||||
logger.debug(
|
||||
f"Validation artifacts on process {accelerator.process_index}: {list(artifacts.keys())}",
|
||||
main_process_only=False,
|
||||
)
|
||||
|
||||
for key, value in list(artifacts.items()):
|
||||
artifact_type = value["type"]
|
||||
artifact_value = value["value"]
|
||||
if artifact_type not in ["image", "video"] or artifact_value is None:
|
||||
continue
|
||||
|
||||
extension = "png" if artifact_type == "image" else "mp4"
|
||||
filename = f"validation-{step}-{accelerator.process_index}-{prompt_filename}.{extension}"
|
||||
filename = os.path.join(self.args.output_dir, filename)
|
||||
|
||||
if artifact_type == "image":
|
||||
logger.debug(f"Saving image to {filename}")
|
||||
artifact_value.save(filename)
|
||||
artifact_value = wandb.Image(filename)
|
||||
elif artifact_type == "video":
|
||||
logger.debug(f"Saving video to {filename}")
|
||||
export_to_video(artifact_value, filename, fps=15)
|
||||
artifact_value = wandb.Video(filename, caption=prompt)
|
||||
|
||||
all_processes_artifacts.append(artifact_value)
|
||||
|
||||
all_artifacts = gather_object(all_processes_artifacts)
|
||||
|
||||
if accelerator.is_main_process:
|
||||
for tracker in accelerator.trackers:
|
||||
if tracker.name == "wandb":
|
||||
tracker.log({"validation": all_artifacts}, step=step)
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
free_memory()
|
||||
memory_statistics = get_memory_statistics()
|
||||
logger.info(f"Memory after validation end: {json.dumps(memory_statistics, indent=4)}")
|
||||
self.transformer.train()
|
||||
|
||||
def evaluate(self) -> None:
|
||||
logger.info("Starting evaluation")
|
||||
# TODO: implement metrics for evaluation
|
||||
|
||||
def _init_distributed(self) -> None:
|
||||
logging_dir = Path(self.args.output_dir, self.args.logging_dir)
|
||||
project_config = ProjectConfiguration(project_dir=self.args.output_dir, logging_dir=logging_dir)
|
||||
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||||
init_process_group_kwargs = InitProcessGroupKwargs(
|
||||
backend="nccl", timeout=timedelta(seconds=self.args.nccl_timeout)
|
||||
)
|
||||
mixed_precision = "no" if torch.backends.mps.is_available() else self.args.mixed_precision
|
||||
report_to = None if self.args.report_to.lower() == "none" else self.args.report_to
|
||||
|
||||
accelerator = Accelerator(
|
||||
project_config=project_config,
|
||||
gradient_accumulation_steps=self.args.gradient_accumulation_steps,
|
||||
mixed_precision=mixed_precision,
|
||||
log_with=report_to,
|
||||
kwargs_handlers=[ddp_kwargs, init_process_group_kwargs],
|
||||
)
|
||||
|
||||
# Disable AMP for MPS.
|
||||
if torch.backends.mps.is_available():
|
||||
accelerator.native_amp = False
|
||||
|
||||
self.state.accelerator = accelerator
|
||||
|
||||
if self.args.seed is not None:
|
||||
self.state.seed = self.args.seed
|
||||
set_seed(self.args.seed)
|
||||
|
||||
def _init_logging(self) -> None:
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=FINETRAINERS_LOG_LEVEL,
|
||||
)
|
||||
if self.state.accelerator.is_local_main_process:
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
|
||||
logger.info("Initialized FineTrainers")
|
||||
logger.info(self.state.accelerator.state, main_process_only=False)
|
||||
|
||||
def _init_directories_and_repositories(self) -> None:
|
||||
if self.state.accelerator.is_main_process:
|
||||
self.args.output_dir = Path(self.args.output_dir)
|
||||
self.args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.state.output_dir = self.args.output_dir
|
||||
|
||||
if self.args.push_to_hub:
|
||||
repo_id = self.args.hub_model_id or Path(self.args.output_dir).name
|
||||
self.state.repo_id = create_repo(token=self.args.hub_token, name=repo_id).repo_id
|
||||
|
||||
def _model_config_call(self, fn, kwargs):
|
||||
accepted_kwargs = inspect.signature(fn).parameters.keys()
|
||||
kwargs = {k: v for k, v in kwargs.items() if k in accepted_kwargs}
|
||||
return fn(**kwargs)
|
||||
@@ -0,0 +1,5 @@
|
||||
from .file_utils import find_files, delete_files
|
||||
from .diffusion_utils import resolution_dependant_timestep_flow_shift, default_flow_shift
|
||||
from .memory_utils import get_memory_statistics, bytes_to_gigabytes, free_memory, make_contiguous
|
||||
from .torch_utils import unwrap_model
|
||||
from .optimizer_utils import get_optimizer, gradient_norm, max_gradient
|
||||
@@ -0,0 +1,30 @@
|
||||
import torch
|
||||
|
||||
|
||||
# Default values copied from https://github.com/huggingface/diffusers/blob/8957324363d8b239d82db4909fbf8c0875683e3d/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py#L47
|
||||
def resolution_dependant_timestep_flow_shift(
|
||||
latents: torch.Tensor,
|
||||
sigmas: torch.Tensor,
|
||||
base_image_seq_len: int = 256,
|
||||
max_image_seq_len: int = 4096,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
) -> torch.Tensor:
|
||||
image_or_video_sequence_length = 0
|
||||
if latents.ndim == 4:
|
||||
image_or_video_sequence_length = latents.shape[2] * latents.shape[3]
|
||||
elif latents.ndim == 5:
|
||||
image_or_video_sequence_length = latents.shape[2] * latents.shape[3] * latents.shape[4]
|
||||
else:
|
||||
raise ValueError(f"Expected 4D or 5D tensor, got {latents.ndim}D tensor")
|
||||
|
||||
m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len)
|
||||
b = base_shift - m * base_image_seq_len
|
||||
mu = m * image_or_video_sequence_length + b
|
||||
sigmas = default_flow_shift(latents, sigmas, shift=mu)
|
||||
return sigmas
|
||||
|
||||
|
||||
def default_flow_shift(sigmas: torch.Tensor, shift: float = 1.0) -> torch.Tensor:
|
||||
sigmas = (sigmas * shift) / (1 + (shift - 1) * sigmas)
|
||||
return sigmas
|
||||
@@ -0,0 +1,44 @@
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
|
||||
logger = logging.getLogger("finetrainers")
|
||||
logger.setLevel(os.environ.get("FINETRAINERS_LOG_LEVEL", "INFO"))
|
||||
|
||||
|
||||
def find_files(dir: Union[str, Path], prefix: str = "checkpoint") -> List[str]:
|
||||
if not isinstance(dir, Path):
|
||||
dir = Path(dir)
|
||||
if not dir.exists():
|
||||
return []
|
||||
checkpoints = os.listdir(dir.as_posix())
|
||||
checkpoints = [c for c in checkpoints if c.startswith(prefix)]
|
||||
checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
|
||||
return checkpoints
|
||||
|
||||
|
||||
def delete_files(dirs: Union[str, List[str], Path, List[Path]]) -> None:
|
||||
if not isinstance(dirs, list):
|
||||
dirs = [dirs]
|
||||
dirs = [Path(d) if isinstance(d, str) else d for d in dirs]
|
||||
logger.info(f"Deleting files: {dirs}")
|
||||
for dir in dirs:
|
||||
if not dir.exists():
|
||||
continue
|
||||
shutil.rmtree(dir, ignore_errors=True)
|
||||
|
||||
|
||||
def string_to_filename(s: str) -> str:
|
||||
return (
|
||||
s.replace(" ", "-")
|
||||
.replace("/", "-")
|
||||
.replace(":", "-")
|
||||
.replace(".", "-")
|
||||
.replace(",", "-")
|
||||
.replace(";", "-")
|
||||
.replace("!", "-")
|
||||
.replace("?", "-")
|
||||
)
|
||||
@@ -0,0 +1,59 @@
|
||||
import gc
|
||||
import logging
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
import torch
|
||||
from accelerate.logging import get_logger
|
||||
|
||||
|
||||
logger = get_logger("finetrainers")
|
||||
|
||||
|
||||
def get_memory_statistics(precision: int = 3) -> Dict[str, Any]:
|
||||
memory_allocated = None
|
||||
memory_reserved = None
|
||||
max_memory_allocated = None
|
||||
max_memory_reserved = None
|
||||
|
||||
if torch.cuda.is_available():
|
||||
device = torch.cuda.current_device()
|
||||
memory_allocated = torch.cuda.memory_allocated(device)
|
||||
memory_reserved = torch.cuda.memory_reserved(device)
|
||||
max_memory_allocated = torch.cuda.max_memory_allocated(device)
|
||||
max_memory_reserved = torch.cuda.max_memory_reserved(device)
|
||||
|
||||
elif torch.mps.is_available():
|
||||
memory_allocated = torch.mps.current_allocated_memory()
|
||||
|
||||
else:
|
||||
logger.warning("No CUDA, MPS, or ROCm device found. Memory statistics are not available.")
|
||||
|
||||
return {
|
||||
"memory_allocated": round(bytes_to_gigabytes(memory_allocated), ndigits=precision),
|
||||
"memory_reserved": round(bytes_to_gigabytes(memory_reserved), ndigits=precision),
|
||||
"max_memory_allocated": round(bytes_to_gigabytes(max_memory_allocated), ndigits=precision),
|
||||
"max_memory_reserved": round(bytes_to_gigabytes(max_memory_reserved), ndigits=precision),
|
||||
}
|
||||
|
||||
|
||||
def bytes_to_gigabytes(x: int) -> float:
|
||||
if x is not None:
|
||||
return x / 1024**3
|
||||
|
||||
|
||||
def free_memory() -> None:
|
||||
if torch.cuda.is_available():
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
# TODO(aryan): handle non-cuda devices
|
||||
|
||||
|
||||
def make_contiguous(x: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
if isinstance(x, torch.Tensor):
|
||||
return x.contiguous()
|
||||
elif isinstance(x, dict):
|
||||
return {k: make_contiguous(v) for k, v in x.items()}
|
||||
else:
|
||||
return x
|
||||
@@ -0,0 +1,178 @@
|
||||
import inspect
|
||||
import logging
|
||||
|
||||
from accelerate.logging import get_logger
|
||||
import torch
|
||||
|
||||
logger = get_logger("finetrainers")
|
||||
|
||||
|
||||
def get_optimizer(
|
||||
params_to_optimize,
|
||||
optimizer_name: str = "adam",
|
||||
learning_rate: float = 1e-3,
|
||||
beta1: float = 0.9,
|
||||
beta2: float = 0.95,
|
||||
beta3: float = 0.98,
|
||||
epsilon: float = 1e-8,
|
||||
weight_decay: float = 1e-4,
|
||||
prodigy_decouple: bool = False,
|
||||
prodigy_use_bias_correction: bool = False,
|
||||
prodigy_safeguard_warmup: bool = False,
|
||||
use_8bit: bool = False,
|
||||
use_4bit: bool = False,
|
||||
use_torchao: bool = False,
|
||||
use_deepspeed: bool = False,
|
||||
use_cpu_offload_optimizer: bool = False,
|
||||
offload_gradients: bool = False,
|
||||
) -> torch.optim.Optimizer:
|
||||
optimizer_name = optimizer_name.lower()
|
||||
|
||||
# Use DeepSpeed optimzer
|
||||
if use_deepspeed:
|
||||
from accelerate.utils import DummyOptim
|
||||
|
||||
return DummyOptim(
|
||||
params_to_optimize,
|
||||
lr=learning_rate,
|
||||
betas=(beta1, beta2),
|
||||
eps=epsilon,
|
||||
weight_decay=weight_decay,
|
||||
)
|
||||
|
||||
if use_8bit and use_4bit:
|
||||
raise ValueError("Cannot set both `use_8bit` and `use_4bit` to True.")
|
||||
|
||||
if (use_torchao and (use_8bit or use_4bit)) or use_cpu_offload_optimizer:
|
||||
try:
|
||||
import torchao
|
||||
|
||||
torchao.__version__
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"To use optimizers from torchao, please install the torchao library: `USE_CPP=0 pip install torchao`."
|
||||
)
|
||||
|
||||
if not use_torchao and use_4bit:
|
||||
raise ValueError("4-bit Optimizers are only supported with torchao.")
|
||||
|
||||
# Optimizer creation
|
||||
supported_optimizers = ["adam", "adamw", "prodigy", "came"]
|
||||
if optimizer_name not in supported_optimizers:
|
||||
logger.warning(
|
||||
f"Unsupported choice of optimizer: {optimizer_name}. Supported optimizers include {supported_optimizers}. Defaulting to `AdamW`."
|
||||
)
|
||||
optimizer_name = "adamw"
|
||||
|
||||
if (use_8bit or use_4bit) and optimizer_name not in ["adam", "adamw"]:
|
||||
raise ValueError("`use_8bit` and `use_4bit` can only be used with the Adam and AdamW optimizers.")
|
||||
|
||||
if use_8bit:
|
||||
try:
|
||||
import bitsandbytes as bnb
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
|
||||
)
|
||||
|
||||
if optimizer_name == "adamw":
|
||||
if use_torchao:
|
||||
from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit
|
||||
|
||||
optimizer_class = AdamW8bit if use_8bit else AdamW4bit if use_4bit else torch.optim.AdamW
|
||||
else:
|
||||
optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
|
||||
|
||||
init_kwargs = {
|
||||
"betas": (beta1, beta2),
|
||||
"eps": epsilon,
|
||||
"weight_decay": weight_decay,
|
||||
}
|
||||
|
||||
elif optimizer_name == "adam":
|
||||
if use_torchao:
|
||||
from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit
|
||||
|
||||
optimizer_class = Adam8bit if use_8bit else Adam4bit if use_4bit else torch.optim.Adam
|
||||
else:
|
||||
optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
|
||||
|
||||
init_kwargs = {
|
||||
"betas": (beta1, beta2),
|
||||
"eps": epsilon,
|
||||
"weight_decay": weight_decay,
|
||||
}
|
||||
|
||||
elif optimizer_name == "prodigy":
|
||||
try:
|
||||
import prodigyopt
|
||||
except ImportError:
|
||||
raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
|
||||
|
||||
optimizer_class = prodigyopt.Prodigy
|
||||
|
||||
if learning_rate <= 0.1:
|
||||
logger.warning(
|
||||
"Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
|
||||
)
|
||||
|
||||
init_kwargs = {
|
||||
"lr": learning_rate,
|
||||
"betas": (beta1, beta2),
|
||||
"beta3": beta3,
|
||||
"eps": epsilon,
|
||||
"weight_decay": weight_decay,
|
||||
"decouple": prodigy_decouple,
|
||||
"use_bias_correction": prodigy_use_bias_correction,
|
||||
"safeguard_warmup": prodigy_safeguard_warmup,
|
||||
}
|
||||
|
||||
elif optimizer_name == "came":
|
||||
try:
|
||||
import came_pytorch
|
||||
except ImportError:
|
||||
raise ImportError("To use CAME, please install the came-pytorch library: `pip install came-pytorch`")
|
||||
|
||||
optimizer_class = came_pytorch.CAME
|
||||
|
||||
init_kwargs = {
|
||||
"lr": learning_rate,
|
||||
"eps": (1e-30, 1e-16),
|
||||
"betas": (beta1, beta2, beta3),
|
||||
"weight_decay": weight_decay,
|
||||
}
|
||||
|
||||
if use_cpu_offload_optimizer:
|
||||
from torchao.prototype.low_bit_optim import CPUOffloadOptimizer
|
||||
|
||||
if "fused" in inspect.signature(optimizer_class.__init__).parameters:
|
||||
init_kwargs.update({"fused": True})
|
||||
|
||||
optimizer = CPUOffloadOptimizer(
|
||||
params_to_optimize, optimizer_class=optimizer_class, offload_gradients=offload_gradients, **init_kwargs
|
||||
)
|
||||
else:
|
||||
optimizer = optimizer_class(params_to_optimize, **init_kwargs)
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
def gradient_norm(parameters):
|
||||
norm = 0
|
||||
for param in parameters:
|
||||
if param.grad is None:
|
||||
continue
|
||||
local_norm = param.grad.detach().data.norm(2)
|
||||
norm += local_norm.item() ** 2
|
||||
norm = norm**0.5
|
||||
return norm
|
||||
|
||||
|
||||
def max_gradient(parameters):
|
||||
max_grad_value = float("-inf")
|
||||
for param in parameters:
|
||||
if param.grad is None:
|
||||
continue
|
||||
local_max_grad = param.grad.detach().data.abs().max()
|
||||
max_grad_value = max(max_grad_value, local_max_grad.item())
|
||||
return max_grad_value
|
||||
@@ -0,0 +1,8 @@
|
||||
from accelerate import Accelerator
|
||||
from diffusers.utils.torch_utils import is_compiled_module
|
||||
|
||||
|
||||
def unwrap_model(accelerator: Accelerator, model):
|
||||
model = accelerator.unwrap_model(model)
|
||||
model = model._orig_mod if is_compiled_module(model) else model
|
||||
return model
|
||||
Reference in New Issue
Block a user