Is your feature request related to a problem? Please describe.
I want to change the time_embed_dim of UNet2DModel, but it is hard coded as time_embed_dim = block_out_channels[0] * 4 in the __init__ function.

Describe the solution you'd like.
Make time_embedding_dim a parameter of the __init__ function, with the default value of None. Use time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 in the function body.

Describe alternatives you've considered.
N/A.

Additional context.
The same thing in UNet2DConditionModel can be changed via the parameter time_embedding_dim of its __init__ function.