finetrainers/args.py

import argparse
import sys
from typing import Any, Dict, List, Optional, Tuple

import torch

from .constants import DEFAULT_IMAGE_RESOLUTION_BUCKETS, DEFAULT_VIDEO_RESOLUTION_BUCKETS
from .models import SUPPORTED_MODEL_CONFIGS


class Args:
    r"""
    The arguments for the finetrainers training script.

    For helpful information about arguments, run `python train.py --help`.

    TODO(aryan): add `python train.py --recommend_configs --model_name <model_name>` to recommend
    good training configs for a model after extensive testing.
    TODO(aryan): add `python train.py --memory_requirements --model_name <model_name>` to show
    memory requirements per model, per training type with sensible training settings.

    MODEL ARGUMENTS
    ---------------
    model_name (`str`):
        Name of model to train. To get a list of models, run `python train.py --list_models`.
    pretrained_model_name_or_path (`str`):
        Path to pretrained model or model identifier from https://huggingface.co/models. The model should be
        loadable based on specified `model_name`.
    revision (`str`, defaults to `None`):
        If provided, the model will be loaded from a specific branch of the model repository.
    variant (`str`, defaults to `None`):
        Variant of model weights to use. Some models provide weight variants, such as `fp16`, to reduce disk
        storage requirements.
    cache_dir (`str`, defaults to `None`):
        The directory where the downloaded models and datasets will be stored, or loaded from.
    text_encoder_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
        Data type for the text encoder when generating text embeddings.
    text_encoder_2_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
        Data type for the text encoder 2 when generating text embeddings.
    text_encoder_3_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
        Data type for the text encoder 3 when generating text embeddings.
    transformer_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
        Data type for the transformer model.
    vae_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
        Data type for the VAE model.
    layerwise_upcasting_modules (`List[str]`, defaults to `[]`):
        Modules that should have fp8 storage weights but higher precision computation. Choose between ['transformer'].
    layerwise_upcasting_storage_dtype (`torch.dtype`, defaults to `float8_e4m3fn`):
        Data type for the layerwise upcasting storage. Choose between ['float8_e4m3fn', 'float8_e5m2'].
    layerwise_upcasting_skip_modules_pattern (`List[str]`, defaults to `["patch_embed", "pos_embed", "x_embedder", "context_embedder", "^proj_in$", "^proj_out$", "norm"]`):
        Modules to skip for layerwise upcasting. Layers such as normalization and modulation, when casted to fp8 precision
        naively (as done in layerwise upcasting), can lead to poorer training and inference quality. We skip these layers
        by default, and recommend adding more layers to the default list based on the model architecture.

    DATASET ARGUMENTS
    -----------------
    data_root (`str`):
        A folder containing the training data.
    dataset_file (`str`, defaults to `None`):
        Path to a CSV/JSON/JSONL file containing metadata for training. This should be provided if you're not using
        a directory dataset format containing a simple `prompts.txt` and `videos.txt`/`images.txt` for example.
    video_column (`str`):
        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.
    caption_column (`str`):
        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.
    id_token (`str`, defaults to `None`):
        Identifier token appended to the start of each prompt if provided. This is useful for LoRA-type training.
    image_resolution_buckets (`List[Tuple[int, int]]`, defaults to `None`):
        Resolution buckets for images. This should be a list of integer tuples, where each tuple represents the
        resolution (height, width) of the image. All images will be resized to the nearest bucket resolution.
    video_resolution_buckets (`List[Tuple[int, int, int]]`, defaults to `None`):
        Resolution buckets for videos. This should be a list of integer tuples, where each tuple represents the
        resolution (num_frames, height, width) of the video. All videos will be resized to the nearest bucket
        resolution.
    video_reshape_mode (`str`, defaults to `None`):
        All input videos are reshaped to this mode. Choose between ['center', 'random', 'none'].
        TODO(aryan): We don't support this.
    caption_dropout_p (`float`, defaults to `0.00`):
        Probability of dropout for the caption tokens. This is useful to improve the unconditional generation
        quality of the model.
    caption_dropout_technique (`str`, defaults to `empty`):
        Technique to use for caption dropout. Choose between ['empty', 'zero']. Some models apply caption dropout
        by setting the prompt condition to an empty string, while others zero-out the text embedding tensors.
    precompute_conditions (`bool`, defaults to `False`):
        Whether or not to precompute the conditionings for the model. This is useful for faster training, and
        reduces the memory requirements.
    remove_common_llm_caption_prefixes (`bool`, defaults to `False`):
        Whether or not to remove common LLM caption prefixes. This is useful for improving the quality of the
        generated text.

    DATALOADER_ARGUMENTS
    --------------------
    See https://pytorch.org/docs/stable/data.html for more information.

    dataloader_num_workers (`int`, defaults to `0`):
        Number of subprocesses to use for data loading. `0` means that the data will be loaded in a blocking manner
        on the main process.
    pin_memory (`bool`, defaults to `False`):
        Whether or not to use the pinned memory setting in PyTorch dataloader. This is useful for faster data loading.

    DIFFUSION ARGUMENTS
    -------------------
    flow_resolution_shifting (`bool`, defaults to `False`):
        Resolution-dependent shifting of timestep schedules.
        [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206).
        TODO(aryan): We don't support this yet.
    flow_base_seq_len (`int`, defaults to `256`):
        Base number of tokens for images/video when applying resolution-dependent shifting.
    flow_max_seq_len (`int`, defaults to `4096`):
        Maximum number of tokens for images/video when applying resolution-dependent shifting.
    flow_base_shift (`float`, defaults to `0.5`):
        Base shift for timestep schedules when applying resolution-dependent shifting.
    flow_max_shift (`float`, defaults to `1.15`):
        Maximum shift for timestep schedules when applying resolution-dependent shifting.
    flow_shift (`float`, defaults to `1.0`):
        Instead of training with uniform/logit-normal sigmas, shift them as (shift * sigma) / (1 + (shift - 1) * sigma).
        Setting it higher is helpful when trying to train models for high-resolution generation or to produce better
        samples in lower number of inference steps.
    flow_weighting_scheme (`str`, defaults to `none`):
        We default to the "none" weighting scheme for uniform sampling and uniform loss.
        Choose between ['sigma_sqrt', 'logit_normal', 'mode', 'cosmap', 'none'].
    flow_logit_mean (`float`, defaults to `0.0`):
        Mean to use when using the `'logit_normal'` weighting scheme.
    flow_logit_std (`float`, defaults to `1.0`):
        Standard deviation to use when using the `'logit_normal'` weighting scheme.
    flow_mode_scale (`float`, defaults to `1.29`):
        Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.

    TRAINING ARGUMENTS
    ------------------
    training_type (`str`, defaults to `None`):
        Type of training to perform. Choose between ['lora'].
    seed (`int`, defaults to `42`):
        A seed for reproducible training.
    batch_size (`int`, defaults to `1`):
        Per-device batch size.
    train_epochs (`int`, defaults to `1`):
        Number of training epochs.
    train_steps (`int`, defaults to `None`):
        Total number of training steps to perform. If provided, overrides `train_epochs`.
    rank (`int`, defaults to `128`):
        The rank for LoRA matrices.
    lora_alpha (`float`, defaults to `64`):
        The lora_alpha to compute scaling factor (lora_alpha / rank) for LoRA matrices.
    target_modules (`List[str]`, defaults to `["to_k", "to_q", "to_v", "to_out.0"]`):
        The target modules for LoRA. Make sure to modify this based on the model.
    gradient_accumulation_steps (`int`, defaults to `1`):
        Number of gradients steps to accumulate before performing an optimizer step.
    gradient_checkpointing (`bool`, defaults to `False`):
        Whether or not to use gradient/activation checkpointing to save memory at the expense of slower
        backward pass.
    checkpointing_steps (`int`, defaults to `500`):
        Save a checkpoint of the training state every X training steps. 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`.
    checkpointing_limit (`int`, defaults to `None`):
        Max number of checkpoints to store.
    resume_from_checkpoint (`str`, defaults to `None`):
        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.

    OPTIMIZER ARGUMENTS
    -------------------
    optimizer (`str`, defaults to `adamw`):
        The optimizer type to use. Choose between ['adam', 'adamw'].
    use_8bit_bnb (`bool`, defaults to `False`):
        Whether to use 8bit variant of the `optimizer` using `bitsandbytes`.
    lr (`float`, defaults to `1e-4`):
        Initial learning rate (after the potential warmup period) to use.
    scale_lr (`bool`, defaults to `False`):
        Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.
    lr_scheduler (`str`, defaults to `cosine_with_restarts`):
        The scheduler type to use. Choose between ['linear', 'cosine', 'cosine_with_restarts', 'polynomial',
        'constant', 'constant_with_warmup'].
    lr_warmup_steps (`int`, defaults to `500`):
        Number of steps for the warmup in the lr scheduler.
    lr_num_cycles (`int`, defaults to `1`):
        Number of hard resets of the lr in cosine_with_restarts scheduler.
    lr_power (`float`, defaults to `1.0`):
        Power factor of the polynomial scheduler.
    beta1 (`float`, defaults to `0.9`):
    beta2 (`float`, defaults to `0.95`):
    beta3 (`float`, defaults to `0.999`):
    weight_decay (`float`, defaults to `0.0001`):
        Penalty for large weights in the model.
    epsilon (`float`, defaults to `1e-8`):
        Small value to avoid division by zero in the optimizer.
    max_grad_norm (`float`, defaults to `1.0`):
        Maximum gradient norm to clip the gradients.

    VALIDATION ARGUMENTS
    --------------------
    validation_prompts (`List[str]`, defaults to `None`):
        List of prompts to use for validation. If not provided, a random prompt will be selected from the training
        dataset.
    validation_images (`List[str]`, defaults to `None`):
        List of image paths to use for validation.
    validation_videos (`List[str]`, defaults to `None`):
        List of video paths to use for validation.
    validation_heights (`List[int]`, defaults to `None`):
        List of heights for the validation videos.
    validation_widths (`List[int]`, defaults to `None`):
        List of widths for the validation videos.
    validation_num_frames (`List[int]`, defaults to `None`):
        List of number of frames for the validation videos.
    num_validation_videos_per_prompt (`int`, defaults to `1`):
        Number of videos to use for validation per prompt.
    validation_every_n_epochs (`int`, defaults to `None`):
        Perform validation every `n` training epochs.
    validation_every_n_steps (`int`, defaults to `None`):
        Perform validation every `n` training steps.
    enable_model_cpu_offload (`bool`, defaults to `False`):
        Whether or not to offload different modeling components to CPU during validation.
    validation_frame_rate (`int`, defaults to `25`):
        Frame rate to use for the validation videos. This value is defaulted to 25, as used in LTX Video pipeline.

    MISCELLANEOUS ARGUMENTS
    -----------------------
    tracker_name (`str`, defaults to `finetrainers`):
        Name of the tracker/project to use for logging training metrics.
    push_to_hub (`bool`, defaults to `False`):
        Whether or not to push the model to the Hugging Face Hub.
    hub_token (`str`, defaults to `None`):
        The API token to use for pushing the model to the Hugging Face Hub.
    hub_model_id (`str`, defaults to `None`):
        The model identifier to use for pushing the model to the Hugging Face Hub.
    output_dir (`str`, defaults to `None`):
        The directory where the model checkpoints and logs will be stored.
    logging_dir (`str`, defaults to `logs`):
        The directory where the logs will be stored.
    allow_tf32 (`bool`, defaults to `False`):
        Whether or not to allow the use of TF32 matmul on compatible hardware.
    nccl_timeout (`int`, defaults to `1800`):
        Timeout for the NCCL communication.
    report_to (`str`, defaults to `wandb`):
        The name of the logger to use for logging training metrics. Choose between ['wandb'].
    """

    # 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
    text_encoder_dtype: torch.dtype = torch.bfloat16
    text_encoder_2_dtype: torch.dtype = torch.bfloat16
    text_encoder_3_dtype: torch.dtype = torch.bfloat16
    transformer_dtype: torch.dtype = torch.bfloat16
    vae_dtype: torch.dtype = torch.bfloat16
    layerwise_upcasting_modules: List[str] = []
    layerwise_upcasting_storage_dtype: torch.dtype = torch.float8_e4m3fn
    layerwise_upcasting_skip_modules_pattern: List[str] = [
        "patch_embed",
        "pos_embed",
        "x_embedder",
        "context_embedder",
        "time_embed",
        "^proj_in$",
        "^proj_out$",
        "norm",
    ]

    # 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"
    precompute_conditions: bool = False
    remove_common_llm_caption_prefixes: bool = False

    # Dataloader arguments
    dataloader_num_workers: int = 0
    pin_memory: bool = False

    # Diffusion arguments
    flow_resolution_shifting: bool = False
    flow_base_seq_len: int = 256
    flow_max_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
    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"
    use_8bit_bnb: bool = False
    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
    validation_frame_rate: int = 25

    # 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,
                "text_encoder_dtype": self.text_encoder_dtype,
                "text_encoder_2_dtype": self.text_encoder_2_dtype,
                "text_encoder_3_dtype": self.text_encoder_3_dtype,
                "transformer_dtype": self.transformer_dtype,
                "vae_dtype": self.vae_dtype,
                "layerwise_upcasting_modules": self.layerwise_upcasting_modules,
                "layerwise_upcasting_storage_dtype": self.layerwise_upcasting_storage_dtype,
                "layerwise_upcasting_skip_modules_pattern": self.layerwise_upcasting_skip_modules_pattern,
            },
            "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,
                "caption_dropout_technique": self.caption_dropout_technique,
                "precompute_conditions": self.precompute_conditions,
                "remove_common_llm_caption_prefixes": self.remove_common_llm_caption_prefixes,
            },
            "dataloader_arguments": {
                "dataloader_num_workers": self.dataloader_num_workers,
                "pin_memory": self.pin_memory,
            },
            "diffusion_arguments": {
                "flow_resolution_shifting": self.flow_resolution_shifting,
                "flow_base_seq_len": self.flow_base_seq_len,
                "flow_max_seq_len": self.flow_max_seq_len,
                "flow_base_shift": self.flow_base_shift,
                "flow_max_shift": self.flow_max_shift,
                "flow_shift": self.flow_shift,
                "flow_weighting_scheme": self.flow_weighting_scheme,
                "flow_logit_mean": self.flow_logit_mean,
                "flow_logit_std": self.flow_logit_std,
                "flow_mode_scale": self.flow_mode_scale,
            },
            "training_arguments": {
                "training_type": self.training_type,
                "seed": self.seed,
                "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,
                "use_8bit_bnb": self.use_8bit_bnb,
                "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,
                "validation_frame_rate": self.validation_frame_rate,
            },
            "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,
            },
        }


# TODO(aryan): handle more informative messages
_IS_ARGUMENTS_REQUIRED = "--list_models" not in sys.argv


def parse_arguments() -> Args:
    parser = argparse.ArgumentParser()

    if _IS_ARGUMENTS_REQUIRED:
        _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)
    else:
        _add_helper_arguments(parser)

        args = parser.parse_args()
        _display_helper_messages(args)
        sys.exit(0)


def validate_args(args: Args):
    _validated_model_args(args)
    _validate_training_args(args)
    _validate_validation_args(args)


def _add_model_arguments(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--model_name",
        type=str,
        required=True,
        choices=list(SUPPORTED_MODEL_CONFIGS.keys()),
        help="Name of model to train.",
    )
    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        required=True,
        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.",
    )
    parser.add_argument("--text_encoder_dtype", type=str, default="bf16", help="Data type for the text encoder.")
    parser.add_argument("--text_encoder_2_dtype", type=str, default="bf16", help="Data type for the text encoder 2.")
    parser.add_argument("--text_encoder_3_dtype", type=str, default="bf16", help="Data type for the text encoder 3.")
    parser.add_argument("--transformer_dtype", type=str, default="bf16", help="Data type for the transformer model.")
    parser.add_argument("--vae_dtype", type=str, default="bf16", help="Data type for the VAE model.")
    parser.add_argument(
        "--layerwise_upcasting_modules",
        type=str,
        default=[],
        nargs="+",
        choices=["transformer"],
        help="Modules that should have fp8 storage weights but higher precision computation.",
    )
    parser.add_argument(
        "--layerwise_upcasting_storage_dtype",
        type=str,
        default="float8_e4m3fn",
        choices=["float8_e4m3fn", "float8_e5m2"],
        help="Data type for the layerwise upcasting storage.",
    )
    parser.add_argument(
        "--layerwise_upcasting_skip_modules_pattern",
        type=str,
        default=["patch_embed", "pos_embed", "x_embedder", "context_embedder", "^proj_in$", "^proj_out$", "norm"],
        nargs="+",
        help="Modules to skip for layerwise upcasting.",
    )


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,
        required=True,
        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.",
    )
    parser.add_argument(
        "--precompute_conditions",
        action="store_true",
        help="Whether or not to precompute the conditionings for the model.",
    )
    parser.add_argument(
        "--remove_common_llm_caption_prefixes",
        action="store_true",
        help="Whether or not to remove common LLM caption prefixes.",
    )


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-dependent shifting of timestep schedules.",
    )
    parser.add_argument(
        "--flow_base_seq_len",
        type=int,
        default=256,
        help="Base image/video sequence length for the diffusion model.",
    )
    parser.add_argument(
        "--flow_max_seq_len",
        type=int,
        default=4096,
        help="Maximum image/video sequence length for the diffusion model.",
    )
    parser.add_argument(
        "--flow_base_shift",
        type=float,
        default=0.5,
        help="Base shift as described in [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206)",
    )
    parser.add_argument(
        "--flow_max_shift",
        type=float,
        default=1.15,
        help="Maximum shift as described in [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206)",
    )
    parser.add_argument(
        "--flow_shift",
        type=float,
        default=1.0,
        help="Shift value to use for the flow matching timestep schedule.",
    )
    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="Standard deviation 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,
        choices=["lora", "full-finetune"],
        required=True,
        help="Type of training to perform. Choose between ['lora', 'full-finetune']",
    )
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument(
        "--batch_size",
        type=int,
        default=1,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument("--train_epochs", type=int, default=1, help="Number of training epochs.")
    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",
        help="Whether or not to use VAE slicing for saving memory.",
    )
    parser.add_argument(
        "--enable_tiling",
        action="store_true",
        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",
        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(
        "--use_8bit_bnb",
        action="store_true",
        help=("Whether to use 8bit variant of the `--optimizer` using `bitsandbytes`."),
    )
    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(
        "--validation_frame_rate",
        type=int,
        default=25,
        help="Frame rate to use for the validation videos.",
    )
    parser.add_argument(
        "--enable_model_cpu_offload",
        action="store_true",
        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="finetrainers-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",
        choices=["none", "wandb"],
        help="The integration to report the results and logs to.",
    )


def _add_helper_arguments(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--list_models",
        action="store_true",
        help="List all the supported models.",
    )


_DTYPE_MAP = {
    "bf16": torch.bfloat16,
    "fp16": torch.float16,
    "fp32": torch.float32,
    "float8_e4m3fn": torch.float8_e4m3fn,
    "float8_e5m2": torch.float8_e5m2,
}


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
    result_args.text_encoder_dtype = _DTYPE_MAP[args.text_encoder_dtype]
    result_args.text_encoder_2_dtype = _DTYPE_MAP[args.text_encoder_2_dtype]
    result_args.text_encoder_3_dtype = _DTYPE_MAP[args.text_encoder_3_dtype]
    result_args.transformer_dtype = _DTYPE_MAP[args.transformer_dtype]
    result_args.vae_dtype = _DTYPE_MAP[args.vae_dtype]
    result_args.layerwise_upcasting_modules = args.layerwise_upcasting_modules
    result_args.layerwise_upcasting_storage_dtype = _DTYPE_MAP[args.layerwise_upcasting_storage_dtype]
    result_args.layerwise_upcasting_skip_modules_pattern = args.layerwise_upcasting_skip_modules_pattern

    # 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
    result_args.caption_dropout_technique = args.caption_dropout_technique
    result_args.precompute_conditions = args.precompute_conditions
    result_args.remove_common_llm_caption_prefixes = args.remove_common_llm_caption_prefixes

    # 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_base_seq_len = args.flow_base_seq_len
    result_args.flow_max_seq_len = args.flow_max_seq_len
    result_args.flow_base_shift = args.flow_base_shift
    result_args.flow_max_shift = args.flow_max_shift
    result_args.flow_shift = args.flow_shift
    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.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.target_modules = args.target_modules
    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.use_8bit_bnb = args.use_8bit_bnb
    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
    result_args.validation_frame_rate = args.validation_frame_rate

    # 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 _validated_model_args(args: Args):
    if args.training_type == "full-finetune":
        assert (
            "transformer" not in args.layerwise_upcasting_modules
        ), "Layerwise upcasting is not supported for full-finetune training"


def _validate_training_args(args: Args):
    if args.training_type == "lora":
        assert args.rank is not None, "Rank is required for LoRA training"
        assert args.lora_alpha is not None, "LoRA alpha is required for LoRA training"
        assert (
            args.target_modules is not None and len(args.target_modules) > 0
        ), "Target modules are required for LoRA training"


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"


def _display_helper_messages(args: argparse.Namespace):
    if args.list_models:
        print("Supported models:")
        for index, model_name in enumerate(SUPPORTED_MODEL_CONFIGS.keys()):
            print(f"  {index + 1}. {model_name}")