




import inspect
import unittest

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

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 transformers import EfficientNetForImageClassification, EfficientNetModel
 from transformers.models.efficientnet.modeling_efficientnet import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor


class       SCREAMING_SNAKE_CASE_							:



 """simple docstring"""


 def __init__(				self		:	Optional[int]				,    lowerCAmelCase		:	List[Any]				,    lowerCAmelCase		:	int=13				,    lowerCAmelCase		:	Union[str, Any]=32				,    lowerCAmelCase		:	List[Any]=3				,    lowerCAmelCase		:	Tuple=[3, 3, 5]				,    lowerCAmelCase		:	str=[32, 16, 24]				,    lowerCAmelCase		:	Dict=[16, 24, 40]				,    lowerCAmelCase		:	Optional[Any]=[1, 1, 2]				,    lowerCAmelCase		:	int=[1, 1, 2]				,    lowerCAmelCase		:	Any=[1, 6, 6]				,    lowerCAmelCase		:	str=True				,    lowerCAmelCase		:	Tuple=True				,    lowerCAmelCase		:	Dict=37				,    lowerCAmelCase		:	Any="gelu"				,    lowerCAmelCase		:	List[str]=10				,    )  ->					Optional[Any]:







      """simple docstring"""







      __UpperCamelCase				:  List[Any]								=			parent
      __UpperCamelCase				:  Any								=			batch_size
      __UpperCamelCase				:  List[str]								=			image_size
      __UpperCamelCase				:  int								=			num_channels
      __UpperCamelCase				:  Tuple								=			kernel_sizes
      __UpperCamelCase				:  Union[str, Any]								=			in_channels
      __UpperCamelCase				:  int								=			out_channels
      __UpperCamelCase				:  List[str]								=			strides
      __UpperCamelCase				:  Optional[Any]								=			num_block_repeats
      __UpperCamelCase				:  Tuple								=			expand_ratios
      __UpperCamelCase				:  Optional[int]								=			is_training
      __UpperCamelCase				:  Dict								=			hidden_act
      __UpperCamelCase				:  Optional[int]								=			num_labels
      __UpperCamelCase				:  List[Any]								=			use_labels






 def 				lowerCamelCase__    (				self		:	Optional[Any]  )  ->					Tuple:







      """simple docstring"""







      __UpperCamelCase				:  Tuple								=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

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

      __UpperCamelCase				:  int								=			self.get_config()
      return config, pixel_values, labels






 def 				lowerCamelCase__    (				self		:	Tuple  )  ->					Any:







      """simple docstring"""







      return EfficientNetConfig(
          num_channels=self.num_channels				,    kernel_sizes=self.kernel_sizes				,    in_channels=self.in_channels				,    out_channels=self.out_channels				,    strides=self.strides				,    num_block_repeats=self.num_block_repeats				,    expand_ratios=self.expand_ratios				,    hidden_act=self.hidden_act				,    num_labels=self.num_labels				,    )






 def 				lowerCamelCase__    (				self		:	Optional[int]				,    lowerCAmelCase		:	int				,    lowerCAmelCase		:	Tuple				,    lowerCAmelCase		:	Optional[Any]  )  ->					Union[str, Any]:







      """simple docstring"""







      __UpperCamelCase				:  int								=			EfficientNetModel(config=lowerCAmelCase  )
      model.to(lowerCAmelCase  )
      model.eval()
      __UpperCamelCase				:  Union[str, Any]								=			model(lowerCAmelCase  )
      # expected last hidden states: B, C, H // 4, W // 4
      self.parent.assertEqual(
          result.last_hidden_state.shape				,    (self.batch_size, config.hidden_dim, self.image_size // 4, self.image_size // 4)				,    )






 def 				lowerCamelCase__    (				self		:	int				,    lowerCAmelCase		:	Union[str, Any]				,    lowerCAmelCase		:	Optional[Any]				,    lowerCAmelCase		:	str  )  ->					Optional[int]:







      """simple docstring"""







      __UpperCamelCase				:  List[Any]								=			EfficientNetForImageClassification(lowerCAmelCase  )
      model.to(lowerCAmelCase  )
      model.eval()
      __UpperCamelCase				:  Optional[int]								=			model(lowerCAmelCase				,    labels=lowerCAmelCase  )
      self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.num_labels)  )






 def 				lowerCamelCase__    (				self		:	Optional[Any]  )  ->					Optional[int]:







      """simple docstring"""







      __UpperCamelCase				:  Optional[int]								=			self.prepare_config_and_inputs()
      __UpperCamelCase				:  str								=			config_and_inputs
      __UpperCamelCase				:  Optional[Any]								=			{"""pixel_values""": pixel_values}
      return config, inputs_dict






@require_torch
class       SCREAMING_SNAKE_CASE_							(					_UpperCamelCase  ,     _UpperCamelCase  ,     unittest.TestCase      ):



 """simple docstring"""


 __magic_name__      :					List[Any]			  =		(EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
 __magic_name__      :					Tuple			  =		(
     {'feature-extraction': EfficientNetModel, 'image-classification': EfficientNetForImageClassification}
     if is_torch_available()
     else {}
 )

 __magic_name__      :					Union[str, Any]			  =		False
 __magic_name__      :					Dict			  =		False
 __magic_name__      :					str			  =		False
 __magic_name__      :					Union[str, Any]			  =		False
 __magic_name__      :					int			  =		False






 def 				lowerCamelCase__    (				self		:	Dict  )  ->					Optional[int]:







      """simple docstring"""







      __UpperCamelCase				:  List[str]								=			EfficientNetModelTester(self  )
      __UpperCamelCase				:  Optional[Any]								=			ConfigTester(
          self				,    config_class=lowerCAmelCase				,    has_text_modality=lowerCAmelCase				,    hidden_size=37  )






 def 				lowerCamelCase__    (				self		:	List[str]  )  ->					int:







      """simple docstring"""







      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 				lowerCamelCase__    (				self		:	Optional[Any]  )  ->					Dict:







      """simple docstring"""







      return






 @unittest.skip(reason="""EfficientNet does not use inputs_embeds"""  )
 def 				lowerCamelCase__    (				self		:	str  )  ->					List[str]:







      """simple docstring"""







      pass






 @unittest.skip(reason="""EfficientNet does not support input and output embeddings"""  )
 def 				lowerCamelCase__    (				self		:	Optional[Any]  )  ->					Any:







      """simple docstring"""







      pass






 @unittest.skip(reason="""EfficientNet does not use feedforward chunking"""  )
 def 				lowerCamelCase__    (				self		:	Tuple  )  ->					Optional[int]:







      """simple docstring"""







      pass






 def 				lowerCamelCase__    (				self		:	Tuple  )  ->					List[str]:







      """simple docstring"""







      __UpperCamelCase				:  Dict								=			self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           __UpperCamelCase				:  List[str]								=			model_class(lowerCAmelCase  )
           __UpperCamelCase				:  str								=			inspect.signature(model.forward  )
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           __UpperCamelCase				:  Tuple								=			[*signature.parameters.keys()]

           __UpperCamelCase				:  Union[str, Any]								=			["""pixel_values"""]
           self.assertListEqual(arg_names[:1]				,    lowerCAmelCase  )






 def 				lowerCamelCase__    (				self		:	Any  )  ->					Optional[int]:







      """simple docstring"""







      __UpperCamelCase				:  Optional[int]								=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCAmelCase  )






 def 				lowerCamelCase__    (				self		:	Tuple  )  ->					List[Any]:







      """simple docstring"""







      def check_hidden_states_output(lowerCAmelCase		:	Optional[Any]				,    lowerCAmelCase		:	str				,    lowerCAmelCase		:	Any  ):
           __UpperCamelCase				:  int								=			model_class(lowerCAmelCase  )
           model.to(lowerCAmelCase  )
           model.eval()

           with torch.no_grad():
                __UpperCamelCase				:  Tuple								=			model(**self._prepare_for_class(lowerCAmelCase				,    lowerCAmelCase  )  )

           __UpperCamelCase				:  Any								=			outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
           __UpperCamelCase				:  Union[str, Any]								=			sum(config.num_block_repeats  ) * 4
           self.assertEqual(len(lowerCAmelCase  )				,    lowerCAmelCase  )

           # EfficientNet'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 // 2, self.model_tester.image_size // 2]				,    )

      __UpperCamelCase				:  Union[str, Any]								=			self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           __UpperCamelCase				:  Union[str, Any]								=			True
           check_hidden_states_output(lowerCAmelCase				,    lowerCAmelCase				,    lowerCAmelCase  )

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           __UpperCamelCase				:  Optional[Any]								=			True

           check_hidden_states_output(lowerCAmelCase				,    lowerCAmelCase				,    lowerCAmelCase  )






 def 				lowerCamelCase__    (				self		:	List[str]  )  ->					List[Any]:







      """simple docstring"""







      __UpperCamelCase				:  Optional[Any]								=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase  )






 @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
 def 				lowerCamelCase__    (				self		:	Optional[int]  )  ->					Dict:







      """simple docstring"""







      pass






 @slow
 def 				lowerCamelCase__    (				self		:	str  )  ->					Any:







      """simple docstring"""







      for model_name in EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           __UpperCamelCase				:  Union[str, Any]								=			EfficientNetModel.from_pretrained(lowerCAmelCase  )
           self.assertIsNotNone(lowerCAmelCase  )






 @is_pipeline_test
 @require_vision
 @slow
 def 				lowerCamelCase__    (				self		:	List[str]  )  ->					Union[str, Any]:







      """simple docstring"""







      super().test_pipeline_image_classification()






def 							A__   ()    ->	Union[str, Any]:
     __UpperCamelCase				:  List[str]								=			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""							)
     return image







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



 """simple docstring"""


 @cached_property
 def 				lowerCamelCase__    (				self		:	List[str]  )  ->					Union[str, Any]:







      """simple docstring"""







      return AutoImageProcessor.from_pretrained("""google/efficientnet-b7"""  ) if is_vision_available() else None






 @slow
 def 				lowerCamelCase__    (				self		:	List[str]  )  ->					Dict:







      """simple docstring"""







      __UpperCamelCase				:  Optional[Any]								=			EfficientNetForImageClassification.from_pretrained("""google/efficientnet-b7"""  ).to(lowerCAmelCase  )

      __UpperCamelCase				:  Any								=			self.default_image_processor
      __UpperCamelCase				:  Any								=			prepare_img()
      __UpperCamelCase				:  Optional[Any]								=			image_processor(images=lowerCAmelCase				,    return_tensors="""pt"""  ).to(lowerCAmelCase  )

      # forward pass
      with torch.no_grad():
           __UpperCamelCase				:  Tuple								=			model(**lowerCAmelCase  )

      # verify the logits
      __UpperCamelCase				:  Optional[int]								=			torch.Size((1, 1000)  )
      self.assertEqual(outputs.logits.shape				,    lowerCAmelCase  )

      __UpperCamelCase				:  int								=			torch.tensor([0.00_01, 0.00_02, 0.00_02]  ).to(lowerCAmelCase  )
      self.assertTrue(torch.allclose(outputs.logits[0, :3]				,    lowerCAmelCase				,    atol=1E-4  )  )












import argparse
import os

import torch
from torchvision.datasets.utils import download_url

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel


a__		    =    {512: '''DiT-XL-2-512x512.pt''', 256: '''DiT-XL-2-256x256.pt'''}





def 							A__   (snake_case      :		List[Any]							)    ->	Dict:
					__UpperCamelCase				:  Optional[int]								=			F'''pretrained_models/{model_name}'''
					if not os.path.isfile(snake_case							):
										os.makedirs("""pretrained_models"""		,  exist_ok=snake_case							)
										__UpperCamelCase				:  Optional[int]								=			F'''https://dl.fbaipublicfiles.com/DiT/models/{model_name}'''
										download_url(snake_case		,  """pretrained_models"""							)
					__UpperCamelCase				:  Union[str, Any]								=			torch.load(snake_case		,  map_location=lambda snake_case		,  snake_case							: storage							)
					return model





def 							A__   (snake_case      :		Optional[Any]							)    ->	Union[str, Any]:
					__UpperCamelCase				:  Any								=			download_model(pretrained_models[args.image_size]							)

					__UpperCamelCase				:  Optional[Any]								=			state_dict["""x_embedder.proj.weight"""]
					__UpperCamelCase				:  List[str]								=			state_dict["""x_embedder.proj.bias"""]
					state_dict.pop("""x_embedder.proj.weight"""							)
					state_dict.pop("""x_embedder.proj.bias"""							)

					for depth in range(28							):
										__UpperCamelCase				:  Tuple								=			state_dict[
										    """t_embedder.mlp.0.weight"""
										]
										__UpperCamelCase				:  Optional[Any]								=			state_dict[
										    """t_embedder.mlp.0.bias"""
										]
										__UpperCamelCase				:  Any								=			state_dict[
										    """t_embedder.mlp.2.weight"""
										]
										__UpperCamelCase				:  Optional[Any]								=			state_dict[
										    """t_embedder.mlp.2.bias"""
										]
										__UpperCamelCase				:  Optional[int]								=			state_dict[
										    """y_embedder.embedding_table.weight"""
										]

										__UpperCamelCase				:  List[str]								=			state_dict[
										    F'''blocks.{depth}.adaLN_modulation.1.weight'''
										]
										__UpperCamelCase				:  Union[str, Any]								=			state_dict[
										    F'''blocks.{depth}.adaLN_modulation.1.bias'''
										]

										__UpperCamelCase			, __UpperCamelCase			, __UpperCamelCase				:  str								=			torch.chunk(state_dict[F'''blocks.{depth}.attn.qkv.weight''']		,  3		,  dim=0							)
										__UpperCamelCase			, __UpperCamelCase			, __UpperCamelCase				:  Any								=			torch.chunk(state_dict[F'''blocks.{depth}.attn.qkv.bias''']		,  3		,  dim=0							)

										__UpperCamelCase				:  Optional[Any]								=			q
										__UpperCamelCase				:  Tuple								=			q_bias
										__UpperCamelCase				:  List[str]								=			k
										__UpperCamelCase				:  Tuple								=			k_bias
										__UpperCamelCase				:  List[Any]								=			v
										__UpperCamelCase				:  Dict								=			v_bias

										__UpperCamelCase				:  Tuple								=			state_dict[
										    F'''blocks.{depth}.attn.proj.weight'''
										]
										__UpperCamelCase				:  Union[str, Any]								=			state_dict[F'''blocks.{depth}.attn.proj.bias''']

										__UpperCamelCase				:  List[Any]								=			state_dict[F'''blocks.{depth}.mlp.fc1.weight''']
										__UpperCamelCase				:  List[str]								=			state_dict[F'''blocks.{depth}.mlp.fc1.bias''']
										__UpperCamelCase				:  Optional[Any]								=			state_dict[F'''blocks.{depth}.mlp.fc2.weight''']
										__UpperCamelCase				:  int								=			state_dict[F'''blocks.{depth}.mlp.fc2.bias''']

										state_dict.pop(F'''blocks.{depth}.attn.qkv.weight'''							)
										state_dict.pop(F'''blocks.{depth}.attn.qkv.bias'''							)
										state_dict.pop(F'''blocks.{depth}.attn.proj.weight'''							)
										state_dict.pop(F'''blocks.{depth}.attn.proj.bias'''							)
										state_dict.pop(F'''blocks.{depth}.mlp.fc1.weight'''							)
										state_dict.pop(F'''blocks.{depth}.mlp.fc1.bias'''							)
										state_dict.pop(F'''blocks.{depth}.mlp.fc2.weight'''							)
										state_dict.pop(F'''blocks.{depth}.mlp.fc2.bias'''							)
										state_dict.pop(F'''blocks.{depth}.adaLN_modulation.1.weight'''							)
										state_dict.pop(F'''blocks.{depth}.adaLN_modulation.1.bias'''							)

					state_dict.pop("""t_embedder.mlp.0.weight"""							)
					state_dict.pop("""t_embedder.mlp.0.bias"""							)
					state_dict.pop("""t_embedder.mlp.2.weight"""							)
					state_dict.pop("""t_embedder.mlp.2.bias"""							)
					state_dict.pop("""y_embedder.embedding_table.weight"""							)

					__UpperCamelCase				:  str								=			state_dict["""final_layer.adaLN_modulation.1.weight"""]
					__UpperCamelCase				:  List[str]								=			state_dict["""final_layer.adaLN_modulation.1.bias"""]
					__UpperCamelCase				:  Union[str, Any]								=			state_dict["""final_layer.linear.weight"""]
					__UpperCamelCase				:  Any								=			state_dict["""final_layer.linear.bias"""]

					state_dict.pop("""final_layer.linear.weight"""							)
					state_dict.pop("""final_layer.linear.bias"""							)
					state_dict.pop("""final_layer.adaLN_modulation.1.weight"""							)
					state_dict.pop("""final_layer.adaLN_modulation.1.bias"""							)

					# DiT XL/2
					__UpperCamelCase				:  Optional[Any]								=			TransformeraDModel(
					    sample_size=args.image_size // 8		,  num_layers=28		,  attention_head_dim=72		,  in_channels=4		,  out_channels=8		,  patch_size=2		,  attention_bias=snake_case		,  num_attention_heads=16		,  activation_fn="""gelu-approximate"""		,  num_embeds_ada_norm=10_00		,  norm_type="""ada_norm_zero"""		,  norm_elementwise_affine=snake_case		,  )
					transformer.load_state_dict(snake_case		,  strict=snake_case							)

					__UpperCamelCase				:  Union[str, Any]								=			DDIMScheduler(
					    num_train_timesteps=10_00		,  beta_schedule="""linear"""		,  prediction_type="""epsilon"""		,  clip_sample=snake_case		,  )

					__UpperCamelCase				:  Optional[int]								=			AutoencoderKL.from_pretrained(args.vae_model							)

					__UpperCamelCase				:  str								=			DiTPipeline(transformer=snake_case		,  vae=snake_case		,  scheduler=snake_case							)

					if args.save:
										pipeline.save_pretrained(args.checkpoint_path							)


if __name__ == "__main__":
	a__		    =    argparse.ArgumentParser()

	parser.add_argument(
	    '''--image_size''',
	    default=256,
	    type=int,
	    required=False,
	    help='''Image size of pretrained model, either 256 or 512.''',
	)
	parser.add_argument(
	    '''--vae_model''',
	    default='''stabilityai/sd-vae-ft-ema''',
	    type=str,
	    required=False,
	    help='''Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.''',
	)
	parser.add_argument(
	    '''--save''', default=True, type=bool, required=False, help='''Whether to save the converted pipeline or not.'''
	)
	parser.add_argument(
	    '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the output pipeline.'''
	)

	a__		    =    parser.parse_args()
	main(args)






