





'''simple docstring'''




import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

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,
    assert_mean_pixel_difference,
)


enable_full_determinism()
class 			_lowerCAmelCase		(			A__							,     A__							,     A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   StableUnCLIPPipeline
     snake_case_						      =   TEXT_TO_IMAGE_PARAMS
     snake_case_						      =   TEXT_TO_IMAGE_BATCH_PARAMS
     snake_case_						      =   TEXT_TO_IMAGE_IMAGE_PARAMS
     snake_case_						      =   TEXT_TO_IMAGE_IMAGE_PARAMS

     # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
     snake_case_						      =   False



     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[str]:
        snake_case 						=   32
        snake_case 						=   embedder_hidden_size

        # prior components

        torch.manual_seed(0						)
        snake_case 						=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

        torch.manual_seed(0						)
        snake_case 						=   CLIPTextModelWithProjection(
            CLIPTextConfig(
                bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=__snake_case   ,       projection_dim=__snake_case   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       pad_token_id=1   ,       vocab_size=10_00   ,       )						)

        torch.manual_seed(0						)
        snake_case 						=   PriorTransformer(
            num_attention_heads=2   ,       attention_head_dim=12   ,       embedding_dim=__snake_case   ,       num_layers=1   ,       )

        torch.manual_seed(0						)
        snake_case 						=   DDPMScheduler(
            variance_type="""fixed_small_log"""   ,       prediction_type="""sample"""   ,       num_train_timesteps=10_00   ,       clip_sample=__snake_case   ,       clip_sample_range=5.0   ,       beta_schedule="""squaredcos_cap_v2"""   ,       )

        # regular denoising components

        torch.manual_seed(0						)
        snake_case 						=   StableUnCLIPImageNormalizer(embedding_dim=__snake_case						)
        snake_case 						=   DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""						)

        torch.manual_seed(0						)
        snake_case 						=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

        torch.manual_seed(0						)
        snake_case 						=   CLIPTextModel(
            CLIPTextConfig(
                bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=__snake_case   ,       projection_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       pad_token_id=1   ,       vocab_size=10_00   ,       )						)

        torch.manual_seed(0						)
        snake_case 						=   UNetaDConditionModel(
            sample_size=32   ,       in_channels=4   ,       out_channels=4   ,       down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")   ,       up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")   ,       block_out_channels=(32, 64)   ,       attention_head_dim=(2, 4)   ,       class_embed_type="""projection"""   ,       projection_class_embeddings_input_dim=embedder_projection_dim * 2   ,       cross_attention_dim=__snake_case   ,       layers_per_block=1   ,       upcast_attention=__snake_case   ,       use_linear_projection=__snake_case   ,       )

        torch.manual_seed(0						)
        snake_case 						=   DDIMScheduler(
            beta_schedule="""scaled_linear"""   ,       beta_start=0.0_00_85   ,       beta_end=0.0_12   ,       prediction_type="""v_prediction"""   ,       set_alpha_to_one=__snake_case   ,       steps_offset=1   ,       )

        torch.manual_seed(0						)
        snake_case 						=   AutoencoderKL()

        snake_case 						=   {
            # prior components
            """prior_tokenizer""": prior_tokenizer,
            """prior_text_encoder""": prior_text_encoder,
            """prior""": prior,
            """prior_scheduler""": prior_scheduler,
            # image noising components
            """image_normalizer""": image_normalizer,
            """image_noising_scheduler""": image_noising_scheduler,
            # regular denoising components
            """tokenizer""": tokenizer,
            """text_encoder""": text_encoder,
            """unet""": unet,
            """scheduler""": scheduler,
            """vae""": vae,
        }

        return components



     def      lowerCAmelCase			(   self							:			Union[str, Any]   ,       __snake_case							:			int   ,       __snake_case							:			List[Any]=0						)->	List[Any]:
        if str(__snake_case						).startswith("""mps"""						):
           snake_case 						=   torch.manual_seed(__snake_case						)
        else:
           snake_case 						=   torch.Generator(device=__snake_case						).manual_seed(__snake_case						)
        snake_case 						=   {
            """prompt""": """A painting of a squirrel eating a burger""",
            """generator""": generator,
            """num_inference_steps""": 2,
            """prior_num_inference_steps""": 2,
            """output_type""": """numpy""",
        }
        return inputs



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Optional[int]:
        snake_case 						=   torch_device == """cpu"""

        self._test_attention_slicing_forward_pass(test_max_difference=__snake_case						)



     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[Any]:
        snake_case 						=   torch_device in ["""cpu""", """mps"""]

        self._test_inference_batch_single_identical(test_max_difference=__snake_case						)




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






     """simple docstring"""






     def      lowerCAmelCase			(   self							:			Dict						)->	Tuple:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



     def      lowerCAmelCase			(   self							:			Optional[int]						)->	Optional[Any]:
        snake_case 						=   load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""						)

        snake_case 						=   StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""   ,       torch_dtype=torch.floataa						)
        pipe.to(__snake_case						)
        pipe.set_progress_bar_config(disable=__snake_case						)
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        snake_case 						=   torch.Generator(device="""cpu"""						).manual_seed(0						)
        snake_case 						=   pipe("""anime turle"""   ,       generator=__snake_case   ,       output_type="""np"""						)

        snake_case 						=   output.images[0]

        assert image.shape == (7_68, 7_68, 3)

        assert_mean_pixel_difference(__snake_case   ,       __snake_case						)



     def      lowerCAmelCase			(   self							:			int						)->	str:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        snake_case 						=   StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""   ,       torch_dtype=torch.floataa						)
        snake_case 						=   pipe.to(__snake_case						)
        pipe.set_progress_bar_config(disable=__snake_case						)
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        snake_case 						=   pipe(
            """anime turtle"""   ,       prior_num_inference_steps=2   ,       num_inference_steps=2   ,       output_type="""np"""   ,       )

        snake_case 						=   torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 10**9








'''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 transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


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





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[int]      ,    __lowerCAmelCase       :							List[Any]=False		)		->  List[str]:
   snake_case 						=   []
   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''')		)

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ("""cls_token""", """vit.embeddings.cls_token"""),
           ("""patch_embed.proj.weight""", """vit.embeddings.patch_embeddings.projection.weight"""),
           ("""patch_embed.proj.bias""", """vit.embeddings.patch_embeddings.projection.bias"""),
           ("""pos_embed""", """vit.embeddings.position_embeddings"""),
       ]		)

   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"
      snake_case 						=   [(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"""),
          ]		)

   return rename_keys





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Dict      ,    __lowerCAmelCase       :							Optional[Any]      ,    __lowerCAmelCase       :							int=False		)		->  Union[str, Any]:
   for i in range(config.num_hidden_layers		):
      if base_model:
         snake_case 						=   """"""
      else:
         snake_case 						=   """vit."""
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      snake_case 						=   state_dict.pop(F'''blocks.{i}.attn.qkv.weight'''		)
      snake_case 						=   state_dict.pop(F'''blocks.{i}.attn.qkv.bias'''		)
      # next, add query, keys and values (in that order) to the state dict
      snake_case 						=   in_proj_weight[
          : config.hidden_size, :
      ]
      snake_case 						=   in_proj_bias[: config.hidden_size]
      snake_case 						=   in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      snake_case 						=   in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      snake_case 						=   in_proj_weight[
          -config.hidden_size :, :
      ]
      snake_case 						=   in_proj_bias[-config.hidden_size :]





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]		)		->  List[str]:
   snake_case 						=   ["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(__lowerCAmelCase      ,    __lowerCAmelCase		)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[Any]      ,    __lowerCAmelCase       :							int      ,    __lowerCAmelCase       :							Tuple		)		->  int:
   snake_case 						=   dct.pop(__lowerCAmelCase		)
   snake_case 						=   val





def 	__lowerCamelCase			(      )		->  Optional[Any]:
   snake_case 						=   """http://images.cocodataset.org/val2017/000000039769.jpg"""
   snake_case 						=   Image.open(requests.get(__lowerCAmelCase      ,    stream=__lowerCAmelCase		).raw		)
   return im





@torch.no_grad()
def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]      ,    __lowerCAmelCase       :							Any		)		->  Dict:
   snake_case 						=   ViTConfig()
   snake_case 						=   False
   # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
   if vit_name[-5:] == "in21k":
      snake_case 						=   True
      snake_case 						=   int(vit_name[-12:-10]		)
      snake_case 						=   int(vit_name[-9:-6]		)
   else:
      snake_case 						=   10_00
      snake_case 						=   """huggingface/label-files"""
      snake_case 						=   """imagenet-1k-id2label.json"""
      snake_case 						=   json.load(open(hf_hub_download(__lowerCAmelCase      ,    __lowerCAmelCase      ,    repo_type="""dataset"""		)      ,    """r"""		)		)
      snake_case 						=   {int(__lowerCAmelCase		): v for k, v in idalabel.items()}
      snake_case 						=   idalabel
      snake_case 						=   {v: k for k, v in idalabel.items()}
      snake_case 						=   int(vit_name[-6:-4]		)
      snake_case 						=   int(vit_name[-3:]		)
   # size of the architecture
   if "deit" in vit_name:
      if vit_name[9:].startswith("""tiny"""		):
         snake_case 						=   1_92
         snake_case 						=   7_68
         snake_case 						=   12
         snake_case 						=   3
      elif vit_name[9:].startswith("""small"""		):
         snake_case 						=   3_84
         snake_case 						=   15_36
         snake_case 						=   12
         snake_case 						=   6
      else:
         pass
   else:
      if vit_name[4:].startswith("""small"""		):
         snake_case 						=   7_68
         snake_case 						=   23_04
         snake_case 						=   8
         snake_case 						=   8
      elif vit_name[4:].startswith("""base"""		):
         pass
      elif vit_name[4:].startswith("""large"""		):
         snake_case 						=   10_24
         snake_case 						=   40_96
         snake_case 						=   24
         snake_case 						=   16
      elif vit_name[4:].startswith("""huge"""		):
         snake_case 						=   12_80
         snake_case 						=   51_20
         snake_case 						=   32
         snake_case 						=   16

    # load original model from timm
   snake_case 						=   timm.create_model(__lowerCAmelCase      ,    pretrained=__lowerCAmelCase		)
   timm_model.eval()

   # load state_dict of original model, remove and rename some keys
   snake_case 						=   timm_model.state_dict()
   if base_model:
      remove_classification_head_(__lowerCAmelCase		)
   snake_case 						=   create_rename_keys(__lowerCAmelCase      ,    __lowerCAmelCase		)
   for src, dest in rename_keys:
      rename_key(__lowerCAmelCase      ,    __lowerCAmelCase      ,    __lowerCAmelCase		)
   read_in_q_k_v(__lowerCAmelCase      ,    __lowerCAmelCase      ,    __lowerCAmelCase		)

   # load HuggingFace model
   if vit_name[-5:] == "in21k":
      snake_case 						=   ViTModel(__lowerCAmelCase		).eval()
   else:
      snake_case 						=   ViTForImageClassification(__lowerCAmelCase		).eval()
   model.load_state_dict(__lowerCAmelCase		)

   # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
   if "deit" in vit_name:
      snake_case 						=   DeiTImageProcessor(size=config.image_size		)
   else:
      snake_case 						=   ViTImageProcessor(size=config.image_size		)
   snake_case 						=   image_processor(images=prepare_img()      ,    return_tensors="""pt"""		)
   snake_case 						=   encoding["""pixel_values"""]
   snake_case 						=   model(__lowerCAmelCase		)

   if base_model:
      snake_case 						=   timm_model.forward_features(__lowerCAmelCase		)
      assert timm_pooled_output.shape == outputs.pooler_output.shape
      assert torch.allclose(__lowerCAmelCase      ,    outputs.pooler_output      ,    atol=1e-3		)
   else:
      snake_case 						=   timm_model(__lowerCAmelCase		)
      assert timm_logits.shape == outputs.logits.shape
      assert torch.allclose(__lowerCAmelCase      ,    outputs.logits      ,    atol=1e-3		)

   Path(__lowerCAmelCase		).mkdir(exist_ok=__lowerCAmelCase		)
   print(F'''Saving model {vit_name} to {pytorch_dump_folder_path}'''		)
   model.save_pretrained(__lowerCAmelCase		)
   print(F'''Saving image processor to {pytorch_dump_folder_path}'''		)
   image_processor.save_pretrained(__lowerCAmelCase		)


if __name__ == "__main__":
 _SCREAMING_SNAKE_CASE							     =			argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--vit_name",
     default="vit_base_patch16_224",
     type=str,
     help="Name of the 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."
 )

 _SCREAMING_SNAKE_CASE							     =			parser.parse_args()
 convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)

