



import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_									=							logging.get_logger(__name__)

UpperCAmelCase_									=							{
    "b0": efficientnet.EfficientNetBa,
    "b1": efficientnet.EfficientNetBa,
    "b2": efficientnet.EfficientNetBa,
    "b3": efficientnet.EfficientNetBa,
    "b4": efficientnet.EfficientNetBa,
    "b5": efficientnet.EfficientNetBa,
    "b6": efficientnet.EfficientNetBa,
    "b7": efficientnet.EfficientNetBa,
}

UpperCAmelCase_									=							{
    "b0": {
        "hidden_dim": 1_2_8_0,
        "width_coef": 1.0,
        "depth_coef": 1.0,
        "image_size": 2_2_4,
        "dropout_rate": 0.2,
        "dw_padding": [],
    },
    "b1": {
        "hidden_dim": 1_2_8_0,
        "width_coef": 1.0,
        "depth_coef": 1.1,
        "image_size": 2_4_0,
        "dropout_rate": 0.2,
        "dw_padding": [1_6],
    },
    "b2": {
        "hidden_dim": 1_4_0_8,
        "width_coef": 1.1,
        "depth_coef": 1.2,
        "image_size": 2_6_0,
        "dropout_rate": 0.3,
        "dw_padding": [5, 8, 1_6],
    },
    "b3": {
        "hidden_dim": 1_5_3_6,
        "width_coef": 1.2,
        "depth_coef": 1.4,
        "image_size": 3_0_0,
        "dropout_rate": 0.3,
        "dw_padding": [5, 1_8],
    },
    "b4": {
        "hidden_dim": 1_7_9_2,
        "width_coef": 1.4,
        "depth_coef": 1.8,
        "image_size": 3_8_0,
        "dropout_rate": 0.4,
        "dw_padding": [6],
    },
    "b5": {
        "hidden_dim": 2_0_4_8,
        "width_coef": 1.6,
        "depth_coef": 2.2,
        "image_size": 4_5_6,
        "dropout_rate": 0.4,
        "dw_padding": [1_3, 2_7],
    },
    "b6": {
        "hidden_dim": 2_3_0_4,
        "width_coef": 1.8,
        "depth_coef": 2.6,
        "image_size": 5_2_8,
        "dropout_rate": 0.5,
        "dw_padding": [3_1],
    },
    "b7": {
        "hidden_dim": 2_5_6_0,
        "width_coef": 2.0,
        "depth_coef": 3.1,
        "image_size": 6_0_0,
        "dropout_rate": 0.5,
        "dw_padding": [1_8],
    },
}
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Optional[Any]				)->Any:
      _lowerCAmelCase 			=  EfficientNetConfig()
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''hidden_dim''']
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''width_coef''']
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''depth_coef''']
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''image_size''']
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''dropout_rate''']
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''dw_padding''']

      _lowerCAmelCase 			=  '''huggingface/label-files'''
      _lowerCAmelCase 			=  '''imagenet-1k-id2label.json'''
      _lowerCAmelCase 			=  1_0_0_0
      _lowerCAmelCase 			=  json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				repo_type='''dataset'''				)     ,				'''r'''				)				)
      _lowerCAmelCase 			=  {int(_SCREAMING_SNAKE_CASE				): v for k, v in idalabel.items()}

      _lowerCAmelCase 			=  idalabel
      _lowerCAmelCase 			=  {v: k for k, v in idalabel.items()}
      return config
def 							UpperCAmelCase__		(				)->str:
      _lowerCAmelCase 			=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
      _lowerCAmelCase 			=  Image.open(requests.get(_SCREAMING_SNAKE_CASE     ,				stream=_SCREAMING_SNAKE_CASE				).raw				)
      return im
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Optional[int]				)->int:
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''image_size''']
      _lowerCAmelCase 			=  EfficientNetImageProcessor(
          size={'''height''': size, '''width''': size}     ,				image_mean=[0.485, 0.456, 0.406]     ,				image_std=[0.4785_3944, 0.473_2864, 0.4743_4163]     ,				do_center_crop=_SCREAMING_SNAKE_CASE     ,				)
      return preprocessor
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Dict				)->str:
      _lowerCAmelCase 			=  [v.split('''_'''				)[0].split('''block'''				)[1] for v in original_param_names if v.startswith('''block'''				)]
      _lowerCAmelCase 			=  sorted(set(_SCREAMING_SNAKE_CASE				)				)
      _lowerCAmelCase 			=  len(_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  {b: str(_SCREAMING_SNAKE_CASE				) for b, i in zip(_SCREAMING_SNAKE_CASE     ,				range(_SCREAMING_SNAKE_CASE				)				)}

      _lowerCAmelCase 			=  []
      rename_keys.append(('''stem_conv/kernel:0''', '''embeddings.convolution.weight''')				)
      rename_keys.append(('''stem_bn/gamma:0''', '''embeddings.batchnorm.weight''')				)
      rename_keys.append(('''stem_bn/beta:0''', '''embeddings.batchnorm.bias''')				)
      rename_keys.append(('''stem_bn/moving_mean:0''', '''embeddings.batchnorm.running_mean''')				)
      rename_keys.append(('''stem_bn/moving_variance:0''', '''embeddings.batchnorm.running_var''')				)

      for b in block_names:
            _lowerCAmelCase 			=  block_name_mapping[b]
            rename_keys.append((f'''block{b}_expand_conv/kernel:0''', f'''encoder.blocks.{hf_b}.expansion.expand_conv.weight''')				)
            rename_keys.append((f'''block{b}_expand_bn/gamma:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.weight''')				)
            rename_keys.append((f'''block{b}_expand_bn/beta:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.bias''')				)
            rename_keys.append(
                (f'''block{b}_expand_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.running_mean''')				)
            rename_keys.append(
                (f'''block{b}_expand_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.running_var''')				)
            rename_keys.append(
                (f'''block{b}_dwconv/depthwise_kernel:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight''')				)
            rename_keys.append((f'''block{b}_bn/gamma:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight''')				)
            rename_keys.append((f'''block{b}_bn/beta:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias''')				)
            rename_keys.append(
                (f'''block{b}_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean''')				)
            rename_keys.append(
                (f'''block{b}_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var''')				)

            rename_keys.append((f'''block{b}_se_reduce/kernel:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.reduce.weight''')				)
            rename_keys.append((f'''block{b}_se_reduce/bias:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.reduce.bias''')				)
            rename_keys.append((f'''block{b}_se_expand/kernel:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.expand.weight''')				)
            rename_keys.append((f'''block{b}_se_expand/bias:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.expand.bias''')				)
            rename_keys.append(
                (f'''block{b}_project_conv/kernel:0''', f'''encoder.blocks.{hf_b}.projection.project_conv.weight''')				)
            rename_keys.append((f'''block{b}_project_bn/gamma:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.weight''')				)
            rename_keys.append((f'''block{b}_project_bn/beta:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.bias''')				)
            rename_keys.append(
                (f'''block{b}_project_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.running_mean''')				)
            rename_keys.append(
                (f'''block{b}_project_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.running_var''')				)

      rename_keys.append(('''top_conv/kernel:0''', '''encoder.top_conv.weight''')				)
      rename_keys.append(('''top_bn/gamma:0''', '''encoder.top_bn.weight''')				)
      rename_keys.append(('''top_bn/beta:0''', '''encoder.top_bn.bias''')				)
      rename_keys.append(('''top_bn/moving_mean:0''', '''encoder.top_bn.running_mean''')				)
      rename_keys.append(('''top_bn/moving_variance:0''', '''encoder.top_bn.running_var''')				)

      _lowerCAmelCase 			=  {}
      for item in rename_keys:
            if item[0] in original_param_names:
                  _lowerCAmelCase 			=  '''efficientnet.''' + item[1]

      _lowerCAmelCase 			=  '''classifier.weight'''
      _lowerCAmelCase 			=  '''classifier.bias'''
      return key_mapping
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Dict     ,				_SCREAMING_SNAKE_CASE   :	Dict     ,				_SCREAMING_SNAKE_CASE   :	Tuple				)->Dict:
      for key, value in tf_params.items():
            if "normalization" in key:
                  continue

            _lowerCAmelCase 			=  key_mapping[key]
            if "_conv" in key and "kernel" in key:
                  _lowerCAmelCase 			=  torch.from_numpy(_SCREAMING_SNAKE_CASE				).permute(3     ,				2     ,				0     ,				1				)
            elif "depthwise_kernel" in key:
                  _lowerCAmelCase 			=  torch.from_numpy(_SCREAMING_SNAKE_CASE				).permute(2     ,				3     ,				0     ,				1				)
            elif "kernel" in key:
                  _lowerCAmelCase 			=  torch.from_numpy(np.transpose(_SCREAMING_SNAKE_CASE				)				)
            else:
                  _lowerCAmelCase 			=  torch.from_numpy(_SCREAMING_SNAKE_CASE				)

            # Replace HF parameters with original TF model parameters
            assert hf_params[hf_key].shape == new_hf_value.shape
            hf_params[hf_key].copy_(_SCREAMING_SNAKE_CASE				)
@torch.no_grad()
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	List[Any]     ,				_SCREAMING_SNAKE_CASE   :	int     ,				_SCREAMING_SNAKE_CASE   :	str     ,				_SCREAMING_SNAKE_CASE   :	List[Any]				)->Optional[Any]:
      _lowerCAmelCase 			=  model_classes[model_name](
          include_top=_SCREAMING_SNAKE_CASE     ,				weights='''imagenet'''     ,				input_tensor=_SCREAMING_SNAKE_CASE     ,				input_shape=_SCREAMING_SNAKE_CASE     ,				pooling=_SCREAMING_SNAKE_CASE     ,				classes=1_0_0_0     ,				classifier_activation='''softmax'''     ,				)

      _lowerCAmelCase 			=  original_model.trainable_variables
      _lowerCAmelCase 			=  original_model.non_trainable_variables
      _lowerCAmelCase 			=  {param.name: param.numpy() for param in tf_params}
      for param in tf_non_train_params:
            _lowerCAmelCase 			=  param.numpy()
      _lowerCAmelCase 			=  list(tf_params.keys()				)

      # Load HuggingFace model
      _lowerCAmelCase 			=  get_efficientnet_config(_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  EfficientNetForImageClassification(_SCREAMING_SNAKE_CASE				).eval()
      _lowerCAmelCase 			=  hf_model.state_dict()

      # Create src-to-dst parameter name mapping dictionary
      print('''Converting parameters...'''				)
      _lowerCAmelCase 			=  rename_keys(_SCREAMING_SNAKE_CASE				)
      replace_params(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE				)

      # Initialize preprocessor and preprocess input image
      _lowerCAmelCase 			=  convert_image_processor(_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  preprocessor(images=prepare_img()     ,				return_tensors='''pt'''				)

      # HF model inference
      hf_model.eval()
      with torch.no_grad():
            _lowerCAmelCase 			=  hf_model(**_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  outputs.logits.detach().numpy()

      # Original model inference
      _lowerCAmelCase 			=  False
      _lowerCAmelCase 			=  CONFIG_MAP[model_name]['''image_size''']
      _lowerCAmelCase 			=  prepare_img().resize((image_size, image_size)     ,				resample=PIL.Image.NEAREST				)
      _lowerCAmelCase 			=  image.img_to_array(_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  np.expand_dims(_SCREAMING_SNAKE_CASE     ,				axis=0				)
      _lowerCAmelCase 			=  original_model.predict(_SCREAMING_SNAKE_CASE				)

      # Check whether original and HF model outputs match  -> np.allclose
      assert np.allclose(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				atol=1e-3				), "The predicted logits are not the same."
      print('''Model outputs match!'''				)

      if save_model:
            # Create folder to save model
            if not os.path.isdir(_SCREAMING_SNAKE_CASE				):
                  os.mkdir(_SCREAMING_SNAKE_CASE				)
            # Save converted model and image processor
            hf_model.save_pretrained(_SCREAMING_SNAKE_CASE				)
            preprocessor.save_pretrained(_SCREAMING_SNAKE_CASE				)

      if push_to_hub:
            # Push model and image processor to hub
            print(f'''Pushing converted {model_name} to the hub...'''				)
            _lowerCAmelCase 			=  f'''efficientnet-{model_name}'''
            preprocessor.push_to_hub(_SCREAMING_SNAKE_CASE				)
            hf_model.push_to_hub(_SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
    UpperCAmelCase_									=							argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--model_name",
        default="b0",
        type=str,
        help="Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default="hf_model",
        type=str,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument("--save_model", action="store_true", help="Save model to local")
    parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")

    UpperCAmelCase_									=							parser.parse_args()
    convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)




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 		UpperCAmelCase						(						snake_case_       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     DiTPipeline
   SCREAMING_SNAKE_CASE__							=     CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
   SCREAMING_SNAKE_CASE__							=     PipelineTesterMixin.required_optional_params - {
       '''latents''',
       '''num_images_per_prompt''',
       '''callback''',
       '''callback_steps''',
   }
   SCREAMING_SNAKE_CASE__							=     CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
   SCREAMING_SNAKE_CASE__							=     False





   def 	__lowerCAmelCase							(	self			):
         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=_lowerCAmelCase       ,     activation_fn='''gelu-approximate'''       ,     num_embeds_ada_norm=1_000       ,     norm_type='''ada_norm_zero'''       ,     norm_elementwise_affine=_lowerCAmelCase       ,     )
         _lowerCAmelCase 			=  AutoencoderKL()
         _lowerCAmelCase 			=  DDIMScheduler()
         _lowerCAmelCase 			=  {'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
         return components





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase=0			):
         if str(_lowerCAmelCase			).startswith('''mps'''			):
               _lowerCAmelCase 			=  torch.manual_seed(_lowerCAmelCase			)
         else:
               _lowerCAmelCase 			=  torch.Generator(device=_lowerCAmelCase			).manual_seed(_lowerCAmelCase			)
         _lowerCAmelCase 			=  {
             '''class_labels''': [1],
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''output_type''': '''numpy''',
         }
         return inputs





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''cpu'''

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

         _lowerCAmelCase 			=  self.get_dummy_inputs(_lowerCAmelCase			)
         _lowerCAmelCase 			=  pipe(**_lowerCAmelCase			).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(_lowerCAmelCase       ,     1E-3			)





   def 	__lowerCAmelCase							(	self			):
         self._test_inference_batch_single_identical(relax_max_difference=_lowerCAmelCase       ,     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			):
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3			)






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





   def 	__lowerCAmelCase							(	self			):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





   def 	__lowerCAmelCase							(	self			):
         _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(_lowerCAmelCase			)

         _lowerCAmelCase 			=  pipe(_lowerCAmelCase       ,     generator=_lowerCAmelCase       ,     num_inference_steps=40       ,     output_type='''np'''			).images

         for word, image in zip(_lowerCAmelCase       ,     _lowerCAmelCase			):
               _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			):
         _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(_lowerCAmelCase			)

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  pipe(_lowerCAmelCase       ,     generator=_lowerCAmelCase       ,     num_inference_steps=25       ,     output_type='''np'''			).images

         for word, image in zip(_lowerCAmelCase       ,     _lowerCAmelCase			):
               _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