def convert_hifigan_checkpoint(checkpoint, config):
    """
    Takes a state dict and config, and returns a converted HiFiGAN vocoder checkpoint.
    """
    # extract state dict for vocoder
    vocoder_state_dict = extract_sub_model(checkpoint, key_prefix="first_stage_model.vocoder.")

    # fix upsampler keys, everything else is correct already
    for i in range(len(config.upsample_rates)):
        vocoder_state_dict[f"upsampler.{i}.weight"] = vocoder_state_dict.pop(f"ups.{i}.weight")
        vocoder_state_dict[f"upsampler.{i}.bias"] = vocoder_state_dict.pop(f"ups.{i}.bias")

    if not config.normalize_before:
        # if we don't set normalize_before then these variables are unused, so we set them to their initialised values
        vocoder_state_dict["mean"] = torch.zeros(config.model_in_dim)
        vocoder_state_dict["scale"] = torch.ones(config.model_in_dim)

    return vocoder_state_dict

    Takes a state dict and config, and returns a converted HiFiGAN vocoder checkpoint.
    
scripts/convert_original_audioldm2_to_diffusers.py