835 lines
30 KiB
Python
835 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
|
|
"""The arguments of FastVideo Inference."""
|
|
|
|
import argparse
|
|
import dataclasses
|
|
from contextlib import contextmanager
|
|
from dataclasses import field
|
|
from typing import Any, Callable, List, Optional, Tuple
|
|
|
|
from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
|
|
from fastvideo.v1.logger import init_logger
|
|
from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def preprocess_text(prompt: str) -> str:
|
|
return prompt
|
|
|
|
|
|
def postprocess_text(output: Any) -> Any:
|
|
raise NotImplementedError
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class FastVideoArgs:
|
|
# Model and path configuration
|
|
model_path: str
|
|
|
|
# Cache strategy
|
|
cache_strategy: str = "none"
|
|
|
|
# Distributed executor backend
|
|
distributed_executor_backend: str = "mp"
|
|
|
|
inference_mode: bool = True # if False == training mode
|
|
|
|
# HuggingFace specific parameters
|
|
trust_remote_code: bool = False
|
|
revision: Optional[str] = None
|
|
|
|
# Parallelism
|
|
num_gpus: int = 1
|
|
tp_size: int = -1
|
|
sp_size: int = -1
|
|
hsdp_replicate_dim: int = 1
|
|
hsdp_shard_dim: int = -1
|
|
dist_timeout: Optional[int] = None # timeout for torch.distributed
|
|
|
|
# Video generation parameters
|
|
embedded_cfg_scale: float = 6.0
|
|
flow_shift: Optional[float] = None
|
|
|
|
output_type: str = "pil"
|
|
|
|
# DiT configuration
|
|
dit_config: DiTConfig = field(default_factory=DiTConfig)
|
|
precision: str = "bf16"
|
|
use_cpu_offload: bool = True
|
|
use_fsdp_inference: bool = True
|
|
|
|
# VAE configuration
|
|
vae_precision: str = "fp16"
|
|
vae_tiling: bool = True # Might change in between forward passes
|
|
vae_sp: bool = False # Might change in between forward passes
|
|
# vae_scale_factor: Optional[int] = None # Deprecated
|
|
vae_config: VAEConfig = field(default_factory=VAEConfig)
|
|
|
|
# Image encoder configuration
|
|
image_encoder_precision: str = "fp32"
|
|
image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
|
|
|
|
# Text encoder configuration
|
|
DEFAULT_TEXT_ENCODER_PRECISIONS = (
|
|
"fp16",
|
|
# "fp16",
|
|
)
|
|
text_encoder_precisions: Tuple[str, ...] = field(
|
|
default_factory=lambda: FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS)
|
|
text_encoder_configs: Tuple[EncoderConfig, ...] = field(
|
|
default_factory=lambda: (EncoderConfig(), ))
|
|
preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
|
|
default_factory=lambda: (preprocess_text, ))
|
|
postprocess_text_funcs: Tuple[Callable[[Any], Any], ...] = field(
|
|
default_factory=lambda: (postprocess_text, ))
|
|
|
|
# STA parameters
|
|
STA_mode: Optional[str] = None
|
|
skip_time_steps: int = 15
|
|
# LoRA parameters
|
|
lora_path: Optional[str] = None
|
|
lora_nickname: Optional[
|
|
str] = "default" # for swapping adapters in the pipeline
|
|
lora_target_names: Optional[List[
|
|
str]] = None # can restrict list of layers to adapt, e.g. ["q_proj"]
|
|
|
|
# STA parameters
|
|
mask_strategy_file_path: Optional[str] = None
|
|
enable_torch_compile: bool = False
|
|
|
|
disable_autocast: bool = False
|
|
|
|
# StepVideo specific parameters
|
|
pos_magic: Optional[str] = None
|
|
neg_magic: Optional[str] = None
|
|
timesteps_scale: Optional[bool] = None
|
|
|
|
# Logging
|
|
log_level: str = "info"
|
|
|
|
@property
|
|
def training_mode(self) -> bool:
|
|
return not self.inference_mode
|
|
|
|
def __post_init__(self):
|
|
self.check_fastvideo_args()
|
|
|
|
@staticmethod
|
|
def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
|
|
# Model and path configuration
|
|
parser.add_argument(
|
|
"--model-path",
|
|
type=str,
|
|
help=
|
|
"The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
|
|
)
|
|
parser.add_argument(
|
|
"--dit-weight",
|
|
type=str,
|
|
help="Path to the DiT model weights",
|
|
)
|
|
parser.add_argument(
|
|
"--model-dir",
|
|
type=str,
|
|
help="Directory containing StepVideo model",
|
|
)
|
|
|
|
# distributed_executor_backend
|
|
parser.add_argument(
|
|
"--distributed-executor-backend",
|
|
type=str,
|
|
choices=["mp"],
|
|
default=FastVideoArgs.distributed_executor_backend,
|
|
help="The distributed executor backend to use",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--inference-mode",
|
|
action=StoreBoolean,
|
|
default=FastVideoArgs.inference_mode,
|
|
help="Whether to use inference mode",
|
|
)
|
|
|
|
# HuggingFace specific parameters
|
|
parser.add_argument(
|
|
"--trust-remote-code",
|
|
action=StoreBoolean,
|
|
default=FastVideoArgs.trust_remote_code,
|
|
help="Trust remote code when loading HuggingFace models",
|
|
)
|
|
parser.add_argument(
|
|
"--revision",
|
|
type=str,
|
|
default=FastVideoArgs.revision,
|
|
help=
|
|
"The specific model version to use (can be a branch name, tag name, or commit id)",
|
|
)
|
|
|
|
# Parallelism
|
|
parser.add_argument(
|
|
"--num-gpus",
|
|
type=int,
|
|
default=FastVideoArgs.num_gpus,
|
|
help="The number of GPUs to use.",
|
|
)
|
|
parser.add_argument(
|
|
"--tensor-parallel-size",
|
|
"--tp-size",
|
|
type=int,
|
|
default=FastVideoArgs.tp_size,
|
|
help="The tensor parallelism size.",
|
|
)
|
|
parser.add_argument(
|
|
"--sequence-parallel-size",
|
|
"--sp-size",
|
|
type=int,
|
|
default=FastVideoArgs.sp_size,
|
|
help="The sequence parallelism size.",
|
|
)
|
|
parser.add_argument(
|
|
"--hsdp-replicate-dim",
|
|
type=int,
|
|
default=FastVideoArgs.hsdp_replicate_dim,
|
|
help="The data parallelism size.",
|
|
)
|
|
parser.add_argument(
|
|
"--hsdp-shard-dim",
|
|
type=int,
|
|
default=FastVideoArgs.hsdp_shard_dim,
|
|
help="The data parallelism shards.",
|
|
)
|
|
parser.add_argument(
|
|
"--dist-timeout",
|
|
type=int,
|
|
default=FastVideoArgs.dist_timeout,
|
|
help="Set timeout for torch.distributed initialization.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--embedded-cfg-scale",
|
|
type=float,
|
|
default=FastVideoArgs.embedded_cfg_scale,
|
|
help="Embedded CFG scale",
|
|
)
|
|
parser.add_argument(
|
|
"--flow-shift",
|
|
"--shift",
|
|
type=float,
|
|
default=FastVideoArgs.flow_shift,
|
|
help="Flow shift parameter",
|
|
)
|
|
parser.add_argument(
|
|
"--output-type",
|
|
type=str,
|
|
default=FastVideoArgs.output_type,
|
|
choices=["pil"],
|
|
help="Output type for the generated video",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--precision",
|
|
type=str,
|
|
default=FastVideoArgs.precision,
|
|
choices=["fp32", "fp16", "bf16"],
|
|
help="Precision for the model",
|
|
)
|
|
|
|
# VAE configuration
|
|
parser.add_argument(
|
|
"--vae-precision",
|
|
type=str,
|
|
default=FastVideoArgs.vae_precision,
|
|
choices=["fp32", "fp16", "bf16"],
|
|
help="Precision for VAE",
|
|
)
|
|
parser.add_argument(
|
|
"--vae-tiling",
|
|
action=StoreBoolean,
|
|
default=FastVideoArgs.vae_tiling,
|
|
help="Enable VAE tiling",
|
|
)
|
|
parser.add_argument(
|
|
"--vae-sp",
|
|
action=StoreBoolean,
|
|
help="Enable VAE spatial parallelism",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--text-encoder-precisions",
|
|
nargs="+",
|
|
type=str,
|
|
default=FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS,
|
|
choices=["fp32", "fp16", "bf16"],
|
|
help="Precision for each text encoder",
|
|
)
|
|
|
|
# Image encoder config
|
|
parser.add_argument(
|
|
"--image-encoder-precision",
|
|
type=str,
|
|
default=FastVideoArgs.image_encoder_precision,
|
|
choices=["fp32", "fp16", "bf16"],
|
|
help="Precision for image encoder",
|
|
)
|
|
|
|
# STA parameters
|
|
parser.add_argument(
|
|
"--STA-mode",
|
|
type=str,
|
|
default=FastVideoArgs.STA_mode,
|
|
choices=[
|
|
"STA_inference", "STA_searching", "STA_tuning",
|
|
"STA_tuning_cfg", None
|
|
],
|
|
help="STA mode",
|
|
)
|
|
parser.add_argument(
|
|
"--skip-time-steps",
|
|
type=int,
|
|
default=FastVideoArgs.skip_time_steps,
|
|
help="Number of time steps to warmup (full attention) for STA",
|
|
)
|
|
parser.add_argument(
|
|
"--mask-strategy-file-path",
|
|
type=str,
|
|
help="Path to mask strategy JSON file for STA",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-torch-compile",
|
|
action=StoreBoolean,
|
|
help=
|
|
"Use torch.compile for speeding up STA inference without teacache",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--use-cpu-offload",
|
|
action=StoreBoolean,
|
|
help=
|
|
"Use CPU offload for model inference. Enable if run out of memory with FSDP.",
|
|
)
|
|
parser.add_argument(
|
|
"--use-fsdp-inference",
|
|
action=StoreBoolean,
|
|
help=
|
|
"Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--disable-autocast",
|
|
action=StoreBoolean,
|
|
help=
|
|
"Disable autocast for denoising loop and vae decoding in pipeline sampling",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--pos_magic",
|
|
type=str,
|
|
default=FastVideoArgs.pos_magic,
|
|
help="Positive magic prompt for sampling",
|
|
)
|
|
parser.add_argument(
|
|
"--neg_magic",
|
|
type=str,
|
|
default=FastVideoArgs.neg_magic,
|
|
help="Negative magic prompt for sampling",
|
|
)
|
|
parser.add_argument(
|
|
"--timesteps_scale",
|
|
type=bool,
|
|
default=FastVideoArgs.timesteps_scale,
|
|
help="Bool for applying scheduler scale in set_timesteps",
|
|
)
|
|
|
|
# Logging
|
|
parser.add_argument(
|
|
"--log-level",
|
|
type=str,
|
|
default=FastVideoArgs.log_level,
|
|
help="The logging level of all loggers.",
|
|
)
|
|
|
|
# Add VAE configuration arguments
|
|
from fastvideo.v1.configs.models.vaes.base import VAEConfig
|
|
VAEConfig.add_cli_args(parser)
|
|
|
|
# Add DiT configuration arguments
|
|
from fastvideo.v1.configs.models.dits.base import DiTConfig
|
|
DiTConfig.add_cli_args(parser)
|
|
|
|
return parser
|
|
|
|
@classmethod
|
|
def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
|
|
args.tp_size = args.tensor_parallel_size
|
|
args.sp_size = args.sequence_parallel_size
|
|
args.flow_shift = getattr(args, "shift", args.flow_shift)
|
|
|
|
# Get all fields from the dataclass
|
|
attrs = [attr.name for attr in dataclasses.fields(cls)]
|
|
|
|
# Create a dictionary of attribute values, with defaults for missing attributes
|
|
kwargs = {}
|
|
for attr in attrs:
|
|
# Handle renamed attributes or those with multiple CLI names
|
|
if attr == 'tp_size' and hasattr(args, 'tensor_parallel_size'):
|
|
kwargs[attr] = args.tensor_parallel_size
|
|
elif attr == 'sp_size' and hasattr(args, 'sequence_parallel_size'):
|
|
kwargs[attr] = args.sequence_parallel_size
|
|
elif attr == 'flow_shift' and hasattr(args, 'shift'):
|
|
kwargs[attr] = args.shift
|
|
# Use getattr with default value from the dataclass for potentially missing attributes
|
|
else:
|
|
default_value = getattr(cls, attr, None)
|
|
value = getattr(args, attr, default_value)
|
|
if value is not None:
|
|
kwargs[attr] = value
|
|
|
|
return cls(**kwargs)
|
|
|
|
def check_fastvideo_args(self) -> None:
|
|
"""Validate inference arguments for consistency"""
|
|
if not self.inference_mode:
|
|
assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
|
|
assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
|
|
assert self.sp_size != -1, "sp_size must be set for training"
|
|
|
|
if self.tp_size == -1:
|
|
self.tp_size = self.num_gpus
|
|
if self.sp_size == -1:
|
|
self.sp_size = self.num_gpus
|
|
if self.hsdp_shard_dim == -1:
|
|
self.hsdp_shard_dim = self.num_gpus
|
|
|
|
assert self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0, "num_gpus must >= and be divisible by sp_size"
|
|
assert self.hsdp_replicate_dim <= self.num_gpus and self.num_gpus % self.hsdp_replicate_dim == 0, "num_gpus must >= and be divisible by hsdp_replicate_dim"
|
|
assert self.hsdp_shard_dim <= self.num_gpus and self.num_gpus % self.hsdp_shard_dim == 0, "num_gpus must >= and be divisible by hsdp_shard_dim"
|
|
|
|
if self.num_gpus < max(self.tp_size, self.sp_size):
|
|
self.num_gpus = max(self.tp_size, self.sp_size)
|
|
|
|
if self.tp_size != self.sp_size:
|
|
raise ValueError(
|
|
f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"
|
|
)
|
|
|
|
# Validate VAE spatial parallelism with VAE tiling
|
|
if self.vae_sp and not self.vae_tiling:
|
|
raise ValueError(
|
|
"Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
|
|
)
|
|
|
|
if len(self.text_encoder_configs) != len(self.text_encoder_precisions):
|
|
raise ValueError(
|
|
f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})"
|
|
)
|
|
|
|
if len(self.text_encoder_configs) != len(self.preprocess_text_funcs):
|
|
raise ValueError(
|
|
f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
|
|
)
|
|
|
|
if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs):
|
|
raise ValueError(
|
|
f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
|
|
)
|
|
|
|
if self.enable_torch_compile and self.num_gpus > 1:
|
|
logger.warning(
|
|
"Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"
|
|
)
|
|
self.enable_torch_compile = False
|
|
|
|
|
|
_current_fastvideo_args = None
|
|
|
|
|
|
def prepare_fastvideo_args(argv: List[str]) -> FastVideoArgs:
|
|
"""
|
|
Prepare the inference arguments from the command line arguments.
|
|
|
|
Args:
|
|
argv: The command line arguments. Typically, it should be `sys.argv[1:]`
|
|
to ensure compatibility with `parse_args` when no arguments are passed.
|
|
|
|
Returns:
|
|
The inference arguments.
|
|
"""
|
|
parser = FlexibleArgumentParser()
|
|
FastVideoArgs.add_cli_args(parser)
|
|
raw_args = parser.parse_args(argv)
|
|
fastvideo_args = FastVideoArgs.from_cli_args(raw_args)
|
|
global _current_fastvideo_args
|
|
_current_fastvideo_args = fastvideo_args
|
|
return fastvideo_args
|
|
|
|
|
|
@contextmanager
|
|
def set_current_fastvideo_args(fastvideo_args: FastVideoArgs):
|
|
"""
|
|
Temporarily set the current fastvideo config.
|
|
Used during model initialization.
|
|
We save the current fastvideo config in a global variable,
|
|
so that all modules can access it, e.g. custom ops
|
|
can access the fastvideo config to determine how to dispatch.
|
|
"""
|
|
global _current_fastvideo_args
|
|
old_fastvideo_args = _current_fastvideo_args
|
|
try:
|
|
_current_fastvideo_args = fastvideo_args
|
|
yield
|
|
finally:
|
|
_current_fastvideo_args = old_fastvideo_args
|
|
|
|
|
|
def get_current_fastvideo_args() -> FastVideoArgs:
|
|
if _current_fastvideo_args is None:
|
|
# in ci, usually when we test custom ops/modules directly,
|
|
# we don't set the fastvideo config. In that case, we set a default
|
|
# config.
|
|
# TODO(will): may need to handle this for CI.
|
|
raise ValueError("Current fastvideo args is not set.")
|
|
return _current_fastvideo_args
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TrainingArgs(FastVideoArgs):
|
|
"""
|
|
Training arguments. Inherits from FastVideoArgs and adds training-specific
|
|
arguments. If there are any conflicts, the training arguments will take
|
|
precedence.
|
|
"""
|
|
data_path: str = ""
|
|
dataloader_num_workers: int = 0
|
|
num_height: int = 0
|
|
num_width: int = 0
|
|
num_frames: int = 0
|
|
|
|
train_batch_size: int = 0
|
|
num_latent_t: int = 0
|
|
group_frame: bool = False
|
|
group_resolution: bool = False
|
|
|
|
# text encoder & vae & diffusion model
|
|
pretrained_model_name_or_path: str = ""
|
|
dit_model_name_or_path: str = ""
|
|
cache_dir: str = ""
|
|
|
|
# diffusion setting
|
|
ema_decay: float = 0.0
|
|
ema_start_step: int = 0
|
|
cfg: float = 0.0
|
|
precondition_outputs: bool = False
|
|
|
|
# validation & logs
|
|
validation_prompt_dir: str = ""
|
|
validation_sampling_steps: str = ""
|
|
validation_guidance_scale: str = ""
|
|
validation_steps: float = 0.0
|
|
log_validation: bool = False
|
|
tracker_project_name: str = ""
|
|
wandb_run_name: str = ""
|
|
seed: Optional[int] = None
|
|
|
|
# output
|
|
output_dir: str = ""
|
|
checkpoints_total_limit: int = 0
|
|
checkpointing_steps: int = 0
|
|
resume_from_checkpoint: bool = False
|
|
|
|
# optimizer & scheduler
|
|
num_train_epochs: int = 0
|
|
max_train_steps: int = 0
|
|
gradient_accumulation_steps: int = 0
|
|
learning_rate: float = 0.0
|
|
scale_lr: bool = False
|
|
lr_scheduler: str = "constant"
|
|
lr_warmup_steps: int = 0
|
|
max_grad_norm: float = 0.0
|
|
gradient_checkpointing: bool = False
|
|
selective_checkpointing: float = 0.0
|
|
allow_tf32: bool = False
|
|
mixed_precision: str = ""
|
|
train_sp_batch_size: int = 0
|
|
fsdp_sharding_startegy: str = ""
|
|
|
|
weighting_scheme: str = ""
|
|
logit_mean: float = 0.0
|
|
logit_std: float = 1.0
|
|
mode_scale: float = 0.0
|
|
|
|
num_euler_timesteps: int = 0
|
|
lr_num_cycles: int = 0
|
|
lr_power: float = 0.0
|
|
not_apply_cfg_solver: bool = False
|
|
distill_cfg: float = 0.0
|
|
scheduler_type: str = ""
|
|
linear_quadratic_threshold: float = 0.0
|
|
linear_range: float = 0.0
|
|
weight_decay: float = 0.0
|
|
use_ema: bool = False
|
|
multi_phased_distill_schedule: str = ""
|
|
pred_decay_weight: float = 0.0
|
|
pred_decay_type: str = ""
|
|
hunyuan_teacher_disable_cfg: bool = False
|
|
|
|
# master_weight_type
|
|
master_weight_type: str = ""
|
|
|
|
# For fast checking in LoRA pipeline
|
|
training_mode: bool = True
|
|
|
|
@classmethod
|
|
def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
|
|
# Get all fields from the dataclass
|
|
attrs = [attr.name for attr in dataclasses.fields(cls)]
|
|
|
|
# Create a dictionary of attribute values, with defaults for missing attributes
|
|
kwargs = {}
|
|
for attr in attrs:
|
|
# Handle renamed attributes or those with multiple CLI names
|
|
if attr == 'tp_size' and hasattr(args, 'tensor_parallel_size'):
|
|
kwargs[attr] = args.tensor_parallel_size
|
|
elif attr == 'sp_size' and hasattr(args, 'sequence_parallel_size'):
|
|
kwargs[attr] = args.sequence_parallel_size
|
|
elif attr == 'flow_shift' and hasattr(args, 'shift'):
|
|
kwargs[attr] = args.shift
|
|
else:
|
|
default_value = getattr(cls, attr, None)
|
|
if getattr(args, attr, default_value) is not None:
|
|
kwargs[attr] = getattr(args, attr, default_value)
|
|
|
|
return cls(**kwargs)
|
|
|
|
@staticmethod
|
|
def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
|
|
parser.add_argument("--data-path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to parquet files")
|
|
parser.add_argument("--dataloader-num-workers",
|
|
type=int,
|
|
required=True,
|
|
help="Number of workers for dataloader")
|
|
parser.add_argument("--num-height",
|
|
type=int,
|
|
required=True,
|
|
help="Number of heights")
|
|
parser.add_argument("--num-width",
|
|
type=int,
|
|
required=True,
|
|
help="Number of widths")
|
|
parser.add_argument("--num-frames",
|
|
type=int,
|
|
required=True,
|
|
help="Number of frames")
|
|
|
|
# Training batch and model configuration
|
|
parser.add_argument("--train-batch-size",
|
|
type=int,
|
|
required=True,
|
|
help="Training batch size")
|
|
parser.add_argument("--num-latent-t",
|
|
type=int,
|
|
required=True,
|
|
help="Number of latent time steps")
|
|
parser.add_argument("--group-frame",
|
|
action=StoreBoolean,
|
|
help="Whether to group frames during training")
|
|
parser.add_argument("--group-resolution",
|
|
action=StoreBoolean,
|
|
help="Whether to group resolutions during training")
|
|
|
|
# Model paths
|
|
parser.add_argument("--pretrained-model-name-or-path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to pretrained model or model name")
|
|
parser.add_argument("--dit-model-name-or-path",
|
|
type=str,
|
|
required=False,
|
|
help="Path to DiT model or model name")
|
|
parser.add_argument("--cache-dir",
|
|
type=str,
|
|
help="Directory to cache models")
|
|
|
|
# Diffusion settings
|
|
parser.add_argument("--ema-decay",
|
|
type=float,
|
|
default=0.999,
|
|
help="EMA decay rate")
|
|
parser.add_argument("--ema-start-step",
|
|
type=int,
|
|
default=0,
|
|
help="Step to start EMA")
|
|
parser.add_argument("--cfg",
|
|
type=float,
|
|
help="Classifier-free guidance scale")
|
|
parser.add_argument(
|
|
"--precondition-outputs",
|
|
action=StoreBoolean,
|
|
help="Whether to precondition the outputs of the model")
|
|
|
|
# Validation and logging
|
|
parser.add_argument("--validation-prompt-dir",
|
|
type=str,
|
|
help="Directory containing validation prompts")
|
|
parser.add_argument("--validation-sampling-steps",
|
|
type=str,
|
|
help="Validation sampling steps")
|
|
parser.add_argument("--validation-guidance-scale",
|
|
type=str,
|
|
help="Validation guidance scale")
|
|
parser.add_argument("--validation-steps",
|
|
type=float,
|
|
help="Number of validation steps")
|
|
parser.add_argument("--log-validation",
|
|
action=StoreBoolean,
|
|
help="Whether to log validation results")
|
|
parser.add_argument("--tracker-project-name",
|
|
type=str,
|
|
help="Project name for tracking")
|
|
parser.add_argument("--wandb-run-name",
|
|
type=str,
|
|
help="Run name for wandb")
|
|
parser.add_argument("--seed",
|
|
type=int,
|
|
default=42,
|
|
help="Seed for deterministic training")
|
|
|
|
# Output configuration
|
|
parser.add_argument("--output-dir",
|
|
type=str,
|
|
required=True,
|
|
help="Output directory for checkpoints and logs")
|
|
parser.add_argument("--checkpoints-total-limit",
|
|
type=int,
|
|
help="Maximum number of checkpoints to keep")
|
|
parser.add_argument("--checkpointing-steps",
|
|
type=int,
|
|
help="Steps between checkpoints")
|
|
parser.add_argument("--resume-from-checkpoint",
|
|
type=str,
|
|
help="Path to checkpoint to resume from")
|
|
parser.add_argument("--logging-dir",
|
|
type=str,
|
|
help="Directory for logging")
|
|
|
|
# Training configuration
|
|
parser.add_argument("--num-train-epochs",
|
|
type=int,
|
|
help="Number of training epochs")
|
|
parser.add_argument("--max-train-steps",
|
|
type=int,
|
|
help="Maximum number of training steps")
|
|
parser.add_argument("--gradient-accumulation-steps",
|
|
type=int,
|
|
help="Number of steps to accumulate gradients")
|
|
parser.add_argument("--learning-rate",
|
|
type=float,
|
|
required=True,
|
|
help="Learning rate")
|
|
parser.add_argument("--scale-lr",
|
|
action=StoreBoolean,
|
|
help="Whether to scale learning rate")
|
|
parser.add_argument("--lr-scheduler",
|
|
type=str,
|
|
default="constant",
|
|
help="Learning rate scheduler type")
|
|
parser.add_argument("--lr-warmup-steps",
|
|
type=int,
|
|
default=10,
|
|
help="Number of warmup steps for learning rate")
|
|
parser.add_argument("--max-grad-norm",
|
|
type=float,
|
|
help="Maximum gradient norm")
|
|
parser.add_argument("--gradient-checkpointing",
|
|
action=StoreBoolean,
|
|
help="Whether to use gradient checkpointing")
|
|
parser.add_argument("--selective-checkpointing",
|
|
type=float,
|
|
help="Selective checkpointing threshold")
|
|
parser.add_argument("--allow-tf32",
|
|
action=StoreBoolean,
|
|
help="Whether to allow TF32")
|
|
parser.add_argument("--mixed-precision",
|
|
type=str,
|
|
help="Mixed precision training type")
|
|
parser.add_argument("--train-sp-batch-size",
|
|
type=int,
|
|
help="Training spatial parallelism batch size")
|
|
|
|
parser.add_argument("--fsdp-sharding-strategy",
|
|
type=str,
|
|
help="FSDP sharding strategy")
|
|
|
|
parser.add_argument(
|
|
"--weighting_scheme",
|
|
type=str,
|
|
default="uniform",
|
|
choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
|
|
)
|
|
parser.add_argument(
|
|
"--logit_mean",
|
|
type=float,
|
|
default=0.0,
|
|
help="mean to use when using the `'logit_normal'` weighting scheme.",
|
|
)
|
|
parser.add_argument(
|
|
"--logit_std",
|
|
type=float,
|
|
default=1.0,
|
|
help="std to use when using the `'logit_normal'` weighting scheme.",
|
|
)
|
|
parser.add_argument(
|
|
"--mode_scale",
|
|
type=float,
|
|
default=1.29,
|
|
help=
|
|
"Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
|
|
)
|
|
|
|
# Additional training parameters
|
|
parser.add_argument("--num-euler-timesteps",
|
|
type=int,
|
|
help="Number of Euler timesteps")
|
|
parser.add_argument("--lr-num-cycles",
|
|
type=int,
|
|
help="Number of learning rate cycles")
|
|
parser.add_argument("--lr-power",
|
|
type=float,
|
|
help="Learning rate power")
|
|
parser.add_argument("--not-apply-cfg-solver",
|
|
action=StoreBoolean,
|
|
help="Whether to not apply CFG solver")
|
|
parser.add_argument("--distill-cfg",
|
|
type=float,
|
|
help="Distillation CFG scale")
|
|
parser.add_argument("--scheduler-type", type=str, help="Scheduler type")
|
|
parser.add_argument("--linear-quadratic-threshold",
|
|
type=float,
|
|
help="Linear quadratic threshold")
|
|
parser.add_argument("--linear-range", type=float, help="Linear range")
|
|
parser.add_argument("--weight-decay", type=float, help="Weight decay")
|
|
parser.add_argument("--use-ema",
|
|
action=StoreBoolean,
|
|
help="Whether to use EMA")
|
|
parser.add_argument("--multi-phased-distill-schedule",
|
|
type=str,
|
|
help="Multi-phased distillation schedule")
|
|
parser.add_argument("--pred-decay-weight",
|
|
type=float,
|
|
help="Prediction decay weight")
|
|
parser.add_argument("--pred-decay-type",
|
|
type=str,
|
|
help="Prediction decay type")
|
|
parser.add_argument("--hunyuan-teacher-disable-cfg",
|
|
action=StoreBoolean,
|
|
help="Whether to disable CFG for Hunyuan teacher")
|
|
parser.add_argument("--master-weight-type",
|
|
type=str,
|
|
help="Master weight type")
|
|
|
|
return parser
|