# this file is only for batch size search and is not used for training

# Define dataset
dataset = dict(
    type="VariableVideoTextDataset",
    data_path=None,
    num_frames=None,
    frame_interval=3,
    image_size=(None, None),
    transform_name="resize_crop",
)

# bucket config format:
# 1. { resolution: {num_frames: (prob, batch_size)} }, in this case batch_size is ignored when searching
# 2. { resolution: {num_frames: (prob, (max_batch_size, ))} }, batch_size is searched in the range [batch_size_start, max_batch_size), batch_size_start is configured via CLI
# 3. { resolution: {num_frames: (prob, (min_batch_size, max_batch_size))} }, batch_size is searched in the range [min_batch_size, max_batch_size)
# 4. { resolution: {num_frames: (prob, (min_batch_size, max_batch_size, step_size))} }, batch_size is searched in the range [min_batch_size, max_batch_size) with step_size (grid search)
# 5. { resolution: {num_frames: (0.0, None)} }, this bucket will not be used

bucket_config = {
    # == manual search ==
    # "240p": {128: (1.0, 2)}, # 4.28s/it
    # "240p": {64: (1.0, 4)},
    # "240p": {32: (1.0, 8)},  # 4.6s/it
    # "240p": {16: (1.0, 16)},  # 4.6s/it
    # "480p": {16: (1.0, 4)},  # 4.6s/it
    # "720p": {16: (1.0, 2)},  # 5.89s/it
    # "256": {1: (1.0, 256)},  # 4.5s/it
    # "512": {1: (1.0, 96)}, # 4.7s/it
    # "512": {1: (1.0, 128)}, # 6.3s/it
    # "480p": {1: (1.0, 50)},  # 4.0s/it
    # "1024": {1: (1.0, 32)},  # 6.8s/it
    # "1024": {1: (1.0, 20)}, # 4.3s/it
    # "1080p": {1: (1.0, 16)}, # 8.6s/it
    # "1080p": {1: (1.0, 8)},  # 4.4s/it
    # == stage 2 ==
    # "240p": {
    #     16: (1.0, (2, 32)),
    #     32: (1.0, (2, 16)),
    #     64: (1.0, (2, 8)),
    #     128: (1.0, (2, 6)),
    # },
    # "256": {1: (1.0, (128, 300))},
    # "512": {1: (0.5, (64, 128))},
    # "480p": {1: (0.4, (32, 128)), 16: (0.4, (2, 32)), 32: (0.0, None)},
    # "720p": {16: (0.1, (2, 16)), 32: (0.0, None)},  # No examples now
    # "1024": {1: (0.3, (8, 64))},
    # "1080p": {1: (0.3, (2, 32))},
    # == stage 3 ==
    "720p": {1: (20, 40), 32: (0.5, (2, 4)), 64: (0.5, (1, 1))},
}


# Define acceleration
num_workers = 4
num_bucket_build_workers = 16
dtype = "bf16"
grad_checkpoint = True
plugin = "zero2"
sp_size = 1

# Define model
model = dict(
    type="STDiT2-XL/2",
    from_pretrained=None,
    input_sq_size=512,  # pretrained model is trained on 512x512
    qk_norm=True,
    enable_flashattn=True,
    enable_layernorm_kernel=True,
)
vae = dict(
    type="VideoAutoencoderKL",
    from_pretrained="stabilityai/sd-vae-ft-ema",
    micro_batch_size=4,
    local_files_only=True,
)
text_encoder = dict(
    type="t5",
    from_pretrained="DeepFloyd/t5-v1_1-xxl",
    model_max_length=200,
    shardformer=True,
    local_files_only=True,
)
scheduler = dict(
    type="iddpm",
    timestep_respacing="",
)

# Others
seed = 42
outputs = "outputs"
wandb = False

epochs = 1000
log_every = 10
ckpt_every = 1000
load = None

batch_size = None
lr = 2e-5
grad_clip = 1.0