def get_autoencoder_kl_config(block_out_channels=None, norm_num_groups=None):
    block_out_channels = block_out_channels or [2, 4]
    norm_num_groups = norm_num_groups or 2
    init_dict = {
        "block_out_channels": block_out_channels,
        "in_channels": 3,
        "out_channels": 3,
        "down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
        "up_block_types": ["UpDecoderBlock2D"] * len(block_out_channels),
        "latent_channels": 4,
        "norm_num_groups": norm_num_groups,
    }
    return init_dict


def get_asym_autoencoder_kl_config(block_out_channels=None, norm_num_groups=None):
    block_out_channels = block_out_channels or [2, 4]
    norm_num_groups = norm_num_groups or 2
    init_dict = {
        "in_channels": 3,
        "out_channels": 3,
        "down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
        "down_block_out_channels": block_out_channels,
        "layers_per_down_block": 1,
        "up_block_types": ["UpDecoderBlock2D"] * len(block_out_channels),
        "up_block_out_channels": block_out_channels,
        "layers_per_up_block": 1,
        "act_fn": "silu",
        "latent_channels": 4,
        "norm_num_groups": norm_num_groups,
        "sample_size": 32,
        "scaling_factor": 0.18215,
    }
    return init_dict


def get_autoencoder_tiny_config(block_out_channels=None):
    block_out_channels = (len(block_out_channels) * [32]) if block_out_channels is not None else [32, 32]
    init_dict = {
        "in_channels": 3,
        "out_channels": 3,
        "encoder_block_out_channels": block_out_channels,
        "decoder_block_out_channels": block_out_channels,
        "num_encoder_blocks": [b // min(block_out_channels) for b in block_out_channels],
        "num_decoder_blocks": [b // min(block_out_channels) for b in reversed(block_out_channels)],
    }
    return init_dict


def get_consistency_vae_config(block_out_channels=None, norm_num_groups=None):
    block_out_channels = block_out_channels or [2, 4]
    norm_num_groups = norm_num_groups or 2
    return {
        "encoder_block_out_channels": block_out_channels,
        "encoder_in_channels": 3,
        "encoder_out_channels": 4,
        "encoder_down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels),
        "decoder_add_attention": False,
        "decoder_block_out_channels": block_out_channels,
        "decoder_down_block_types": ["ResnetDownsampleBlock2D"] * len(block_out_channels),
        "decoder_downsample_padding": 1,
        "decoder_in_channels": 7,
        "decoder_layers_per_block": 1,
        "decoder_norm_eps": 1e-05,
        "decoder_norm_num_groups": norm_num_groups,
        "encoder_norm_num_groups": norm_num_groups,
        "decoder_num_train_timesteps": 1024,
        "decoder_out_channels": 6,
        "decoder_resnet_time_scale_shift": "scale_shift",
        "decoder_time_embedding_type": "learned",
        "decoder_up_block_types": ["ResnetUpsampleBlock2D"] * len(block_out_channels),
        "scaling_factor": 1,
        "latent_channels": 4,
    }


def get_autoencoder_oobleck_config(block_out_channels=None):
    init_dict = {
        "encoder_hidden_size": 12,
        "decoder_channels": 12,
        "decoder_input_channels": 6,
        "audio_channels": 2,
        "downsampling_ratios": [2, 4],
        "channel_multiples": [1, 2],
    }
    return init_dict