

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()
__UpperCAmelCase								=			logging.get_logger(__name__)





def      snake_case_				(__A					:						List[Any] ,       __A					:						Dict=False					)      ->     Dict:
   __lowerCAmelCase  :      str    =						[]

   # 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"
      __lowerCAmelCase  :      str    =						[(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      snake_case_				(__A					:						Optional[int] ,       __A					:						str ,       __A					:						Any=False					)      ->     Union[str, Any]:





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





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





def      snake_case_				(__A					:						List[Any] ,       __A					:						int ,       __A					:						Any					)      ->     Union[str, Any]:
   __lowerCAmelCase  :      Optional[int]    =						dct.pop(SCREAMING_SNAKE_CASE__					)
   __lowerCAmelCase  :      Dict    =						val





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





@torch.no_grad()
def      snake_case_				(__A					:						Any ,       __A					:						int ,       __A					:						Dict=False					)      ->     Tuple:
   __lowerCAmelCase  :      Dict    =						BitConfig(
       global_padding="""same""" ,       layer_type="""bottleneck""" ,       depths=(3, 4, 9) ,       out_features=["""stage3"""] ,       embedding_dynamic_padding=SCREAMING_SNAKE_CASE__ ,       )
   __lowerCAmelCase  :      Optional[int]    =						ViTHybridConfig(backbone_config=SCREAMING_SNAKE_CASE__ ,       image_size=3_8_4 ,       num_labels=1_0_0_0					)
   __lowerCAmelCase  :      Tuple    =						False

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

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

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

   # load HuggingFace model
   if vit_name[-5:] == "in21k":
      __lowerCAmelCase  :      Union[str, Any]    =						ViTHybridModel(SCREAMING_SNAKE_CASE__					).eval()
   else:
      __lowerCAmelCase  :      Any    =						ViTHybridForImageClassification(SCREAMING_SNAKE_CASE__					).eval()
   model.load_state_dict(SCREAMING_SNAKE_CASE__					)

   # create image processor
   __lowerCAmelCase  :      int    =						create_transform(**resolve_data_config({} ,       model=SCREAMING_SNAKE_CASE__					)					)
   __lowerCAmelCase  :      Union[str, Any]    =						transform.transforms

   __lowerCAmelCase  :      List[str]    =						{
       """bilinear""": PILImageResampling.BILINEAR,
       """bicubic""": PILImageResampling.BICUBIC,
       """nearest""": PILImageResampling.NEAREST,
   }

   __lowerCAmelCase  :      Union[str, Any]    =						ViTHybridImageProcessor(
       do_resize=SCREAMING_SNAKE_CASE__ ,       size={"""shortest_edge""": timm_transforms[0].size} ,       resample=pillow_resamplings[timm_transforms[0].interpolation.value] ,       do_center_crop=SCREAMING_SNAKE_CASE__ ,       crop_size={"""height""": timm_transforms[1].size[0], """width""": timm_transforms[1].size[1]} ,       do_normalize=SCREAMING_SNAKE_CASE__ ,       image_mean=timm_transforms[-1].mean.tolist() ,       image_std=timm_transforms[-1].std.tolist() ,       )

   __lowerCAmelCase  :      int    =						prepare_img()
   __lowerCAmelCase  :      Dict    =						transform(SCREAMING_SNAKE_CASE__					).unsqueeze(0					)
   __lowerCAmelCase  :      str    =						processor(SCREAMING_SNAKE_CASE__ ,       return_tensors="""pt"""					).pixel_values

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

   # verify logits
   with torch.no_grad():
      __lowerCAmelCase  :      str    =						model(SCREAMING_SNAKE_CASE__					)
      __lowerCAmelCase  :      int    =						outputs.logits

   print("""Predicted class:""" ,       logits.argmax(-1					).item()					)
   if base_model:
      __lowerCAmelCase  :      int    =						timm_model.forward_features(SCREAMING_SNAKE_CASE__					)
      assert timm_pooled_output.shape == outputs.pooler_output.shape
      assert torch.allclose(SCREAMING_SNAKE_CASE__ ,       outputs.pooler_output ,       atol=1e-3					)
   else:
      __lowerCAmelCase  :      Tuple    =						timm_model(SCREAMING_SNAKE_CASE__					)
      assert timm_logits.shape == outputs.logits.shape
      assert torch.allclose(SCREAMING_SNAKE_CASE__ ,       outputs.logits ,       atol=1e-3					)
   print("""Looks ok!"""					)

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

   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__":
   __UpperCAmelCase								=			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."""
   )

   __UpperCAmelCase								=			parser.parse_args()
   convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)





import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()



class SCREAMING_SNAKE_CASE     (  a_       ,							a_       ,							unittest.TestCase			):
  """simple docstring"""






  lowerCamelCase   :      Tuple										=AutoencoderKL
  lowerCamelCase   :      Tuple										="sample"
  lowerCamelCase   :      Dict										=1e-2







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      str    =						4
     __lowerCAmelCase  :      Dict    =						3
     __lowerCAmelCase  :      Optional[Any]    =						(32, 32)

     __lowerCAmelCase  :      Union[str, Any]    =						floats_tensor((batch_size, num_channels) + sizes   ).to(lowerCAmelCase   )

     return {"sample": image}







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any   )							->     int:






     """simple docstring"""







     return (3, 32, 32)







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   List[str]   )							->     str:






     """simple docstring"""







     return (3, 32, 32)







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






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						{
         """block_out_channels""": [32, 64],
         """in_channels""": 3,
         """out_channels""": 3,
         """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
         """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
         """latent_channels""": 4,
     }
     __lowerCAmelCase  :      Optional[int]    =						self.dummy_input
     return init_dict, inputs_dict







  def 	SCREAMING_SNAKE_CASE   (      self			:   int   )							->     str:






     """simple docstring"""







     pass







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     Union[str, Any]:






     """simple docstring"""







     pass







  @unittest.skipIf(torch_device == """mps"""	,							"""Gradient checkpointing skipped on MPS"""   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase   ,__lowerCAmelCase  :      str    =						self.prepare_init_args_and_inputs_for_common()
     __lowerCAmelCase  :      Dict    =						self.model_class(**lowerCAmelCase   )
     model.to(lowerCAmelCase   )

     assert not model.is_gradient_checkpointing and model.training

     __lowerCAmelCase  :      str    =						model(**lowerCAmelCase   ).sample
     # run the backwards pass on the model. For backwards pass, for simplicity purpose,
     # we won't calculate the loss and rather backprop on out.sum()
     model.zero_grad()

     __lowerCAmelCase  :      Any    =						torch.randn_like(lowerCAmelCase   )
     __lowerCAmelCase  :      str    =						(out - labels).mean()
     loss.backward()

     # re-instantiate the model now enabling gradient checkpointing
     __lowerCAmelCase  :      List[str]    =						self.model_class(**lowerCAmelCase   )
     # clone model
     model_a.load_state_dict(model.state_dict()   )
     model_a.to(lowerCAmelCase   )
     model_a.enable_gradient_checkpointing()

     assert model_a.is_gradient_checkpointing and model_a.training

     __lowerCAmelCase  :      Any    =						model_a(**lowerCAmelCase   ).sample
     # run the backwards pass on the model. For backwards pass, for simplicity purpose,
     # we won't calculate the loss and rather backprop on out.sum()
     model_a.zero_grad()
     __lowerCAmelCase  :      Dict    =						(out_a - labels).mean()
     loss_a.backward()

     # compare the output and parameters gradients
     self.assertTrue((loss - loss_a).abs() < 1e-5   )
     __lowerCAmelCase  :      int    =						dict(model.named_parameters()   )
     __lowerCAmelCase  :      Union[str, Any]    =						dict(model_a.named_parameters()   )
     for name, param in named_params.items():
        self.assertTrue(torch_all_close(param.grad.data	,							named_params_a[name].grad.data	,							atol=5e-5   )   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   str   )							->     List[str]:






     """simple docstring"""







     __lowerCAmelCase   ,__lowerCAmelCase  :      List[Any]    =						AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""	,							output_loading_info=lowerCAmelCase   )
     self.assertIsNotNone(lowerCAmelCase   )
     self.assertEqual(len(loading_info["""missing_keys"""]   )	,							0   )

     model.to(lowerCAmelCase   )
     __lowerCAmelCase  :      int    =						model(**self.dummy_input   )

     assert image is not None, "Make sure output is not None"







  def 	SCREAMING_SNAKE_CASE   (      self			:   Any   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""   )
     __lowerCAmelCase  :      Optional[Any]    =						model.to(lowerCAmelCase   )
     model.eval()

     if torch_device == "mps":
        __lowerCAmelCase  :      List[Any]    =						torch.manual_seed(0   )
     else:
        __lowerCAmelCase  :      Any    =						torch.Generator(device=lowerCAmelCase   ).manual_seed(0   )

     __lowerCAmelCase  :      Optional[int]    =						torch.randn(
         1	,							model.config.in_channels	,							model.config.sample_size	,							model.config.sample_size	,							generator=torch.manual_seed(0   )	,							)
     __lowerCAmelCase  :      Optional[int]    =						image.to(lowerCAmelCase   )
     with torch.no_grad():
        __lowerCAmelCase  :      Union[str, Any]    =						model(lowerCAmelCase	,							sample_posterior=lowerCAmelCase	,							generator=lowerCAmelCase   ).sample

     __lowerCAmelCase  :      Dict    =						output[0, -1, -3:, -3:].flatten().cpu()

     # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
     # the expected output slices are not the same for CPU and GPU.
     if torch_device == "mps":
        __lowerCAmelCase  :      List[str]    =						torch.tensor(
            [
                -4.0078e-01,
                -3.8323e-04,
                -1.2681e-01,
                -1.1462e-01,
                2.0095e-01,
                1.0893e-01,
                -8.8247e-02,
                -3.0361e-01,
                -9.8644e-03,
            ]   )
     elif torch_device == "cpu":
        __lowerCAmelCase  :      Union[str, Any]    =						torch.tensor(
            [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]   )
     else:
        __lowerCAmelCase  :      Tuple    =						torch.tensor(
            [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]   )

     self.assertTrue(torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							rtol=1e-2   )   )






@slow
class SCREAMING_SNAKE_CASE     (  unittest.TestCase			):
  """simple docstring"""






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






     """simple docstring"""







     return f'''gaussian_noise_s={seed}_shape={'_'.join([str(lowerCAmelCase   ) for s in shape]   )}.npy'''







  def 	SCREAMING_SNAKE_CASE   (      self			:   int   )							->     Union[str, Any]:






     """simple docstring"""







     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







  def 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   Union[str, Any]=0	,							lowerCAmelCase			:   Any=(4, 3, 5_12, 5_12)	,							lowerCAmelCase			:   Any=False   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						torch.floataa if fpaa else torch.floataa
     __lowerCAmelCase  :      Optional[int]    =						torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCAmelCase	,							lowerCAmelCase   )   )   ).to(lowerCAmelCase   ).to(lowerCAmelCase   )
     return image







  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]	,							lowerCAmelCase			:   Optional[Any]="CompVis/stable-diffusion-v1-4"	,							lowerCAmelCase			:   int=False   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						"""fp16""" if fpaa else None
     __lowerCAmelCase  :      List[str]    =						torch.floataa if fpaa else torch.floataa

     __lowerCAmelCase  :      Dict    =						AutoencoderKL.from_pretrained(
         lowerCAmelCase	,							subfolder="""vae"""	,							torch_dtype=lowerCAmelCase	,							revision=lowerCAmelCase	,							)
     model.to(lowerCAmelCase   ).eval()

     return model







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple	,							lowerCAmelCase			:   Tuple=0   )							->     Tuple:






     """simple docstring"""







     if torch_device == "mps":
        return torch.manual_seed(lowerCAmelCase   )
     return torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )







  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
          [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
          # fmt: on
      ]   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   Dict   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[int]    =						self.get_sd_image(lowerCAmelCase   )
     __lowerCAmelCase  :      List[str]    =						self.get_generator(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[Any]    =						model(lowerCAmelCase	,							generator=lowerCAmelCase	,							sample_posterior=lowerCAmelCase   ).sample

     assert sample.shape == image.shape

     __lowerCAmelCase  :      Any    =						sample[-1, -2:, -2:, :2].flatten().float().cpu()
     __lowerCAmelCase  :      List[str]    =						torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice   )

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=3e-3   )







  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
          [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
          # fmt: on
      ]   )
  @require_torch_gpu
  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Optional[int]   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						self.get_sd_vae_model(fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      Tuple    =						self.get_sd_image(lowerCAmelCase	,							fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      Optional[int]    =						self.get_generator(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Dict    =						model(lowerCAmelCase	,							generator=lowerCAmelCase	,							sample_posterior=lowerCAmelCase   ).sample

     assert sample.shape == image.shape

     __lowerCAmelCase  :      List[str]    =						sample[-1, -2:, :2, -2:].flatten().float().cpu()
     __lowerCAmelCase  :      Optional[int]    =						torch.tensor(lowerCAmelCase   )

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=1e-2   )







  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
          [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
          # fmt: on
      ]   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   Any	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Any   )							->     str:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[int]    =						self.get_sd_image(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      List[Any]    =						model(lowerCAmelCase   ).sample

     assert sample.shape == image.shape

     __lowerCAmelCase  :      Union[str, Any]    =						sample[-1, -2:, -2:, :2].flatten().float().cpu()
     __lowerCAmelCase  :      str    =						torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice   )

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=3e-3   )







  @parameterized.expand(
      [
          # fmt: off
          [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
          [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
          # fmt: on
      ]   )
  @require_torch_gpu
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   int   )							->     str:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[Any]    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[Any]    =						model.decode(lowerCAmelCase   ).sample

     assert list(sample.shape   ) == [3, 3, 5_12, 5_12]

     __lowerCAmelCase  :      Union[str, Any]    =						sample[-1, -2:, :2, -2:].flatten().cpu()
     __lowerCAmelCase  :      Tuple    =						torch.tensor(lowerCAmelCase   )

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=1e-3   )







  @parameterized.expand(
      [
          # fmt: off
          [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
          [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
          # fmt: on
      ]   )
  @require_torch_gpu
  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Any   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						self.get_sd_vae_model(fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      str    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)	,							fpaa=lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Dict    =						model.decode(lowerCAmelCase   ).sample

     assert list(sample.shape   ) == [3, 3, 5_12, 5_12]

     __lowerCAmelCase  :      Any    =						sample[-1, -2:, :2, -2:].flatten().float().cpu()
     __lowerCAmelCase  :      Union[str, Any]    =						torch.tensor(lowerCAmelCase   )

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=5e-3   )







  @parameterized.expand([(13,), (16,), (27,)]   )
  @require_torch_gpu
  @unittest.skipIf(not is_xformers_available()	,							reason="""xformers is not required when using PyTorch 2.0."""   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]	,							lowerCAmelCase			:   Any   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						self.get_sd_vae_model(fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      Union[str, Any]    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)	,							fpaa=lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Union[str, Any]    =						model.decode(lowerCAmelCase   ).sample

     model.enable_xformers_memory_efficient_attention()
     with torch.no_grad():
        __lowerCAmelCase  :      int    =						model.decode(lowerCAmelCase   ).sample

     assert list(sample.shape   ) == [3, 3, 5_12, 5_12]

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=1e-1   )







  @parameterized.expand([(13,), (16,), (37,)]   )
  @require_torch_gpu
  @unittest.skipIf(not is_xformers_available()	,							reason="""xformers is not required when using PyTorch 2.0."""   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict	,							lowerCAmelCase			:   List[Any]   )							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[int]    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[Any]    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[Any]    =						model.decode(lowerCAmelCase   ).sample

     model.enable_xformers_memory_efficient_attention()
     with torch.no_grad():
        __lowerCAmelCase  :      Tuple    =						model.decode(lowerCAmelCase   ).sample

     assert list(sample.shape   ) == [3, 3, 5_12, 5_12]

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=1e-2   )







  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
          [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
          # fmt: on
      ]   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   List[Any]   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      List[str]    =						self.get_sd_image(lowerCAmelCase   )
     __lowerCAmelCase  :      Any    =						self.get_generator(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[int]    =						model.encode(lowerCAmelCase   ).latent_dist
        __lowerCAmelCase  :      Union[str, Any]    =						dist.sample(generator=lowerCAmelCase   )

     assert list(sample.shape   ) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

     __lowerCAmelCase  :      Any    =						sample[0, -1, -3:, -3:].flatten().cpu()
     __lowerCAmelCase  :      int    =						torch.tensor(lowerCAmelCase   )

     __lowerCAmelCase  :      str    =						3e-3 if torch_device != """mps""" else 1e-2
     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=lowerCAmelCase   )


