


"""simple docstring"""


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()
A	:  Optional[int]            =     logging.get_logger("transformers.models.encodec")

A	:  str            =     {
    "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",
}
A	:  Optional[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",
}
A	:  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",
}
A	:  Optional[Any]            =     {
    "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",
}
A	:  int            =     {
    "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",
}
A	:  Optional[int]            =     {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
A	:  Union[str, Any]            =     {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
A	:  Tuple            =     []
A	:  Any            =     []


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			for attribute in key.split("." ):
						__lowerCAmelCase							       =     getattr(_UpperCamelCase       ,					_UpperCamelCase )

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

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


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




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


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			__lowerCAmelCase							       =     []

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

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

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

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

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

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


@torch.no_grad()
def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase=None       ,					_UpperCamelCase=None       ,					):

			'''simple docstring'''




			if config_path is not None:
						__lowerCAmelCase							       =     EncodecConfig.from_pretrained(_UpperCamelCase )
			else:
						__lowerCAmelCase							       =     EncodecConfig()

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

			__lowerCAmelCase							       =     EncodecModel(_UpperCamelCase )

			__lowerCAmelCase							       =     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(_UpperCamelCase )

			__lowerCAmelCase							       =     torch.load(_UpperCamelCase )
			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
						__lowerCAmelCase							       =     original_checkpoint["best_state"]
			recursively_load_weights(_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase )
			model.save_pretrained(_UpperCamelCase )

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


if __name__ == "__main__":
	A	:  str            =     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."
	)

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








"""simple docstring"""


import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
@slow
@require_torch_gpu
class     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							def      snake_case							(	self      ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     StableDiffusionKDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"      )
										__lowerCAmelCase							       =     sd_pipe.to(__a      )
										sd_pipe.set_progress_bar_config(disable=__a      )

										sd_pipe.set_scheduler("sample_euler"      )

										__lowerCAmelCase							       =     "A painting of a squirrel eating a burger"
										__lowerCAmelCase							       =     torch.manual_seed(0      )
										__lowerCAmelCase							       =     sd_pipe([prompt]       ,  generator=__a       ,  guidance_scale=9.0       ,  num_inference_steps=20       ,  output_type="np"      )

										__lowerCAmelCase							       =     output.images

										__lowerCAmelCase							       =     image[0, -3:, -3:, -1]

										assert image.shape == (1, 5_12, 5_12, 3)
										__lowerCAmelCase							       =     np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]      )

										assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"      )
										__lowerCAmelCase							       =     sd_pipe.to(__a      )
										sd_pipe.set_progress_bar_config(disable=__a      )

										sd_pipe.set_scheduler("sample_euler"      )

										__lowerCAmelCase							       =     "A painting of a squirrel eating a burger"
										__lowerCAmelCase							       =     torch.manual_seed(0      )
										__lowerCAmelCase							       =     sd_pipe([prompt]       ,  generator=__a       ,  guidance_scale=9.0       ,  num_inference_steps=20       ,  output_type="np"      )

										__lowerCAmelCase							       =     output.images

										__lowerCAmelCase							       =     image[0, -3:, -3:, -1]

										assert image.shape == (1, 5_12, 5_12, 3)
										__lowerCAmelCase							       =     np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]      )

										assert np.abs(image_slice.flatten() - expected_slice      ).max() < 5e-1



							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"      )
										__lowerCAmelCase							       =     sd_pipe.to(__a      )
										sd_pipe.set_progress_bar_config(disable=__a      )

										sd_pipe.set_scheduler("sample_dpmpp_2m"      )

										__lowerCAmelCase							       =     "A painting of a squirrel eating a burger"
										__lowerCAmelCase							       =     torch.manual_seed(0      )
										__lowerCAmelCase							       =     sd_pipe(
										    [prompt]       ,  generator=__a       ,  guidance_scale=7.5       ,  num_inference_steps=15       ,  output_type="np"       ,  use_karras_sigmas=__a       ,  )

										__lowerCAmelCase							       =     output.images

										__lowerCAmelCase							       =     image[0, -3:, -3:, -1]

										assert image.shape == (1, 5_12, 5_12, 3)
										__lowerCAmelCase							       =     np.array(
										    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]      )

										assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2




