

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__								=		logging.get_logger("""transformers.models.encodec""")

lowerCAmelCase__								=		{
    """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__								=		{
    """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__								=		{
    """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__								=		{
    """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__								=		{
    """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__								=		{
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
lowerCAmelCase__								=		{
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
lowerCAmelCase__								=		[]
lowerCAmelCase__								=		[]






def 					lowerCamelCase_     (							UpperCAmelCase_							:     Tuple							,       UpperCAmelCase_							:     Optional[int]							,       UpperCAmelCase_							:     List[str]							,       UpperCAmelCase_							:     Optional[int]							,       UpperCAmelCase_							:     Union[str, Any]	):




	'''simple docstring'''




	for attribute in key.split('.'	):
		_UpperCamelCase   :		Optional[int]			=						getattr(__UpperCamelCase							,       __UpperCamelCase	)

	if weight_type is not None:
		_UpperCamelCase   :		List[Any]			=						getattr(__UpperCamelCase							,       __UpperCamelCase	).shape
	else:
		_UpperCamelCase   :		Tuple			=						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   :		List[Any]			=						value
	elif weight_type == "weight_g":
		_UpperCamelCase   :		Dict			=						value
	elif weight_type == "weight_v":
		_UpperCamelCase   :		List[str]			=						value
	elif weight_type == "bias":
		_UpperCamelCase   :		Any			=						value
	elif weight_type == "running_mean":
		_UpperCamelCase   :		str			=						value
	elif weight_type == "running_var":
		_UpperCamelCase   :		Dict			=						value
	elif weight_type == "num_batches_tracked":
		_UpperCamelCase   :		Union[str, Any]			=						value
	elif weight_type == "weight_ih_l0":
		_UpperCamelCase   :		Any			=						value
	elif weight_type == "weight_hh_l0":
		_UpperCamelCase   :		List[Any]			=						value
	elif weight_type == "bias_ih_l0":
		_UpperCamelCase   :		Optional[Any]			=						value
	elif weight_type == "bias_hh_l0":
		_UpperCamelCase   :		Any			=						value
	elif weight_type == "weight_ih_l1":
		_UpperCamelCase   :		Any			=						value
	elif weight_type == "weight_hh_l1":
		_UpperCamelCase   :		Optional[Any]			=						value
	elif weight_type == "bias_ih_l1":
		_UpperCamelCase   :		int			=						value
	elif weight_type == "bias_hh_l1":
		_UpperCamelCase   :		Union[str, Any]			=						value
	else:
		_UpperCamelCase   :		Tuple			=						value

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






def 					lowerCamelCase_     (							UpperCAmelCase_							:     int							,       UpperCAmelCase_							:     Any	):




	'''simple docstring'''




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






def 					lowerCamelCase_     (							UpperCAmelCase_							:     Tuple							,       UpperCAmelCase_							:     str							,       UpperCAmelCase_							:     Tuple	):




	'''simple docstring'''




	_UpperCamelCase   :		Dict			=						[]

	if model_name == "encodec_24khz" or "encodec_32khz":
		_UpperCamelCase   :		str			=						MAPPING_24K
	elif model_name == "encodec_48khz":
		_UpperCamelCase   :		str			=						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

		_UpperCamelCase   :		List[str]			=						False
		for key, mapped_key in MAPPING.items():
			if "*" in key:
				_UpperCamelCase	,		_UpperCamelCase   :		str			=						key.split('.*.'	)
				if prefix in name and suffix in name:
					_UpperCamelCase   :		Tuple			=						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   :		str			=						True
				if "*" in mapped_key:
					_UpperCamelCase   :		List[str]			=						name.split(__UpperCamelCase	)[0].split('.'	)[-2]
					_UpperCamelCase   :		Any			=						mapped_key.replace('*'							,       __UpperCamelCase	)
				if "weight_g" in name:
					_UpperCamelCase   :		Union[str, Any]			=						'weight_g'
				elif "weight_v" in name:
					_UpperCamelCase   :		Dict			=						'weight_v'
				elif "weight_ih_l0" in name:
					_UpperCamelCase   :		str			=						'weight_ih_l0'
				elif "weight_hh_l0" in name:
					_UpperCamelCase   :		str			=						'weight_hh_l0'
				elif "bias_ih_l0" in name:
					_UpperCamelCase   :		Tuple			=						'bias_ih_l0'
				elif "bias_hh_l0" in name:
					_UpperCamelCase   :		Dict			=						'bias_hh_l0'
				elif "weight_ih_l1" in name:
					_UpperCamelCase   :		Dict			=						'weight_ih_l1'
				elif "weight_hh_l1" in name:
					_UpperCamelCase   :		Tuple			=						'weight_hh_l1'
				elif "bias_ih_l1" in name:
					_UpperCamelCase   :		List[str]			=						'bias_ih_l1'
				elif "bias_hh_l1" in name:
					_UpperCamelCase   :		Union[str, Any]			=						'bias_hh_l1'
				elif "bias" in name:
					_UpperCamelCase   :		str			=						'bias'
				elif "weight" in name:
					_UpperCamelCase   :		Union[str, Any]			=						'weight'
				elif "running_mean" in name:
					_UpperCamelCase   :		Optional[Any]			=						'running_mean'
				elif "running_var" in name:
					_UpperCamelCase   :		Optional[int]			=						'running_var'
				elif "num_batches_tracked" in name:
					_UpperCamelCase   :		str			=						'num_batches_tracked'
				else:
					_UpperCamelCase   :		Tuple			=						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_							:     Optional[Any]							,       UpperCAmelCase_							:     Any							,       UpperCAmelCase_							:     Optional[Any]							,       UpperCAmelCase_							:     Optional[int]=None							,       UpperCAmelCase_							:     Optional[int]=None							,       ):




	'''simple docstring'''




	if config_path is not None:
		_UpperCamelCase   :		Union[str, Any]			=						EncodecConfig.from_pretrained(__UpperCamelCase	)
	else:
		_UpperCamelCase   :		Any			=						EncodecConfig()

	if model_name == "encodec_24khz":
		pass  # config is already correct
	elif model_name == "encodec_32khz":
		_UpperCamelCase   :		Optional[int]			=						[8, 5, 4, 4]
		_UpperCamelCase   :		List[Any]			=						[2.2]
		_UpperCamelCase   :		Any			=						6_4
		_UpperCamelCase   :		Optional[int]			=						3_2_0_0_0
		_UpperCamelCase   :		Union[str, Any]			=						2_0_4_8
		_UpperCamelCase   :		Dict			=						False
		_UpperCamelCase   :		List[Any]			=						False
		_UpperCamelCase   :		List[Any]			=						False
	elif model_name == "encodec_48khz":
		_UpperCamelCase   :		List[str]			=						[8, 5, 4, 2]
		_UpperCamelCase   :		List[str]			=						[3.0, 6.0, 1_2.0, 2_4.0]
		_UpperCamelCase   :		List[str]			=						4_8_0_0_0
		_UpperCamelCase   :		Any			=						2
		_UpperCamelCase   :		List[Any]			=						False
		_UpperCamelCase   :		Dict			=						'time_group_norm'
		_UpperCamelCase   :		Any			=						True
		_UpperCamelCase   :		str			=						1.0
		_UpperCamelCase   :		List[str]			=						0.0_1
	else:
		raise ValueError(F'''Unknown model name: {model_name}'''	)

	_UpperCamelCase   :		Tuple			=						EncodecModel(__UpperCamelCase	)

	_UpperCamelCase   :		Tuple			=						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	)

	_UpperCamelCase   :		List[str]			=						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
		_UpperCamelCase   :		List[str]			=						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__":
							lowerCAmelCase__								=		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__								=		parser.parse_args()
							convert_checkpoint(
							    args.model,
							    args.checkpoint_path,
							    args.pytorch_dump_folder_path,
							    args.config_path,
							    args.push_to_hub,
							)







import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
							from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image







class 	lowercase    :


		"""simple docstring"""




		def __init__(				self	,       __snake_case	,       __snake_case=3	,       __snake_case=32	,       __snake_case=3	,       __snake_case=10	,       __snake_case=[8, 16, 32, 64]	,       __snake_case=[1, 1, 2, 1]	,       __snake_case=True	,       __snake_case=True	,       __snake_case="relu"	,       __snake_case=3	,       __snake_case=None	,       __snake_case=["stage2", "stage3", "stage4"]	,       __snake_case=[2, 3, 4]	,       __snake_case=1	,       ):
			_UpperCamelCase   :		List[Any]			=						parent
			_UpperCamelCase   :		Dict			=						batch_size
			_UpperCamelCase   :		Optional[int]			=						image_size
			_UpperCamelCase   :		str			=						num_channels
			_UpperCamelCase   :		Optional[Any]			=						embeddings_size
			_UpperCamelCase   :		Tuple			=						hidden_sizes
			_UpperCamelCase   :		Dict			=						depths
			_UpperCamelCase   :		str			=						is_training
			_UpperCamelCase   :		Optional[int]			=						use_labels
			_UpperCamelCase   :		str			=						hidden_act
			_UpperCamelCase   :		Optional[int]			=						num_labels
			_UpperCamelCase   :		Optional[int]			=						scope
			_UpperCamelCase   :		Tuple			=						len(__snake_case)
			_UpperCamelCase   :		Dict			=						out_features
			_UpperCamelCase   :		Union[str, Any]			=						out_indices
			_UpperCamelCase   :		int			=						num_groups


		def 		A__					(				self):
			_UpperCamelCase   :		Tuple			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

			_UpperCamelCase   :		str			=						None
			if self.use_labels:
				_UpperCamelCase   :		Optional[int]			=						ids_tensor([self.batch_size]	,       self.num_labels)

			_UpperCamelCase   :		str			=						self.get_config()

			return config, pixel_values, labels


		def 		A__					(				self):
			return BitConfig(
			    num_channels=self.num_channels	,       embeddings_size=self.embeddings_size	,       hidden_sizes=self.hidden_sizes	,       depths=self.depths	,       hidden_act=self.hidden_act	,       num_labels=self.num_labels	,       out_features=self.out_features	,       out_indices=self.out_indices	,       num_groups=self.num_groups	,       )


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		str			=						BitModel(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[Any]			=						model(__snake_case)
			self.parent.assertEqual(
			    result.last_hidden_state.shape	,       (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)	,       )


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		Dict			=						self.num_labels
			_UpperCamelCase   :		Dict			=						BitForImageClassification(__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Dict			=						model(__snake_case	,       labels=__snake_case)
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.num_labels))


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		Optional[Any]			=						BitBackbone(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		List[Any]			=						model(__snake_case)

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps)	,       len(config.out_features))
			self.parent.assertListEqual(list(result.feature_maps[0].shape)	,       [self.batch_size, self.hidden_sizes[1], 4, 4])

			# verify channels
			self.parent.assertEqual(len(model.channels)	,       len(config.out_features))
			self.parent.assertListEqual(model.channels	,       config.hidden_sizes[1:])

			# verify backbone works with out_features=None
			_UpperCamelCase   :		Any			=						None
			_UpperCamelCase   :		str			=						BitBackbone(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Any			=						model(__snake_case)

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps)	,       1)
			self.parent.assertListEqual(list(result.feature_maps[0].shape)	,       [self.batch_size, self.hidden_sizes[-1], 1, 1])

			# verify channels
			self.parent.assertEqual(len(model.channels)	,       1)
			self.parent.assertListEqual(model.channels	,       [config.hidden_sizes[-1]])






		def 		A__					(				self):
			_UpperCamelCase   :		Optional[int]			=						self.prepare_config_and_inputs()
			_UpperCamelCase	,		_UpperCamelCase	,		_UpperCamelCase   :		int			=						config_and_inputs
			_UpperCamelCase   :		int			=						{'pixel_values': pixel_values}
			return config, inputs_dict







@require_torch
class 	lowercase    (				_lowercase    ,							_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
		a__					       =  (
		    {"feature-extraction": BitModel, "image-classification": BitForImageClassification}
		    if is_torch_available()
		    else {}
		)

		a__					       =  False
		a__					       =  False
		a__					       =  False
		a__					       =  False
		a__					       =  False


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						BitModelTester(self)
			_UpperCamelCase   :		Optional[Any]			=						ConfigTester(self	,       config_class=__snake_case	,       has_text_modality=__snake_case)


		def 		A__					(				self):
			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()


		def 		A__					(				self):
			return


		@unittest.skip(reason='Bit does not output attentions')
		def 		A__					(				self):
			pass


		@unittest.skip(reason='Bit does not use inputs_embeds')
		def 		A__					(				self):
			pass


		@unittest.skip(reason='Bit does not support input and output embeddings')
		def 		A__					(				self):
			pass


		def 		A__					(				self):
			_UpperCamelCase	,		_UpperCamelCase   :		List[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				_UpperCamelCase   :		int			=						model_class(__snake_case)
				_UpperCamelCase   :		List[Any]			=						inspect.signature(model.forward)
				# signature.parameters is an OrderedDict => so arg_names order is deterministic
				_UpperCamelCase   :		Optional[int]			=						[*signature.parameters.keys()]

				_UpperCamelCase   :		List[Any]			=						['pixel_values']
				self.assertListEqual(arg_names[:1]	,       __snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Tuple			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_backbone(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase	,		_UpperCamelCase   :		Optional[int]			=						self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				_UpperCamelCase   :		Union[str, Any]			=						model_class(config=__snake_case)
				for name, module in model.named_modules():
					if isinstance(__snake_case	,       (nn.BatchNormad, nn.GroupNorm)):
						self.assertTrue(
						    torch.all(module.weight == 1)	,       msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''	,       )
						self.assertTrue(
						    torch.all(module.bias == 0)	,       msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''	,       )


		def 		A__					(				self):
			def check_hidden_states_output(__snake_case	,       __snake_case	,       __snake_case):
				_UpperCamelCase   :		str			=						model_class(__snake_case)
				model.to(__snake_case)
				model.eval()

				with torch.no_grad():
					_UpperCamelCase   :		Union[str, Any]			=						model(**self._prepare_for_class(__snake_case	,       __snake_case))

				_UpperCamelCase   :		List[Any]			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

				_UpperCamelCase   :		str			=						self.model_tester.num_stages
				self.assertEqual(len(__snake_case)	,       expected_num_stages + 1)

				# Bit's feature maps are of shape (batch_size, num_channels, height, width)
				self.assertListEqual(
				    list(hidden_states[0].shape[-2:])	,       [self.model_tester.image_size // 4, self.model_tester.image_size // 4]	,       )

			_UpperCamelCase	,		_UpperCamelCase   :		int			=						self.model_tester.prepare_config_and_inputs_for_common()
			_UpperCamelCase   :		List[str]			=						['preactivation', 'bottleneck']
			for model_class in self.all_model_classes:
				for layer_type in layers_type:
					_UpperCamelCase   :		Any			=						layer_type
					_UpperCamelCase   :		Tuple			=						True
					check_hidden_states_output(__snake_case	,       __snake_case	,       __snake_case)

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					_UpperCamelCase   :		List[str]			=						True

					check_hidden_states_output(__snake_case	,       __snake_case	,       __snake_case)


		@unittest.skip(reason='Bit does not use feedforward chunking')
		def 		A__					(				self):
			pass


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*__snake_case)






		@slow
		def 		A__					(				self):
			for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
				_UpperCamelCase   :		Optional[Any]			=						BitModel.from_pretrained(__snake_case)
				self.assertIsNotNone(__snake_case)






def 					lowerCamelCase_     (							)						->				Optional[int]:




	'''simple docstring'''




	_UpperCamelCase   :		List[Any]			=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
	return image







@require_torch
@require_vision
class 	lowercase    (				unittest.TestCase							):


		"""simple docstring"""




		@cached_property
		def 		A__					(				self):
			return (
			    BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]) if is_vision_available() else None
			)






		@slow
		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(__snake_case)

			_UpperCamelCase   :		str			=						self.default_image_processor
			_UpperCamelCase   :		List[str]			=						prepare_img()
			_UpperCamelCase   :		int			=						image_processor(images=__snake_case	,       return_tensors='pt').to(__snake_case)

			# forward pass
			with torch.no_grad():
				_UpperCamelCase   :		Any			=						model(**__snake_case)

			# verify the logits
			_UpperCamelCase   :		Dict			=						torch.Size((1, 10_00))
			self.assertEqual(outputs.logits.shape	,       __snake_case)

			_UpperCamelCase   :		Optional[int]			=						torch.tensor([[-0.6_5_2_6, -0.5_2_6_3, -1.4_3_9_8]]).to(__snake_case)

			self.assertTrue(torch.allclose(outputs.logits[0, :3]	,       __snake_case	,       atol=1e-4))







@require_torch
class 	lowercase    (				_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  (BitBackbone,) if is_torch_available() else ()
		a__					       =  BitConfig

		a__					       =  False


		def 		A__					(				self):
			_UpperCamelCase   :		List[str]			=						BitModelTester(self)




