





import argparse

import torch

from transformers import (
    EncodecConfig,
    EncodecFeatureExtractor,
    EncodecModel,
    logging,
)


# checkpoints downloaded from:
# https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th
# https://huggingface.co/facebook/musicgen-small/resolve/main/compression_state_dict.bin
# https://dl.fbaipublicfiles.com/encodec/v0/encodec_48khz-7e698e3e.th


logging.set_verbosity_info()
UpperCamelCase__   : Dict					  = logging.get_logger("""transformers.models.encodec""")

UpperCamelCase__   : Optional[int]					  = {
    """quantizer.vq.layers.*._codebook.inited""": """quantizer.layers.*.codebook.inited""",
    """quantizer.vq.layers.*._codebook.cluster_size""": """quantizer.layers.*.codebook.cluster_size""",
    """quantizer.vq.layers.*._codebook.embed""": """quantizer.layers.*.codebook.embed""",
    """quantizer.vq.layers.*._codebook.embed_avg""": """quantizer.layers.*.codebook.embed_avg""",
}
UpperCamelCase__   : List[Any]					  = {
    """encoder.model.0.conv.conv""": """encoder.layers.0.conv""",
    """encoder.model.1.block.1.conv.conv""": """encoder.layers.1.block.1.conv""",
    """encoder.model.1.block.3.conv.conv""": """encoder.layers.1.block.3.conv""",
    """encoder.model.1.shortcut.conv.conv""": """encoder.layers.1.shortcut.conv""",
    """encoder.model.3.conv.conv""": """encoder.layers.3.conv""",
    """encoder.model.4.block.1.conv.conv""": """encoder.layers.4.block.1.conv""",
    """encoder.model.4.block.3.conv.conv""": """encoder.layers.4.block.3.conv""",
    """encoder.model.4.shortcut.conv.conv""": """encoder.layers.4.shortcut.conv""",
    """encoder.model.6.conv.conv""": """encoder.layers.6.conv""",
    """encoder.model.7.block.1.conv.conv""": """encoder.layers.7.block.1.conv""",
    """encoder.model.7.block.3.conv.conv""": """encoder.layers.7.block.3.conv""",
    """encoder.model.7.shortcut.conv.conv""": """encoder.layers.7.shortcut.conv""",
    """encoder.model.9.conv.conv""": """encoder.layers.9.conv""",
    """encoder.model.10.block.1.conv.conv""": """encoder.layers.10.block.1.conv""",
    """encoder.model.10.block.3.conv.conv""": """encoder.layers.10.block.3.conv""",
    """encoder.model.10.shortcut.conv.conv""": """encoder.layers.10.shortcut.conv""",
    """encoder.model.12.conv.conv""": """encoder.layers.12.conv""",
    """encoder.model.13.lstm""": """encoder.layers.13.lstm""",
    """encoder.model.15.conv.conv""": """encoder.layers.15.conv""",
}
UpperCamelCase__   : Union[str, Any]					  = {
    """encoder.model.0.conv.norm""": """encoder.layers.0.norm""",
    """encoder.model.1.block.1.conv.norm""": """encoder.layers.1.block.1.norm""",
    """encoder.model.1.block.3.conv.norm""": """encoder.layers.1.block.3.norm""",
    """encoder.model.1.shortcut.conv.norm""": """encoder.layers.1.shortcut.norm""",
    """encoder.model.3.conv.norm""": """encoder.layers.3.norm""",
    """encoder.model.4.block.1.conv.norm""": """encoder.layers.4.block.1.norm""",
    """encoder.model.4.block.3.conv.norm""": """encoder.layers.4.block.3.norm""",
    """encoder.model.4.shortcut.conv.norm""": """encoder.layers.4.shortcut.norm""",
    """encoder.model.6.conv.norm""": """encoder.layers.6.norm""",
    """encoder.model.7.block.1.conv.norm""": """encoder.layers.7.block.1.norm""",
    """encoder.model.7.block.3.conv.norm""": """encoder.layers.7.block.3.norm""",
    """encoder.model.7.shortcut.conv.norm""": """encoder.layers.7.shortcut.norm""",
    """encoder.model.9.conv.norm""": """encoder.layers.9.norm""",
    """encoder.model.10.block.1.conv.norm""": """encoder.layers.10.block.1.norm""",
    """encoder.model.10.block.3.conv.norm""": """encoder.layers.10.block.3.norm""",
    """encoder.model.10.shortcut.conv.norm""": """encoder.layers.10.shortcut.norm""",
    """encoder.model.12.conv.norm""": """encoder.layers.12.norm""",
    """encoder.model.15.conv.norm""": """encoder.layers.15.norm""",
}
UpperCamelCase__   : int					  = {
    """decoder.model.0.conv.conv""": """decoder.layers.0.conv""",
    """decoder.model.1.lstm""": """decoder.layers.1.lstm""",
    """decoder.model.3.convtr.convtr""": """decoder.layers.3.conv""",
    """decoder.model.4.block.1.conv.conv""": """decoder.layers.4.block.1.conv""",
    """decoder.model.4.block.3.conv.conv""": """decoder.layers.4.block.3.conv""",
    """decoder.model.4.shortcut.conv.conv""": """decoder.layers.4.shortcut.conv""",
    """decoder.model.6.convtr.convtr""": """decoder.layers.6.conv""",
    """decoder.model.7.block.1.conv.conv""": """decoder.layers.7.block.1.conv""",
    """decoder.model.7.block.3.conv.conv""": """decoder.layers.7.block.3.conv""",
    """decoder.model.7.shortcut.conv.conv""": """decoder.layers.7.shortcut.conv""",
    """decoder.model.9.convtr.convtr""": """decoder.layers.9.conv""",
    """decoder.model.10.block.1.conv.conv""": """decoder.layers.10.block.1.conv""",
    """decoder.model.10.block.3.conv.conv""": """decoder.layers.10.block.3.conv""",
    """decoder.model.10.shortcut.conv.conv""": """decoder.layers.10.shortcut.conv""",
    """decoder.model.12.convtr.convtr""": """decoder.layers.12.conv""",
    """decoder.model.13.block.1.conv.conv""": """decoder.layers.13.block.1.conv""",
    """decoder.model.13.block.3.conv.conv""": """decoder.layers.13.block.3.conv""",
    """decoder.model.13.shortcut.conv.conv""": """decoder.layers.13.shortcut.conv""",
    """decoder.model.15.conv.conv""": """decoder.layers.15.conv""",
}
UpperCamelCase__   : Tuple					  = {
    """decoder.model.0.conv.norm""": """decoder.layers.0.norm""",
    """decoder.model.3.convtr.norm""": """decoder.layers.3.norm""",
    """decoder.model.4.block.1.conv.norm""": """decoder.layers.4.block.1.norm""",
    """decoder.model.4.block.3.conv.norm""": """decoder.layers.4.block.3.norm""",
    """decoder.model.4.shortcut.conv.norm""": """decoder.layers.4.shortcut.norm""",
    """decoder.model.6.convtr.norm""": """decoder.layers.6.norm""",
    """decoder.model.7.block.1.conv.norm""": """decoder.layers.7.block.1.norm""",
    """decoder.model.7.block.3.conv.norm""": """decoder.layers.7.block.3.norm""",
    """decoder.model.7.shortcut.conv.norm""": """decoder.layers.7.shortcut.norm""",
    """decoder.model.9.convtr.norm""": """decoder.layers.9.norm""",
    """decoder.model.10.block.1.conv.norm""": """decoder.layers.10.block.1.norm""",
    """decoder.model.10.block.3.conv.norm""": """decoder.layers.10.block.3.norm""",
    """decoder.model.10.shortcut.conv.norm""": """decoder.layers.10.shortcut.norm""",
    """decoder.model.12.convtr.norm""": """decoder.layers.12.norm""",
    """decoder.model.13.block.1.conv.norm""": """decoder.layers.13.block.1.norm""",
    """decoder.model.13.block.3.conv.norm""": """decoder.layers.13.block.3.norm""",
    """decoder.model.13.shortcut.conv.norm""": """decoder.layers.13.shortcut.norm""",
    """decoder.model.15.conv.norm""": """decoder.layers.15.norm""",
}
UpperCamelCase__   : Optional[Any]					  = {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
UpperCamelCase__   : Union[str, Any]					  = {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
UpperCamelCase__   : List[str]					  = []
UpperCamelCase__   : List[Any]					  = []




def    A_(			A				,      A				,      A				,      A				,      A   ):
 for attribute in key.split("""."""   ):
  UpperCAmelCase_								=     getattr(A				,      A   )

 if weight_type is not None:
  UpperCAmelCase_								=     getattr(A				,      A   ).shape
 else:
  UpperCAmelCase_								=     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":
  UpperCAmelCase_								=     value
 elif weight_type == "weight_g":
  UpperCAmelCase_								=     value
 elif weight_type == "weight_v":
  UpperCAmelCase_								=     value
 elif weight_type == "bias":
  UpperCAmelCase_								=     value
 elif weight_type == "running_mean":
  UpperCAmelCase_								=     value
 elif weight_type == "running_var":
  UpperCAmelCase_								=     value
 elif weight_type == "num_batches_tracked":
  UpperCAmelCase_								=     value
 elif weight_type == "weight_ih_l0":
  UpperCAmelCase_								=     value
 elif weight_type == "weight_hh_l0":
  UpperCAmelCase_								=     value
 elif weight_type == "bias_ih_l0":
  UpperCAmelCase_								=     value
 elif weight_type == "bias_hh_l0":
  UpperCAmelCase_								=     value
 elif weight_type == "weight_ih_l1":
  UpperCAmelCase_								=     value
 elif weight_type == "weight_hh_l1":
  UpperCAmelCase_								=     value
 elif weight_type == "bias_ih_l1":
  UpperCAmelCase_								=     value
 elif weight_type == "bias_hh_l1":
  UpperCAmelCase_								=     value
 else:
  UpperCAmelCase_								=     value

 logger.info(f"""{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}."""   )




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




def    A_(			A				,      A				,      A   ):
 UpperCAmelCase_								=     []

 if model_name == "encodec_24khz" or "encodec_32khz":
  UpperCAmelCase_								=     MAPPING_24K
 elif model_name == "encodec_48khz":
  UpperCAmelCase_								=     MAPPING_48K
 else:
  raise ValueError(f"""Unsupported model: {model_name}"""   )

 for name, value in orig_dict.items():
  if should_ignore(A				,      A   ):
   logger.info(f"""{name} was ignored"""   )
   continue

  UpperCAmelCase_								=     False
  for key, mapped_key in MAPPING.items():
   if "*" in key:
    UpperCAmelCase_						,	UpperCAmelCase_								=     key.split(""".*."""   )
    if prefix in name and suffix in name:
     UpperCAmelCase_								=     suffix

   if key in name:
    # HACK otherwise .embed gets initialized with .embed_avg too
    if key.endswith("""embed"""   ) and name.endswith("""embed_avg"""   ):
     continue

    UpperCAmelCase_								=     True
    if "*" in mapped_key:
     UpperCAmelCase_								=     name.split(A   )[0].split("""."""   )[-2]
     UpperCAmelCase_								=     mapped_key.replace("""*"""				,      A   )
    if "weight_g" in name:
     UpperCAmelCase_								=     """weight_g"""
    elif "weight_v" in name:
     UpperCAmelCase_								=     """weight_v"""
    elif "weight_ih_l0" in name:
     UpperCAmelCase_								=     """weight_ih_l0"""
    elif "weight_hh_l0" in name:
     UpperCAmelCase_								=     """weight_hh_l0"""
    elif "bias_ih_l0" in name:
     UpperCAmelCase_								=     """bias_ih_l0"""
    elif "bias_hh_l0" in name:
     UpperCAmelCase_								=     """bias_hh_l0"""
    elif "weight_ih_l1" in name:
     UpperCAmelCase_								=     """weight_ih_l1"""
    elif "weight_hh_l1" in name:
     UpperCAmelCase_								=     """weight_hh_l1"""
    elif "bias_ih_l1" in name:
     UpperCAmelCase_								=     """bias_ih_l1"""
    elif "bias_hh_l1" in name:
     UpperCAmelCase_								=     """bias_hh_l1"""
    elif "bias" in name:
     UpperCAmelCase_								=     """bias"""
    elif "weight" in name:
     UpperCAmelCase_								=     """weight"""
    elif "running_mean" in name:
     UpperCAmelCase_								=     """running_mean"""
    elif "running_var" in name:
     UpperCAmelCase_								=     """running_var"""
    elif "num_batches_tracked" in name:
     UpperCAmelCase_								=     """num_batches_tracked"""
    else:
     UpperCAmelCase_								=     None
    set_recursively(A				,      A				,      A				,      A				,      A   )
   continue
  if not is_used:
   unused_weights.append(A   )

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




@torch.no_grad()
def    A_(			A				,      A				,      A				,      A=None				,      A=None				,      ):
 if config_path is not None:
  UpperCAmelCase_								=     EncodecConfig.from_pretrained(A   )
 else:
  UpperCAmelCase_								=     EncodecConfig()

 if model_name == "encodec_24khz":
  pass  # config is already correct
 elif model_name == "encodec_32khz":
  UpperCAmelCase_								=     [8, 5, 4, 4]
  UpperCAmelCase_								=     [2.2]
  UpperCAmelCase_								=     64
  UpperCAmelCase_								=     32000
  UpperCAmelCase_								=     2048
  UpperCAmelCase_								=     False
  UpperCAmelCase_								=     False
  UpperCAmelCase_								=     False
 elif model_name == "encodec_48khz":
  UpperCAmelCase_								=     [8, 5, 4, 2]
  UpperCAmelCase_								=     [3.0, 6.0, 12.0, 24.0]
  UpperCAmelCase_								=     48000
  UpperCAmelCase_								=     2
  UpperCAmelCase_								=     False
  UpperCAmelCase_								=     """time_group_norm"""
  UpperCAmelCase_								=     True
  UpperCAmelCase_								=     1.0
  UpperCAmelCase_								=     0.01
 else:
  raise ValueError(f"""Unknown model name: {model_name}"""   )

 UpperCAmelCase_								=     EncodecModel(A   )

 UpperCAmelCase_								=     EncodecFeatureExtractor(
     feature_size=config.audio_channels				,      sampling_rate=config.sampling_rate				,      chunk_length_s=config.chunk_length_s				,      overlap=config.overlap				,      )
 feature_extractor.save_pretrained(A   )

 UpperCAmelCase_								=     torch.load(A   )
 if "best_state" in original_checkpoint:
  # we might have a training state saved, in which case discard the yaml results and just retain the weights
  UpperCAmelCase_								=     original_checkpoint["""best_state"""]
 recursively_load_weights(A				,      A				,      A   )
 model.save_pretrained(A   )

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


if __name__ == "__main__":
    UpperCamelCase__   : Dict					  = argparse.ArgumentParser()
    parser.add_argument(
        """--model""",
        default="""encodec_24khz""",
        type=str,
        help="""The model to convert. Should be one of 'encodec_24khz', 'encodec_32khz', 'encodec_48khz'.""",
    )
    parser.add_argument("""--checkpoint_path""", required=True, default=None, type=str, help="""Path to original checkpoint""")
    parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
    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."""
    )

    UpperCamelCase__   : Tuple					  = parser.parse_args()
    convert_checkpoint(
        args.model,
        args.checkpoint_path,
        args.pytorch_dump_folder_path,
        args.config_path,
        args.push_to_hub,
    )









from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import (
    BaseOutput,
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_onnx_available,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)







@dataclass
class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					lowerCamelCase     :     Union[List[PIL.Image.Image], np.ndarray]
					lowerCamelCase     :     Optional[List[bool]]


try:
				if not (is_transformers_available() and is_torch_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
				from .pipeline_cycle_diffusion import CycleDiffusionPipeline
				from .pipeline_stable_diffusion import StableDiffusionPipeline
				from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
				from .pipeline_stable_diffusion_imgaimg import StableDiffusionImgaImgPipeline
				from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
				from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
				from .pipeline_stable_diffusion_instruct_pixapix import StableDiffusionInstructPixaPixPipeline
				from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
				from .pipeline_stable_diffusion_ldmad import StableDiffusionLDMaDPipeline
				from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
				from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
				from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
				from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
				from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
				from .pipeline_stable_unclip import StableUnCLIPPipeline
				from .pipeline_stable_unclip_imgaimg import StableUnCLIPImgaImgPipeline
				from .safety_checker import StableDiffusionSafetyChecker
				from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

try:
				if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
else:
				from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline


try:
				if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.26.0""")):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import (
				    StableDiffusionDepthaImgPipeline,
				    StableDiffusionDiffEditPipeline,
				    StableDiffusionPixaPixZeroPipeline,
				)
else:
				from .pipeline_stable_diffusion_depthaimg import StableDiffusionDepthaImgPipeline
				from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
				from .pipeline_stable_diffusion_pixapix_zero import StableDiffusionPixaPixZeroPipeline


try:
				if not (
				    is_torch_available()
				    and is_transformers_available()
				    and is_k_diffusion_available()
				    and is_k_diffusion_version(""">=""", """0.0.12""")
				):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
				from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline

try:
				if not (is_transformers_available() and is_onnx_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_onnx_objects import *  # noqa F403
else:
				from .pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline, StableDiffusionOnnxPipeline
				from .pipeline_onnx_stable_diffusion_imgaimg import OnnxStableDiffusionImgaImgPipeline
				from .pipeline_onnx_stable_diffusion_inpaint import OnnxStableDiffusionInpaintPipeline
				from .pipeline_onnx_stable_diffusion_inpaint_legacy import OnnxStableDiffusionInpaintPipelineLegacy
				from .pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

if is_transformers_available() and is_flax_available():
				import flax







				@flax.struct.dataclass
				class   _UpperCamelCase  (       A_  ):







									'''simple docstring'''

									lowerCamelCase     :     np.ndarray
									lowerCamelCase     :     List[bool]


				from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
				from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline
				from .pipeline_flax_stable_diffusion_imgaimg import FlaxStableDiffusionImgaImgPipeline
				from .pipeline_flax_stable_diffusion_inpaint import FlaxStableDiffusionInpaintPipeline
				from .safety_checker_flax import FlaxStableDiffusionSafetyChecker


