(bug) audioldm ckpt to diffuser script is not working at medium or large size ckpt.

[circuluspibo]

HI. i tried to convert diffusers model from audioldm.
small size ckpt like as audioldm-full-s-v2.ckpt audioldm-s-text-ft.ckpt are well,
but middle or large size are not woking (audioldm-m-text-ft.ckpt, audioldm-m-full.ckpt)below

shape torch.Size([960]) from checkpoint, the shape in current model is
torch.Size([640]).
size mismatch for mid_block.resnets.1.conv2.weight: copying a param with
shape torch.Size([960, 960, 3, 3]) from checkpoint, the shape in current model
is torch.Size([640, 640, 3, 3]).
size mismatch for mid_block.resnets.1.conv2.bias: copying a param with
shape torch.Size([960]) from checkpoint, the shape in current model is
torch.Size([640]).
size mismatch for conv_norm_out.weight: copying a param with shape
torch.Size([192]) from checkpoint, the shape in current model is
torch.Size([128]).
size mismatch for conv_norm_out.bias: copying a param with shape
torch.Size([192]) from checkpoint, the shape in current model is
torch.Size([128]).
size mismatch for conv_out.weight: copying a param with shape
torch.Size([8, 192, 3, 3]) from checkpoint, the shape in current model is
torch.Size([8, 128, 3, 3]).
(base) circulus@circulus:~/git/diffusers/scr

[sayakpaul]

Cc: @patrickvonplaten @williamberman

[patrickvonplaten]

cc @sanchit-gandhi here

[circuluspibo]

@sanchit-gandhi please check this article :)

[sanchit-gandhi]

Hey @circuluspibo - if you check the docstrings, you'll see that the conversion script defaults to using the config for the small checkpoint, and that the recommendation is to specify a new config if converting checkpoints of different dimensionality. If we check the config for the original models, we see that for the medium checkpoints we need to increase the UNet model_channels, and for the large checkpoints we need to increase both the model_channels and num_head_channels. It's precisely these config attributes that are giving the shape mismatch you've reported!

To make it easier to convert these checkpoints, I've opened a PR that let's you specify explicitly model_channels and num_head_channels - you should pass these two arguments to the conversion script when converting the medium and large checkpoints:

Small:

• model_channels = None (leave to default to original config value of 128)
• num_head_channels = None (leave to default to original config value of 32)

Medium:

• model_channels = 192
• num_head_channels = None (leave to default to original config value of 32)

Large:

• model_channels = 256
• num_head_channels = 64

All of the recommended audioldm-X-full checkpoints are available on the HF Hub under the CVSSP org: https://huggingface.co/cvssp So you only need to convert these checkpoints if you've done any additional fine-tuning.