Lightstar/venv/Lib/site-packages/transformers/models/dac/convert_dac_checkpoint.py

# coding=utf-8
# Copyright 2024 Descript and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import fnmatch
import re

import torch

from transformers import (
    DacConfig,
    DacFeatureExtractor,
    DacModel,
    logging,
)


# checkpoints downloaded using:
# pip install descript-audio-codec
# python3 -m dac download # downloads the default 44kHz variant
# python3 -m dac download --model_type 44khz # downloads the 44kHz variant
# python3 -m dac download --model_type 24khz # downloads the 24kHz variant
# python3 -m dac download --model_type 16khz # downloads the 16kHz variant
# More informations: https://github.com/descriptinc/descript-audio-codec/tree/main

logging.set_verbosity_info()
logger = logging.get_logger("transformers.models.dac")


def match_pattern(string, pattern):
    # Split the pattern into parts
    pattern_parts = pattern.split(".")
    string_parts = string.split(".")

    pattern_block_count = string_block_count = 0

    for part in pattern_parts:
        if part.startswith("block"):
            pattern_block_count += 1

    for part in string_parts:
        if part.startswith("block"):
            string_block_count += 1

    return fnmatch.fnmatch(string, pattern) and string_block_count == pattern_block_count


TOP_LEVEL_KEYS = []
IGNORE_KEYS = []


MAPPING_ENCODER = {
    "encoder.block.0": ["encoder.conv1"],
    "encoder.block.5": ["encoder.snake1"],
    "encoder.block.6": ["encoder.conv2"],
    "encoder.block.*.block.*.block.0".replace("*", r"\d+"): ["encoder.block", "res_unit", "snake1"],
    "encoder.block.*.block.*.block.1".replace("*", r"\d+"): ["encoder.block", "res_unit", "conv1"],
    "encoder.block.*.block.*.block.2".replace("*", r"\d+"): ["encoder.block", "res_unit", "snake2"],
    "encoder.block.*.block.*.block.3".replace("*", r"\d+"): ["encoder.block", "res_unit", "conv2"],
    "encoder.block.*.block.3".replace("*", r"\d+"): ["encoder.block", "snake1"],
    "encoder.block.*.block.4".replace("*", r"\d+"): ["encoder.block", "conv1"],
}

MAPPING_QUANTIZER = {
    "quantizer.quantizers.*": ["quantizer.quantizers.*"],
}

MAPPING_DECODER = {
    "decoder.model.0": ["decoder.conv1"],
    "decoder.model.5": ["decoder.snake1"],
    "decoder.model.6": ["decoder.conv2"],
    "decoder.model.*.block.0".replace("*", r"\d+"): ["decoder.block", "snake1"],
    "decoder.model.*.block.1".replace("*", r"\d+"): ["decoder.block", "conv_t1"],
    "decoder.model.*.block.*.block.0".replace("*", r"\d+"): ["decoder.block", "res_unit", "snake1"],
    "decoder.model.*.block.*.block.1".replace("*", r"\d+"): ["decoder.block", "res_unit", "conv1"],
    "decoder.model.*.block.*.block.2".replace("*", r"\d+"): ["decoder.block", "res_unit", "snake2"],
    "decoder.model.*.block.*.block.3".replace("*", r"\d+"): ["decoder.block", "res_unit", "conv2"],
}


MAPPING = {
    **MAPPING_ENCODER,
    **MAPPING_QUANTIZER,
    **MAPPING_DECODER,
}


def set_recursively(hf_pointer, key, value, full_name, weight_type):
    for attribute in key.split("."):
        hf_pointer = getattr(hf_pointer, attribute)

    if weight_type is not None:
        hf_shape = getattr(hf_pointer, weight_type).shape
    else:
        hf_shape = hf_pointer.shape

    if hf_shape != value.shape:
        raise ValueError(
            f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
            f" {value.shape} for {full_name}"
        )

    if weight_type == "weight":
        hf_pointer.weight.data = value
    elif weight_type == "weight_g":
        hf_pointer.weight_g.data = value
    elif weight_type == "weight_v":
        hf_pointer.weight_v.data = value
    elif weight_type == "bias":
        hf_pointer.bias.data = value
    elif weight_type == "alpha":
        hf_pointer.alpha.data = value
    logger.info(f"{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}.")


def should_ignore(name, ignore_keys):
    for key in ignore_keys:
        if key.endswith(".*"):
            if name.startswith(key[:-1]):
                return True
        elif ".*." in key:
            prefix, suffix = key.split(".*.")
            if prefix in name and suffix in name:
                return True
        elif key in name:
            return True
    return False


def recursively_load_weights(orig_dict, hf_model, model_name):
    unused_weights = []

    if model_name not in ["dac_16khz", "dac_24khz", "dac_44khz"]:
        raise ValueError(f"Unsupported model: {model_name}")

    for name, value in orig_dict.items():
        is_used = False
        for key, mapped_key in MAPPING.items():
            regex = re.compile(key)
            if regex.search(name):
                if len(mapped_key) == 1:
                    if mapped_key[0][0] == "q":
                        mapped_key = ".".join(name.split(".")[:-1])
                    else:
                        mapped_key = mapped_key[0]
                elif len(mapped_key) == 3:
                    integers = re.findall(r"\b\d+\b", name)
                    if mapped_key[0][0] == "d":
                        mapped_key = "{}.{}.{}{}.{}".format(
                            mapped_key[0],
                            str(int(integers[0]) - 1),
                            mapped_key[1],
                            str(int(integers[1]) - 1),
                            mapped_key[2],
                        )
                    else:
                        mapped_key = "{}.{}.{}{}.{}".format(
                            mapped_key[0],
                            str(int(integers[0]) - 1),
                            mapped_key[1],
                            str(int(integers[1]) + 1),
                            mapped_key[2],
                        )
                elif len(mapped_key) == 2:
                    integers = re.findall(r"\b\d+\b", name)
                    mapped_key = "{}.{}.{}".format(mapped_key[0], str(int(integers[0]) - 1), mapped_key[1])

                is_used = True
                if "weight_g" in name:
                    weight_type = "weight_g"
                elif "weight_v" in name:
                    weight_type = "weight_v"
                elif "bias" in name:
                    weight_type = "bias"
                elif "alpha" in name:
                    weight_type = "alpha"
                elif "weight" in name:
                    weight_type = "weight"
                set_recursively(hf_model, mapped_key, value, name, weight_type)

        if not is_used:
            unused_weights.append(name)

    print(list(set(unused_weights)))

    logger.warning(f"Unused weights: {unused_weights}")


@torch.no_grad()
def convert_checkpoint(
    model_name,
    checkpoint_path,
    pytorch_dump_folder_path,
    sample_rate=16000,
    repo_id=None,
):
    model_dict = torch.load(checkpoint_path, "cpu")

    config = DacConfig()

    metadata = model_dict["metadata"]["kwargs"]
    config.encoder_hidden_size = metadata["encoder_dim"]
    config.downsampling_ratios = metadata["encoder_rates"]
    config.codebook_size = metadata["codebook_size"]
    config.n_codebooks = metadata["n_codebooks"]
    config.codebook_dim = metadata["codebook_dim"]
    config.decoder_hidden_size = metadata["decoder_dim"]
    config.upsampling_ratios = metadata["decoder_rates"]
    config.quantizer_dropout = float(metadata["quantizer_dropout"])
    config.sampling_rate = sample_rate

    model = DacModel(config)
    feature_extractor = DacFeatureExtractor()
    feature_extractor.sampling_rate = sample_rate

    original_checkpoint = model_dict["state_dict"]

    model.apply_weight_norm()
    recursively_load_weights(original_checkpoint, model, model_name)
    model.remove_weight_norm()

    model.save_pretrained(pytorch_dump_folder_path)

    if repo_id:
        print("Pushing to the hub...")
        feature_extractor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model",
        default="dac_44khz",
        type=str,
        help="The model to convert. Should be one of 'dac_16khz', 'dac_24khz', 'dac_44khz'.",
    )
    parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
    parser.add_argument(
        "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
    )
    parser.add_argument("--sample_rate", default=None, type=str, help="Sample rate used by DacFeatureExtractor")
    args = parser.parse_args()

    convert_checkpoint(
        args.model, args.checkpoint_path, args.pytorch_dump_folder_path, args.sample_rate, args.push_to_hub
    )