


from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import tensorflow as tf

	from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
	from PIL import Image

	from transformers import ViTImageProcessor




class       lowercase_					:





				def __init__(				self       ,  lowercase_       ,  lowercase_=13       ,  lowercase_=30       ,  lowercase_=2       ,  lowercase_=3       ,  lowercase_=True       ,  lowercase_=True       ,  lowercase_=32       ,  lowercase_=2       ,  lowercase_=4       ,  lowercase_=37       ,  lowercase_="gelu"       ,  lowercase_=0.1       ,  lowercase_=0.1       ,  lowercase_=10       ,  lowercase_=0.02       ,  lowercase_=3       ,  lowercase_=0.6       ,  lowercase_=None       ,  )  -> Dict:
											a__      =parent
											a__      =batch_size
											a__      =image_size
											a__      =patch_size
											a__      =num_channels
											a__      =is_training
											a__      =use_labels
											a__      =hidden_size
											a__      =num_hidden_layers
											a__      =num_attention_heads
											a__      =intermediate_size
											a__      =hidden_act
											a__      =hidden_dropout_prob
											a__      =attention_probs_dropout_prob
											a__      =type_sequence_label_size
											a__      =initializer_range
											a__      =mask_ratio
											a__      =scope

											# in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
											# (we add 1 for the [CLS] token)
											a__      =(image_size // patch_size) ** 2
											a__      =int(math.ceil((1 - mask_ratio) * (num_patches + 1)))





				def 						__UpperCamelCase    (				self)  -> List[Any]:
											a__      =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

											a__      =None
											if self.use_labels:
																		a__      =ids_tensor([self.batch_size]       ,  self.type_sequence_label_size)

											a__      =self.get_config()

											return config, pixel_values, labels





				def 						__UpperCamelCase    (				self)  -> Dict:
											return ViTMAEConfig(
											    image_size=self.image_size       ,  patch_size=self.patch_size       ,  num_channels=self.num_channels       ,  hidden_size=self.hidden_size       ,  num_hidden_layers=self.num_hidden_layers       ,  num_attention_heads=self.num_attention_heads       ,  intermediate_size=self.intermediate_size       ,  decoder_hidden_size=self.hidden_size       ,  decoder_num_hidden_layers=self.num_hidden_layers       ,  decoder_num_attention_heads=self.num_attention_heads       ,  decoder_intermediate_size=self.intermediate_size       ,  hidden_act=self.hidden_act       ,  hidden_dropout_prob=self.hidden_dropout_prob       ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,  is_decoder=lowercase_       ,  initializer_range=self.initializer_range       ,  mask_ratio=self.mask_ratio       ,  )





				def 						__UpperCamelCase    (				self       ,  lowercase_       ,  lowercase_       ,  lowercase_)  -> Any:
											a__      =TFViTMAEModel(config=lowercase_)
											a__      =model(lowercase_       ,  training=lowercase_)
											self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size))





				def 						__UpperCamelCase    (				self       ,  lowercase_       ,  lowercase_       ,  lowercase_)  -> Optional[int]:
											a__      =TFViTMAEForPreTraining(lowercase_)
											a__      =model(lowercase_       ,  training=lowercase_)
											# expected sequence length = num_patches
											a__      =(self.image_size // self.patch_size) ** 2
											a__      =self.patch_size**2 * self.num_channels
											self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, num_patches, expected_num_channels))

											# test greyscale images
											a__      =1
											a__      =TFViTMAEForPreTraining(lowercase_)

											a__      =floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
											a__      =model(lowercase_       ,  training=lowercase_)
											a__      =self.patch_size**2
											self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, num_patches, expected_num_channels))





				def 						__UpperCamelCase    (				self)  -> List[str]:
											a__      =self.prepare_config_and_inputs()
											((a__)    ,     (a__)    ,     (a__))      =config_and_inputs
											a__      ={'pixel_values': pixel_values}
											return config, inputs_dict







@require_tf
class       lowercase_					(lowercase__      ,					lowercase__      ,					unittest.TestCase   ):
				snake_case	    =(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
				snake_case	    ={'feature-extraction': TFViTMAEModel} if is_tf_available() else {}

				snake_case	    =False
				snake_case	    =False
				snake_case	    =False
				snake_case	    =False





				def 						__UpperCamelCase    (				self)  -> Dict:
											a__      =TFViTMAEModelTester(self)
											a__      =ConfigTester(self       ,  config_class=lowercase_       ,  has_text_modality=lowercase_       ,  hidden_size=37)





				def 						__UpperCamelCase    (				self)  -> Tuple:
											self.config_tester.run_common_tests()





				@unittest.skip(reason='ViTMAE does not use inputs_embeds')
				def 						__UpperCamelCase    (				self)  -> Optional[int]:
											pass





				def 						__UpperCamelCase    (				self)  -> List[Any]:
											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																		a__      =model_class(lowercase_)
																		self.assertIsInstance(model.get_input_embeddings()       ,  (tf.keras.layers.Layer))
																		a__      =model.get_output_embeddings()
																		self.assertTrue(x is None or isinstance(lowercase_       ,  tf.keras.layers.Layer))





				def 						__UpperCamelCase    (				self)  -> Dict:
											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																		a__      =model_class(lowercase_)
																		a__      =inspect.signature(model.call)
																		# signature.parameters is an OrderedDict => so arg_names order is deterministic
																		a__      =[*signature.parameters.keys()]

																		a__      =['pixel_values']
																		self.assertListEqual(arg_names[:1]       ,  lowercase_)





				def 						__UpperCamelCase    (				self)  -> List[Any]:
											a__      =self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*lowercase_)





				def 						__UpperCamelCase    (				self)  -> Union[str, Any]:
											a__      =self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_pretraining(*lowercase_)





				def 						__UpperCamelCase    (				self)  -> Any:
											# make the mask reproducible
											np.random.seed(2)

											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											a__      =int((config.image_size // config.patch_size) ** 2)
											a__      =np.random.uniform(size=(self.model_tester.batch_size, num_patches))

											for model_class in self.all_model_classes:
																		a__      =model_class(lowercase_)
																		a__      =self._prepare_for_class(lowercase_       ,  lowercase_)

																		a__      =model(lowercase_       ,  noise=lowercase_)

																		a__      =copy.deepcopy(self._prepare_for_class(lowercase_       ,  lowercase_))
																		a__      =model(**lowercase_       ,  noise=lowercase_)
																		a__      =outputs_dict[0].numpy()
																		a__      =outputs_keywords[0].numpy()

																		self.assertLess(np.sum(np.abs(output_dict - output_keywords))       ,  1e-6)





				def 						__UpperCamelCase    (				self)  -> Any:
											# make the mask reproducible
											np.random.seed(2)

											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											a__      =int((config.image_size // config.patch_size) ** 2)
											a__      =np.random.uniform(size=(self.model_tester.batch_size, num_patches))

											def prepare_numpy_arrays(lowercase_):
																		a__      ={}
																		for k, v in inputs_dict.items():
																									if tf.is_tensor(lowercase_):
																																a__      =v.numpy()
																									else:
																																a__      =np.array(lowercase_)

																		return inputs_np_dict

											for model_class in self.all_model_classes:
																		a__      =model_class(lowercase_)

																		a__      =self._prepare_for_class(lowercase_       ,  lowercase_)
																		a__      =prepare_numpy_arrays(lowercase_)

																		a__      =model(lowercase_       ,  noise=lowercase_)
																		a__      =model(**lowercase_       ,  noise=lowercase_)
																		self.assert_outputs_same(lowercase_       ,  lowercase_)





				def 						__UpperCamelCase    (				self       ,  lowercase_       ,  lowercase_       ,  lowercase_)  -> Union[str, Any]:
											# make masks reproducible
											np.random.seed(2)

											a__      =int((tf_model.config.image_size // tf_model.config.patch_size) ** 2)
											a__      =np.random.uniform(size=(self.model_tester.batch_size, num_patches))
											a__      =tf.constant(lowercase_)

											# Add `noise` argument.
											# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
											a__      =tf_noise

											super().check_pt_tf_models(lowercase_       ,  lowercase_       ,  lowercase_)





				def 						__UpperCamelCase    (				self)  -> List[Any]:
											# make mask reproducible
											np.random.seed(2)

											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											a__      ={
											    module_member
											    for model_class in self.all_model_classes
											    for module in (import_module(model_class.__module__),)
											    for module_member_name in dir(lowercase_)
											    if module_member_name.endswith('MainLayer')
											    # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
											    and module_member_name[: -len('MainLayer')] == model_class.__name__[: -len('Model')]
											    for module_member in (getattr(lowercase_       ,  lowercase_),)
											    if isinstance(lowercase_       ,  lowercase_)
											    and tf.keras.layers.Layer in module_member.__bases__
											    and getattr(lowercase_       ,  '_keras_serializable'       ,  lowercase_)
											}

											a__      =int((config.image_size // config.patch_size) ** 2)
											a__      =np.random.uniform(size=(self.model_tester.batch_size, num_patches))
											a__      =tf.convert_to_tensor(lowercase_)
											inputs_dict.update({'noise': noise})

											for main_layer_class in tf_main_layer_classes:
																		a__      =main_layer_class(lowercase_)

																		a__      ={
																		    name: tf.keras.Input(tensor.shape[1:]       ,  dtype=tensor.dtype) for name, tensor in inputs_dict.items()
																		}

																		a__      =tf.keras.Model(lowercase_       ,  outputs=main_layer(lowercase_))
																		a__      =model(lowercase_)

																		with tempfile.TemporaryDirectory() as tmpdirname:
																									a__      =os.path.join(lowercase_       ,  'keras_model.h5')
																									model.save(lowercase_)
																									a__      =tf.keras.models.load_model(
																									    lowercase_       ,  custom_objects={main_layer_class.__name__: main_layer_class})
																									assert isinstance(lowercase_       ,  tf.keras.Model)
																									a__      =model(lowercase_)
																									self.assert_outputs_same(lowercase_       ,  lowercase_)





				@slow
				def 						__UpperCamelCase    (				self)  -> Union[str, Any]:
											# make mask reproducible
											np.random.seed(2)

											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											a__      =int((config.image_size // config.patch_size) ** 2)
											a__      =np.random.uniform(size=(self.model_tester.batch_size, num_patches))

											for model_class in self.all_model_classes:
																		a__      =model_class(lowercase_)
																		a__      =self._prepare_for_class(lowercase_       ,  lowercase_)
																		a__      =model(lowercase_       ,  noise=lowercase_)

																		if model_class.__name__ == "TFViTMAEModel":
																									a__      =outputs.last_hidden_state.numpy()
																									a__      =0
																		else:
																									a__      =outputs.logits.numpy()
																									a__      =0

																		with tempfile.TemporaryDirectory() as tmpdirname:
																									model.save_pretrained(lowercase_       ,  saved_model=lowercase_)
																									a__      =model_class.from_pretrained(lowercase_)
																									a__      =model(lowercase_       ,  noise=lowercase_)

																									if model_class.__name__ == "TFViTMAEModel":
																																a__      =after_outputs['last_hidden_state'].numpy()
																																a__      =0
																									else:
																																a__      =after_outputs['logits'].numpy()
																																a__      =0

																									a__      =np.amax(np.abs(out_a - out_a))
																									self.assertLessEqual(lowercase_       ,  1e-5)





				def 						__UpperCamelCase    (				self)  -> Any:
											# make mask reproducible
											np.random.seed(2)

											a__    ,     a__      =self.model_tester.prepare_config_and_inputs_for_common()

											a__      =int((config.image_size // config.patch_size) ** 2)
											a__      =np.random.uniform(size=(self.model_tester.batch_size, num_patches))

											for model_class in self.all_model_classes:
																		a__      =model_class(lowercase_)
																		a__      =self._prepare_for_class(lowercase_       ,  lowercase_)

																		a__      =model(lowercase_       ,  noise=lowercase_)
																		a__      =model.get_config()
																		# make sure that returned config is jsonifiable, which is required by keras
																		json.dumps(lowercase_)
																		a__      =model_class.from_config(model.get_config())
																		# make sure it also accepts a normal config
																		a__      =model_class.from_config(model.config)
																		a__      =new_model(lowercase_)  # Build model
																		new_model.set_weights(model.get_weights())
																		a__      =new_model(lowercase_       ,  noise=lowercase_)

																		self.assert_outputs_same(lowercase_       ,  lowercase_)





				@unittest.skip(
				    reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.')
				def 						__UpperCamelCase    (				self)  -> Any:
											pass





				@unittest.skip(reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load')
				def 						__UpperCamelCase    (				self)  -> List[str]:
											pass





				@slow
				def 						__UpperCamelCase    (				self)  -> int:
											a__      =TFViTMAEModel.from_pretrained('google/vit-base-patch16-224')
											self.assertIsNotNone(lowercase_)







def   _lowercase(					):
							a__      =Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
							return image







@require_tf
@require_vision
class       lowercase_					(unittest.TestCase   ):





				@cached_property
				def 						__UpperCamelCase    (				self)  -> Dict:
											return ViTImageProcessor.from_pretrained('facebook/vit-mae-base') if is_vision_available() else None





				@slow
				def 						__UpperCamelCase    (				self)  -> Optional[int]:
											# make random mask reproducible across the PT and TF model
											np.random.seed(2)

											a__      =TFViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base')

											a__      =self.default_image_processor
											a__      =prepare_img()
											a__      =image_processor(images=lowercase_       ,  return_tensors='tf')

											# prepare a noise vector that will be also used for testing the TF model
											# (this way we can ensure that the PT and TF models operate on the same inputs)
											a__      =ViTMAEConfig()
											a__      =int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2)
											a__      =np.random.uniform(size=(1, num_patches))

											# forward pass
											a__      =model(**lowercase_       ,  noise=lowercase_)

											# verify the logits
											a__      =tf.convert_to_tensor([1, 196, 768])
											self.assertEqual(outputs.logits.shape       ,  lowercase_)

											a__      =tf.convert_to_tensor(
											    [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]])

											tf.debugging.assert_near(outputs.logits[0, :3, :3]       ,  lowercase_       ,  atol=1e-4)





import argparse

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_controlnet_from_original_ckpt


if __name__ == "__main__":
	_lowerCAmelCase:						int   				=       argparse.ArgumentParser()

	parser.add_argument(
	    '--checkpoint_path', default=None, type=str, required=True, help='Path to the checkpoint to convert.'
	)
	parser.add_argument(
	    '--original_config_file',
	    type=str,
	    required=True,
	    help='The YAML config file corresponding to the original architecture.',
	)
	parser.add_argument(
	    '--num_in_channels',
	    default=None,
	    type=int,
	    help='The number of input channels. If `None` number of input channels will be automatically inferred.',
	)
	parser.add_argument(
	    '--image_size',
	    default=512,
	    type=int,
	    help=(
	        'The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2'
	        ' Base. Use 768 for Stable Diffusion v2.'
	    ),
	)
	parser.add_argument(
	    '--extract_ema',
	    action='store_true',
	    help=(
	        'Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights'
	        ' or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield'
	        ' higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning.'
	    ),
	)
	parser.add_argument(
	    '--upcast_attention',
	    action='store_true',
	    help=(
	        'Whether the attention computation should always be upcasted. This is necessary when running stable'
	        ' diffusion 2.1.'
	    ),
	)
	parser.add_argument(
	    '--from_safetensors',
	    action='store_true',
	    help='If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.',
	)
	parser.add_argument(
	    '--to_safetensors',
	    action='store_true',
	    help='Whether to store pipeline in safetensors format or not.',
	)
	parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the output model.')
	parser.add_argument('--device', type=str, help='Device to use (e.g. cpu, cuda:0, cuda:1, etc.)')






	def   _lowercase(					__a				:						Optional[Any]     ):
								if string == "True":
															return True
								elif string == "False":
															return False
								else:
															raise ValueError(f"""could not parse string as bool {string}"""     )

	parser.add_argument(
	    '--use_linear_projection', help='Override for use linear projection', required=False, type=parse_bool
	)

	parser.add_argument('--cross_attention_dim', help='Override for cross attention_dim', required=False, type=int)

	_lowerCAmelCase:						str   				=       parser.parse_args()

	_lowerCAmelCase:						Tuple   				=       download_controlnet_from_original_ckpt(
	    checkpoint_path=args.checkpoint_path,
	    original_config_file=args.original_config_file,
	    image_size=args.image_size,
	    extract_ema=args.extract_ema,
	    num_in_channels=args.num_in_channels,
	    upcast_attention=args.upcast_attention,
	    from_safetensors=args.from_safetensors,
	    device=args.device,
	    use_linear_projection=args.use_linear_projection,
	    cross_attention_dim=args.cross_attention_dim,
	)







	controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)

