




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

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



def        a__					(							snake_case__			,					snake_case__			,					snake_case__			,					snake_case__			,					snake_case__							)		->   Optional[int]:
			for attribute in key.split("""."""							):
						lowerCamelCase      =   getattr(snake_case__			,					snake_case__							)

			if weight_type is not None:
						lowerCamelCase      =   getattr(snake_case__			,					snake_case__							).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        a__					(							snake_case__			,					snake_case__							)		->   List[str]:
			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        a__					(							snake_case__			,					snake_case__			,					snake_case__							)		->   Optional[Any]:
			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(snake_case__			,					snake_case__							):
									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(snake_case__							)[0].split("""."""							)[-2]
															lowerCamelCase      =   mapped_key.replace("""*"""			,					snake_case__							)
												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(snake_case__			,					snake_case__			,					snake_case__			,					snake_case__			,					snake_case__							)
									continue
						if not is_used:
									unused_weights.append(snake_case__							)

			logger.warning(F'Unused weights: {unused_weights}'							)



@torch.no_grad()
def        a__					(							snake_case__			,					snake_case__			,					snake_case__			,					snake_case__=None			,					snake_case__=None			,					)		->   Union[str, Any]:
			if config_path is not None:
						lowerCamelCase      =   EncodecConfig.from_pretrained(snake_case__							)
			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_20_00
						lowerCamelCase      =   20_48
						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_80_00
						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(snake_case__							)

			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(snake_case__							)

			lowerCamelCase      =   torch.load(snake_case__							)
			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(snake_case__			,					snake_case__			,					snake_case__							)
			model.save_pretrained(snake_case__							)

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


if __name__ == "__main__":
					lowerCAmelCase				:     List[Any]					=				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."""
					)

					lowerCAmelCase				:     List[Any]					=				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 AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class        __magic_name__				(			UpperCAmelCase__		,   unittest.TestCase							):






				'''simple docstring'''
				__UpperCamelCase         =		DiTPipeline
				__UpperCamelCase         =		CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
				__UpperCamelCase         =		PipelineTesterMixin.required_optional_params - {
				    "latents",
				    "num_images_per_prompt",
				    "callback",
				    "callback_steps",
				}
				__UpperCamelCase         =		CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
				__UpperCamelCase         =		False





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							torch.manual_seed(0					)
							lowerCamelCase      =   TransformeraDModel(
							    sample_size=16						,					num_layers=2						,					patch_size=4						,					attention_head_dim=8						,					num_attention_heads=2						,					in_channels=4						,					out_channels=8						,					attention_bias=_a						,					activation_fn="""gelu-approximate"""						,					num_embeds_ada_norm=1_000						,					norm_type="""ada_norm_zero"""						,					norm_elementwise_affine=_a						,					)
							lowerCamelCase      =   AutoencoderKL()
							lowerCamelCase      =   DDIMScheduler()
							lowerCamelCase      =   {"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
							return components





				def   _lowerCAmelCase							(			self						,					_a						,					_a=0					):





							"""simple docstring"""
							if str(_a					).startswith("""mps"""					):
										lowerCamelCase      =   torch.manual_seed(_a					)
							else:
										lowerCamelCase      =   torch.Generator(device=_a					).manual_seed(_a					)
							lowerCamelCase      =   {
							    """class_labels""": [1],
							    """generator""": generator,
							    """num_inference_steps""": 2,
							    """output_type""": """numpy""",
							}
							return inputs





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   """cpu"""

							lowerCamelCase      =   self.get_dummy_components()
							lowerCamelCase      =   self.pipeline_class(**_a					)
							pipe.to(_a					)
							pipe.set_progress_bar_config(disable=_a					)

							lowerCamelCase      =   self.get_dummy_inputs(_a					)
							lowerCamelCase      =   pipe(**_a					).images
							lowerCamelCase      =   image[0, -3:, -3:, -1]

							self.assertEqual(image.shape						,					(1, 16, 16, 3)					)
							lowerCamelCase      =   np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]					)
							lowerCamelCase      =   np.abs(image_slice.flatten() - expected_slice					).max()
							self.assertLessEqual(_a						,					1e-3					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							self._test_inference_batch_single_identical(relax_max_difference=_a						,					expected_max_diff=1e-3					)





				@unittest.skipIf(
				    torch_device != """cuda""" or not is_xformers_available()						,					reason="""XFormers attention is only available with CUDA and `xformers` installed"""						,					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3					)






@require_torch_gpu
@slow
class        __magic_name__				(			unittest.TestCase							):






				'''simple docstring'''





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   torch.manual_seed(0					)

							lowerCamelCase      =   DiTPipeline.from_pretrained("""facebook/DiT-XL-2-256"""					)
							pipe.to("""cuda"""					)

							lowerCamelCase      =   ["""vase""", """umbrella""", """white shark""", """white wolf"""]
							lowerCamelCase      =   pipe.get_label_ids(_a					)

							lowerCamelCase      =   pipe(_a						,					generator=_a						,					num_inference_steps=40						,					output_type="""np"""					).images

							for word, image in zip(_a						,					_a					):
										lowerCamelCase      =   load_numpy(
										    f'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'					)
										assert np.abs((expected_image - image).max()					) < 1e-2





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""					)
							lowerCamelCase      =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config					)
							pipe.to("""cuda"""					)

							lowerCamelCase      =   ["""vase""", """umbrella"""]
							lowerCamelCase      =   pipe.get_label_ids(_a					)

							lowerCamelCase      =   torch.manual_seed(0					)
							lowerCamelCase      =   pipe(_a						,					generator=_a						,					num_inference_steps=25						,					output_type="""np"""					).images

							for word, image in zip(_a						,					_a					):
										lowerCamelCase      =   load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
										    f'/dit/{word}_512.npy'					)

										assert np.abs((expected_image - image).max()					) < 1e-1



