Function main

train_code/train_cogvideox_motion_FrameINO.py:549–1218

(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"]  
    enable_slicing = config["enable_slicing"]
    enable_tiling = config["enable_tiling"]
    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"]
    variant = config["variant"]
    lr_scheduler = config["lr_scheduler"]
    use_rotary_positional_embeddings = config["use_rotary_positional_embeddings"],
    use_learned_positional_embeddings = config["use_learned_positional_embeddings"]
    validation_step = config["validation_step"]
    first_iter_validation = config["first_iter_validation"]
    num_inference_steps = config["num_inference_steps"]

    # 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"])