Function main

train_code/train_wan_motion_FrameINO.py:677–1333

(config, args)

def main(config, args):

    # Process args
    debug = args.debug
    use_8BitAdam = args.use_8BitAdam


    # Read Frequently Used config
    resume_from_checkpoint = config["resume_from_checkpoint"]
    output_folder = config["output_folder"]
    experiment_name = config["experiment_name"]
    mixed_precision = config["mixed_precision"]
    report_to = config["report_to"]
    seed = config["seed"]
    base_model_path = config["base_model_path"]
    pretrained_transformer_path = config["pretrained_transformer_path"]
    download_folder_path = config["download_folder_path"]
    train_csv_relative_path = config["train_csv_relative_path"]
    validation_csv_relative_path = config["validation_csv_relative_path"]
    gradient_checkpointing = config["gradient_checkpointing"]         
    learning_rate = config["learning_rate"]
    train_batch_size = config["train_batch_size"]
    dataloader_num_workers = config["dataloader_num_workers"] if not debug else 1       # In debug mode, only has 1 worker
    gradient_accumulation_steps = config["gradient_accumulation_steps"]
    max_train_steps = config["max_train_steps"]
    lr_warmup_steps = config["lr_warmup_steps"]
    checkpointing_steps = config["checkpointing_steps"]
    scale_lr = config["scale_lr"]
    checkpoints_total_limit = config["checkpoints_total_limit"]
    revision = config["revision"]
    lr_scheduler = config["lr_scheduler"]
    validation_step = config["validation_step"]
    first_iter_validation = config["first_iter_validation"]
    num_inference_steps = config["num_inference_steps"]
    max_grad_norm = config["max_grad_norm"]


    # Organize
    use_FrameIn = True      


    # Value Check
    if max_train_steps is None:
        print("max_train_steps must be set")
        os.exit(0)

    if torch.backends.mps.is_available() and mixed_precision == "bf16":
        # 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."
        )

    output_dir = os.path.join(output_folder, experiment_name)
    logging_dir = Path(output_dir, config["logging_name"])

    accelerator_project_config = ProjectConfiguration(project_dir = output_dir, logging_dir = logging_dir)
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)        # HACK: his find_unused_parameters is said to increase the computation cost
    init_kwargs = InitProcessGroupKwargs(backend="nccl", timeout=timedelta(seconds=config["nccl_timeout"]))