



'''simple docstring'''






import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase:      Union[str, Any]		 =		logging.get_logger(__name__)
def    lowerCamelCase__ (		_A				,    _A=False       ):
  a		:      Optional[int]       =  []

  # fmt: off
  # stem:
  rename_keys.append(('cls_token', 'vit.embeddings.cls_token')       )
  rename_keys.append(('pos_embed', 'vit.embeddings.position_embeddings')       )

  rename_keys.append(('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight')       )
  rename_keys.append(('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias')       )

  # backbone
  rename_keys.append(('patch_embed.backbone.stem.conv.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight')       )
  rename_keys.append(('patch_embed.backbone.stem.norm.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight')       )
  rename_keys.append(('patch_embed.backbone.stem.norm.bias', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias')       )

  for stage_idx in range(len(config.backbone_config.depths       )       ):
    for layer_idx in range(config.backbone_config.depths[stage_idx]       ):
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight""")       )
      rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias""")       )

    rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight""")       )
    rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight""")       )
    rename_keys.append((f"""patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias""", f"""vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias""")       )

  # transformer encoder
  for i in range(config.num_hidden_layers       ):
    # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
    rename_keys.append((f"""blocks.{i}.norm1.weight""", f"""vit.encoder.layer.{i}.layernorm_before.weight""")       )
    rename_keys.append((f"""blocks.{i}.norm1.bias""", f"""vit.encoder.layer.{i}.layernorm_before.bias""")       )
    rename_keys.append((f"""blocks.{i}.attn.proj.weight""", f"""vit.encoder.layer.{i}.attention.output.dense.weight""")       )
    rename_keys.append((f"""blocks.{i}.attn.proj.bias""", f"""vit.encoder.layer.{i}.attention.output.dense.bias""")       )
    rename_keys.append((f"""blocks.{i}.norm2.weight""", f"""vit.encoder.layer.{i}.layernorm_after.weight""")       )
    rename_keys.append((f"""blocks.{i}.norm2.bias""", f"""vit.encoder.layer.{i}.layernorm_after.bias""")       )
    rename_keys.append((f"""blocks.{i}.mlp.fc1.weight""", f"""vit.encoder.layer.{i}.intermediate.dense.weight""")       )
    rename_keys.append((f"""blocks.{i}.mlp.fc1.bias""", f"""vit.encoder.layer.{i}.intermediate.dense.bias""")       )
    rename_keys.append((f"""blocks.{i}.mlp.fc2.weight""", f"""vit.encoder.layer.{i}.output.dense.weight""")       )
    rename_keys.append((f"""blocks.{i}.mlp.fc2.bias""", f"""vit.encoder.layer.{i}.output.dense.bias""")       )

  if base_model:
    # layernorm + pooler
    rename_keys.extend(
        [
            ('norm.weight', 'layernorm.weight'),
            ('norm.bias', 'layernorm.bias'),
            ('pre_logits.fc.weight', 'pooler.dense.weight'),
            ('pre_logits.fc.bias', 'pooler.dense.bias'),
        ]       )

    # if just the base model, we should remove "vit" from all keys that start with "vit"
    a		:      int       =  [(pair[0], pair[1][4:]) if pair[1].startswith('vit'       ) else pair for pair in rename_keys]
  else:
    # layernorm + classification head
    rename_keys.extend(
        [
            ('norm.weight', 'vit.layernorm.weight'),
            ('norm.bias', 'vit.layernorm.bias'),
            ('head.weight', 'classifier.weight'),
            ('head.bias', 'classifier.bias'),
        ]       )
  # fmt: on

  return rename_keys
def    lowerCamelCase__ (		_A				,    _A				,    _A=False       ):
  for i in range(config.num_hidden_layers       ):
    if base_model:
      a		:      Any       =  ''
    else:
      a		:      Optional[Any]       =  'vit.'
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    a		:      str       =  state_dict.pop(f"""blocks.{i}.attn.qkv.weight"""       )
    a		:      List[Any]       =  state_dict.pop(f"""blocks.{i}.attn.qkv.bias"""       )
    # next, add query, keys and values (in that order) to the state dict
    a		:      List[Any]       =  in_proj_weight[
        : config.hidden_size, :
    ]
    a		:      Tuple       =  in_proj_bias[: config.hidden_size]
    a		:      Union[str, Any]       =  in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    a		:      Any       =  in_proj_bias[
        config.hidden_size : config.hidden_size * 2
    ]
    a		:      Dict       =  in_proj_weight[
        -config.hidden_size :, :
    ]
    a		:      int       =  in_proj_bias[-config.hidden_size :]
def    lowerCamelCase__ (		_A       ):
  a		:      str       =  ['head.weight', 'head.bias']
  for k in ignore_keys:
    state_dict.pop(_A				,    _A       )
def    lowerCamelCase__ (		_A				,    _A				,    _A       ):
  a		:      str       =  dct.pop(_A       )
  a		:      str       =  val
def    lowerCamelCase__ (		):
  a		:      Tuple       =  'http://images.cocodataset.org/val2017/000000039769.jpg'
  a		:      Optional[int]       =  Image.open(requests.get(_A				,    stream=_A       ).raw       )
  return im
@torch.no_grad()
def    lowerCamelCase__ (		_A				,    _A				,    _A=False       ):
  a		:      Any       =  BitConfig(
      global_padding='same'				,    layer_type='bottleneck'				,    depths=(3, 4, 9)				,    out_features=['stage3']				,    embedding_dynamic_padding=_A				,    )
  a		:      Tuple       =  ViTHybridConfig(backbone_config=_A				,    image_size=384				,    num_labels=1000       )
  a		:      List[str]       =  False

  # load original model from timm
  a		:      List[Any]       =  timm.create_model(_A				,    pretrained=_A       )
  timm_model.eval()

  # load state_dict of original model, remove and rename some keys
  a		:      List[str]       =  timm_model.state_dict()
  if base_model:
    remove_classification_head_(_A       )
  a		:      int       =  create_rename_keys(_A				,    _A       )
  for src, dest in rename_keys:
    rename_key(_A				,    _A				,    _A       )
  read_in_q_k_v(_A				,    _A				,    _A       )

  a		:      Union[str, Any]       =  'huggingface/label-files'
  a		:      Optional[int]       =  'imagenet-1k-id2label.json'
  a		:      str       =  json.load(open(hf_hub_download(_A				,    _A				,    repo_type='dataset'       )				,    'r'       )       )
  a		:      Dict       =  {int(_A       ): v for k, v in idalabel.items()}
  a		:      Union[str, Any]       =  idalabel
  a		:      Dict       =  {v: k for k, v in idalabel.items()}

  # load HuggingFace model
  if vit_name[-5:] == "in21k":
    a		:      str       =  ViTHybridModel(_A       ).eval()
  else:
    a		:      List[str]       =  ViTHybridForImageClassification(_A       ).eval()
  model.load_state_dict(_A       )

  # create image processor
  a		:      Optional[int]       =  create_transform(**resolve_data_config({}				,    model=_A       )       )
  a		:      Dict       =  transform.transforms

  a		:      Optional[Any]       =  {
      'bilinear': PILImageResampling.BILINEAR,
      'bicubic': PILImageResampling.BICUBIC,
      'nearest': PILImageResampling.NEAREST,
  }

  a		:      Optional[Any]       =  ViTHybridImageProcessor(
      do_resize=_A				,    size={'shortest_edge': timm_transforms[0].size}				,    resample=pillow_resamplings[timm_transforms[0].interpolation.value]				,    do_center_crop=_A				,    crop_size={'height': timm_transforms[1].size[0], 'width': timm_transforms[1].size[1]}				,    do_normalize=_A				,    image_mean=timm_transforms[-1].mean.tolist()				,    image_std=timm_transforms[-1].std.tolist()				,    )

  a		:      List[str]       =  prepare_img()
  a		:      Union[str, Any]       =  transform(_A       ).unsqueeze(0       )
  a		:      Optional[int]       =  processor(_A				,    return_tensors='pt'       ).pixel_values

  # verify pixel values
  assert torch.allclose(_A				,    _A       )

  # verify logits
  with torch.no_grad():
    a		:      Any       =  model(_A       )
    a		:      Any       =  outputs.logits

  print('Predicted class:'				,    logits.argmax(-1       ).item()       )
  if base_model:
    a		:      List[str]       =  timm_model.forward_features(_A       )
    assert timm_pooled_output.shape == outputs.pooler_output.shape
    assert torch.allclose(_A				,    outputs.pooler_output				,    atol=1E-3       )
  else:
    a		:      Union[str, Any]       =  timm_model(_A       )
    assert timm_logits.shape == outputs.logits.shape
    assert torch.allclose(_A				,    outputs.logits				,    atol=1E-3       )
  print('Looks ok!'       )

  if pytorch_dump_folder_path is not None:
    Path(_A       ).mkdir(exist_ok=_A       )
    print(f"""Saving model {vit_name} to {pytorch_dump_folder_path}"""       )
    model.save_pretrained(_A       )
    print(f"""Saving processor to {pytorch_dump_folder_path}"""       )
    processor.save_pretrained(_A       )

  if push_to_hub:
    print(f"""Pushing model and processor to the hub {vit_name}"""       )
    model.push_to_hub(f"""ybelkada/{vit_name}"""       )
    processor.push_to_hub(f"""ybelkada/{vit_name}"""       )


if __name__ == "__main__":
     lowerCAmelCase:      Optional[Any]		 =		argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--vit_name',
         default='vit_base_r50_s16_384',
         type=str,
         help='Name of the hybrid ViT timm model you\'d like to convert.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )
     parser.add_argument(
         '--push_to_hub', action='store_true', help='Whether to upload the model to the HuggingFace hub.'
     )

     lowerCAmelCase:      List[str]		 =		parser.parse_args()
     convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''simple docstring'''






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


lowerCAmelCase:      Any		 =		False

@skip_mps
class 							a__(      lowerCamelCase__	,				lowerCamelCase__	,				lowerCamelCase__	,				unittest.TestCase			):
    lowercase__				=					StableDiffusionAttendAndExcitePipeline
    lowercase__				=					False
    lowercase__				=					TEXT_TO_IMAGE_PARAMS
    lowercase__				=					TEXT_TO_IMAGE_BATCH_PARAMS.union({"""token_indices"""}			)
    lowercase__				=					TEXT_TO_IMAGE_IMAGE_PARAMS
    lowercase__				=					TEXT_TO_IMAGE_IMAGE_PARAMS







    @classmethod
    def     lowercase_ (   cls				:						Union[str, Any]			):
      super().setUpClass()
      torch.use_deterministic_algorithms(__snake_case			)







    @classmethod
    def     lowercase_ (   cls				:						Optional[Any]			):
      super().tearDownClass()
      torch.use_deterministic_algorithms(__snake_case			)







    def     lowercase_ (   self				:						List[Any]			):
      torch.manual_seed(0			)
      a		:      List[Any]       =  UNetaDConditionModel(
          block_out_channels=(32, 64)		,    layers_per_block=1		,    sample_size=32		,    in_channels=4		,    out_channels=4		,    down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,    up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,    cross_attention_dim=32		,    attention_head_dim=(2, 4)		,    use_linear_projection=__snake_case		,    )
      a		:      Optional[int]       =  DDIMScheduler(
          beta_start=0.00085		,    beta_end=0.012		,    beta_schedule='scaled_linear'		,    clip_sample=__snake_case		,    set_alpha_to_one=__snake_case		,    )
      torch.manual_seed(0			)
      a		:      str       =  AutoencoderKL(
          block_out_channels=[32, 64]		,    in_channels=3		,    out_channels=3		,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']		,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']		,    latent_channels=4		,    sample_size=1_28		,    )
      torch.manual_seed(0			)
      a		:      str       =  CLIPTextConfig(
          bos_token_id=0		,    eos_token_id=2		,    hidden_size=32		,    intermediate_size=37		,    layer_norm_eps=1e-0_5		,    num_attention_heads=4		,    num_hidden_layers=5		,    pad_token_id=1		,    vocab_size=10_00		,    hidden_act='gelu'		,    projection_dim=5_12		,    )
      a		:      Any       =  CLIPTextModel(__snake_case			)
      a		:      Tuple       =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

      a		:      List[str]       =  {
          'unet': unet,
          'scheduler': scheduler,
          'vae': vae,
          'text_encoder': text_encoder,
          'tokenizer': tokenizer,
          'safety_checker': None,
          'feature_extractor': None,
      }

      return components







    def     lowercase_ (   self				:						str		,    __snake_case				:						Tuple		,    __snake_case				:						Optional[int]=0			):
      if str(__snake_case			).startswith('mps'			):
        a		:      Any       =  torch.manual_seed(__snake_case			)
      else:
        a		:      Any       =  torch.Generator(device=__snake_case			).manual_seed(__snake_case			)
      a		:      int       =  {
          'prompt': 'a cat and a frog',
          'token_indices': [2, 5],
          'generator': generator,
          'num_inference_steps': 1,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
          'max_iter_to_alter': 2,
          'thresholds': {0: 0.7},
      }
      return inputs







    def     lowercase_ (   self				:						List[Any]			):
      a		:      Union[str, Any]       =  'cpu'

      a		:      Any       =  self.get_dummy_components()
      a		:      List[str]       =  self.pipeline_class(**__snake_case			)
      pipe.to(__snake_case			)
      pipe.set_progress_bar_config(disable=__snake_case			)

      a		:      Any       =  self.get_dummy_inputs(__snake_case			)
      a		:      int       =  pipe(**__snake_case			).images
      a		:      int       =  image[0, -3:, -3:, -1]

      self.assertEqual(image.shape		,    (1, 64, 64, 3)			)
      a		:      str       =  np.array(
          [0.63905364, 0.62897307, 0.48599017, 0.5133624, 0.5550048, 0.45769516, 0.50326973, 0.5023139, 0.45384496]			)
      a		:      int       =  np.abs(image_slice.flatten() - expected_slice			).max()
      self.assertLessEqual(__snake_case		,    1e-3			)







    def     lowercase_ (   self				:						Dict			):
      super().test_cpu_offload_forward_pass(expected_max_diff=5e-4			)







    def     lowercase_ (   self				:						Tuple			):
      # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
      self._test_inference_batch_consistent(batch_sizes=[1, 2]			)







    def     lowercase_ (   self				:						Union[str, Any]			):
      self._test_inference_batch_single_identical(batch_size=2		,    expected_max_diff=7e-4			)







    def     lowercase_ (   self				:						Tuple			):
      super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3			)







    def     lowercase_ (   self				:						str			):
      super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4			)







    def     lowercase_ (   self				:						Any			):
      super().test_save_load_local(expected_max_difference=5e-4			)







    def     lowercase_ (   self				:						List[Any]			):
      super().test_save_load_optional_components(expected_max_difference=4e-4			)

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







    @classmethod
    def     lowercase_ (   cls				:						Union[str, Any]			):
      super().setUpClass()
      torch.use_deterministic_algorithms(__snake_case			)







    @classmethod
    def     lowercase_ (   cls				:						Union[str, Any]			):
      super().tearDownClass()
      torch.use_deterministic_algorithms(__snake_case			)







    def     lowercase_ (   self				:						Union[str, Any]			):
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()







    def     lowercase_ (   self				:						Optional[Any]			):
      a		:      List[Any]       =  torch.manual_seed(51			)

      a		:      Dict       =  StableDiffusionAttendAndExcitePipeline.from_pretrained(
          'CompVis/stable-diffusion-v1-4'		,    safety_checker=__snake_case		,    torch_dtype=torch.floataa			)
      pipe.to('cuda'			)

      a		:      Optional[Any]       =  'a painting of an elephant with glasses'
      a		:      Any       =  [5, 7]

      a		:      Tuple       =  pipe(
          prompt=__snake_case		,    token_indices=__snake_case		,    guidance_scale=7.5		,    generator=__snake_case		,    num_inference_steps=5		,    max_iter_to_alter=5		,    output_type='numpy'		,    ).images[0]

      a		:      str       =  load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy'			)
      assert np.abs((expected_image - image).max()			) < 5e-1